fix(data-warehouse): Fixed row counts being missing (#23726)

PostHog · Jul 16, 2024 · b58c853 · b58c853
1 parent e963e0a
commit b58c853
Showing 1 changed file with 10 additions and 2 deletions.
diff --git a/posthog/temporal/data_imports/pipelines/pipeline.py b/posthog/temporal/data_imports/pipelines/pipeline.py
@@ -120,7 +120,10 @@ def _run(self) -> dict[str, int]:
                         if e.exception.args[0] != "Generic error: No data source supplied to write command.":
                             raise
 
-                row_counts = pipeline.last_trace.last_normalize_info.row_counts
+                if pipeline.last_trace.last_normalize_info is not None:
+                    row_counts = pipeline.last_trace.last_normalize_info.row_counts
+                else:
+                    row_counts = {}
                 # Remove any DLT tables from the counts
                 filtered_rows = dict(filter(lambda pair: not pair[0].startswith("_dlt"), row_counts.items()))
                 counts = Counter(filtered_rows)
@@ -148,7 +151,12 @@ def _run(self) -> dict[str, int]:
                 if isinstance(e.exception, DeltaError):
                     if e.exception.args[0] != "Generic error: No data source supplied to write command.":
                         raise
-            row_counts = pipeline.last_trace.last_normalize_info.row_counts
+
+            if pipeline.last_trace.last_normalize_info is not None:
+                row_counts = pipeline.last_trace.last_normalize_info.row_counts
+            else:
+                row_counts = {}
+
             filtered_rows = dict(filter(lambda pair: not pair[0].startswith("_dlt"), row_counts.items()))
             counts = Counter(filtered_rows)
             total_counts = total_counts + counts