format and tests

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

Author: tchow-zlai

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

@@ -82,7 +82,8 @@ class DelegatingBigQueryMetastoreCatalog extends CatalogExtension {
 
   override def listTables(namespace: Array[String]): Array[Identifier] = icebergCatalog.listTables(namespace)
 
-  override def loadTable(ident: Identifier): Table = {
+  override def loadTable(rawIdent: Identifier): Table = {
+    val ident = Identifier.of(rawIdent.namespace.flatMap(_.split("\\.")), rawIdent.name)
     Try { icebergCatalog.loadTable(ident) }
       .recover {
         case _ => {
@@ -113,7 +114,8 @@ class DelegatingBigQueryMetastoreCatalog extends CatalogExtension {
                               Map(TableCatalog.PROP_EXTERNAL -> "true", TableCatalog.PROP_LOCATION -> uri))
             }
             case _: StandardTableDefinition => {
-              val connectorTable = connectorCatalog.loadTable(ident)
+              // Hack because there's a bug in the BigQueryCatalog where they ignore the projectId.
+              val connectorTable = connectorCatalog.loadTable(Identifier.of(Array(tId.getDataset), tId.getTable))
               DelegatingTable(connectorTable, Map(TableCatalog.PROP_EXTERNAL -> "false"))
             }
             case _ => throw new IllegalStateException(s"Cannot support table of type: ${table.getFriendlyName}")

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

@@ -28,7 +28,7 @@ class BigQueryCatalogTest extends AnyFlatSpec with MockitoSugar {
       Map(
         "spark.chronon.table.format_provider.class" -> classOf[GcpFormatProvider].getName,
         "hive.metastore.uris" -> "thrift://localhost:9083",
-        "spark.chronon.partition.column" -> "c",
+        "spark.chronon.partition.column" -> "ds",
         "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,
@@ -98,13 +98,14 @@ 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.
-    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(Set("ds"), partitions.flatMap(_.keys).toSet)
+    val externalPartitions = tableUtils.partitions("data.checkouts_parquet_partitioned")
+     assertEquals(Seq("2023-11-30"), externalPartitions)
+    val nativePartitions = tableUtils.partitions("data.purchases")
+    assertEquals(Set(20231118, 20231122, 20231125, 20231102, 20231123, 20231119, 20231130, 20231101, 20231117, 20231110, 20231108, 20231112, 20231115, 20231116, 20231113, 20231104, 20231103, 20231106, 20231121, 20231124, 20231128, 20231109, 20231127, 20231129, 20231126, 20231114, 20231107, 20231111, 20231120, 20231105).map(_.toString), nativePartitions.toSet)
+
+    val df = tableUtils.loadTable("`canary-443022.data`.purchases")
+    df.show
+
   }
 
   it should "kryo serialization for ResolvingFileIO" in {