pyworkflow-engine 0.1.36__tar.gz → 0.1.37__tar.gz

{pyworkflow_engine-0.1.36 → pyworkflow_engine-0.1.37}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: pyworkflow-engine
-Version: 0.1.36
+Version: 0.1.37
 Summary: A Python implementation of durable, event-sourced workflows inspired by Vercel Workflow
 Author: PyWorkflow Contributors
 License: MIT

{pyworkflow_engine-0.1.36 → pyworkflow_engine-0.1.37}/pyproject.toml

@@ -7,7 +7,7 @@ packages = [{include = "pyworkflow"}]
 
 [project]
 name = "pyworkflow-engine"
-version = "0.1.36"
+version = "0.1.37"
 description = "A Python implementation of durable, event-sourced workflows inspired by Vercel Workflow"
 readme = "README.md"
 requires-python = ">=3.11"

{pyworkflow_engine-0.1.36 → pyworkflow_engine-0.1.37}/pyworkflow/__init__.py

@@ -29,7 +29,7 @@ Quick Start:
     >>> run_id = await start(my_workflow, "Alice")
 """
 
-__version__ = "0.1.36"
+__version__ = "0.1.37"
 
 # Configuration
 from pyworkflow.config import (

{pyworkflow_engine-0.1.36 → pyworkflow_engine-0.1.37}/pyworkflow/celery/app.py

@@ -13,6 +13,7 @@ garbage collector and Celery's saferepr module. It does not affect functionality
 """
 
 import os
+from datetime import timedelta
 from typing import Any
 
 from celery import Celery
@@ -343,8 +344,14 @@ def create_celery_app(
         # Monitoring
         "worker_send_task_events": True,
         "task_send_sent_event": True,
-        # Beat scheduler (for sleep resumption)
-        "beat_schedule": {},
+        # Beat scheduler (for sleep resumption and periodic tasks)
+        "beat_schedule": {
+            "pyworkflow-data-retention": {
+                "task": "pyworkflow.run_data_retention",
+                "schedule": timedelta(hours=24),
+                "options": {"queue": "pyworkflow.default"},
+            },
+        },
         # Logging
         "worker_log_format": "[%(asctime)s: %(levelname)s/%(processName)s] %(message)s",
         "worker_task_log_format": "[%(asctime)s: %(levelname)s/%(processName)s] [%(task_name)s(%(task_id)s)] %(message)s",

{pyworkflow_engine-0.1.36 → pyworkflow_engine-0.1.37}/pyworkflow/celery/tasks.py

@@ -2749,3 +2749,45 @@ async def _handle_continue_as_new_celery(
     )
 
     return new_run_id
+
+
+@celery_app.task(
+    name="pyworkflow.run_data_retention",
+    base=SingletonWorkflowTask,
+    bind=True,
+    queue="pyworkflow.default",
+    release_lock_on_failure=True,
+    lock_expiry=7200,  # 2h safety net
+)
+def run_data_retention_task(self: SingletonWorkflowTask) -> dict[str, Any]:
+    """
+    Periodic task: delete workflow runs older than data_retention_days.
+
+    Singleton (only one instance runs at a time). Skips if data_retention_days
+    is not configured.
+    """
+    from datetime import timedelta
+
+    from pyworkflow.config import get_config
+    from pyworkflow.storage.config import storage_to_config
+
+    config = get_config()
+    if config.data_retention_days is None:
+        logger.debug("Data retention not configured; skipping.")
+        return {"deleted": 0, "skipped": True}
+
+    storage = config.storage
+    if storage is None:
+        logger.warning("No storage configured; skipping data retention.")
+        return {"deleted": 0, "skipped": True}
+
+    cutoff = datetime.now(UTC) - timedelta(days=config.data_retention_days)
+    logger.info(
+        "Running data retention: deleting runs updated before {}",
+        cutoff.isoformat(),
+    )
+    storage_config = storage_to_config(storage)
+    backend = _get_storage_backend(storage_config)
+    count = run_async(backend.delete_old_runs(cutoff))
+    logger.info("Data retention complete: deleted {} runs", count)
+    return {"deleted": count, "skipped": False}

{pyworkflow_engine-0.1.36 → pyworkflow_engine-0.1.37}/pyworkflow/config.py

@@ -18,7 +18,7 @@ Usage:
     ... )
 
 Environment Variables:
