[ch] Refactor/unify s3 lambda, improve s3 backfill script (#5730)

Refactor the s3 replicator lambda to reduce dup code and have better
error handling (upload to same database so we don't have to create a new
one each time)

Tested via:
* randomly selecting a key for merges, merge_bases,
rerun_disabled_tests, external_contribution_stats and checking that it
could be uploaded
* backfilled inductor_torchao_perf_stats and torchbench_userbenchmark
* rebackfilled queue_times_historical and removed dups
* raise an exception to see if the error handler wrote to the database
(it did)

Reused logic from lambda replicator in s3 backfill script

Author: clee2000

aws/lambda/clickhouse-replicator-s3/lambda_function.py

@@ -99,6 +99,8 @@ def get_skipped_failure_parser_helper(name, type, field_to_check_for_existence):
         ) as {name}
         """
 
+    # Cannot use general_adapter due to custom field for now()::DateTime64(9)
+    # time_inserted
     query = f"""
     insert into {table}
     select
@@ -158,14 +160,7 @@ def get_skipped_failure_parser_helper(name, type, field_to_check_for_existence):
     try:
         get_clickhouse_client().query(query)
     except Exception as e:
-        if "Expected not greater than" in str(e):
-            get_clickhouse_client().query(
-                f"insert into errors.{table}_ingest_errors values ('{key}', 'file is too large?')"
-            )
-        else:
-            get_clickhouse_client().query(
-                f"insert into errors.{table}_ingest_errors values ('{key}', '{json.dumps(str(e))}')"
-            )
+        log_failure_to_clickhouse(table, bucket, key, e)
 
 
 def rerun_disabled_tests_adapter(table, bucket, key):
@@ -180,26 +175,7 @@ def rerun_disabled_tests_adapter(table, bucket, key):
     `workflow_run_attempt` Int64
     """
 
-    url = f"https://{bucket}.s3.amazonaws.com/{encode_url_component(key)}"
-
-    def get_insert_query(compression):
-        return f"""
-        insert into {table}
-        select *, ('{bucket}', '{key}')
-        from s3('{url}', 'JSONEachRow', '{schema}', '{compression}')
-        """
-
-    try:
-        get_clickhouse_client().query(get_insert_query("gzip"))
-    except Exception as e:
-        if "Expected not greater than" in str(e):
-            get_clickhouse_client().query(
-                f"insert into errors.{table}_ingest_errors values ('{key}', 'file is too large?')"
-            )
-        else:
-            get_clickhouse_client().query(
-                f"insert into errors.{table}_ingest_errors values ('{key}', '{json.dumps(str(e))}')"
-            )
+    general_adapter(table, bucket, key, schema, ["gzip"], "JSONEachRow")
 
 
 def handle_test_run_summary(table, bucket, key) -> None:
@@ -219,6 +195,7 @@ def handle_test_run_summary(table, bucket, key) -> None:
     """
     url = f"https://{bucket}.s3.amazonaws.com/{encode_url_component(key)}"
 
+    # Cannot use general_adapter due to custom field for now()::DateTime64(9)
     def get_insert_query(compression):
         return f"""
         insert into {table} SETTINGS async_insert=1, wait_for_async_insert=1
@@ -229,14 +206,7 @@ def get_insert_query(compression):
     try:
         get_clickhouse_client().query(get_insert_query("gzip"))
     except Exception as e:
-        if "Expected not greater than" in str(e):
-            get_clickhouse_client().query(
-                f"insert into errors.{table}_ingest_errors values ('{key}', 'file is too large?')"
-            )
-        else:
-            get_clickhouse_client().query(
-                f"insert into errors.{table}_ingest_errors values ('{key}', '{json.dumps(str(e))}')"
-            )
+        log_failure_to_clickhouse(table, bucket, key, e)
 
 
 def merges_adapter(table, bucket, key) -> None:
@@ -263,21 +233,7 @@ def merges_adapter(table, bucket, key) -> None:
     `unstable_checks` Array(Array(String))
     """
 
-    url = f"https://{bucket}.s3.amazonaws.com/{encode_url_component(key)}"
-
-    def get_insert_query(compression):
-        return f"""
-        insert into {table}
-        select *, ('{bucket}', '{key}')
-        from s3('{url}', 'JSONEachRow', '{schema}', '{compression}')
-        """
-
-    try:
-        get_clickhouse_client().query(get_insert_query("none"))
-    except Exception as e:
-        get_clickhouse_client().query(
-            f"insert into errors.{table}_ingest_errors values ('{key}', '{json.dumps(str(e))}')"
-        )
+    general_adapter(table, bucket, key, schema, ["none"], "JSONEachRow")
 
 
 def merge_bases_adapter(table, bucket, key) -> None:
@@ -289,19 +245,7 @@ def merge_bases_adapter(table, bucket, key) -> None:
     `sha` String
     """
 
-    url = f"https://{bucket}.s3.amazonaws.com/{encode_url_component(key)}"
-
-    def get_insert_query(compression):
-        return f"""
-        insert into {table}
-        select *, ('{bucket}', '{key}')
-        from s3('{url}', 'JSONEachRow', '{schema}', '{compression}')
-        """
-
-    try:
-        get_clickhouse_client().query(get_insert_query("gzip"))
-    except:
-        get_clickhouse_client().query(get_insert_query("none"))
+    general_adapter(table, bucket, key, schema, ["gzip", "none"], "JSONEachRow")
 
 
 def queue_times_historical_adapter(table, bucket, key) -> None:
@@ -311,20 +255,7 @@ def queue_times_historical_adapter(table, bucket, key) -> None:
     `count` Int64,
     `time` DateTime64(9)
     """
