Fix workers

PostHog · Dec 20, 2024 · 12c5461 · 12c5461
1 parent f55dca1
commit 12c5461
Show file tree

Hide file tree

Showing 2 changed files with 23 additions and 18 deletions.
diff --git a/posthog/tasks/tasks.py b/posthog/tasks/tasks.py
@@ -730,9 +730,9 @@ def calculate_decide_usage() -> None:
 
     ph_client = get_ph_client()
 
-    capture_decide_usage_for_all_teams(ph_client)
-
-    ph_client.shutdown()
+    if ph_client:
+        capture_decide_usage_for_all_teams(ph_client)
+        ph_client.shutdown()
 
 
 @shared_task(ignore_result=True)

diff --git a/posthog/tasks/warehouse.py b/posthog/tasks/warehouse.py
@@ -35,23 +35,25 @@ def capture_workspace_rows_synced_by_team(team_id: int) -> None:
     team.external_data_workspace_last_synced_at = now
 
     for job in ExternalDataJob.objects.filter(team_id=team_id, created_at__gte=begin).order_by("created_at").all():
-        ph_client.capture(
-            team_id,
-            "$data_sync_job_completed",
-            {
-                "team_id": team_id,
-                "workspace_id": team.external_data_workspace_id,
-                "count": job.rows_synced,
-                "start_time": job.created_at,
-                "job_id": str(job.pk),
-            },
-        )
+        if ph_client:
+            ph_client.capture(
+                team_id,
+                "$data_sync_job_completed",
+                {
+                    "team_id": team_id,
+                    "workspace_id": team.external_data_workspace_id,
+                    "count": job.rows_synced,
+                    "start_time": job.created_at,
+                    "job_id": str(job.pk),
+                },
+            )
 
         team.external_data_workspace_last_synced_at = job.created_at
 
     team.save()
 
-    ph_client.shutdown()
+    if ph_client:
+        ph_client.shutdown()
 
 
 @shared_task(ignore_result=True)
@@ -66,14 +68,17 @@ def validate_data_warehouse_table_columns(team_id: int, table_id: str) -> None:
             table.columns[column]["valid"] = table.validate_column_type(column)
         table.save()
 
-        ph_client.capture(team_id, "validate_data_warehouse_table_columns succeeded")
+        if ph_client:
+            ph_client.capture(team_id, "validate_data_warehouse_table_columns succeeded")
     except Exception as e:
         logger.exception(
             f"validate_data_warehouse_table_columns raised an exception for table: {table_id}",
             exc_info=e,
             team_id=team_id,
         )
 
-        ph_client.capture(team_id, "validate_data_warehouse_table_columns errored")
+        if ph_client:
+            ph_client.capture(team_id, "validate_data_warehouse_table_columns errored")
     finally:
-        ph_client.shutdown()
+        if ph_client:
+            ph_client.shutdown()