fix(data-warehouse): Use database async wrapper (#25971)

PostHog · Nov 4, 2024 · 571add2 · 571add2
1 parent 0011ba8
commit 571add2
Show file tree

Hide file tree

Showing 4 changed files with 11 additions and 10 deletions.
diff --git a/posthog/temporal/data_imports/workflow_activities/create_job_model.py b/posthog/temporal/data_imports/workflow_activities/create_job_model.py
@@ -7,7 +7,7 @@
 # TODO: remove dependency
 
 from posthog.warehouse.external_data_source.jobs import (
-    create_external_data_job,
+    acreate_external_data_job,
 )
 from posthog.warehouse.models.external_data_schema import (
     ExternalDataSchema,
@@ -27,7 +27,7 @@ async def create_external_data_job_model_activity(inputs: CreateExternalDataJobM
     logger = await bind_temporal_worker_logger(team_id=inputs.team_id)
 
     try:
-        job = await sync_to_async(create_external_data_job)(
+        job = await acreate_external_data_job(
             team_id=inputs.team_id,
             external_data_source_id=inputs.source_id,
             external_data_schema_id=inputs.schema_id,

diff --git a/posthog/temporal/tests/data_imports/test_end_to_end.py b/posthog/temporal/tests/data_imports/test_end_to_end.py
@@ -908,9 +908,9 @@ def get_jobs():
         return list(jobs)
 
     with mock.patch(
-        "posthog.temporal.data_imports.workflow_activities.create_job_model.create_external_data_job",
-    ) as mock_list_limited_team_attributes:
-        mock_list_limited_team_attributes.side_effect = Exception("Ruhoh!")
+        "posthog.temporal.data_imports.workflow_activities.create_job_model.acreate_external_data_job",
+    ) as acreate_external_data_job:
+        acreate_external_data_job.side_effect = Exception("Ruhoh!")
 
         with pytest.raises(Exception):
             await _execute_run(workflow_id, inputs, stripe_customer["data"])

diff --git a/posthog/temporal/tests/external_data/test_external_data_job.py b/posthog/temporal/tests/external_data/test_external_data_job.py
@@ -25,7 +25,7 @@
     SyncNewSchemasActivityInputs,
     sync_new_schemas_activity,
 )
-from posthog.warehouse.external_data_source.jobs import create_external_data_job
+from posthog.warehouse.external_data_source.jobs import acreate_external_data_job
 from posthog.warehouse.models import (
     get_latest_run_if_exists,
     ExternalDataJob,
@@ -238,7 +238,7 @@ async def test_update_external_job_activity(activity_environment, team, **kwargs
         should_sync=True,
     )
 
-    new_job = await sync_to_async(create_external_data_job)(
+    new_job = await acreate_external_data_job(
         team_id=team.id,
         external_data_source_id=new_source.pk,
         workflow_id=activity_environment.info.workflow_id,
@@ -283,7 +283,7 @@ async def test_update_external_job_activity_with_retryable_error(activity_enviro
         should_sync=True,
     )
 
-    new_job = await sync_to_async(create_external_data_job)(
+    new_job = await acreate_external_data_job(
         team_id=team.id,
         external_data_source_id=new_source.pk,
         workflow_id=activity_environment.info.workflow_id,
@@ -329,7 +329,7 @@ async def test_update_external_job_activity_with_non_retryable_error(activity_en
         should_sync=True,
     )
 
-    new_job = await sync_to_async(create_external_data_job)(
+    new_job = await acreate_external_data_job(
         team_id=team.id,
         external_data_source_id=new_source.pk,
         workflow_id=activity_environment.info.workflow_id,

diff --git a/posthog/warehouse/external_data_source/jobs.py b/posthog/warehouse/external_data_source/jobs.py
@@ -9,7 +9,8 @@ def get_external_data_source(team_id: str, external_data_source_id: str) -> Exte
     return ExternalDataSource.objects.get(team_id=team_id, id=external_data_source_id)
 
 
-def create_external_data_job(
+@database_sync_to_async
+def acreate_external_data_job(
     external_data_source_id: UUID,
     external_data_schema_id: UUID,
     workflow_id: str,