feat(data-warehouse): external data job rewrite

latest_migrations.manifest

@@ -5,7 +5,7 @@ contenttypes: 0002_remove_content_type_name
 ee: 0016_rolemembership_organization_member
 otp_static: 0002_throttling
 otp_totp: 0002_auto_20190420_0723
-posthog: 0400_datawarehousetable_row_count
+posthog: 0401_externaldatajob_schema
 sessions: 0001_initial
 social_django: 0010_uid_db_index
 two_factor: 0007_auto_20201201_1019

posthog/migrations/0401_externaldatajob_schema.py

@@ -0,0 +1,20 @@
+# Generated by Django 4.1.13 on 2024-04-10 15:12
+
+from django.db import migrations, models
+import django.db.models.deletion
+
+
+class Migration(migrations.Migration):
+    dependencies = [
+        ("posthog", "0400_datawarehousetable_row_count"),
+    ]
+
+    operations = [
+        migrations.AddField(
+            model_name="externaldatajob",
+            name="schema",
+            field=models.ForeignKey(
+                null=True, on_delete=django.db.models.deletion.CASCADE, to="posthog.externaldataschema"
+            ),
+        ),
+    ]

posthog/temporal/common/client.py

@@ -49,3 +49,16 @@ async def sync_connect() -> Client:
         settings.TEMPORAL_CLIENT_KEY,
     )
     return client
+
+
+async def async_connect() -> Client:
+    """Asynchronous connect to Temporal and return a Client."""
+    client = await connect(
+        settings.TEMPORAL_HOST,
+        settings.TEMPORAL_PORT,
+        settings.TEMPORAL_NAMESPACE,
+        settings.TEMPORAL_CLIENT_ROOT_CA,
+        settings.TEMPORAL_CLIENT_CERT,
+        settings.TEMPORAL_CLIENT_KEY,
+    )
+    return client

posthog/temporal/common/schedule.py

@@ -12,6 +12,15 @@ async def create_schedule(temporal: Client, id: str, schedule: Schedule, trigger
     )
 
 
+async def a_create_schedule(temporal: Client, id: str, schedule: Schedule, trigger_immediately: bool = False):
+    """Async create a Temporal Schedule."""
+    return await temporal.create_schedule(
+        id=id,
+        schedule=schedule,
+        trigger_immediately=trigger_immediately,
+    )
+
+
 @async_to_sync
 async def update_schedule(temporal: Client, id: str, schedule: Schedule) -> None:
     """Update a Temporal Schedule."""
@@ -25,6 +34,18 @@ async def updater(_: ScheduleUpdateInput) -> ScheduleUpdate:
     )
 
 
+async def a_update_schedule(temporal: Client, id: str, schedule: Schedule) -> None:
+    """Async update a Temporal Schedule."""
+    handle = temporal.get_schedule_handle(id)
+
+    async def updater(_: ScheduleUpdateInput) -> ScheduleUpdate:
+        return ScheduleUpdate(schedule=schedule)
+
+    return await handle.update(
+        updater=updater,
+    )
+
+
 @async_to_sync
 async def unpause_schedule(temporal: Client, schedule_id: str, note: str | None = None) -> None:
     """Unpause a Temporal Schedule."""
@@ -39,6 +60,12 @@ async def delete_schedule(temporal: Client, schedule_id: str) -> None:
     await handle.delete()
 
 
+async def a_delete_schedule(temporal: Client, schedule_id: str) -> None:
+    """Async delete a Temporal Schedule."""
+    handle = temporal.get_schedule_handle(schedule_id)
+    await handle.delete()
+
+
 @async_to_sync
 async def describe_schedule(temporal: Client, schedule_id: str):
     """Describe a Temporal Schedule."""
@@ -55,6 +82,21 @@ async def pause_schedule(temporal: Client, schedule_id: str, note: str | None =
 
 @async_to_sync
 async def trigger_schedule(temporal: Client, schedule_id: str, note: str | None = None) -> None:
-    """Pause a Temporal Schedule."""
+    """Trigger a Temporal Schedule."""
     handle = temporal.get_schedule_handle(schedule_id)
     await handle.trigger()
+
+
+async def a_trigger_schedule(temporal: Client, schedule_id: str, note: str | None = None) -> None:
+    """Trigger a Temporal Schedule."""
+    handle = temporal.get_schedule_handle(schedule_id)
+    await handle.trigger()
+
+
+async def a_schedule_exists(temporal: Client, schedule_id: str) -> bool:
+    """Check whether a schedule exists."""
+    try:
+        await temporal.get_schedule_handle(schedule_id).describe()
+        return True
+    except:
+        return False

posthog/temporal/data_imports/__init__.py

@@ -1,18 +1,20 @@
 from posthog.temporal.data_imports.external_data_job import (
     ExternalDataJobWorkflow,
-    create_external_data_job_model,
+    create_external_data_job_model_activity,
     create_source_templates,
-    run_external_data_job,
+    import_data_activity,
     update_external_data_job_model,
     validate_schema_activity,
+    check_schedule_activity,
 )
 
 WORKFLOWS = [ExternalDataJobWorkflow]
 
 ACTIVITIES = [
-    create_external_data_job_model,
+    create_external_data_job_model_activity,
     update_external_data_job_model,
-    run_external_data_job,
+    import_data_activity,
     validate_schema_activity,
     create_source_templates,
+    check_schedule_activity,
 ]