-    PYWORKFLOW_STORAGE_TYPE: Storage backend type (file, memory, sqlite, postgres, mysql)
+    PYWORKFLOW_STORAGE_TYPE: Storage backend type (file, memory, sqlite, postgres, mysql, citus)
     PYWORKFLOW_STORAGE_PATH: Path for file/sqlite backends
     PYWORKFLOW_POSTGRES_HOST: PostgreSQL host
     PYWORKFLOW_POSTGRES_PORT: PostgreSQL port
@@ -33,6 +33,7 @@ Environment Variables:
     PYWORKFLOW_CELERY_BROKER: Celery broker URL
     PYWORKFLOW_CELERY_RESULT_BACKEND: Celery result backend URL
     PYWORKFLOW_RUNTIME: Default runtime (local, celery)
+    PYWORKFLOW_DATA_RETENTION_DAYS: Days to retain completed/failed/cancelled runs (unset = keep forever)
 """
 
 import os
@@ -58,9 +59,9 @@ def _load_env_storage_config() -> dict[str, Any] | None:
 
     storage_type = storage_type.lower()
 
-    if storage_type == "postgres":
+    if storage_type in ("postgres", "citus"):
         return {
-            "type": "postgres",
+            "type": storage_type,
             "host": os.getenv("PYWORKFLOW_POSTGRES_HOST", "localhost"),
             "port": int(os.getenv("PYWORKFLOW_POSTGRES_PORT", "5432")),
             "user": os.getenv("PYWORKFLOW_POSTGRES_USER", "pyworkflow"),
@@ -161,6 +162,9 @@ class PyWorkflowConfig:
     celery_broker: str | None = None
     aws_region: str | None = None
 
+    # Data retention policy
+    data_retention_days: int | None = None  # None = keep forever
+
     # Event limit settings (WARNING: Do not modify unless you understand the implications)
     # These limits prevent runaway workflows from consuming excessive resources
     event_soft_limit: int = 10_000  # Log warning at this count
@@ -196,11 +200,21 @@ def _config_from_env_and_yaml() -> PyWorkflowConfig:
     celery_config = yaml_config.get("celery", {})
     celery_broker = os.getenv("PYWORKFLOW_CELERY_BROKER") or celery_config.get("broker")
 
+    # Data retention: env var > yaml > None (keep forever)
+    retention_env = os.getenv("PYWORKFLOW_DATA_RETENTION_DAYS")
+    if retention_env is not None:
+        data_retention_days: int | None = int(retention_env)
+    elif yaml_config.get("retention_days") is not None:
+        data_retention_days = int(yaml_config["retention_days"])
+    else:
+        data_retention_days = None
+
     return PyWorkflowConfig(
         default_runtime=runtime,
         default_durable=durable,
         storage=storage,
         celery_broker=celery_broker,
+        data_retention_days=data_retention_days,
     )
 
 
@@ -348,11 +362,16 @@ def configure_from_yaml(path: str | Path, discover: bool = True) -> None:
     celery_config = yaml_config.get("celery", {})
     celery_broker = celery_config.get("broker")
 
+    # Data retention
+    retention_raw = yaml_config.get("retention_days")
+    data_retention_days: int | None = int(retention_raw) if retention_raw is not None else None
+
     _config = PyWorkflowConfig(
         default_runtime=runtime,
         default_durable=durable,
         storage=storage,
         celery_broker=celery_broker,
+        data_retention_days=data_retention_days,
     )
     _config_loaded_from_yaml = True
 

{pyworkflow_engine-0.1.36 → pyworkflow_engine-0.1.37}/pyworkflow/storage/__init__.py

@@ -47,6 +47,12 @@ try:
 except ImportError:
     MySQLStorageBackend = None  # type: ignore
 
+# Citus distributed PostgreSQL backend - optional import (requires asyncpg + Citus extension)
+try:
+    from pyworkflow.storage.citus import CitusStorageBackend
+except ImportError:
+    CitusStorageBackend = None  # type: ignore
+
 __all__ = [
     "StorageBackend",
     "FileStorageBackend",
@@ -56,6 +62,7 @@ __all__ = [
     "DynamoDBStorageBackend",
     "CassandraStorageBackend",
     "MySQLStorageBackend",
+    "CitusStorageBackend",
     "WorkflowRun",
     "StepExecution",
     "Hook",

{pyworkflow_engine-0.1.36 → pyworkflow_engine-0.1.37}/pyworkflow/storage/base.py

@@ -700,6 +700,22 @@ class StorageBackend(ABC):
         """
         pass
 
