bq client native

Author: tchow-zlai

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

@@ -12,21 +12,28 @@ import com.google.cloud.bigquery.ExternalTableDefinition
 import com.google.cloud.bigquery.StandardTableDefinition
 
 case class GCPFormatProvider(sparkSession: SparkSession) extends FormatProvider {
-  lazy val bigQueryClient = BigQueryOptions.getDefaultInstance.getService
+  lazy val bqOptions = BigQueryOptions.getDefaultInstance
+  lazy val bigQueryClient = bqOptions.getService
   def readFormat(tableName: String): Format = {
 
     val btTableIdentifier: TableId = BigQueryUtil.parseTableId(tableName)
+
+    // Order of Precedence for Default Project
+    // Explicitly configured project in code (e.g., setProjectId()).
+    // GOOGLE_CLOUD_PROJECT environment variable.
+    // project_id from the ADC service account JSON file.
+    // Active project in the gcloud CLI configuration.
+    // No default project: An error will occur if no project ID is available.
+
     val unshadedTI: BTableId =
-      BTableId.of(btTableIdentifier.getProject, btTableIdentifier.getDataset, btTableIdentifier.getTable)
+      BTableId.of(bqOptions.getProjectId, btTableIdentifier.getDataset, btTableIdentifier.getTable)
 
     val tableOpt = Option(bigQueryClient.getTable(unshadedTI))
-
     tableOpt match {
       case Some(table) => {
-        table.getDefinition match {
-          case ExternalTableDefinition => BQuery(unshadedTI.getProject)
-          case StandardTableDefinition => GCS(unshadedTI.getProject)
-        }
+        if (table.getDefinition.isInstanceOf[ExternalTableDefinition]) GCS(unshadedTI.getProject)
+        else if (table.getDefinition.isInstanceOf[StandardTableDefinition]) BQuery(unshadedTI.getProject)
+        else throw new IllegalStateException(s"Cannot support table of type: ${table.getDefinition}")
       }
       case None => Hive
     }

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

@@ -48,7 +48,7 @@ class BigQueryCatalogTest extends AnyFunSuite with MockitoSugar {
     table.show
   }
 
-  ignore("integration testing bigquery partitions") {
+  test("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: