Changes to fix all the test failures with Bazel

.github/image/Dockerfile

@@ -19,6 +19,10 @@ ENV PATH=$PATH:$JAVA_HOME/bin
 RUN curl -L "https://github.com/sbt/sbt/releases/download/v1.8.2/sbt-1.8.2.tgz" | tar -xz -C /usr/local
 ENV PATH="/usr/local/sbt/bin:${PATH}"
 
+# bazel
+RUN curl -fsSL "https://github.com/bazelbuild/bazelisk/releases/download/v1.18.0/bazelisk-linux-amd64" -o /usr/local/bin/bazel
+ENV PATH="/usr/local/bin/bazel:${PATH}"
+
 # thrift
 ARG THRIFT_VERSION=0.21.0
 RUN apk add --no-cache \

.github/workflows/test_scala_no_spark.yaml

@@ -48,8 +48,9 @@ jobs:
 
       - name: Run Flink tests
         run: |
-          export SBT_OPTS="-Xmx24G -Xms4G --add-opens=java.base/sun.nio.ch=ALL-UNNAMED"
-          sbt "++ 2.12.18 flink/test"
+#          export SBT_OPTS="-Xmx24G -Xms4G --add-opens=java.base/sun.nio.ch=ALL-UNNAMED"
+#          sbt "++ 2.12.18 flink/test"
+          bazel test //flink:test
 
       - name: Run Async KV Store writer test
         run: |

.github/workflows/test_scala_spark.yaml

@@ -56,8 +56,9 @@ jobs:
 
       - name: Run spark join tests
         run: |
-          export SBT_OPTS="-Xmx8G -Xms2G --add-opens=java.base/sun.nio.ch=ALL-UNNAMED"
-          sbt "spark/testOnly -- -n jointest"
+#          export SBT_OPTS="-Xmx8G -Xms2G --add-opens=java.base/sun.nio.ch=ALL-UNNAMED"
+#          sbt "spark/testOnly -- -n jointest"
+          bazel test //spark:test_test_suite_src_test_scala_ai_chronon_spark_test_JoinUtilsTest.scala
 
   table_utils_delta_format_spark_tests:
     runs-on: ubuntu-latest

.ijwb/.bazelproject

@@ -0,0 +1,22 @@
+directories:
+  # Add the directories you want added as source here
+  # By default, we've added your entire workspace ('.')
+  .
+
+# Automatically includes all relevant targets under the 'directories' above
+derive_targets_from_directories: true
+
+targets:
+  # If source code isn't resolving, add additional targets that compile it here
+
+additional_languages:
+  # Uncomment any additional languages you want supported
+  # android
+  # dart
+  # go
+  # javascript
+  # kotlin
+  python
+  scala
+  typescript
+  java

BUILD.bazel

@@ -0,0 +1 @@
+package(default_visibility = ["//visibility:public"])

WORKSPACE

@@ -23,22 +23,29 @@ http_archive(
         "https://github.com/bazelbuild/rules_java/releases/download/7.0.6/rules_java-7.0.6.tar.gz",
     ],
 )
-load("@rules_java//java:repositories.bzl", "rules_java_dependencies", "rules_java_toolchains")
+
+load("@rules_java//java:repositories.bzl", "remote_jdk17_repos", "rules_java_dependencies", "rules_java_toolchains")
+
 rules_java_dependencies()
+
 rules_java_toolchains()
-load("@rules_java//java:repositories.bzl", "remote_jdk17_repos")
+
 remote_jdk17_repos()
 
 # For JVM support
 http_archive(
     name = "rules_jvm_external",
-    strip_prefix = "rules_jvm_external-6.6",
     sha256 = "3afe5195069bd379373528899c03a3072f568d33bd96fe037bd43b1f590535e7",
-    url = "https://github.com/bazel-contrib/rules_jvm_external/releases/download/6.6/rules_jvm_external-6.6.tar.gz"
+    strip_prefix = "rules_jvm_external-6.6",
+    url = "https://github.com/bazel-contrib/rules_jvm_external/releases/download/6.6/rules_jvm_external-6.6.tar.gz",
 )
