fix: properly detect bigquery catalog (#629)

## Summary

- Now that we have configurable catalogs, we should rely on that to
determine whether we're interacting with BigQuery.

## Checklist
- [ ] Added Unit Tests
- [ ] Covered by existing CI
- [ ] Integration tested
- [ ] Documentation update
<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit

## Summary by CodeRabbit

- **New Features**
- Enhanced cloud integration logic to accurately determine data table
catalogs, ensuring more reliable data format retrieval.
- Improved error handling, providing clearer feedback when encountering
issues with data format detection.
- Introduced a new method for converting table names into identifiers,
enhancing catalog detection functionality.
- Added a method for retrieving the catalog name based on table names,
improving the overall functionality.

- **Bug Fixes**
	- Corrected a typo in the error message for the `NoSuchTableException`.

- **Tests**
- Added a test case to verify the functionality of catalog detection in
the `GcpFormatProvider` class, increasing test coverage.
- Introduced a new test case for validating catalog detection based on
various input strings in the `TableUtilsTest` class.
- Enhanced test coverage for format detection scenarios in the
`BigQueryCatalogTest` class.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->

<!-- av pr metadata
This information is embedded by the av CLI when creating PRs to track
the status of stacks when using Aviator. Please do not delete or edit
this section of the PR.
```
{"parent":"main","parentHead":"","trunk":"main"}
```
-->

---------

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

Author: tchow-zlai

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

@@ -21,6 +21,7 @@ import org.apache.spark.sql.execution.datasources.v2.parquet.ParquetTable
 import org.apache.spark.sql.types.StructType
 import org.apache.spark.sql.util.CaseInsensitiveStringMap
 import org.apache.spark.sql.catalyst.analysis.NoSuchTableException
+import org.apache.iceberg.gcp.bigquery.BigQueryMetastoreCatalog
 
 import java.util
 import scala.jdk.CollectionConverters._
@@ -89,6 +90,8 @@ class DelegatingBigQueryMetastoreCatalog extends TableCatalog with SupportsNames
   private var catalogName: String =
     null // This corresponds to `spark_catalog in `spark.sql.catalog.spark_catalog`. This is necessary for spark to correctly choose which implementation to use.
 
+  private var catalogProps: Map[String, String] = Map.empty[String, String]
+
   override def listNamespaces: Array[Array[String]] = icebergCatalog.listNamespaces()
 
   override def listNamespaces(namespace: Array[String]): Array[Array[String]] = icebergCatalog.listNamespaces(namespace)
@@ -114,13 +117,27 @@ class DelegatingBigQueryMetastoreCatalog extends TableCatalog with SupportsNames
   override def listTables(namespace: Array[String]): Array[Identifier] = icebergCatalog.listTables(namespace)
 
   override def loadTable(rawIdent: Identifier): Table = {
-    val ident = Identifier.of(rawIdent.namespace.flatMap(_.split("\\.")), rawIdent.name)
-    Try { icebergCatalog.loadTable(ident) }
+    // Remove the catalog segment. We've already consumed it, now it's time to figure out the namespace.
+    val identNoCatalog = Identifier.of(
+      rawIdent.namespace.flatMap(_.split("\\.")).toList match {
+        case catalog :: namespace :: Nil => Array(namespace)
+        case namespace :: Nil            => Array(namespace)
+      },
+      rawIdent.name
+    )
+    Try {
+      val icebergSparkTable = icebergCatalog.loadTable(identNoCatalog)
+      DelegatingTable(icebergSparkTable,
+                      additionalProperties =
+                        Map(TableCatalog.PROP_EXTERNAL -> "false", TableCatalog.PROP_PROVIDER -> "ICEBERG"))
+    }
       .recover {
         case _ => {
-          val tId = ident.namespace().toList match {
-            case database :: Nil            => TableId.of(database, ident.name())
-            case project :: database :: Nil => TableId.of(project, database, ident.name())
+          val project =
+            catalogProps.getOrElse(BigQueryMetastoreCatalog.PROPERTIES_KEY_GCP_PROJECT, bqOptions.getProjectId)
+          val tId = identNoCatalog.namespace().toList match {
+            case database :: Nil            => TableId.of(project, database, identNoCatalog.name())
+            case catalog :: database :: Nil => TableId.of(project, database, identNoCatalog.name())
             case Nil =>
               throw new IllegalArgumentException(s"Table identifier namespace ${rawIdent} must have at least one part.")
           }
@@ -143,7 +160,9 @@ class DelegatingBigQueryMetastoreCatalog extends TableCatalog with SupportsNames
                                                 None,
                                                 classOf[ParquetFileFormat])
               DelegatingTable(fileBasedTable,
-                              Map(TableCatalog.PROP_EXTERNAL -> "true", TableCatalog.PROP_LOCATION -> uri))
+                              Map(TableCatalog.PROP_EXTERNAL -> "true",
+                                  TableCatalog.PROP_LOCATION -> uri,
+                                  TableCatalog.PROP_PROVIDER -> "PARQUET"))
             }
             case _: StandardTableDefinition => {
               //todo(tchow): Support partitioning
@@ -153,13 +172,14 @@ class DelegatingBigQueryMetastoreCatalog extends TableCatalog with SupportsNames
               val connectorTable = connectorCatalog.loadTable(Identifier.of(Array(tId.getDataset), tId.getTable))
               // ideally it should be the below:
               // val connectorTable = connectorCatalog.loadTable(ident)
-              DelegatingTable(connectorTable, Map(TableCatalog.PROP_EXTERNAL -> "false"))
+              DelegatingTable(connectorTable,
+                              Map(TableCatalog.PROP_EXTERNAL -> "false", TableCatalog.PROP_PROVIDER -> "BIGQUERY"))
             }
             case _ => throw new IllegalStateException(s"Cannot support table of type: ${table.getFriendlyName}")
           }
         }
       }
