rebase

Author: tchow-zlai

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

@@ -1,6 +1,61 @@
 package ai.chronon.integrations.cloud_gcp.test
 
+import ai.chronon.integrations.cloud_gcp.GcpFormatProvider
+import ai.chronon.integrations.cloud_gcp.BQuery
+import ai.chronon.spark.{SparkSessionBuilder, TableUtils}
 import org.scalatest.funsuite.AnyFunSuite
 import org.scalatestplus.mockito.MockitoSugar
 
-class BigQueryCatalogTest extends AnyFunSuite with MockitoSugar {}
+class BigQueryCatalogTest extends AnyFunSuite with MockitoSugar {
+
+  lazy val spark: SparkSession = SparkSessionBuilder.build(
+    "BigQuerySparkTest",
+    local = true,
+    additionalConfig = Some(
+      Map(
+        "spark.chronon.table.format_provider.class" -> classOf[GcpFormatProvider].getName,
+        "hive.metastore.uris" -> "thrift://localhost:9083",
+        "spark.chronon.partition.column" -> "c",
+        "spark.hadoop.fs.gs.impl" -> classOf[GoogleHadoopFileSystem].getName,
+        "spark.hadoop.fs.AbstractFileSystem.gs.impl" -> classOf[GoogleHadoopFS].getName,
+        "spark.hadoop.google.cloud.auth.service.account.enable" -> true.toString,
+        "spark.hadoop.fs.gs.impl" -> classOf[GoogleHadoopFileSystem].getName
+      ))
+  )
+  lazy val tableUtils: TableUtils = TableUtils(spark)
+
+  test("hive uris are set") {
+    assertEquals("thrift://localhost:9083", spark.sqlContext.getConf("hive.metastore.uris"))
+  }
+
+  test("verify dynamic classloading of GCP providers") {
+
+    assertTrue(tableUtils.tableReadFormat("data.sample_native") match {
+      case BQuery(_) => true
+      case _         => false
+    })
+  }
+
+  test("integration testing bigquery load table") {
+    val externalTable = "data.checkouts_parquet"
+    val table = tableUtils.loadTable(externalTable)
+    val partitioned = tableUtils.isPartitioned(externalTable)
+    val database = tableUtils.createDatabase("test_database")
+    val allParts = tableUtils.allPartitions(externalTable)
+    table.show
+  }
+
+  ignore("integration testing bigquery partitions") {
+    // 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)
+  }
+}