wip

Author: tchow-zlai

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

@@ -42,6 +42,15 @@ class BigQueryCatalogTest extends AnyFunSuite with MockitoSugar {
     })
   }
 
+  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:

spark/src/main/scala/ai/chronon/spark/TableUtils.scala

@@ -141,16 +141,24 @@ case class TableUtils(sparkSession: SparkSession) {
       rdd
     }
 
-  def tableExists(tableName: String): Boolean = sparkSession.catalog.tableExists(tableName)
+  // Needs provider
+  def tableExists(tableName: String): Boolean = {
+    sparkSession.catalog.tableExists(tableName)
+  }
 
-  def loadTable(tableName: String): DataFrame = sparkSession.table(tableName)
+  // Needs provider
+  def loadTable(tableName: String): DataFrame = {
+    sparkSession.table(tableName)
+  }
 
+  // Needs provider
   def isPartitioned(tableName: String): Boolean = {
     // TODO: use proper way to detect if a table is partitioned or not
     val schema = getSchemaFromTable(tableName)
     schema.fieldNames.contains(partitionColumn)
   }
 
+  // Needs provider
   def createDatabase(database: String): Boolean = {
     try {
       val command = s"CREATE DATABASE IF NOT EXISTS $database"
@@ -168,6 +176,7 @@ case class TableUtils(sparkSession: SparkSession) {
 
   def tableReadFormat(tableName: String): Format = tableFormatProvider.readFormat(tableName)
 
+  // Needs provider
   // return all specified partition columns in a table in format of Map[partitionName, PartitionValue]
   def allPartitions(tableName: String, partitionColumnsFilter: Seq[String] = Seq.empty): Seq[Map[String, String]] = {
     if (!tableExists(tableName)) return Seq.empty[Map[String, String]]
@@ -182,6 +191,7 @@ case class TableUtils(sparkSession: SparkSession) {
     }
   }
 
+  // Needs provider
   def partitions(tableName: String, subPartitionsFilter: Map[String, String] = Map.empty): Seq[String] = {
     if (!tableExists(tableName)) return Seq.empty[String]
     val format = tableReadFormat(tableName)
@@ -222,11 +232,13 @@ case class TableUtils(sparkSession: SparkSession) {
     }
   }
 
+  // Needs provider
   def getSchemaFromTable(tableName: String): StructType = {
     sparkSession.sql(s"SELECT * FROM $tableName LIMIT 1").schema
   }
 
   // method to check if a user has access to a table
+  // Needs provider
   def checkTablePermission(tableName: String,
                            fallbackPartition: String =
                              partitionSpec.before(partitionSpec.at(System.currentTimeMillis()))): Boolean = {
@@ -252,12 +264,15 @@ case class TableUtils(sparkSession: SparkSession) {
     }
   }
 
+  // Needs provider
   def lastAvailablePartition(tableName: String, subPartitionFilters: Map[String, String] = Map.empty): Option[String] =
     partitions(tableName, subPartitionFilters).reduceOption((x, y) => Ordering[String].max(x, y))
 
+  // Needs provider
   def firstAvailablePartition(tableName: String, subPartitionFilters: Map[String, String] = Map.empty): Option[String] =
     partitions(tableName, subPartitionFilters).reduceOption((x, y) => Ordering[String].min(x, y))
 
+  // Needs provider
   def insertPartitions(df: DataFrame,
                        tableName: String,
                        tableProperties: Map[String, String] = null,
@@ -351,6 +366,7 @@ case class TableUtils(sparkSession: SparkSession) {
     }
   }
 
+  // Needs provider
   def insertUnPartitioned(df: DataFrame,
                           tableName: String,
                           tableProperties: Map[String, String] = null,
@@ -412,6 +428,7 @@ case class TableUtils(sparkSession: SparkSession) {
     }.get
   }
 
+  // Needs provider
   private def repartitionAndWriteInternal(df: DataFrame,
                                           tableName: String,
                                           saveMode: SaveMode,
@@ -488,6 +505,7 @@ case class TableUtils(sparkSession: SparkSession) {
     }
   }
 
+  // Needs provider
   private def createTableSql(tableName: String,
                              schema: StructType,
                              partitionColumns: Seq[String],
@@ -526,6 +544,7 @@ case class TableUtils(sparkSession: SparkSession) {
     Seq(createFragment, partitionFragment, fileFormatString, propertiesFragment).mkString("\n")
   }
 
+  // Needs provider
   private def alterTablePropertiesSql(tableName: String, properties: Map[String, String]): String = {
     // Only SQL api exists for setting TBLPROPERTIES
     val propertiesString = properties
@@ -612,6 +631,7 @@ case class TableUtils(sparkSession: SparkSession) {
     Some(missingChunks)
   }
 
+  // Needs provider
   def getTableProperties(tableName: String): Option[Map[String, String]] = {
     try {
       val tableId = sparkSession.sessionState.sqlParser.parseTableIdentifier(tableName)
@@ -621,6 +641,7 @@ case class TableUtils(sparkSession: SparkSession) {
     }
   }
 
+  // Needs provider
   def dropTableIfExists(tableName: String): Unit = {
     val command = s"DROP TABLE IF EXISTS $tableName"
     logger.info(s"Dropping table with command: $command")