feat: unit tests for local iteration

[tchow-zlai]

Summary

Checklist

• Added Unit Tests
• Covered by existing CI
• Integration tested
• Documentation update

Summary by CodeRabbit

• New Features

  • Introduced a new BigQuery client for enhanced interaction with BigQuery services.
  • Added functionality for managing partitions in Spark SQL tables through a new utility class.

• Bug Fixes

  • Improved error handling in the database creation process.

• Tests

  • Added a new test class for verifying BigQuery catalog functionality.
  • Updated existing test classes to utilize the new partition management utilities.

• Chores

  • Cleaned up deprecated methods in the TableUtils class.
  • Refactored comments for clarity regarding method dependencies.

[coderabbitai]

Walkthrough

This pull request introduces enhancements to the Google Cloud Platform (GCP) integration in the Chronon project. The changes span multiple files, focusing on dependency management, BigQuery client integration, table utilities, and test infrastructure. The modifications clarify dependency versions, enhance the functionality of format providers, and improve the test framework for GCP-related features.

Changes

File	Change Summary
build.sbt	Updated gcs-connector dependency version to hadoop3-2.2.26
cloud_gcp/src/main/scala/ai/chronon/integrations/cloud_gcp/BigQueryFormat.scala	Added bigQueryClient lazy val, enhanced readFormat and partitions methods
spark/src/main/scala/ai/chronon/spark/TableUtils.scala	Changed to a regular class, added provider comments, removed deprecated methods
spark/src/test/scala/ai/chronon/spark/test/TableTestUtils.scala	New utility class for partition management in tests
Multiple test files	Replaced TableUtils with TableTestUtils in test classes

Possibly related PRs

• Implement stubbed backend APIs for model list, search and time series #22: GCP dependency management in build.sbt.
• Downgrade Play to 2.x to stick to Scala 2.12 across the repo #43: Updates to project dependency configurations.
• Summary upload #50: Enhancements in handling data uploads relevant to GCP.
• Add delta lake integration #51: New variable and dependency for Delta Spark in build.sbt.
• feat: introduce BigQueryFormat #146: Direct connection with BigQuery Format integration.
• chore: small cleanup of TableUtils #153: Cleanup of table utilities, similar to the main PR focus.

Suggested Reviewers

• piyush-zlai
• nikhil-zlai

Poem

🌟 Clouds of Code, Dancing Free
GCP's magic, now set to spree
Connectors spark, utilities shine
BigQuery's rhythm, perfectly fine!
Engineering poetry in motion 🚀

---

🪧 Tips

Chat

There are 3 ways to chat with CodeRabbit:

• Review comments: Directly reply to a review comment made by CodeRabbit. Example:
  • I pushed a fix in commit <commit_id>, please review it.
  • Generate unit testing code for this file.
  • Open a follow-up GitHub issue for this discussion.
• Files and specific lines of code (under the "Files changed" tab): Tag @coderabbitai in a new review comment at the desired location with your query. Examples:
  • @coderabbitai generate unit testing code for this file.
  • @coderabbitai modularize this function.
• PR comments: Tag @coderabbitai in a new PR comment to ask questions about the PR branch. For the best results, please provide a very specific query, as very limited context is provided in this mode. Examples:
  • @coderabbitai gather interesting stats about this repository and render them as a table. Additionally, render a pie chart showing the language distribution in the codebase.
  • @coderabbitai read src/utils.ts and generate unit testing code.
  • @coderabbitai read the files in the src/scheduler package and generate a class diagram using mermaid and a README in the markdown format.
  • @coderabbitai help me debug CodeRabbit configuration file.

Note: Be mindful of the bot's finite context window. It's strongly recommended to break down tasks such as reading entire modules into smaller chunks. For a focused discussion, use review comments to chat about specific files and their changes, instead of using the PR comments.

CodeRabbit Commands (Invoked using PR comments)

• @coderabbitai pause to pause the reviews on a PR.
• @coderabbitai resume to resume the paused reviews.
• @coderabbitai review to trigger an incremental review. This is useful when automatic reviews are disabled for the repository.
• @coderabbitai full review to do a full review from scratch and review all the files again.
• @coderabbitai summary to regenerate the summary of the PR.
• @coderabbitai generate docstrings to generate docstrings for this PR. (Beta)
• @coderabbitai resolve resolve all the CodeRabbit review comments.
• @coderabbitai configuration to show the current CodeRabbit configuration for the repository.
• @coderabbitai help to get help.