-      .getOrElse(throw new NoSuchTableException(f"Tgable: ${ident} not found in bigquery catalog."))
+      .getOrElse(throw new NoSuchTableException(f"Table: ${identNoCatalog} not found in bigquery catalog."))
   }
 
   override def createTable(ident: Identifier,
@@ -187,6 +207,7 @@ class DelegatingBigQueryMetastoreCatalog extends TableCatalog with SupportsNames
     icebergCatalog.initialize(name, options)
     connectorCatalog.initialize(name, options)
     catalogName = name
+    catalogProps = options.asCaseSensitiveMap.asScala.toMap
   }
 
   override def name(): String = catalogName

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

@@ -1,13 +1,13 @@
 package ai.chronon.integrations.cloud_gcp
 import ai.chronon.spark.format.{DefaultFormatProvider, Format, Iceberg}
 import com.google.cloud.bigquery._
-import com.google.cloud.iceberg.bigquery.relocated.com.google.api.services.bigquery.model.TableReference
-import org.apache.iceberg.exceptions.NoSuchIcebergTableException
-import org.apache.iceberg.gcp.bigquery.{BigQueryClient, BigQueryClientImpl}
+import org.apache.iceberg.gcp.bigquery.BigQueryMetastoreCatalog
+import org.apache.iceberg.spark.SparkCatalog
 import org.apache.spark.sql.SparkSession
+import org.apache.spark.sql.connector.catalog.TableCatalog
 
 import scala.jdk.CollectionConverters._
