Improve BQ typed support

build.sbt

@@ -428,6 +428,10 @@ ThisBuild / mimaBinaryIssueFilters ++= Seq(
   // added BQ Json object
   ProblemFilters.exclude[MissingTypesProblem](
     "com.spotify.scio.bigquery.types.package$Json$"
+  ),
+  // dropped custom BigQueryAvroUtilsWrapper
+  ProblemFilters.exclude[MissingClassProblem](
+    "org.apache.beam.sdk.io.gcp.bigquery.BigQueryAvroUtilsWrapper"
   )
 )
 

integration/src/test/scala/com/spotify/scio/bigquery/BigQueryIOIT.scala

@@ -17,47 +17,151 @@
 
 package com.spotify.scio.bigquery
 
+import com.spotify.scio.bigquery.BigQueryTypedTable.Format
+import com.spotify.scio.coders.Coder
+import com.spotify.scio.avro._
 import org.apache.beam.sdk.options._
 import com.spotify.scio.testing._
 import com.spotify.scio.testing.util.ItUtils
+import org.apache.avro.generic.GenericRecord
 
 object BigQueryIOIT {
   @BigQueryType.fromTable("bigquery-public-data:samples.shakespeare")
   class ShakespeareFromTable
 
   @BigQueryType.fromQuery(
-    """
-    SELECT word, word_count FROM `bigquery-public-data.samples.shakespeare` LIMIT 10
-  """
+    """SELECT word, word_count
+      |FROM `bigquery-public-data.samples.shakespeare`
+      |WHERE corpus = 'kinglear'
+      |ORDER BY word_count DESC
+      |LIMIT 5""".stripMargin
   )
   class ShakespeareFromQuery
-
-  val tempLocation: String = ItUtils.gcpTempLocation("bigquery-it")
 }
 
 class BigQueryIOIT extends PipelineSpec {
   import BigQueryIOIT._
   import ItUtils.project
 
+  val tempLocation: String = ItUtils.gcpTempLocation("bigquery-it")
   val options: PipelineOptions = PipelineOptionsFactory
     .fromArgs(s"--project=$project", s"--tempLocation=$tempLocation")
     .create()
 
-  "Select" should "read typed values from a SQL query" in
+  val kinglearTop5: Seq[(String, Long)] = Seq(
+    "the" -> 786L,
+    "I" -> 622L,
+    "and" -> 594L,
+    "of" -> 447L,
+    "to" -> 438L
+  )
+
+  def extractCorpus(r: TableRow): String =
+    r.get("corpus").asInstanceOf[String]
+
+  def extractCorpus(r: GenericRecord): String =
+    r.get("corpus").asInstanceOf[String]
+
+  def extractWordCount(r: TableRow): (String, Long) = {
+    val word = r.get("word").asInstanceOf[String]
+    val count = r.get("word_count").asInstanceOf[String].toLong
+    word -> count
+  }
+
+  def extractWordCount(r: GenericRecord): (String, Long) = {
+    val word = r.get("word").asInstanceOf[String]
+    val count = r.get("word_count").asInstanceOf[String].toLong
+    word -> count
+  }
+
+  "Select" should "read values from a SQL query" in {
+    runWithRealContext(options) { sc =>
+      val query = Query(ShakespeareFromQuery.queryRaw)
+      val scoll = sc.bigQuerySelect(query).map(extractWordCount)
+
+      scoll should containInAnyOrder(kinglearTop5)
+    }
+  }
+
+  it should "read storage values from a SQL query" in {
     runWithRealContext(options) { sc =>
-      val scoll = sc.read(BigQueryTyped[ShakespeareFromQuery])
-      scoll should haveSize(10)
-      scoll should satisfy[ShakespeareFromQuery] {
-        _.forall(_.getClass == classOf[ShakespeareFromQuery])
-      }
+      val query = Query(ShakespeareFromQuery.queryRaw)
+      val scoll = sc
+        .bigQueryStorage(query)
+        .map(extractWordCount)
+
+      scoll should containInAnyOrder(kinglearTop5)
     }
+  }
 
-  "TableRef" should "read typed values from table" in
+  it should "read typed values from a SQL query" in {
     runWithRealContext(options) { sc =>
-      val scoll = sc.read(BigQueryTyped[ShakespeareFromTable])
-      scoll.take(10) should haveSize(10)
-      scoll should satisfy[ShakespeareFromTable] {
-        _.forall(_.getClass == classOf[ShakespeareFromTable])
-      }
+      val scoll = sc
+        .typedBigQuery[ShakespeareFromQuery]()
+        .flatMap { r =>
+          for {
+            w <- r.word
+            c <- r.word_count
+          } yield w -> c
+        }
+
+      scoll should containInAnyOrder(kinglearTop5)
+    }
+  }
+
+  "Table" should "read values from table" in {
+    runWithRealContext(options) { sc =>
+      val table = Table.Spec(ShakespeareFromTable.table)
+      val scoll = sc
+        .bigQueryTable(table)
+        .filter(r => extractCorpus(r) == "kinglear")
+        .map(extractWordCount)
+        .top(5)(Ordering.by(_._2))
+        .flatten
+
+      scoll should containInAnyOrder(kinglearTop5)
+    }
+  }
+
+  it should "read avro values from table" in {
+    runWithRealContext(options) { sc =>
+      // BQ limitation: We can't give an avro reader schema
+      implicit val coder: Coder[GenericRecord] = avroGenericRecordCoder
+      val table = Table.Spec(ShakespeareFromTable.table)
+      val scoll = sc
+        .bigQueryTable[GenericRecord](table, Format.GenericRecord)
+        .filter(r => extractCorpus(r) == "kinglear")
+        .map(extractWordCount)
+        .top(5)(Ordering.by(_._2))
+        .flatten
+
+      scoll should containInAnyOrder(kinglearTop5)
+    }
+  }
+
+  it should "read storage values from table" in {
+    runWithRealContext(options) { sc =>
+      val table = Table.Spec(ShakespeareFromTable.table)
+      val scoll = sc
+        .bigQueryStorage(table)
+        .filter(r => extractCorpus(r) == "kinglear")
+        .map(extractWordCount)
+        .top(5)(Ordering.by(_._2))
+        .flatten
+
+      scoll should containInAnyOrder(kinglearTop5)
+    }
+  }
+
+  it should "read typed values from table" in {
+    runWithRealContext(options) { sc =>
+      val scoll = sc
+        .typedBigQuery[ShakespeareFromTable]()
+        .collect { case r if r.corpus == "kinglear" => r.word -> r.word_count }
+        .top(5)(Ordering.by(_._2))
+        .flatten
+
+      scoll should containInAnyOrder(kinglearTop5)
     }
+  }
 }