Other keywords and placeholders

• Add @coderabbitai ignore anywhere in the PR description to prevent this PR from being reviewed.
• Add @coderabbitai summary to generate the high-level summary at a specific location in the PR description.
• Add @coderabbitai anywhere in the PR title to generate the title automatically.

CodeRabbit Configuration File (.coderabbit.yaml)

• You can programmatically configure CodeRabbit by adding a .coderabbit.yaml file to the root of your repository.
• Please see the configuration documentation for more information.
• If your editor has YAML language server enabled, you can add the path at the top of this file to enable auto-completion and validation: # yaml-language-server: $schema=https://coderabbit.ai/integrations/schema.v2.json

Documentation and Community

• Visit our Documentation for detailed information on how to use CodeRabbit.
• Join our Discord Community to get help, request features, and share feedback.
• Follow us on X/Twitter for updates and announcements.

[nikhil-zlai]

ignore is very cool!

[coderabbitai]

Actionable comments posted: 1

🧹 Nitpick comments (4)

spark/src/test/scala/ai/chronon/spark/test/TableTestUtils.scala (1)

8-11: Consider renaming default param.
“partitionColumn: String = partitionColumn” may shadow the class field and cause confusion.

spark/src/test/scala/ai/chronon/spark/test/TableUtilsTest.scala (1)

84-89: Minor formatting change.
No functional impact.

cloud_gcp/src/test/scala/ai/chronon/integrations/cloud_gcp/test/BigQueryCatalogTest.scala (2)

42-50: Consider adding result assertions.
Relying on table.show might be insufficient for thorough testing.

---

51-63: Ignore block is documented with action steps.
Consider mocking or local stubbing for BigQuery integration.

📜 Review details

Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro (Legacy)

📥 Commits

Reviewing files that changed from the base of the PR and between ea2a1f3 and acde01d.

📒 Files selected for processing (8)

• build.sbt (1 hunks)
• cloud_gcp/src/main/scala/ai/chronon/integrations/cloud_gcp/BigQueryFormat.scala (2 hunks)
• cloud_gcp/src/test/scala/ai/chronon/integrations/cloud_gcp/test/BigQueryCatalogTest.scala (1 hunks)
• spark/src/main/scala/ai/chronon/spark/TableUtils.scala (11 hunks)
• spark/src/test/scala/ai/chronon/spark/test/JoinTest.scala (1 hunks)
• spark/src/test/scala/ai/chronon/spark/test/LabelJoinTest.scala (1 hunks)
• spark/src/test/scala/ai/chronon/spark/test/TableTestUtils.scala (1 hunks)
• spark/src/test/scala/ai/chronon/spark/test/TableUtilsTest.scala (3 hunks)

✅ Files skipped from review due to trivial changes (2)

• spark/src/test/scala/ai/chronon/spark/test/JoinTest.scala
• build.sbt

🔇 Additional comments (30)

spark/src/test/scala/ai/chronon/spark/test/TableTestUtils.scala (3)

12-24: No issues, succinct logic.
Dropping multiple partitions in one command is efficient.

---

25-26: Good fallback.
Logs correctly if table doesn’t exist.

---

29-39: Clever approach to partition range.
Uses Stream and calls dropPartitions.

spark/src/test/scala/ai/chronon/spark/test/TableUtilsTest.scala (2)

49-49: Name change matches new class.
Replacing TableUtils(…) with TableTestUtils(…) is consistent.

---

347-347: Test verification is solid.
Ensures final mapping is correct.

spark/src/test/scala/ai/chronon/spark/test/LabelJoinTest.scala (1)

40-40: Consistent utility reference.
Uses TableTestUtils now.

spark/src/main/scala/ai/chronon/spark/TableUtils.scala (16)

144-147: Simple existence check.
No issues.

---

149-152: Load table is straightforward.
Optimizations might rely on provider.

---

Line range hint 161-178: Database creation.
Handles AlreadyExistsException gracefully.

---

Line range hint 179-193: Retrieves all partitions.
Solid conditional logic.

---

Line range hint 194-198: Primary partitions.
Filters by sub-partitions.

---

235-235: Schema retrieval.
Query-based approach is acceptable.

---

Line range hint 241-258: Permission check.
Attempts to read a sample row.

---

267-267: Retrieves last partition.
No edge-case concerns.

---

271-271: Retrieves first partition.
Similar logic, no issues.

---

Line range hint 275-333: Partition insertion logic is thorough.
Auto-expand and fallback handled well.

---

Line range hint 369-378: Unpartitioned insertion.
Creates or updates.

---

Line range hint 508-529: Table DDL creation.
Handles partition columns properly.

---

547-547: Altering table props.
Concise TBLPROPERTIES handling.

---

Line range hint 634-642: Fetch table props.
Uses catalog metadata, good fallback.

---

644-644: Simple drop.
No concerns.

---

Line range hint 431-460: Repartitioning.
Tune rowCountPerPartition carefully.

✅ Verification successful

Let me gather information about the repartitioning logic to verify the concern.

---

Repartitioning configuration is well-defined and bounded

The repartitioning logic uses:

• Configurable rows per partition via spark.chronon.rowCountPerPartition (default: 1e8)
• Bounded daily file count (10-2000 for cluster, 1-2000 for local)
• Safe division-by-zero handling

🏁 Scripts executed

The following scripts were executed for the analysis:

Error: No shell scripts found

---

Script:

#!/bin/bash
# Search for repartitioning logic and row count per partition configuration
ast-grep --pattern 'val rowCountPerPartition = $_'

# Search for related configuration or constants
rg -A 5 'rowCountPerPartition' 

Length of output: 3455

cloud_gcp/src/test/scala/ai/chronon/integrations/cloud_gcp/test/BigQueryCatalogTest.scala (6)

3-4: No issues with these imports.

---

8-10: Imports for Spark, Hadoop, and JUnit look straightforward.

---

13-13: Class definition is clean and extends AnyFunSuite + MockitoSugar.

---

31-33: Good job verifying hive.metastore.uris.

---

35-40: Well-structured test for dynamic classloading.
Verifies that the table read format is from GCP.

---

15-28: Validate the external services in integration tests.
Ensure the local Hive metastore and GCP credentials are properly configured at runtime.

cloud_gcp/src/main/scala/ai/chronon/integrations/cloud_gcp/BigQueryFormat.scala (2)

130-130: Robust try/finally block usage.
Reverting Spark config ensures minimal side effects.

---

15-15: Initialize BigQuery client carefully.
Ensure credentials are valid when running outside default GCP environments.

[coderabbitai]

💡 Codebase verification

Partition detection needs improvement
The current implementation in isPartitioned only checks if the partition column exists in the schema, which is insufficient. A column's presence doesn't guarantee it's used for partitioning. The table's actual partition information should be retrieved from the table metadata.

• TableUtils.scala: Replace schema-based check with proper metadata query using table format provider's partition information.

🔗 Analysis chain

Partition detection is naive.
Reliable if partitionColumn is consistent.

---

Let me verify the partition detection implementation by searching for related code.

---

Let me search for the schema-related implementation to understand how partitioning is actually determined.

🏁 Scripts executed

The following scripts were executed for the analysis:

Error: No shell scripts found

---

Script:

#!/bin/bash
# Search for partitionColumn definition and usage
rg "partitionColumn" -A 3

# Search for isPartitioned method implementations
ast-grep --pattern 'def isPartitioned($$$) { $$$ }'

# Look for table schema related methods
ast-grep --pattern 'def getSchemaFromTable($$$) { $$$ }'

Length of output: 66094

---

Script:

#!/bin/bash
# Search for getSchemaFromTable implementation
rg "def getSchemaFromTable" -A 10

# Search for table schema related code
rg "schema.*partition" -A 3

Length of output: 11158

[coderabbitai]

Actionable comments posted: 0

🔭 Outside diff range comments (2)

cloud_gcp/src/main/scala/ai/chronon/integrations/cloud_gcp/BigQueryFormat.scala (2)

Line range hint 19-65: Remove commented federation code and add error handling

1. Remove the large commented section
2. Add validation for table identifier parsing

   def readFormat(tableName: String): Format = {
+    try {
       val btTableIdentifier: TableId = BigQueryUtil.parseTableId(tableName)
       // ... rest of the code ...
+    } catch {
+      case e: IllegalArgumentException =>
+        throw new IllegalArgumentException(s"Invalid BigQuery table name: $tableName", e)
+    }
-    /**
-    Using federation
-     ... large commented section ...
-      * */
   }

---

Configuration handling needs improvement in GCSFormat

The GCSFormat class sets the same configuration values but lacks cleanup, which could affect concurrent operations. The suggested helper method is valid.

• GCSFormat.scala: Add configuration cleanup in try-finally block
• BigQueryFormat.scala: Refactor to use the suggested helper method