-
-    url = f"https://{bucket}.s3.amazonaws.com/{encode_url_component(key)}"
-
-    def get_insert_query(compression):
-        return f"""
-        insert into {table}
-        select *, ('{bucket}', '{key}')
-        from s3('{url}', 'JSONEachRow', '{schema}', '{compression}')
-        """
-
-    try:
-        get_clickhouse_client().query(get_insert_query("gzip"))
-    except:
-        get_clickhouse_client().query(get_insert_query("none"))
+    general_adapter(table, bucket, key, schema, ["gzip", "none"], "JSONEachRow")
 
 
 def external_contribution_stats_adapter(table, bucket, key) -> None:
@@ -334,28 +265,22 @@ def external_contribution_stats_adapter(table, bucket, key) -> None:
     `user_count` Int64,
     `users` Array(String)
     """
-    url = f"https://{bucket}.s3.amazonaws.com/{encode_url_component(key)}"
-
-    def get_insert_query(compression):
-        return f"""
-        insert into {table}
-        select *, ('{bucket}', '{key}')
-        from s3('{url}', 'JSONEachRow', '{schema}', '{compression}',
-            extra_credentials(
-                role_arn = 'arn:aws:iam::308535385114:role/clickhouse_role'
-            )
-        )
-        """
-
-    try:
-        get_clickhouse_client().query(get_insert_query("gzip"))
-    except Exception as e:
-        get_clickhouse_client().query(
-            f"insert into errors.gen_errors VALUES ('{table}', '{bucket}', '{key}', '{json.dumps(str(e))}')"
-        )
+    general_adapter(table, bucket, key, schema, ["gzip"], "JSONEachRow")
+
+
+def log_failure_to_clickhouse(table, bucket, key, error) -> None:
+    error = {
+        "table": table,
+        "bucket": bucket,
+        "key": key,
+        "reason": str(error),
+    }
+    get_clickhouse_client().query(
+        f"insert into errors.gen_errors format JSONEachRow {json.dumps(error)}"
+    )
 
 
-def general_adapter(table, bucket, key, schema, compression, format) -> None:
+def general_adapter(table, bucket, key, schema, compressions, format) -> None:
     url = f"https://{bucket}.s3.amazonaws.com/{encode_url_component(key)}"
 
     def get_insert_query(compression):
@@ -370,11 +295,18 @@ def get_insert_query(compression):
         """
 
     try:
-        get_clickhouse_client().query(get_insert_query(compression))
-    except Exception as e:
-        get_clickhouse_client().query(
-            f"insert into errors.gen_errors values ('{table}', '{bucket}', '{key}', '{json.dumps(str(e))}')"
+        exceptions = []
+        for compression in compressions:
+            try:
+                get_clickhouse_client().query(get_insert_query(compression))
+                return
+            except Exception as e:
+                exceptions.append(e)
+        raise Exception(
+            f"Failed to insert into {table} with {[str(x) for x in exceptions]}"
         )
+    except Exception as e:
+        log_failure_to_clickhouse(table, bucket, key, e)
 
 
 def external_aggregated_test_metrics_adapter(table, bucket, key) -> None:
@@ -399,7 +331,7 @@ def external_aggregated_test_metrics_adapter(table, bucket, key) -> None:
     `workflow_name` String,
     `workflow_run_attempt` Int64
     """
-    general_adapter(table, bucket, key, schema, "gzip", "JSONEachRow")
+    general_adapter(table, bucket, key, schema, ["gzip"], "JSONEachRow")
 
 
 def torchao_perf_stats_adapter(table, bucket, key) -> None:
@@ -444,7 +376,7 @@ def torchao_perf_stats_adapter(table, bucket, key) -> None:
     `unique_graphs` String,
     `workflow_id` String
     """
-    general_adapter(table, bucket, key, schema, "none", "CSV")
+    general_adapter(table, bucket, key, schema, ["none"], "CSV")
 
 
 def torchbench_userbenchmark_adapter(table, bucket, key):
@@ -453,7 +385,7 @@ def torchbench_userbenchmark_adapter(table, bucket, key):
     `metrics` String,
     `name` String
     """
-    general_adapter(table, bucket, key, schema, "none", "JSONEachRow")
+    general_adapter(table, bucket, key, schema, ["none"], "JSONEachRow")
 
 
 SUPPORTED_PATHS = {

tools/rockset_migration/dynamo2ch.py

@@ -160,7 +160,7 @@ def upload_to_clickhouse(records, table):
     # Async insert to maybe make insertions more efficient on the ClickHouse side
     # https://clickhouse.com/docs/en/cloud/bestpractices/asynchronous-inserts
     get_clickhouse_client().query(
-        f"INSERT INTO `{table}`  SETTINGS async_insert=1, wait_for_async_insert=1  FORMAT JSONEachRow {body}"
+        f"INSERT INTO {table}  SETTINGS async_insert=1, wait_for_async_insert=1  FORMAT JSONEachRow {body}"
     )