Flink submit updates to allow Etsy jobs to run

[piyush-zlai]

Summary

Some updates to get our Flink jobs running on the Etsy side:

• Configure schema registry via host/port/scheme instead of URL
• Explicitly set the task slots per task manager
• Configure checkpoint directory based on teams.json

Checklist

• Added Unit Tests
• Covered by existing CI
• Integration tested - Kicked off the job on the Etsy cluster and confirmed it's up and running
• Documentation update

Summary by CodeRabbit

• New Features

  • Introduced an enhanced configuration option for Flink job submissions by adding a state URI parameter for improved job state management.
  • Expanded schema registry configuration, enabling greater flexibility with host, port, and scheme settings.

• Chores

  • Adjusted logging levels and refined error messaging to support better troubleshooting.

• Documentation

  • Updated configuration guidance to aid in setting up schema registry integration.

[coderabbitai]

Walkthrough

This pull request updates logging levels, job submission configuration, and schema registry handling across several components. The changes include downgrading a log statement from info to debug in the KV store, modifying Flink job submission methods to incorporate a new state URI parameter, and refactoring schema registry configurations by splitting a single URL into separate host, port, and scheme parameters.

Changes

File(s)	Change Summary
cloud_gcp/src/main/scala/ai/chronon/integrations/cloud_gcp/BigTableKVStoreImpl.scala	Modified logging level in the multiPut method from info to debug.
cloud_gcp/src/main/scala/ai/chronon/integrations/cloud_gcp/DataprocSubmitter.scala, cloud_gcp/src/test/scala/ai/chronon/integrations/cloud_gcp/DataprocSubmitterTest.scala, spark/src/…/JobSubmitter.scala	Updated Flink job submission: added FlinkStateUri extraction, adjusted method signature in buildFlinkJob, required FLINK_STATE_URI env variable, and added a new constant.
flink/src/main/scala/ai/chronon/flink/FlinkJob.scala	Replaced hardcoded "registry_url" check with RegistryHostKey and updated the error message accordingly.
flink/src/main/scala/ai/chronon/flink/SchemaRegistrySchemaProvider.scala, flink/src/test/scala/ai/chronon/flink/test/SchemaRegistrySchemaProviderSpec.scala	Refactored schema registry configuration by replacing a single URL with separate host, port, and scheme fields; updated buildSchemaRegistryClient signature; introduced an object to define constant keys.

Sequence Diagram(s)

sequenceDiagram
    participant Client
    participant DataprocSubmitter
    participant FlinkJobBuilder
    Client->>DataprocSubmitter: submit() with job properties
    DataprocSubmitter->>DataprocSubmitter: Extract FlinkStateUri
    DataprocSubmitter->>FlinkJobBuilder: buildFlinkJob(..., flinkStateUri, ...)
    FlinkJobBuilder-->>DataprocSubmitter: Job.Builder
    DataprocSubmitter-->>Client: Submission response

Loading

sequenceDiagram
    participant FlinkJob
    participant SchemaRegistrySchemaProvider
    participant RegistryClient
    FlinkJob->>SchemaRegistrySchemaProvider: Check for RegistryHostKey
    SchemaRegistrySchemaProvider->>SchemaRegistrySchemaProvider: Build client (scheme, host, port)
    SchemaRegistrySchemaProvider->>RegistryClient: Initialize client
    RegistryClient-->>SchemaRegistrySchemaProvider: Client ready
    SchemaRegistrySchemaProvider-->>FlinkJob: Return schema configuration

Loading

Possibly related PRs

• Summary upload #50: Involves similar modifications to the multiPut functionality, aligning log level updates with other components.

Suggested reviewers

• nikhil-zlai
• david-zlai

Poem

In the code's vibrant glow, we switch to debug light,
Flink jobs now flow with state URI insight.
Registry keys guide the schema's tune,
Each method sings a refactor soon.
CodeRabbit leaps as changes take flight! 🚀
Cheers to code refined, shining ever so bright!

Warning

Review ran into problems

🔥 Problems

GitHub Actions and Pipeline Checks: Resource not accessible by integration - https://docs.github.com/rest/actions/workflow-runs#list-workflow-runs-for-a-repository.

Please grant the required permissions to the CodeRabbit GitHub App under the organization or repository settings.

---

🪧 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.

[coderabbitai]

Actionable comments posted: 0

🧹 Nitpick comments (4)

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

131-135: Consider externalizing memory configuration.

Memory and task slot settings are hardcoded. Consider moving these to configuration files for better flexibility.

---

254-254: Enhance error message.