🔗 Analysis chain

Line range hint 82-130: Add timeout and improve configuration handling

1. Configuration changes could affect concurrent operations
2. Missing timeout for BigQuery queries

+  private def withTemporaryConfig[T](sparkSession: SparkSession)(f: => T): T = {
+    val originalConfig = Map(
+      "viewsEnabled" -> sparkSession.conf.get("viewsEnabled", false.toString),
+      "materializationDataset" -> sparkSession.conf.get("materializationDataset", "")
+    )
+    try {
+      sparkSession.conf.set("viewsEnabled", true)
+      sparkSession.conf.set("materializationDataset", database)
+      f
+    } finally {
+      originalConfig.foreach { case (k, v) => sparkSession.conf.set(k, v) }
+    }
+  }

Run this to verify concurrent access patterns:

🏁 Scripts executed

The following scripts were executed for the analysis:

Script:

#!/bin/bash
# Check for potential concurrent access issues
rg "sparkSession.*conf\.set" --type scala

Length of output: 887

---

Script:

#!/bin/bash
# Check if GCSFormat also handles configuration cleanup
ast-grep --pattern 'try {
  $$$
} finally {
  $$$
}'

Length of output: 14187

🧹 Nitpick comments (4)

cloud_gcp/src/test/scala/ai/chronon/integrations/cloud_gcp/test/BigQueryCatalogTest.scala (2)

25-25: Remove duplicate config.

spark.hadoop.fs.gs.impl is defined twice with the same value.

-        "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,
-        "spark.hadoop.fs.gs.impl" -> classOf[GoogleHadoopFileSystem].getName

Also applies to: 28-28

---

44-51: Add documentation for ignored test.

Add a TODO comment explaining why this test is ignored and what's needed to enable it.

spark/src/main/scala/ai/chronon/spark/TableUtils.scala (1)

Line range hint 144-644: Document provider requirements.

Multiple methods are marked with "Needs provider" without explanation. Add documentation explaining:

• What provider is needed
• Why it's needed
• Implementation plan

cloud_gcp/src/main/scala/ai/chronon/integrations/cloud_gcp/BigQueryFormat.scala (1)

15-17: Consider making BigQuery client injectable

Make client injectable to improve testability.

-  lazy val bigQueryClient = BigQueryOptions.getDefaultInstance.getService
+  protected[cloud_gcp] lazy val bigQueryClient = createBigQueryClient()
+  protected def createBigQueryClient() = BigQueryOptions.getDefaultInstance.getService

📜 Review details

Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro (Legacy)

📥 Commits

Reviewing files that changed from the base of the PR and between acde01d and 3be7d8a.

📒 Files selected for processing (8)

• build.sbt (1 hunks)
• cloud_gcp/src/main/scala/ai/chronon/integrations/cloud_gcp/BigQueryFormat.scala (2 hunks)
• cloud_gcp/src/test/scala/ai/chronon/integrations/cloud_gcp/test/BigQueryCatalogTest.scala (1 hunks)
• spark/src/main/scala/ai/chronon/spark/TableUtils.scala (13 hunks)
• spark/src/test/scala/ai/chronon/spark/test/JoinTest.scala (1 hunks)
• spark/src/test/scala/ai/chronon/spark/test/LabelJoinTest.scala (1 hunks)
• spark/src/test/scala/ai/chronon/spark/test/TableTestUtils.scala (1 hunks)
• spark/src/test/scala/ai/chronon/spark/test/TableUtilsTest.scala (3 hunks)

🚧 Files skipped from review as they are similar to previous changes (5)

• spark/src/test/scala/ai/chronon/spark/test/JoinTest.scala
• build.sbt
• spark/src/test/scala/ai/chronon/spark/test/TableUtilsTest.scala
• spark/src/test/scala/ai/chronon/spark/test/LabelJoinTest.scala
• spark/src/test/scala/ai/chronon/spark/test/TableTestUtils.scala

🔇 Additional comments (4)

cloud_gcp/src/test/scala/ai/chronon/integrations/cloud_gcp/test/BigQueryCatalogTest.scala (2)

33-42: LGTM!

Test methods are well-structured with clear assertions.

---

53-65: LGTM!

Well-documented test with clear setup instructions.

spark/src/main/scala/ai/chronon/spark/TableUtils.scala (2)

154-159: Partition detection needs improvement.

The schema-based check is insufficient for reliable partition detection.

---

799-803: LGTM!

Clean companion object implementation.