+    @abstractmethod
+    async def delete_old_runs(self, older_than: datetime) -> int:
+        """
+        Delete workflow runs in terminal states last updated before `older_than`.
+
+        Terminal states: COMPLETED, FAILED, CANCELLED, CONTINUED_AS_NEW, INTERRUPTED.
+        Associated events, steps, hooks, and cancellation flags are deleted too.
+
+        Args:
+            older_than: Delete runs where updated_at < older_than
+
+        Returns:
+            Number of workflow runs deleted
+        """
+        pass
+
     async def health_check(self) -> bool:
         """
         Check if storage backend is healthy and accessible.

{pyworkflow_engine-0.1.36 → pyworkflow_engine-0.1.37}/pyworkflow/storage/cassandra.py

@@ -1709,6 +1709,72 @@ class CassandraStorageBackend(StorageBackend):
             schedule.updated_at = datetime.now(UTC)
             await self.update_schedule(schedule)
 
+    async def delete_old_runs(self, older_than: datetime) -> int:
+        """Delete terminal runs (and all related data) updated before older_than."""
+        session = self._ensure_connected()
+        terminal = ["completed", "failed", "cancelled", "continued_as_new", "interrupted"]
+        count = 0
+
+        for status_val in terminal:
+            rows = session.execute(
+                SimpleStatement(
+                    "SELECT run_id FROM runs_by_status WHERE status = %s",
+                    consistency_level=self.read_consistency,
+                ),
+                (status_val,),
+            )
+            for row in rows:
+                run_id = row.run_id
+                # Fetch run to check updated_at
+                run_row = session.execute(
+                    SimpleStatement(
+                        "SELECT updated_at FROM workflow_runs WHERE run_id = %s",
+                        consistency_level=self.read_consistency,
+                    ),
+                    (run_id,),
+                ).one()
+                if run_row is None or run_row.updated_at >= older_than:
+                    continue
+                # Delete all related data
+                session.execute(
+                    SimpleStatement(
+                        "DELETE FROM events WHERE run_id = %s",
+                        consistency_level=self.write_consistency,
+                    ),
+                    (run_id,),
+                )
+                session.execute(
+                    SimpleStatement(
+                        "DELETE FROM steps WHERE run_id = %s",
+                        consistency_level=self.write_consistency,
+                    ),
+                    (run_id,),
+                )
+                session.execute(
+                    SimpleStatement(
+                        "DELETE FROM hooks WHERE run_id = %s",
+                        consistency_level=self.write_consistency,
+                    ),
+                    (run_id,),
+                )
+                session.execute(
+                    SimpleStatement(
+                        "DELETE FROM cancellation_flags WHERE run_id = %s",
+                        consistency_level=self.write_consistency,
+                    ),
+                    (run_id,),
+                )
+                session.execute(
+                    SimpleStatement(
+                        "DELETE FROM workflow_runs WHERE run_id = %s",
+                        consistency_level=self.write_consistency,
+                    ),
+                    (run_id,),
+                )
+                count += 1
+
+        return count
+
     # Helper methods for converting Cassandra rows to domain objects
 
     def _row_to_workflow_run(self, row: Any) -> WorkflowRun:

pyworkflow_engine-0.1.37/pyworkflow/storage/citus.py

@@ -0,0 +1,370 @@
+"""
+Citus distributed PostgreSQL storage backend.
+
+Citus is a PostgreSQL extension that shards tables across multiple nodes for
+horizontal scalability. This backend extends PostgresStorageBackend, reusing all
+query logic while overriding schema initialization to create Citus-compatible
+distributed tables.
+
+Distribution strategy:
+- All workflow tables are co-located on `run_id` so that all data for a given
+  workflow run lands on the same Citus shard.
+- Schedules are distributed on `schedule_id` (independent of workflow runs).
+- `schema_versions` becomes a reference table (replicated to all workers).
+
+Citus constraints vs plain PostgreSQL:
+- Primary keys and unique constraints MUST include the distribution column.
+- Cross-shard foreign key constraints are unsupported.
+- Global unique constraints on non-distribution columns cannot be enforced.
+
+See: https://docs.citusdata.com/en/stable/develop/api.html
+"""
+
+import asyncpg
+
+from pyworkflow.storage.migrations import Migration
+from pyworkflow.storage.postgres import PostgresMigrationRunner, PostgresStorageBackend
+
+
+class CitusMigrationRunner(PostgresMigrationRunner):
+    """
+    Citus-specific migration runner.
+
+    Extends PostgresMigrationRunner to handle Citus schema constraints.
+    V2 migration: same step_id backfill as PostgreSQL, but skips creating
+    unique constraints that would conflict with Citus distribution requirements.
+    """
+
+    async def apply_migration(self, migration: Migration) -> None:
+        """Apply a migration with Citus-specific handling."""
+        from datetime import UTC, datetime
+
+        async with self._pool.acquire() as conn, conn.transaction():
+            if migration.version == 2:
+                # V2: Add step_id column to events table
+                # Check if events table exists (fresh databases won't have it yet)
+                table_exists = await conn.fetchrow("""
+                    SELECT EXISTS (
+                        SELECT FROM information_schema.tables
+                        WHERE table_name = 'events'
+                    ) as exists
+                """)
+
+                if table_exists and table_exists["exists"]:
+                    # Use IF NOT EXISTS for idempotency
+                    await conn.execute("""
+                        DO $$
+                        BEGIN
+                            IF NOT EXISTS (
+                                SELECT 1 FROM information_schema.columns
+                                WHERE table_name = 'events' AND column_name = 'step_id'
+                            ) THEN
+                                ALTER TABLE events ADD COLUMN step_id TEXT;
+                            END IF;
+                        END $$
+                    """)
+
+                    # Create index for optimized has_event() queries
+                    await conn.execute("""
+                        CREATE INDEX IF NOT EXISTS idx_events_run_id_step_id_type
+                        ON events(run_id, step_id, type)
+                    """)
+
+                    # Backfill step_id from JSON data
+                    await conn.execute("""
+                        UPDATE events
+                        SET step_id = (data::jsonb)->>'step_id'
+                        WHERE step_id IS NULL
+                          AND (data::jsonb)->>'step_id' IS NOT NULL
+                    """)
+                # NOTE: No UNIQUE constraints added here — Citus requires distribution
+                # column in every unique constraint, which V2 doesn't add.
+            elif migration.up_func:
+                await migration.up_func(conn)
+            elif migration.up_sql and migration.up_sql != "SELECT 1":
+                await conn.execute(migration.up_sql)
+
+            # Record the migration
+            await conn.execute(
+                """
+                INSERT INTO schema_versions (version, applied_at, description)
+                VALUES ($1, $2, $3)
+                """,
+                migration.version,
+                datetime.now(UTC),
+                migration.description,
+            )
+
+
+class CitusStorageBackend(PostgresStorageBackend):
+    """
+    Citus distributed PostgreSQL storage backend.
+
+    Extends PostgresStorageBackend with Citus-specific schema initialization.
+    All query logic is inherited unchanged — only DDL differs to satisfy Citus's
+    distribution column requirements.
+
+    Tables are sharded as follows:
+    - workflow_runs, events, steps, hooks, cancellation_flags: distributed on run_id
+    - schedules: distributed on schedule_id
+    - schema_versions: reference table (replicated to all workers)
+
+    Requirements:
+    - PostgreSQL with the Citus extension installed and loaded
+    - The calling database user must have permission to call Citus functions
+
+    Usage:
+        backend = CitusStorageBackend(
+            host="citus-coordinator",
+            database="pyworkflow",
+        )
+        await backend.connect()
+    """
+
+    async def _initialize_schema(self) -> None:
+        """
+        Create Citus-compatible schema and distribute tables.
+
+        Steps:
+        1. Verify Citus extension is available.
+        2. Create tables with Citus-compatible PKs/constraints (no cross-shard FKs,
+           no cross-shard UNIQUE constraints, composite PKs include distribution col).
+        3. Create indexes.
+        4. Distribute tables (idempotent — already-distributed tables are skipped).
+        5. Run migrations via CitusMigrationRunner.
+        """
+        pool = await self._get_pool()
+
+        # Step 1: Verify Citus extension is present
+        async with pool.acquire() as conn:
+            try:
+                await conn.fetchval("SELECT citus_version()")
+            except asyncpg.UndefinedFunctionError:
+                raise RuntimeError(
+                    "Citus extension is not available on this PostgreSQL server. "
+                    "Install Citus and run `CREATE EXTENSION citus;` before using "
+                    "CitusStorageBackend. See: https://docs.citusdata.com/en/stable/installation/"
+                )
+
+        # Step 2 & 3: Create tables and indexes with Citus-compatible DDL
+        async with pool.acquire() as conn:
+            # schema_versions — created first so CitusMigrationRunner can use it
+            await conn.execute("""
+                CREATE TABLE IF NOT EXISTS schema_versions (
+                    version INTEGER PRIMARY KEY,
+                    applied_at TIMESTAMPTZ NOT NULL,
+                    description TEXT
+                )
+            """)
+
+            # workflow_runs — distributed on run_id
+            # Differences vs postgres:
+            #   - DROP FK on parent_run_id (self-referential FK can't be enforced cross-shard)
+            #   - UNIQUE INDEX on idempotency_key → plain INDEX (global uniqueness unenforceable)
+            await conn.execute("""
+                CREATE TABLE IF NOT EXISTS workflow_runs (
+                    run_id TEXT PRIMARY KEY,
+                    workflow_name TEXT NOT NULL,
+                    status TEXT NOT NULL,
+                    created_at TIMESTAMPTZ NOT NULL,
+                    updated_at TIMESTAMPTZ NOT NULL,
+                    started_at TIMESTAMPTZ,
+                    completed_at TIMESTAMPTZ,
+                    input_args TEXT NOT NULL DEFAULT '[]',
+                    input_kwargs TEXT NOT NULL DEFAULT '{}',
+                    result TEXT,
+                    error TEXT,
+                    idempotency_key TEXT,
+                    max_duration TEXT,
+                    metadata TEXT DEFAULT '{}',
+                    recovery_attempts INTEGER DEFAULT 0,
+                    max_recovery_attempts INTEGER DEFAULT 3,
+                    recover_on_worker_loss BOOLEAN DEFAULT TRUE,
+                    parent_run_id TEXT,
+                    nesting_depth INTEGER DEFAULT 0,
+                    continued_from_run_id TEXT,
+                    continued_to_run_id TEXT
+                )
+            """)
+
+            await conn.execute(
+                "CREATE INDEX IF NOT EXISTS idx_runs_status ON workflow_runs(status)"
+            )
+            await conn.execute(
+                "CREATE INDEX IF NOT EXISTS idx_runs_workflow_name ON workflow_runs(workflow_name)"
+            )
+            await conn.execute(
+                "CREATE INDEX IF NOT EXISTS idx_runs_created_at ON workflow_runs(created_at DESC)"
+            )
+            # Non-unique index: Citus cannot enforce global uniqueness on non-distribution columns
+            await conn.execute(
+                "CREATE INDEX IF NOT EXISTS idx_runs_idempotency_key ON workflow_runs(idempotency_key) WHERE idempotency_key IS NOT NULL"
+            )
+            await conn.execute(
+                "CREATE INDEX IF NOT EXISTS idx_runs_parent_run_id ON workflow_runs(parent_run_id)"
+            )
+
+            # events — distributed on run_id, co-located with workflow_runs
+            # Differences vs postgres:
+            #   - PK: event_id → (run_id, event_id) to include distribution column
+            #   - FK on run_id retained (co-located FK is supported by Citus)
+            await conn.execute("""
+                CREATE TABLE IF NOT EXISTS events (
+                    event_id TEXT NOT NULL,
+                    run_id TEXT NOT NULL,
+                    sequence INTEGER NOT NULL,
+                    type TEXT NOT NULL,
+                    timestamp TIMESTAMPTZ NOT NULL,
+                    data TEXT NOT NULL DEFAULT '{}',
+                    step_id TEXT,
+                    PRIMARY KEY (run_id, event_id)
+                )
+            """)
+
+            await conn.execute(
+                "CREATE INDEX IF NOT EXISTS idx_events_run_id_sequence ON events(run_id, sequence)"
+            )
+            await conn.execute(
+                "CREATE INDEX IF NOT EXISTS idx_events_run_id_type ON events(run_id, type)"
+            )
+            await conn.execute(
+                "CREATE INDEX IF NOT EXISTS idx_events_run_id_step_id_type ON events(run_id, step_id, type)"
+            )
+
+            # steps — distributed on run_id, co-located with workflow_runs
+            # Differences vs postgres:
+            #   - PK: step_id → (run_id, step_id) to include distribution column
+            #   - FK on run_id retained (co-located FK)
+            #   - Extra index on step_id alone for get_step(step_id) scatter-gather queries
+            await conn.execute("""
+                CREATE TABLE IF NOT EXISTS steps (
+                    step_id TEXT NOT NULL,
+                    run_id TEXT NOT NULL,
+                    step_name TEXT NOT NULL,
+                    status TEXT NOT NULL,
+                    created_at TIMESTAMPTZ NOT NULL,
+                    started_at TIMESTAMPTZ,
+                    completed_at TIMESTAMPTZ,
+                    input_args TEXT NOT NULL DEFAULT '[]',
+                    input_kwargs TEXT NOT NULL DEFAULT '{}',
+                    result TEXT,
+                    error TEXT,
+                    retry_count INTEGER DEFAULT 0,
+                    PRIMARY KEY (run_id, step_id)
+                )
+            """)
+
+            await conn.execute("CREATE INDEX IF NOT EXISTS idx_steps_run_id ON steps(run_id)")
+            # Shard-local index to speed up get_step(step_id) scatter-gather queries
+            await conn.execute("CREATE INDEX IF NOT EXISTS idx_steps_step_id ON steps(step_id)")
+
+            # hooks — distributed on run_id, co-located with workflow_runs
+            # Differences vs postgres:
+            #   - UNIQUE on token → plain INDEX (token = run_id:hook_id, collisions impossible)
+            #   - FK on run_id retained (co-located FK)
+            await conn.execute("""
+                CREATE TABLE IF NOT EXISTS hooks (
+                    run_id TEXT NOT NULL,
+                    hook_id TEXT NOT NULL,
+                    token TEXT NOT NULL,
+                    created_at TIMESTAMPTZ NOT NULL,
+                    received_at TIMESTAMPTZ,
+                    expires_at TIMESTAMPTZ,
+                    status TEXT NOT NULL,
+                    payload TEXT,
+                    metadata TEXT DEFAULT '{}',
+                    PRIMARY KEY (run_id, hook_id)
+                )
+            """)
+
+            await conn.execute("CREATE INDEX IF NOT EXISTS idx_hooks_token ON hooks(token)")
+            await conn.execute("CREATE INDEX IF NOT EXISTS idx_hooks_run_id ON hooks(run_id)")
+            await conn.execute("CREATE INDEX IF NOT EXISTS idx_hooks_status ON hooks(status)")
+
+            # schedules — distributed on schedule_id (independent of runs)
+            await conn.execute("""
+                CREATE TABLE IF NOT EXISTS schedules (
+                    schedule_id TEXT PRIMARY KEY,
+                    workflow_name TEXT NOT NULL,
+                    spec TEXT NOT NULL,
+                    spec_type TEXT NOT NULL,
+                    timezone TEXT,
+                    input_args TEXT NOT NULL DEFAULT '[]',
+                    input_kwargs TEXT NOT NULL DEFAULT '{}',
+                    status TEXT NOT NULL,
+                    overlap_policy TEXT NOT NULL,
+                    next_run_time TIMESTAMPTZ,
+                    last_run_time TIMESTAMPTZ,
+                    running_run_ids TEXT DEFAULT '[]',
+                    metadata TEXT DEFAULT '{}',
+                    created_at TIMESTAMPTZ NOT NULL,
+                    updated_at TIMESTAMPTZ NOT NULL,
+                    paused_at TIMESTAMPTZ,
+                    deleted_at TIMESTAMPTZ
+                )
+            """)
+
+            await conn.execute(
+                "CREATE INDEX IF NOT EXISTS idx_schedules_status ON schedules(status)"
+            )
+            await conn.execute(
+                "CREATE INDEX IF NOT EXISTS idx_schedules_next_run_time ON schedules(next_run_time)"
+            )
+            await conn.execute(
+                "CREATE INDEX IF NOT EXISTS idx_schedules_workflow_name ON schedules(workflow_name)"
+            )
+
+            # cancellation_flags — distributed on run_id, co-located with workflow_runs
+            # Differences vs postgres:
+            #   - FK on run_id dropped (cross-shard FK unsupported; co-location ensures correctness)
+            await conn.execute("""
+                CREATE TABLE IF NOT EXISTS cancellation_flags (
+                    run_id TEXT PRIMARY KEY,
+                    created_at TIMESTAMPTZ NOT NULL
+                )
+            """)
+
+        # Step 4: Distribute tables (idempotent — already-distributed tables are skipped)
+        await self._distribute_tables(pool)
+
+        # Step 5: Run migrations via Citus-aware runner
+        runner = CitusMigrationRunner(pool)
+        await runner.run_migrations()
+
+    async def _distribute_tables(self, pool: asyncpg.Pool) -> None:
+        """
+        Call Citus distribution functions for each table.
+
+        Uses pg_dist_partition to detect already-distributed tables so this
+        method is idempotent and safe to call on every connect().
+        """
+        async with pool.acquire() as conn:
+            # Fetch already-distributed table names using relname to avoid
+            # schema-prefix ambiguity (logicalrelid::text is search_path-dependent)
+            rows = await conn.fetch("""
+                SELECT c.relname AS tbl
+                FROM pg_dist_partition dp
+                JOIN pg_class c ON dp.logicalrelid = c.oid
+            """)
+            distributed = {row["tbl"] for row in rows}
+
+            # workflow_runs: anchor table, distribute first
+            if "workflow_runs" not in distributed:
+                await conn.execute("SELECT create_distributed_table('workflow_runs', 'run_id')")
+
+            # Co-located tables: must reference the same distribution column
+            colocated_run_id = ["events", "steps", "hooks", "cancellation_flags"]
+            for table in colocated_run_id:
+                if table not in distributed:
+                    await conn.execute(
+                        f"SELECT create_distributed_table('{table}', 'run_id', "
+                        f"colocate_with => 'workflow_runs')"
+                    )
+
+            # schedules: independent distribution
+            if "schedules" not in distributed:
+                await conn.execute("SELECT create_distributed_table('schedules', 'schedule_id')")
+
+            # schema_versions: reference table (replicated to all workers)
+            if "schema_versions" not in distributed:
+                await conn.execute("SELECT create_reference_table('schema_versions')")