fix(batch-exports): Use existing schema in merge query (#24753)

PostHog · Sep 4, 2024 · a9d32a4 · a9d32a4
1 parent dccb835
commit a9d32a4
Show file tree

Hide file tree

Showing 2 changed files with 19 additions and 2 deletions.
diff --git a/posthog/temporal/batch_exports/bigquery_batch_export.py b/posthog/temporal/batch_exports/bigquery_batch_export.py
@@ -205,6 +205,7 @@ async def amerge_person_tables(
         final_table: bigquery.Table,
         stage_table: bigquery.Table,
         merge_key: collections.abc.Iterable[bigquery.SchemaField],
+        update_fields: collections.abc.Iterable[bigquery.SchemaField] | None = None,
         person_version_key: str = "person_version",
         person_distinct_id_version_key: str = "person_distinct_id_version",
     ):
@@ -221,7 +222,13 @@ async def amerge_person_tables(
         update_clause = ""
         values = ""
         field_names = ""
-        for n, field in enumerate(final_table.schema):
+
+        if not update_fields:
+            update_clause_fields = final_table.schema
+        else:
+            update_clause_fields = update_fields
+
+        for n, field in enumerate(update_clause_fields):
             if n > 0:
                 update_clause += ", "
                 values += ", "
@@ -231,6 +238,9 @@ async def amerge_person_tables(
             field_names += f"`{field.name}`"
             values += f"stage.`{field.name}`"
 
+        if not update_clause:
+            raise ValueError("Empty update clause")
+
         merge_query = f"""
         MERGE `{final_table.full_table_id.replace(":", ".", 1)}` final
         USING `{stage_table.full_table_id.replace(":", ".", 1)}` stage
@@ -481,6 +491,7 @@ async def flush_to_bigquery(
                         final_table=bigquery_table,
                         stage_table=bigquery_stage_table,
                         merge_key=merge_key,
+                        update_fields=schema,
                     )
 
                 return writer.records_total

diff --git a/posthog/temporal/tests/batch_exports/test_bigquery_batch_export_workflow.py b/posthog/temporal/tests/batch_exports/test_bigquery_batch_export_workflow.py
@@ -5,6 +5,7 @@
 import os
 import typing
 import uuid
+import warnings
 
 import pyarrow as pa
 import pytest
@@ -213,7 +214,12 @@ def bigquery_dataset(bigquery_config, bigquery_client) -> typing.Generator[bigqu
 
     yield dataset
 
-    # bigquery_client.delete_dataset(dataset_id, delete_contents=True, not_found_ok=True)
+    try:
+        bigquery_client.delete_dataset(dataset_id, delete_contents=True, not_found_ok=True)
+    except Exception as exc:
+        warnings.warn(
+            f"Failed to clean up dataset: {dataset_id} due to '{exc.__class__.__name__}': {str(exc)}", stacklevel=1
+        )
 
 
 @pytest.fixture