feat(array): split out snapshot creation from task processing (#42488)

Author: joshsny

products/tasks/backend/models.py

@@ -245,19 +245,8 @@ def get_artifact_s3_prefix(self) -> str:
 
     def append_log(self, entries: list[dict]):
         """Append log entries to S3 storage."""
-
-        existing_content = ""
-        is_new_file = False
-        try:
-            existing_content = object_storage.read(self.log_url) or ""
-        except Exception as e:
-            is_new_file = True
-            logger.debug(
-                "task_run.no_existing_logs",
-                task_run_id=str(self.id),
-                log_url=self.log_url,
-                error=str(e),
-            )
+        existing_content = object_storage.read(self.log_url, missing_ok=True) or ""
+        is_new_file = not existing_content
 
         new_lines = "\n".join(json.dumps(entry) for entry in entries)
         content = existing_content + ("\n" if existing_content else "") + new_lines

products/tasks/backend/temporal/__init__.py

@@ -1,32 +1,38 @@
+from .create_snapshot.activities import (
+    cleanup_sandbox as snapshot_cleanup_sandbox,
+    clone_repository as snapshot_clone_repository,
+    create_sandbox as snapshot_create_sandbox,
+    create_snapshot as snapshot_create_snapshot,
+    get_snapshot_context,
+    setup_repository as snapshot_setup_repository,
+)
+from .create_snapshot.workflow import CreateSnapshotForRepositoryWorkflow
 from .process_task.activities import (
-    check_snapshot_exists_for_repository,
-    cleanup_personal_api_key,
     cleanup_sandbox,
-    clone_repository,
-    create_sandbox_from_snapshot,
-    create_snapshot,
     execute_task_in_sandbox,
-    get_sandbox_for_setup,
+    get_sandbox_for_repository,
     get_task_processing_context,
-    setup_repository,
     track_workflow_event,
 )
 from .process_task.workflow import ProcessTaskWorkflow
 
 WORKFLOWS = [
     ProcessTaskWorkflow,
+    CreateSnapshotForRepositoryWorkflow,
 ]
 
 ACTIVITIES = [
+    # process_task activities
     get_task_processing_context,
-    check_snapshot_exists_for_repository,
-    get_sandbox_for_setup,
-    clone_repository,
-    setup_repository,
-    create_snapshot,
-    create_sandbox_from_snapshot,
+    get_sandbox_for_repository,
     execute_task_in_sandbox,
-    cleanup_personal_api_key,
     cleanup_sandbox,
     track_workflow_event,
+    # create_snapshot activities
+    get_snapshot_context,
+    snapshot_create_sandbox,
+    snapshot_clone_repository,
+    snapshot_setup_repository,
+    snapshot_create_snapshot,
+    snapshot_cleanup_sandbox,
 ]

products/tasks/backend/temporal/conftest.py

@@ -5,11 +5,13 @@
 
 from temporalio.testing import ActivityEnvironment
 
-from posthog.models import Integration, Organization, OrganizationMembership, Team, User
+from posthog.models import Integration, OAuthApplication, Organization, OrganizationMembership, Team, User
 from posthog.temporal.common.logger import configure_logger
 
 from products.tasks.backend.models import SandboxSnapshot, Task, TaskRun
 from products.tasks.backend.services.sandbox import Sandbox, SandboxConfig, SandboxTemplate
+from products.tasks.backend.temporal.create_snapshot.activities.get_snapshot_context import SnapshotContext
+from products.tasks.backend.temporal.oauth import ARRAY_APP_CLIENT_ID_DEV
 from products.tasks.backend.temporal.process_task.activities.get_task_processing_context import TaskProcessingContext
 
 
@@ -19,6 +21,22 @@ def activity_environment():
     return ActivityEnvironment()
 
 
+@pytest.fixture
+def array_oauth_app():
+    """Create the Array OAuth application for tests."""
+    app, _ = OAuthApplication.objects.get_or_create(
+        client_id=ARRAY_APP_CLIENT_ID_DEV,
+        defaults={
+            "name": "Array Test App",
+            "client_type": OAuthApplication.CLIENT_PUBLIC,
+            "authorization_grant_type": OAuthApplication.GRANT_AUTHORIZATION_CODE,
+            "redirect_uris": "https://app.posthog.com/callback",
+            "algorithm": "RS256",
+        },
+    )
+    yield app
+
+
 @pytest.fixture
 def organization():
     """A test organization."""
@@ -43,6 +61,12 @@ def team(organization):
     team.delete()
 
 
+@pytest.fixture
+def test_team(team):
+    """Alias for team fixture."""
+    return team
+
+
 @pytest.fixture
 def user(team):
     user = User.objects.create(
@@ -75,7 +99,7 @@ def github_integration(team):
 
 
 @pytest.fixture
-def test_task(team, user, github_integration):
+def test_task(team, user, github_integration, array_oauth_app):
     """Create a test task."""
 
     task = Task.objects.create(
@@ -120,6 +144,16 @@ def task_context(test_task, test_task_run) -> TaskProcessingContext:
     )
 
 
+@pytest.fixture
+def snapshot_context(github_integration, team) -> SnapshotContext:
+    """Create a SnapshotContext for testing."""
+    return SnapshotContext(
+        github_integration_id=github_integration.id,
+        repository="posthog/posthog-js",
+        team_id=team.id,
+    )
+
+
 @pytest.fixture(autouse=True)
 def configure_logger_auto() -> None:
     """Configure logger when running in a Temporal activity environment."""

products/tasks/backend/temporal/create_snapshot/__init__.py

@@ -0,0 +1,41 @@
+from .activities import (
+    CleanupSandboxInput,
+    CloneRepositoryInput,
+    CreateSandboxInput,
+    CreateSandboxOutput,
+    CreateSnapshotInput,
+    GetSnapshotContextInput,
+    SetupRepositoryInput,
+    SnapshotContext,
+    cleanup_sandbox,
+    clone_repository,
+    create_sandbox,
+    create_snapshot,
+    get_snapshot_context,
+    setup_repository,
+)
+from .workflow import (
+    CreateSnapshotForRepositoryInput,
+    CreateSnapshotForRepositoryOutput,
+    CreateSnapshotForRepositoryWorkflow,
+)
+
+__all__ = [
+    "CleanupSandboxInput",
+    "CloneRepositoryInput",
+    "CreateSandboxInput",
+    "CreateSandboxOutput",
+    "CreateSnapshotForRepositoryInput",
+    "CreateSnapshotForRepositoryOutput",
+    "CreateSnapshotForRepositoryWorkflow",
+    "CreateSnapshotInput",
+    "GetSnapshotContextInput",
+    "SetupRepositoryInput",
+    "SnapshotContext",
+    "cleanup_sandbox",
+    "clone_repository",
+    "create_sandbox",
+    "create_snapshot",
+    "get_snapshot_context",
+    "setup_repository",
+]

products/tasks/backend/temporal/create_snapshot/activities/__init__.py

@@ -0,0 +1,23 @@
+from .cleanup_sandbox import CleanupSandboxInput, cleanup_sandbox
+from .clone_repository import CloneRepositoryInput, clone_repository
+from .create_sandbox import CreateSandboxInput, CreateSandboxOutput, create_sandbox
+from .create_snapshot import CreateSnapshotInput, create_snapshot
+from .get_snapshot_context import GetSnapshotContextInput, SnapshotContext, get_snapshot_context
+from .setup_repository import SetupRepositoryInput, setup_repository
+
+__all__ = [
+    "CleanupSandboxInput",
+    "CloneRepositoryInput",
+    "CreateSandboxInput",
+    "CreateSandboxOutput",
+    "CreateSnapshotInput",
+    "GetSnapshotContextInput",
+    "SetupRepositoryInput",
+    "SnapshotContext",
+    "cleanup_sandbox",
+    "clone_repository",
+    "create_sandbox",
+    "create_snapshot",
+    "get_snapshot_context",
+    "setup_repository",
+]

products/tasks/backend/temporal/create_snapshot/activities/cleanup_sandbox.py

@@ -0,0 +1,31 @@
+import logging
+from dataclasses import dataclass
+
+from temporalio import activity
+
+from posthog.temporal.common.utils import asyncify
+
+from products.tasks.backend.services.sandbox import Sandbox
+from products.tasks.backend.temporal.observability import log_activity_execution
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class CleanupSandboxInput:
+    sandbox_id: str
+
+
+@activity.defn
+@asyncify
+def cleanup_sandbox(input: CleanupSandboxInput) -> None:
+    with log_activity_execution(
+        "cleanup_sandbox",
+        sandbox_id=input.sandbox_id,
+    ):
+        try:
+            sandbox = Sandbox.get_by_id(input.sandbox_id)
+            sandbox.destroy()
+        except Exception:
+            # The sandbox has a timeout, and it will eventually terminate if we failed to cleanup
+            pass

products/tasks/backend/temporal/process_task/activities/clone_repository.py renamed to products/tasks/backend/temporal/create_snapshot/activities/clone_repository.py

@@ -6,37 +6,34 @@
 
 from products.tasks.backend.services.sandbox import Sandbox
 from products.tasks.backend.temporal.exceptions import GitHubAuthenticationError, RepositoryCloneError
-from products.tasks.backend.temporal.observability import emit_agent_log, log_activity_execution
+from products.tasks.backend.temporal.observability import log_activity_execution
 
 from ..utils import get_github_token
-from .get_task_processing_context import TaskProcessingContext
+from .get_snapshot_context import SnapshotContext
 
 
 @dataclass
 class CloneRepositoryInput:
-    context: TaskProcessingContext
+    context: SnapshotContext
     sandbox_id: str
 
 
 @activity.defn
 @asyncify
 def clone_repository(input: CloneRepositoryInput) -> str:
-    """Clone repository into sandbox. Idempotent: wipes existing directory. Returns clone logs."""
     ctx = input.context
 
     with log_activity_execution(
         "clone_repository",
         sandbox_id=input.sandbox_id,
         **ctx.to_log_context(),
     ):
-        emit_agent_log(ctx.run_id, "info", f"Cloning repository {ctx.repository}")
-
         try:
             github_token = get_github_token(ctx.github_integration_id)
         except Exception as e:
             raise GitHubAuthenticationError(
                 f"Failed to get GitHub token for integration {ctx.github_integration_id}",
-                {"github_integration_id": ctx.github_integration_id, "task_id": ctx.task_id, "error": str(e)},
+                {"github_integration_id": ctx.github_integration_id, "error": str(e)},
                 cause=e,
             )
 
@@ -50,7 +47,6 @@ def clone_repository(input: CloneRepositoryInput) -> str:
                 {
                     "repository": ctx.repository,
                     "sandbox_id": input.sandbox_id,
-                    "task_id": ctx.task_id,
                     "error": str(e),
                 },
                 cause=e,
@@ -63,10 +59,8 @@ def clone_repository(input: CloneRepositoryInput) -> str:
                     "repository": ctx.repository,
                     "exit_code": result.exit_code,
                     "stderr": result.stderr[:500],
-                    "task_id": ctx.task_id,
                 },
                 cause=RuntimeError(f"Git clone exited with code {result.exit_code}: {result.stderr[:200]}"),
             )
 