+
 load("@rules_jvm_external//:repositories.bzl", "rules_jvm_external_deps")
+
 rules_jvm_external_deps()
+
 load("@rules_jvm_external//:setup.bzl", "rules_jvm_external_setup")
+
 rules_jvm_external_setup()
 
 # For additional rulesets like java_test_suite
@@ -48,9 +55,13 @@ http_archive(
     strip_prefix = "rules_jvm-0.24.0",
     urls = ["https://github.com/bazel-contrib/rules_jvm/releases/download/v0.24.0/rules_jvm-v0.24.0.tar.gz"],
 )
+
 load("@contrib_rules_jvm//:repositories.bzl", "contrib_rules_jvm_deps")
+
 contrib_rules_jvm_deps()
+
 load("@contrib_rules_jvm//:setup.bzl", "contrib_rules_jvm_setup")
+
 contrib_rules_jvm_setup()
 
 # For Scala support
@@ -60,15 +71,37 @@ http_archive(
     strip_prefix = "rules_scala-6.6.0",
     url = "https://github.com/bazelbuild/rules_scala/releases/download/v6.6.0/rules_scala-v6.6.0.tar.gz",
 )
+
 # Initialize Scala with specific version support
 load("@io_bazel_rules_scala//:scala_config.bzl", "scala_config")
+
 scala_config(scala_version = SCALA_VERSION)
+
+load("@io_bazel_rules_scala//scala:scala_maven_import_external.bzl", "scala_maven_import_external")
+
+scala_maven_import_external(
+    name = "scala_compiler_source_2_12_18",
+    artifact = "org.scala-lang:scala-compiler:%s:sources" % SCALA_VERSION,
+    artifact_sha256 = "f79ee80f140218253f2a38c9d73f8a9b552d06afce7a5f61cf08079a388e21df",
+    licenses = ["notice"],
+    server_urls = [
+        "https://repo1.maven.org/maven2",
+        "https://mirror.bazel.build/repo1.maven.org/maven2",
+    ],
+)
+
 load("@io_bazel_rules_scala//scala:scala.bzl", "scala_repositories")
+
 scala_repositories()
+
 load("@io_bazel_rules_scala//scala:toolchains.bzl", "scala_register_toolchains")
+
 scala_register_toolchains()
+
 load("@io_bazel_rules_scala//testing:scalatest.bzl", "scalatest_repositories", "scalatest_toolchain")
+
 scalatest_repositories()
+
 scalatest_toolchain()
 
 # For Protobuf support
@@ -80,10 +113,14 @@ http_archive(
         "https://github.com/bazelbuild/rules_proto/archive/refs/tags/5.3.0-21.7.tar.gz",
     ],
 )
+
 load("@rules_proto//proto:repositories.bzl", "rules_proto_dependencies", "rules_proto_toolchains")
+
 rules_proto_dependencies()
+
 rules_proto_toolchains()
 
 # To load all dependencies used across our modules
 load("//tools/build_rules/dependencies:load_dependencies.bzl", "load_all_dependencies")
-load_all_dependencies()
+
+load_all_dependencies()

aggregator/BUILD.bazel

@@ -17,45 +17,45 @@ scala_library(
     ],
 )
 
