Skip to content

fix: properly detect bigquery catalog #629

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 13 commits into from
Apr 12, 2025
Merged

fix: properly detect bigquery catalog #629

Show file tree
Hide file tree
Changes from 2 commits
Commits
Show all changes
13 commits
Select commit Hold shift + click to select a range
eb08cb1
fix: properly detect bigquery catalog
tchow-zlai Apr 11, 2025
8822ede
update
tchow-zlai Apr 11, 2025
7af162a
wip
tchow-zlai Apr 11, 2025
484368a
update
tchow-zlai Apr 11, 2025
7ace96b
passing tests
tchow-zlai Apr 11, 2025
527333f
passing tests
tchow-zlai Apr 11, 2025
22fae64
ignore
tchow-zlai Apr 11, 2025
b7b44d4
wip
tchow-zlai Apr 11, 2025
1ba3541
wip
tchow-zlai Apr 11, 2025
ef6213a
wip
tchow-zlai Apr 11, 2025
565bdd0
passing tests
tchow-zlai Apr 11, 2025
a423261
exclude
tchow-zlai Apr 11, 2025
a22660f
Merge branch 'main' into tchow/catalog-switch
tchow-zlai Apr 11, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
54 changes: 42 additions & 12 deletions cloud_gcp/src/main/scala/ai/chronon/integrations/cloud_gcp/GcpFormatProvider.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,8 +2,10 @@ 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 com.google.cloud.spark.bigquery.BigQueryCatalog
import org.apache.iceberg.exceptions.NoSuchIcebergTableException
import org.apache.iceberg.gcp.bigquery.{BigQueryClient, BigQueryClientImpl}
import org.apache.iceberg.gcp.bigquery.{BigQueryClient, BigQueryClientImpl, BigQueryMetastoreCatalog}
import org.apache.iceberg.spark.SparkCatalog
import org.apache.spark.sql.SparkSession

import scala.jdk.CollectionConverters._
Expand All @@ -23,19 +25,47 @@ class GcpFormatProvider(override val sparkSession: SparkSession) extends Default
private lazy val icebergClient: BigQueryClient = new BigQueryClientImpl()

override def readFormat(tableName: String): scala.Option[Format] = {
logger.info(s"Retrieving read format for table: ${tableName}")
val parsedCatalog = getCatalog(tableName)
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

💡 Verification agent

🧩 Analysis chain

Missing getCatalog implementation.

Method is called but appears to be missing implementation.

🏁 Script executed:

#!/bin/bash
# Check for getCatalog method implementation
grep -r "def getCatalog" --include="*.scala" .

Length of output: 171

Implement or Inherit getCatalog

The call to getCatalog in GcpFormatProvider.scala (line 23) doesn’t resolve locally. Although a similar method exists in DefaultFormatProvider.scala, it isn’t automatically available here. Please either inherit from a common base that provides the implementation or add a GCP-specific getCatalog method.


if (isBigQueryCatalog(parsedCatalog)) {
logger.info(s"Detected BigQuery catalog: $parsedCatalog")
Try {
val btTableIdentifier = SparkBQUtils.toTableId(tableName)(sparkSession)
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

remind me again but why do we need to convert tableName to a big table identifier?

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

That's just what the bigQueryClient accepts - see the line below.

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

tbh this isn't even entirely correct just yet. We actually need to find the project_id associated with the catalog, and replace that in the table Name before we pass it off to the BQ Client. but will do that in a follow-up.

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

alright I've adjusted this actually to make it work with BigQuery.

val bqTable = bigQueryClient.getTable(btTableIdentifier)
getFormat(bqTable)
} match {
case Success(format) => scala.Option(format)
case Failure(e) =>
throw new IllegalStateException(
s"${tableName} belongs to bigquery catalog ${parsedCatalog} but could not be found",
e)
}
} else {

logger.info(s"Detected non-BigQuery catalog: $parsedCatalog")
super.readFormat(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 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[cloud_gcp] def isBigQueryCatalog(catalog: String): Boolean = {
val cat = sparkSession.sessionState.catalogManager.catalog(catalog)
cat.isInstanceOf[DelegatingBigQueryMetastoreCatalog] || cat
.isInstanceOf[BigQueryCatalog] || (cat.isInstanceOf[SparkCatalog] && cat
.asInstanceOf[SparkCatalog]
.icebergCatalog()
.isInstanceOf[BigQueryMetastoreCatalog])
}

private[cloud_gcp] def getFormat(table: Table): Format = {
Expand Down
11 changes: 11 additions & 0 deletions cloud_gcp/src/test/scala/ai/chronon/integrations/cloud_gcp/BigQueryCatalogTest.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -61,6 +61,17 @@ class BigQueryCatalogTest extends AnyFlatSpec with MockitoSugar {
assertCompiles("classOf[GoogleHadoopFS]")
}

it should "test catalog detection" in {
val fp = FormatProvider.from(spark).asInstanceOf[GcpFormatProvider]
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(""))
}

it should "verify dynamic classloading of GCP providers" in {
assertEquals("thrift://localhost:9083", spark.sqlContext.getConf("hive.metastore.uris"))
assertTrue(FormatProvider.from(spark).isInstanceOf[GcpFormatProvider])
Expand Down