-        # NOTE: git clone returns it's output in stderr
         return result.stderr

products/tasks/backend/temporal/create_snapshot/activities/create_sandbox.py

@@ -0,0 +1,45 @@
+from dataclasses import dataclass
+
+from temporalio import activity
+
+from posthog.temporal.common.utils import asyncify
+
+from products.tasks.backend.services.sandbox import Sandbox, SandboxConfig, SandboxTemplate
+from products.tasks.backend.temporal.observability import log_activity_execution
+
+from ..utils import get_sandbox_name_for_snapshot
+from .get_snapshot_context import SnapshotContext
+
+
+@dataclass
+class CreateSandboxInput:
+    context: SnapshotContext
+
+
+@dataclass
+class CreateSandboxOutput:
+    sandbox_id: str
+
+
+@activity.defn
+@asyncify
+def create_sandbox(input: CreateSandboxInput) -> CreateSandboxOutput:
+    ctx = input.context
+
+    with log_activity_execution(
+        "create_sandbox",
+        **ctx.to_log_context(),
+    ):
+        config = SandboxConfig(
+            name=get_sandbox_name_for_snapshot(ctx.github_integration_id, ctx.repository),
+            template=SandboxTemplate.DEFAULT_BASE,
+            environment_variables={},
+            snapshot_id=None,
+            metadata={"purpose": "snapshot_creation"},
+        )
+
+        sandbox = Sandbox.create(config)
+
+        activity.logger.info(f"Created sandbox {sandbox.id} for snapshot creation")
+
+        return CreateSandboxOutput(sandbox_id=sandbox.id)