+test_deps = [
+    ":lib",
+    "//api:lib",
+    "//api:thrift",
+    # Libraries
+    maven_artifact("org.slf4j:slf4j-api"),
+    maven_artifact("org.apache.logging.log4j:log4j-slf4j-impl"),
+    maven_artifact("com.google.code.gson:gson"),
+    maven_artifact("org.apache.datasketches:datasketches-memory"),
+    maven_artifact("org.apache.datasketches:datasketches-java"),
+    maven_artifact("org.apache.commons:commons-lang3"),
+    maven_artifact("org.apache.commons:commons-math3"),
+    maven_artifact("org.scala-lang.modules:scala-collection-compat_2.12"),
+    # Testing
+    scala_artifact("org.scalatest:scalatest-matchers-core"),
+    scala_artifact("org.scalatest:scalatest-core"),
+    scala_artifact("org.scalatest:scalatest"),
+    scala_artifact("org.scalatest:scalatest-flatspec"),
+    scala_artifact("org.scalatest:scalatest-funsuite"),
+    scala_artifact("org.scalatest:scalatest-shouldmatchers"),
+    scala_artifact("org.scalactic:scalactic"),
+    scala_artifact("org.scalatestplus:mockito-3-4"),
+    scala_artifact("org.mockito:mockito-scala"),
+    maven_artifact("org.mockito:mockito-core"),
+    maven_artifact("org.scalatest:scalatest-compatible"),
+    maven_artifact("junit:junit"),
+    maven_artifact("com.novocode:junit-interface"),
+]
+
 scala_library(
     name = "test-lib",
     srcs = glob(["src/test/**/*.scala"]),
     visibility = ["//visibility:public"],
-    deps = [
-        ":lib",
-        "//api:lib",
-        "//api:thrift",
-        maven_artifact("junit:junit"),
-        maven_artifact("com.novocode:junit-interface"),
-        maven_artifact("org.slf4j:slf4j-api"),
-        maven_artifact("org.apache.logging.log4j:log4j-slf4j-impl"),
-        maven_artifact("com.google.code.gson:gson"),
-        maven_artifact("org.apache.datasketches:datasketches-memory"),
-        maven_artifact("org.apache.datasketches:datasketches-java"),
-        maven_artifact("org.apache.commons:commons-lang3"),
-        maven_artifact("org.apache.commons:commons-math3"),
-        maven_artifact("org.scala-lang.modules:scala-collection-compat_2.12"),
-    ],
+    deps = test_deps,
 )
 
 scala_test_suite(
     name = "test",
     srcs = glob(["src/test/**/*.scala"]),
     visibility = ["//visibility:public"],
-    deps = [
-       ":lib",
-       ":test-lib",
-       "//api:lib",
-       "//api:thrift",
-       maven_artifact("junit:junit"),
-       maven_artifact("com.novocode:junit-interface"),
-       maven_artifact("org.slf4j:slf4j-api"),
-       maven_artifact("org.apache.logging.log4j:log4j-slf4j-impl"),
-       maven_artifact("com.google.code.gson:gson"),
-       maven_artifact("org.apache.datasketches:datasketches-memory"),
-       maven_artifact("org.apache.datasketches:datasketches-java"),
-       maven_artifact("org.apache.commons:commons-lang3"),
-       maven_artifact("org.apache.commons:commons-math3"),
-       maven_artifact("org.scala-lang.modules:scala-collection-compat_2.12"),
-   ],
+    deps = test_deps + [":test-lib"],
 )

api/src/test/scala/ai/chronon/api/test/ExtensionsTest.scala

