From 571add2456c478207bf2bc53eb701ffcf6631087 Mon Sep 17 00:00:00 2001
From: Tom Owers <tom@posthog.com>
Date: Mon, 4 Nov 2024 14:25:27 +0000
Subject: [PATCH] fix(data-warehouse): Use database async wrapper (#25971)

---
 .../data_imports/workflow_activities/create_job_model.py  | 4 ++--
 posthog/temporal/tests/data_imports/test_end_to_end.py    | 6 +++---
 .../tests/external_data/test_external_data_job.py         | 8 ++++----
 posthog/warehouse/external_data_source/jobs.py            | 3 ++-
 4 files changed, 11 insertions(+), 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
index 9cf06b39808b5..2cef81a1ce62c 100644
--- 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
index b9c74c4efefff..9ba969e196cea 100644
--- 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
index 1b5b3d692d5ed..0a59c3aa19a23 100644
--- 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
index 81edddbffaaff..d21210f2ec097 100644
--- 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,