coderabbit good idea

Author: david-zlai

api/py/ai/chronon/repo/run.py

@@ -135,6 +135,7 @@
 ZIPLINE_ONLINE_CLASS_DEFAULT = "ai.chronon.integrations.cloud_gcp.GcpApiImpl"
 ZIPLINE_FLINK_JAR_DEFAULT = "flink-assembly-0.1.0-SNAPSHOT.jar"
 ZIPLINE_DATAPROC_SUBMITTER_JAR = "cloud_gcp_submitter_deploy.jar"
+ZIPLINE_SERVICE_JAR = "service-0.1.0-SNAPSHOT.jar"
 
 ZIPLINE_DIRECTORY = "/tmp/zipline"
 
@@ -864,78 +865,27 @@ def generate_dataproc_submitter_args(user_args: str, job_type: DataprocJobType =
         raise ValueError(f"Invalid job type: {job_type}")
 
 
-def download_dataproc_submitter_jar(destination_dir: str, customer_id: str):
+def download_zipline_jar(destination_dir: str, customer_id: str, jar_name: str):
     bucket_name = f"zipline-artifacts-{customer_id}"
 
-    file_name = ZIPLINE_DATAPROC_SUBMITTER_JAR
-
-    source_blob_name = f"jars/{file_name}"
-    dataproc_jar_destination_path = f"{destination_dir}/{file_name}"
-
-    are_identical = compare_gcs_and_local_file_hashes(bucket_name, source_blob_name,
-                                                      dataproc_jar_destination_path) if os.path.exists(
-        dataproc_jar_destination_path) else False
-
-    if are_identical:
-        print(
-            f"{dataproc_jar_destination_path} matches GCS {bucket_name}/{source_blob_name}")
-    else:
-        print(
-            f"{dataproc_jar_destination_path} does NOT match GCS {bucket_name}/{source_blob_name}")
-        print("Downloading dataproc submitter jar from GCS...")
-        download_gcs_blob(bucket_name, source_blob_name,
-                          dataproc_jar_destination_path)
-
-    return dataproc_jar_destination_path
-
-
-def download_chronon_gcp_jar(destination_dir: str, customer_id: str):
-    bucket_name = f"zipline-artifacts-{customer_id}"
-
-    file_name = ZIPLINE_ONLINE_JAR_DEFAULT
-
-    source_blob_name = f"jars/{file_name}"
-    chronon_gcp_jar_destination_path = f"{destination_dir}/{file_name}"
-
-    are_identical = compare_gcs_and_local_file_hashes(bucket_name, source_blob_name,
-                                                      chronon_gcp_jar_destination_path) if os.path.exists(
-        chronon_gcp_jar_destination_path) else False
-
-    if are_identical:
-        print(
-            f"{chronon_gcp_jar_destination_path} matches GCS {bucket_name}/{source_blob_name}")
-    else:
-        print(
-            f"{chronon_gcp_jar_destination_path} does NOT match GCS {bucket_name}/{source_blob_name}")
-        print("Downloading chronon gcp jar from GCS...")
-        download_gcs_blob(bucket_name, source_blob_name,
-                          chronon_gcp_jar_destination_path)
-    return chronon_gcp_jar_destination_path
-
-
-def download_service_jar(destination_dir: str, customer_id: str):
-    bucket_name = f"zipline-artifacts-{customer_id}"
-
-    file_name = "service-0.1.0-SNAPSHOT.jar"
-
-    source_blob_name = f"jars/{file_name}"
-    service_jar_destination_path = f"{destination_dir}/{file_name}"
+    source_blob_name = f"jars/{jar_name}"
+    destination_path = f"{destination_dir}/{jar_name}"
 
     are_identical = compare_gcs_and_local_file_hashes(bucket_name, source_blob_name,
-                                                      service_jar_destination_path) if os.path.exists(
-        service_jar_destination_path) else False
+                                                      destination_path) if os.path.exists(
+        destination_path) else False
 
     if are_identical:
         print(
-            f"{service_jar_destination_path} matches GCS {bucket_name}/{source_blob_name}")
+            f"{destination_path} matches GCS {bucket_name}/{source_blob_name}")
     else:
         print(
-            f"{service_jar_destination_path} does NOT match GCS {bucket_name}/{source_blob_name}")
-        print("Downloading service jar from GCS...")
+            f"{destination_path} does NOT match GCS {bucket_name}/{source_blob_name}")
+        print(f"Downloading {jar_name} from GCS...")
 
         download_gcs_blob(bucket_name, source_blob_name,
-                          service_jar_destination_path)
-    return service_jar_destination_path
+                          destination_path)
+    return destination_path
 
 
 @retry_decorator(retries=2, backoff=5)
@@ -1085,12 +1035,12 @@ def main(ctx, conf, env, mode, dataproc, ds, app_name, start_ds, end_ds, paralle
     os.makedirs(ZIPLINE_DIRECTORY, exist_ok=True)
 
     if dataproc:
-        jar_path = download_dataproc_submitter_jar(ZIPLINE_DIRECTORY, get_customer_id())
+        jar_path = download_zipline_jar(ZIPLINE_DIRECTORY, get_customer_id(), ZIPLINE_DATAPROC_SUBMITTER_JAR)
     elif chronon_jar:
         jar_path = chronon_jar
     else:
-        service_jar_path = download_service_jar(ZIPLINE_DIRECTORY, get_customer_id())
-        chronon_gcp_jar_path = download_chronon_gcp_jar(ZIPLINE_DIRECTORY, get_customer_id())
+        service_jar_path = download_zipline_jar(ZIPLINE_DIRECTORY, get_customer_id(), ZIPLINE_SERVICE_JAR)
+        chronon_gcp_jar_path = download_zipline_jar(ZIPLINE_DIRECTORY, get_customer_id(), ZIPLINE_ONLINE_JAR_DEFAULT)
         jar_path = f"{service_jar_path}:{chronon_gcp_jar_path}"
 
     Runner(ctx.params, os.path.expanduser(jar_path)).run()