@@ -113,7 +113,7 @@ class ExtensionsTest extends AnyFlatSpec {
   }
 
   it should "group by keys should contain partition column" in {
-    val groupBy = spy(new GroupBy())
+    val groupBy = spy[GroupBy](new GroupBy())
     val baseKeys = List("a", "b")
     val partitionColumn = "ds"
     groupBy.accuracy = Accuracy.SNAPSHOT
@@ -127,7 +127,7 @@ class ExtensionsTest extends AnyFlatSpec {
   }
 
   it should "group by keys should contain time column for temporal accuracy" in {
-    val groupBy = spy(new GroupBy())
+    val groupBy = spy[GroupBy](new GroupBy())
     val baseKeys = List("a", "b")
     val partitionColumn = "ds"
     groupBy.accuracy = Accuracy.TEMPORAL

api/thrift/agent.thrift

@@ -0,0 +1,136 @@
+namespace java ai.chronon.api
+
+
+struct YarnAutoScalingSpec {
+    1: optional i32 minInstances
+    2: optional i32 maxInstances
+    3: optional i32 scaleUpFactor // 1.5x, 2x etc
+    4: optional i32 scaleDownFactor
+    5: optional string cooldownPeriod
+}
+
+// our clusters are created transiently prior to running the job
+struct YarnClusterSpec {
+    1: optional string clusterName
+    2: optional string hostType
+    3: optional i32 hostCount
+
+    // dataproc = x.y.z, emr = x.y.z, etc
+    10: optional string yarnOfferingVersion
+
+    // to access the right data and right back to kvstore
+    20: optional string networkPolicy
+    30: optional YarnAutoScalingSpec autoScalingSpec
+}
+
+enum YarnJobType {
+    SPARK = 0,
+    FLINK = 1,
+}
+
+struct YarnJob {
+    // create transient cluster with this name and runs an app with the same yarn name
+    1: optional string appName
+    2: optional YarnJobType jobType
+
+    10: optional list<string> args
+    11: optional map<string, string> env
+    12: optional map<string, string> conf
+    // creates local file with this name and contents - relative to cwd
+    // contains the groupBy, join, queries etc
+    13: optional map<string, string> fileWithContents
+
+    20: optional string chrononVersion
+    21: optional YarnClusterSpec clusterSpec
+}
+
+struct KvWrite {
+    1: optional string key
+    2: optional string value
+    3: optional string timestamp
+}
+
+// currently used for writing join metadata to kvstore needed prior to fetching joins
+struct KvWriteJob {
+    1: optional string scope // projectId in gcp, account name in aws
+    2: optional string dataset
+    3: optional string table
+    4: optional list<KvWrite> writes
+}
+
+struct PartitionListingJob {
+    1: optional string scope // projectId in gcp, account name in aws
+    2: optional string dataset
+    3: optional string table
+    4: optional string partitionColumn
+    5: optional list<string> extraPartitionFilters
+}
+
+// agent accepts jobs and runs them
+union JobBase {
+    1: YarnJob yarnJob
+    2: KvWriteJob kvWriteJob
+    3: PartitionListingJob partitionListingJob
+}
+
+struct Job {
+    1: optional string jobId
+    2: optional JobBase jobUnion
+    3: optional i32 statusReportInterval
+    4: optional i32 maxRetries
+}
+
+struct JobListGetRequest {
+    // this is only sent on the first request after a start
+    1: optional list<string> existingJobsIds
+}
+
+struct JobListResponse {
+    // controller responds with jobs data plane agent is not aware of
+    1: optional list<Job> jobsToStart
+    2: optional list<string> jobsToStop
+}
+
+enum JobStatusType {
+    PENDING = 0,
+    RUNNING = 1,
+    SUCCEEDED = 2,
+    FAILED = 3,
+    STOPPED = 4
+}
+
+struct ResourceUsage {
+    1: optional i64 vcoreSeconds
+    2: optional i64 megaByteSeconds
+    3: optional i64 cumulativeDiskWriteBytes
+    4: optional i64 cumulativeDiskReadBytes
+}
+
+struct YarnIncrementalJobStatus {
+    // batch / streaming job
+    1: optional map<JobStatusType, i64> statusChangeTimes
+    2: optional ResourceUsage resourceUsage
+    // driver logs - probably only errors and exceptions
+    3: optional list<string> logsSinceLastPush
+}
+
+struct JobInfo {
+    1: optional string jobId
+    2: optional JobStatusType currentStatus
+
+    10: optional YarnIncrementalJobStatus yarnIcrementalStatus
+}
+
+struct DatePartitionRange {
+    1: optional string start
+    2: optional string end
+}
+
+struct PartitionListingPutRequest {
+    1: optional map<PartitionListingJob, list<DatePartitionRange>> partitions
+    2: optional map<PartitionListingJob, string> errors
+}
+
+struct JobInfoPutRequest {
+    1: optional list<JobInfo> jobStatuses
+}