Consider providing more context in the error message about the expected format and purpose of FLINK_STATE_URI.

-        val flinkStateUri = sys.env.getOrElse("FLINK_STATE_URI", throw new Exception("FLINK_STATE_URI not set"))
+        val flinkStateUri = sys.env.getOrElse("FLINK_STATE_URI", 
+          throw new Exception("FLINK_STATE_URI not set. Please set it to the GCS path for storing Flink state and checkpoints."))

flink/src/main/scala/ai/chronon/flink/SchemaRegistrySchemaProvider.scala (1)

36-47: Consider simplifying URL construction.

-  maybePortString match {
-    case Some(portString) =>
-      val registryUrl = s"$schemeString://$registryHost:$portString"
-      new CachedSchemaRegistryClient(registryUrl, CacheCapacity)
-    case None =>
-      val registryUrl = s"$schemeString://$registryHost"
-      new CachedSchemaRegistryClient(registryUrl, CacheCapacity)
-  }
+  val registryUrl = maybePortString.fold(s"$schemeString://$registryHost")(port => s"$schemeString://$registryHost:$port")
+  new CachedSchemaRegistryClient(registryUrl, CacheCapacity)

flink/src/main/scala/ai/chronon/flink/FlinkJob.scala (1)

321-327: Consider using getOrElse for more concise error handling.

-  topicInfo.params.get(RegistryHostKey) match {
-    case Some(_) => new SchemaRegistrySchemaProvider(topicInfo.params)
-    case None =>
-      throw new IllegalArgumentException(
-        s"We only support schema registry based schema lookups. Missing $RegistryHostKey in topic config")
-  }
+  topicInfo.params.get(RegistryHostKey).map(_ => new SchemaRegistrySchemaProvider(topicInfo.params))
+    .getOrElse(throw new IllegalArgumentException(
+      s"We only support schema registry based schema lookups. Missing $RegistryHostKey in topic config"))

📜 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 9356745 and b1f5e80.

📒 Files selected for processing (7)

• cloud_gcp/src/main/scala/ai/chronon/integrations/cloud_gcp/BigTableKVStoreImpl.scala (1 hunks)
• cloud_gcp/src/main/scala/ai/chronon/integrations/cloud_gcp/DataprocSubmitter.scala (5 hunks)
• cloud_gcp/src/test/scala/ai/chronon/integrations/cloud_gcp/DataprocSubmitterTest.scala (3 hunks)
• flink/src/main/scala/ai/chronon/flink/FlinkJob.scala (2 hunks)
• flink/src/main/scala/ai/chronon/flink/SchemaRegistrySchemaProvider.scala (2 hunks)
• flink/src/test/scala/ai/chronon/flink/test/SchemaRegistrySchemaProviderSpec.scala (3 hunks)
• spark/src/main/scala/ai/chronon/spark/JobSubmitter.scala (1 hunks)

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

• spark/src/main/scala/ai/chronon/spark/JobSubmitter.scala
• cloud_gcp/src/main/scala/ai/chronon/integrations/cloud_gcp/BigTableKVStoreImpl.scala

⏰ Context from checks skipped due to timeout of 90000ms (8)

• GitHub Check: table_utils_delta_format_spark_tests
• GitHub Check: mutation_spark_tests
• GitHub Check: other_spark_tests
• GitHub Check: join_spark_tests
• GitHub Check: no_spark_scala_tests
• GitHub Check: fetcher_spark_tests
• GitHub Check: scala_compile_fmt_fix
• GitHub Check: enforce_triggered_workflows

🔇 Additional comments (10)

cloud_gcp/src/test/scala/ai/chronon/integrations/cloud_gcp/DataprocSubmitterTest.scala (3)

5-5: LGTM!

Import aligns with the new state URI functionality.

---

65-68: LGTM!

Clear comments explaining checkpoint and state persistence configuration.

---

87-90: LGTM!

Consistent state URI configuration across test cases.

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

70-73: LGTM!

Proper validation for required state URI parameter.

---

137-138: LGTM!

Consistent state and checkpoint directory configuration.

flink/src/main/scala/ai/chronon/flink/SchemaRegistrySchemaProvider.scala (3)

12-18: LGTM! Clear and comprehensive documentation.

---

22-29: LGTM! Good separation of configuration parameters.

The split configuration allows for more flexible setups, especially with load balancers.

---

73-78: LGTM! Well-organized configuration constants.

flink/src/test/scala/ai/chronon/flink/test/SchemaRegistrySchemaProviderSpec.scala (2)

16-19: LGTM! Test implementation matches production changes.

---

26-27: LGTM! Test configuration updated correctly.