integration/src/test/scala/com/spotify/scio/bigquery/TypedBigQueryIT.scala

@@ -48,7 +48,8 @@ object TypedBigQueryIT {
     timestamp: Instant,
     date: LocalDate,
     time: LocalTime,
-    datetime: LocalDateTime,
+    // BQ DATETIME is problematic with avro: export as 'string(datetime)', load as '(long)local-timestamp-micros'
+    // datetime: LocalDateTime,
     geography: Geography,
     json: Json,
     bigNumeric: BigNumeric
@@ -116,20 +117,32 @@ object TypedBigQueryIT {
 class TypedBigQueryIT extends PipelineSpec with BeforeAndAfterAll {
   import TypedBigQueryIT._
 
+  private val bq = BigQuery.defaultInstance()
+
   override protected def afterAll(): Unit = {
-    val bq = BigQuery.defaultInstance()
     // best effort cleanup
     Try(bq.tables.delete(typedTable.ref))
     Try(bq.tables.delete(tableRowTable.ref))
     Try(bq.tables.delete(avroTable.ref))
   }
 
+  def waitForTable(table: Table.Spec): Unit = {
+    var retries = 0
+    while (!bq.tables.exists(table.ref) && retries < 3) {
+      Thread.sleep(500)
+      retries += 1
+    }
+    if (retries >= 3) throw new RuntimeException(s"Table $table not found")
+  }
+
   "TypedBigQuery" should "handle records as TableRow" in {
     runWithRealContext(options) { sc =>
       sc.parallelize(records)
         .saveAsTypedBigQueryTable(typedTable, createDisposition = CREATE_IF_NEEDED)
     }.waitUntilFinish()
 
+    waitForTable(typedTable)
+
     runWithRealContext(options) { sc =>
       val data = sc.typedBigQuery[Record](typedTable)
       data should containInAnyOrder(records)
@@ -147,15 +160,17 @@ class TypedBigQueryIT extends PipelineSpec with BeforeAndAfterAll {
         )
     }.waitUntilFinish()
 
+    waitForTable(tableRowTable)
+
     runWithRealContext(options) { sc =>
       val data = sc.bigQueryTable(tableRowTable).map(Record.fromTableRow)
       data should containInAnyOrder(records)
     }
   }
 
-  // TODO fix if in beam 2.61
-  ignore should "handle records as avro format" in {
+  it should "handle records as avro format" in {
     implicit val coder: Coder[GenericRecord] = avroGenericRecordCoder(Record.avroSchema)
+
     runWithRealContext(options) { sc =>
       sc.parallelize(records)
         .map(Record.toAvro)
@@ -166,8 +181,11 @@ class TypedBigQueryIT extends PipelineSpec with BeforeAndAfterAll {
         )
     }.waitUntilFinish()
 
+    waitForTable(avroTable)
+
     runWithRealContext(options) { sc =>
-      val data = sc.bigQueryTable(avroTable, Format.GenericRecord).map(Record.fromAvro)
+      val data =
+        sc.bigQueryTable(avroTable, Format.GenericRecordWithLogicalTypes).map(Record.fromAvro)
       data should containInAnyOrder(records)
     }
   }