-import scala.util.{Failure, Success, Try}
+import scala.util.Try
 
 class GcpFormatProvider(override val sparkSession: SparkSession) extends DefaultFormatProvider(sparkSession) {
 
@@ -18,46 +18,29 @@ class GcpFormatProvider(override val sparkSession: SparkSession) extends Default
     * - Active project in the gcloud CLI configuration.
     * - No default project: An error will occur if no project ID is available.
     */
-  private lazy val bqOptions = BigQueryOptions.getDefaultInstance
-  private lazy val bigQueryClient: BigQuery = bqOptions.getService
-  private lazy val icebergClient: BigQueryClient = new BigQueryClientImpl()
 
   override def readFormat(tableName: String): scala.Option[Format] = {
-    logger.info(s"Retrieving read format for table: ${tableName}")
-
-    // order is important here. we want the Hive case where we just check for table in catalog to be last
-    Try {
-      val btTableIdentifier = SparkBQUtils.toTableId(tableName)(sparkSession)
-      val bqTable = bigQueryClient.getTable(btTableIdentifier)
-      getFormat(bqTable)
-    } match {
-      case Success(format) => scala.Option(format)
-      case Failure(e) =>
-        logger.info(s"${tableName} is not a BigQuery table")
-        super.readFormat(tableName)
-    }
-  }
-
-  private[cloud_gcp] def getFormat(table: Table): Format = {
-    table.getDefinition.asInstanceOf[TableDefinition] match {
-      case _: ExternalTableDefinition =>
+    val parsedCatalog = getCatalog(tableName)
+    val identifier = SparkBQUtils.toIdentifier(tableName)(sparkSession)
+    val cat = sparkSession.sessionState.catalogManager.catalog(parsedCatalog)
+    cat match {
+      case delegating: DelegatingBigQueryMetastoreCatalog =>
         Try {
-          val tableRef = new TableReference()
-            .setProjectId(table.getTableId.getProject)
-            .setDatasetId(table.getTableId.getDataset)
-            .setTableId(table.getTableId.getTable)
-
-          icebergClient.getTable(tableRef) // Just try to load it. It'll fail if it's not an iceberg table.
-          Iceberg
-        }.recover {
-          case _: NoSuchIcebergTableException => BigQueryExternal
-          case e: Exception                   => throw e
-        }.get
-
-      case _: StandardTableDefinition => BigQueryNative
-
-      case _ =>
-        throw new IllegalStateException(s"Cannot support table of type: ${table.getFriendlyName}")
+          delegating
+            .loadTable(identifier)
+            .properties
+            .asScala
+            .getOrElse(TableCatalog.PROP_PROVIDER, "")
+            .toUpperCase match {
+            case "ICEBERG"   => Iceberg
+            case "BIGQUERY"  => BigQueryNative
+            case "PARQUET"   => BigQueryExternal
+            case unsupported => throw new IllegalStateException(s"Unsupported provider type: ${unsupported}")
+          }
+        }.toOption
+      case iceberg: SparkCatalog if (iceberg.icebergCatalog().isInstanceOf[BigQueryMetastoreCatalog]) =>
+        scala.Option(Iceberg)
+      case _ => super.readFormat(tableName)
     }
   }
 }

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

@@ -2,6 +2,7 @@ package ai.chronon.integrations.cloud_gcp
 import com.google.cloud.bigquery.connector.common.BigQueryUtil
 import org.apache.spark.sql.SparkSession
 import com.google.cloud.bigquery.TableId
+import org.apache.spark.sql.connector.catalog.Identifier
 
 object SparkBQUtils {
 
@@ -14,4 +15,10 @@ object SparkBQUtils {
       .getOrElse(TableId.of(shadedTid.getDataset, shadedTid.getTable))
   }
 
+  def toIdentifier(tableName: String)(implicit spark: SparkSession): Identifier = {
+    val parseIdentifier = spark.sessionState.sqlParser.parseMultipartIdentifier(tableName).reverse
+    Identifier.of(parseIdentifier.tail.reverse.toArray, parseIdentifier.head)
+
+  }
+
 }

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

@@ -10,6 +10,8 @@ import com.google.cloud.hadoop.fs.gcs.{
   GoogleHadoopFileSystemConfiguration,
   HadoopConfigurationProperty
 }
+import ai.chronon.spark.format.Iceberg
+
 import com.google.cloud.spark.bigquery.SparkBigQueryUtil
 import org.apache.iceberg.gcp.bigquery.{BigQueryMetastoreCatalog => BQMSCatalog}
 import org.apache.iceberg.gcp.gcs.GCSFileIO
@@ -113,6 +115,37 @@ class BigQueryCatalogTest extends AnyFlatSpec with MockitoSugar {
     println(allParts)
   }
 
