com.emarsys.rdb.connector.bigquery.BigQuerySimpleSelect.scala Maven / Gradle / Ivy

Go to download

Show more of this group Show more artifacts with this name
Show all versions of rdb-connector-bigquery_2.12 Show documentation

rdb-connector-bigquery

The newest version!

package com.emarsys.rdb.connector.bigquery

import akka.NotUsed
import akka.stream.scaladsl.Source
import cats.data.EitherT
import com.emarsys.rdb.connector.common.ConnectorResponse
import com.emarsys.rdb.connector.common.defaults.SqlWriter._
import com.emarsys.rdb.connector.common.models.SimpleSelect

import scala.concurrent.duration.FiniteDuration

trait BigQuerySimpleSelect {
  self: BigQueryConnector with BigQueryMetadata =>

  import cats.instances.future._

  override def simpleSelect(
      select: SimpleSelect,
      timeout: FiniteDuration
  ): ConnectorResponse[Source[Seq[String], NotUsed]] = {
    EitherT(listFields(select.table.t)).map { fields =>
      val writer = BigQueryWriter(config, fields)
      import writer._

      bigQueryClient.streamingQuery(select.toSql).completionTimeout(timeout).recoverWithRetries(1, streamErrorHandler)
    }.value
  }

  override protected def runSelectWithGroupLimit(
      select: SimpleSelect,
      groupLimit: Int,
      references: Seq[String],
      timeout: FiniteDuration
  ): ConnectorResponse[Source[Seq[String], NotUsed]] = {
    EitherT(listFields(select.table.t)).map { fields =>
      val writer = BigQueryWriter(config, fields)
      import writer._

      bigQueryClient
        .streamingQuery(select.toSql(simpleSelectWithGroupLimitWriter(references, groupLimit)))
        .completionTimeout(timeout)
        .map(_.dropRight(1))
        .recoverWithRetries(1, streamErrorHandler)
    }.value
  }
}