update test

Co-authored-by: Thomas Chow <[email protected]>

Author: tchow-zlai

cloud_gcp/BUILD.bazel

@@ -35,7 +35,7 @@ shared_deps = [
     maven_artifact("org.threeten:threetenbp"),
     maven_artifact("org.apache.kafka:kafka-clients"),
     maven_artifact("com.google.cloud.spark:spark-3.5-bigquery"),
-    scala_artifact_with_suffix("com.google.cloud.spark:spark-bigquery"),
+    # scala_artifact_with_suffix("com.google.cloud.spark:spark-bigquery"),
     # scala_artifact_with_suffix("com.google.cloud.spark:spark-bigquery-with-dependencies"),
     scala_artifact_with_suffix("org.apache.iceberg:iceberg-spark-runtime-3.5"),
 ]

cloud_gcp/src/main/resources/META-INF/services/org.apache.spark.sql.sources.DataSourceRegister

@@ -1,2 +1 @@
 com.google.cloud.spark.bigquery.v2.Spark35BigQueryTableProvider
-com.google.cloud.spark.bigquery.BigQueryRelationProvider

cloud_gcp/src/main/scala/ai/chronon/integrations/cloud_gcp/BigQueryFormat.scala

@@ -20,7 +20,7 @@ case class BigQueryFormat(project: String, bqClient: BigQuery, override val opti
     extends Format {
   override def name: String = "bigquery"
 
-  private val bqFormat = "com.google.cloud.spark.bigquery.BigQueryRelationProvider"
+  private val bqFormat = "com.google.cloud.spark.bigquery.v2.Spark35BigQueryTableProvider"
 
   override def alterTableProperties(tableName: String,
                                     tableProperties: Map[String, String]): (String => Unit) => Unit = {

cloud_gcp/src/main/scala/ai/chronon/integrations/cloud_gcp/GcpFormatProvider.scala

@@ -63,7 +63,7 @@ case class GcpFormatProvider(sparkSession: SparkSession) extends FormatProvider
       case _ => throw new IllegalStateException(s"Cannot support table of type: ${table.getFriendlyName}")
     }
 
-  private def format(tableName: String): scala.Option[Format] = {
+  private[cloud_gcp] def format(tableName: String): scala.Option[Format] = {
 
     val btTableIdentifier: TableId = BigQueryUtil.parseTableId(tableName)
     val table = scala.Option(bigQueryClient.getTable(btTableIdentifier.getDataset, btTableIdentifier.getTable))

cloud_gcp/src/test/scala/ai/chronon/integrations/cloud_gcp/BigQueryCatalogTest.scala

@@ -98,16 +98,13 @@ class BigQueryCatalogTest extends AnyFlatSpec with MockitoSugar {
 
   it should "integration testing bigquery partitions" ignore {
     // TODO(tchow): This test is ignored because it requires a running instance of the bigquery. Need to figure out stubbing locally.
-    // to run this:
-    //    1. Set up a tunnel to dataproc federation proxy:
-    //       gcloud compute ssh zipline-canary-cluster-m \
-    //        --zone us-central1-c \
-    //        -- -f -N -L 9083:localhost:9083
-    //    2. enable this test and off you go.
     val externalPartitions = tableUtils.partitions("data.checkouts_parquet")
     println(externalPartitions)
     val nativePartitions = tableUtils.partitions("data.sample_native")
     println(nativePartitions)
+    val tblFormat = GcpFormatProvider(spark).format("data.purchases").get
+    val partitions = tblFormat.partitions("data.purchases")(spark)
+    assertEquals(partitions.flatMap(_.keys), Seq("ds"))
   }
 
   it should "kryo serialization for ResolvingFileIO" in {