+  it should "integration testing formats" ignore {
+    val externalTable = "default_iceberg.data.checkouts_parquet"
+    val externalFormat = FormatProvider.from(spark).readFormat(externalTable)
+    assertEquals(Some(BigQueryExternal), externalFormat)
+
+    val externalTableNoCat = "data.checkouts_parquet"
+    val externalFormatNoCat = FormatProvider.from(spark).readFormat(externalTableNoCat)
+    assertEquals(Some(BigQueryExternal), externalFormatNoCat)
+
+    val nativeTable = "default_iceberg.data.checkouts_native"
+    val nativeFormat = FormatProvider.from(spark).readFormat(nativeTable)
+    assertEquals(Some(BigQueryNative), nativeFormat)
+
+    val nativeTableNoCat = "data.checkouts_native"
+    val nativeFormatNoCat = FormatProvider.from(spark).readFormat(nativeTableNoCat)
+    assertEquals(Some(BigQueryNative), nativeFormatNoCat)
+
+    val icebergTable = "default_iceberg.data.quickstart_purchases_davidhan_v1_dev_davidhan"
+    val icebergFormat = FormatProvider.from(spark).readFormat(icebergTable)
+    assertEquals(Some(Iceberg), icebergFormat)
+
+    val icebergTableNoCat = "data.quickstart_purchases_davidhan_v1_dev_davidhan"
+    val icebergFormatNoCat = FormatProvider.from(spark).readFormat(icebergTableNoCat)
+    assertEquals(Some(Iceberg), icebergFormatNoCat)
+
+    val dneTable = "default_iceberg.data.dne"
+    val dneFormat = FormatProvider.from(spark).readFormat(dneTable)
+    assertTrue(dneFormat.isEmpty)
+  }
+
+
   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, set `GOOGLE_APPLICATION_CREDENTIALS=<path_to_application_default_credentials.json>

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

@@ -33,6 +33,6 @@ class GcpFormatProviderTest extends AnyFlatSpec with MockitoSugar {
         .build())
     when(mockTable.getTableId).thenReturn(TableId.of("project", "dataset", "table"))
 
-    val gcsFormat = gcpFormatProvider.getFormat(mockTable)
+    val gcsFormat = gcpFormatProvider.readFormat(tableName)
   }
 }

spark/src/main/scala/ai/chronon/spark/format/DefaultFormatProvider.scala

@@ -22,6 +22,18 @@ class DefaultFormatProvider(val sparkSession: SparkSession) extends FormatProvid
     } else { null })
   }
 
+  def getCatalog(tableName: String): String = {
+    logger.info(s"Retrieving read format for table: ${tableName}")
+    val parsed = sparkSession.sessionState.sqlParser.parseMultipartIdentifier(tableName)
+    val parsedCatalog = parsed.toList match {
+      case catalog :: namespace :: tableName :: Nil => catalog
+      case namespace :: tableName :: Nil            => sparkSession.catalog.currentCatalog()
+      case tableName :: Nil                         => sparkSession.catalog.currentCatalog()
+      case _ => throw new IllegalStateException(s"Invalid table naming convention specified: ${tableName}")
+    }
+    parsedCatalog
+  }
+
   private def isIcebergTable(tableName: String): Boolean =
     Try {
       sparkSession.read.format("iceberg").load(tableName)

spark/src/test/scala/ai/chronon/spark/test/TableUtilsTest.scala

@@ -25,6 +25,10 @@ import org.apache.spark.sql.functions.col
 import org.junit.Assert.{assertEquals, assertTrue}
 import org.scalatest.flatspec.AnyFlatSpec
 
+import ai.chronon.spark.format.FormatProvider
+import ai.chronon.spark.format.DefaultFormatProvider
+import org.apache.spark.sql.catalyst.parser.ParseException
+
 import scala.util.Try
 
 case class TestRecord(ds: String, id: String)
@@ -36,6 +40,7 @@ class SimpleAddUDF extends UDF {
 }
 
 class TableUtilsTest extends AnyFlatSpec {
+
   lazy val spark: SparkSession = SparkSessionBuilder.build("TableUtilsTest", local = true)
   private val tableUtils = TableTestUtils(spark)
   private implicit val partitionSpec: PartitionSpec = tableUtils.partitionSpec
@@ -639,4 +644,15 @@ class TableUtilsTest extends AnyFlatSpec {
     }
   }
 
+  it should "test catalog detection" in {
+    val fp = FormatProvider.from(spark).asInstanceOf[DefaultFormatProvider]
+    assertEquals("catalogA", fp.getCatalog("catalogA.foo.bar"))
+    assertEquals("catalogA", fp.getCatalog("`catalogA`.foo.bar"))
+    assertEquals("spark_catalog", fp.getCatalog("`catalogA.foo`.bar"))
+    assertEquals("spark_catalog", fp.getCatalog("`catalogA.foo.bar`"))
+    assertEquals("spark_catalog", fp.getCatalog("foo.bar"))
+    assertEquals("spark_catalog", fp.getCatalog("bar"))
+    assertThrows[ParseException](fp.getCatalog(""))
+  }
+
 }