pyworkflow-engine 0.2.8__tar.gz → 0.2.10__tar.gz

{pyworkflow_engine-0.2.8 → pyworkflow_engine-0.2.10}/PKG-INFO

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

{pyworkflow_engine-0.2.8 → pyworkflow_engine-0.2.10}/pyproject.toml

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

{pyworkflow_engine-0.2.8 → pyworkflow_engine-0.2.10}/pyworkflow/__init__.py

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

{pyworkflow_engine-0.2.8 → pyworkflow_engine-0.2.10}/pyworkflow/celery/app.py

@@ -61,6 +61,82 @@ def _run_warmup_callable(dotted_path: str) -> None:
         )
 
 
+def _run_startup_migrations() -> None:
+    """Run pending DB migrations before the worker accepts tasks.
+
+    Loads storage configuration from environment variables, creates a throwaway
+    backend, calls ``connect()`` (which triggers ``_initialize_schema()`` →
+    ``run_migrations()``), and disconnects. The pool is bound to a temporary
+    event loop that lives only for this function, so the forked workers below
+    start with no inherited state.
+
+    Behavior:
+    - If ``PYWORKFLOW_STORAGE_TYPE`` is unset, or points at a schema-less backend
+      (``memory`` / ``file``), this is a no-op and tasks fall back to lazy init.
+    - If migrations raise, the exception propagates — the worker fails to start
+      rather than accepting tasks against a schema it can't read. That is the
+      desired failure mode: a restart loop is louder than silent corruption.
+    - Controllable via ``PYWORKFLOW_SKIP_STARTUP_MIGRATIONS=1`` for deployments
+      that intentionally migrate out-of-band (e.g. a dedicated job).
+    """
+    import asyncio
+
+    from loguru import logger as loguru_logger
+
+    if os.getenv("PYWORKFLOW_SKIP_STARTUP_MIGRATIONS", "").lower() in ("1", "true", "yes"):
+        loguru_logger.info(
+            "STARTUP_MIGRATIONS: skipped (PYWORKFLOW_SKIP_STARTUP_MIGRATIONS set)"
+        )
+        return
+
+    from pyworkflow.config import _load_env_storage_config
+
+    config = _load_env_storage_config()
+    if not config:
+        loguru_logger.info(
+            "STARTUP_MIGRATIONS: PYWORKFLOW_STORAGE_TYPE not set — skipping. "
+            "Migrations will run lazily on first task that touches storage."
+        )
+        return
+
+    storage_type = config.get("type")
+    if storage_type in ("memory", "file"):
+        loguru_logger.info(
+            f"STARTUP_MIGRATIONS: skipping for type={storage_type} (schema-less backend)"
+        )
+        return
+
+    loguru_logger.info(
+        f"STARTUP_MIGRATIONS: connecting to {storage_type} to apply pending migrations"
+    )
+
+    async def _run() -> None:
+        # Create the backend directly (bypassing the cache) so that the pool we
+        # open here does not leak into the forked workers' storage cache. The
+        # worker processes will each create their own pool on first task.
+        from pyworkflow.storage.config import _create_storage_backend
+
+        backend = _create_storage_backend(config)
+        try:
+            await backend.connect()
+        finally:
+            try:
+                if hasattr(backend, "disconnect"):
+                    await backend.disconnect()
+            except Exception as e:
+                loguru_logger.warning(f"STARTUP_MIGRATIONS: disconnect failed: {e}")
+
+    try:
+        asyncio.run(_run())
+    except Exception:
+        loguru_logger.opt(exception=True).error(
+            "STARTUP_MIGRATIONS: FAILED — worker will not start"
+        )
+        raise
+
+    loguru_logger.info("STARTUP_MIGRATIONS: completed successfully")
+
+
 def is_sentinel_url(url: str) -> bool:
     """Check if URL uses sentinel:// or sentinel+ssl:// protocol."""
     return url.startswith("sentinel://") or url.startswith("sentinel+ssl://")
@@ -420,6 +496,15 @@ def on_worker_init(**kwargs):
     if warmup_path:
         _run_warmup_callable(warmup_path)
 
+    # Apply pending DB migrations once per pod, BEFORE any task is dispatched.
+    # This replaces the old lazy "first task triggers connect() triggers
+    # migrations" path, which was fragile: backend instances cached across
+    # restarts (or inherited via fork) would skip migrations entirely because
+    # _initialized was True without the migration check ever having run in
+    # this process. Running at startup makes migration state observable in
+    # the pod logs and fails the pod fast if the schema is incompatible.
+    _run_startup_migrations()
+
 
 @worker_process_init.connect
 def on_worker_process_init(**kwargs):

{pyworkflow_engine-0.2.8 → pyworkflow_engine-0.2.10}/pyworkflow/storage/citus.py

@@ -21,6 +21,7 @@ See: https://docs.citusdata.com/en/stable/develop/api.html
 """
 
 import asyncpg
+from loguru import logger
 
 from pyworkflow.storage.migrations import Migration
 from pyworkflow.storage.postgres import PostgresMigrationRunner, PostgresStorageBackend
@@ -39,6 +40,10 @@ class CitusMigrationRunner(PostgresMigrationRunner):
         """Apply a migration with Citus-specific handling."""
         from datetime import UTC, datetime
 
+        logger.info(
+            f"CITUS_MIGRATION: opening transaction for v{migration.version}",
+            version=migration.version,
+        )
         async with self._pool.acquire() as conn, conn.transaction():
             if migration.version == 2:
                 # V2: Add step_id column to events table
@@ -122,14 +127,17 @@ class CitusMigrationRunner(PostgresMigrationRunner):
                 # DDL inside PL/pgSQL DO blocks is executed via SPI and bypasses the
                 # hook, which would leave worker shards missing the column while the
                 # coordinator catalog appears correct.
+                logger.info("CITUS_MIGRATION v5: ALTER stream_subscriptions ADD stream_run_id")
                 await conn.execute(
                     "ALTER TABLE stream_subscriptions "
                     "ADD COLUMN IF NOT EXISTS stream_run_id TEXT NULL"
                 )
+                logger.info("CITUS_MIGRATION v5: CREATE INDEX idx_subscriptions_stream_run")
                 await conn.execute(
                     "CREATE INDEX IF NOT EXISTS idx_subscriptions_stream_run "
                     "ON stream_subscriptions(stream_id, stream_run_id)"
                 )
+                logger.info("CITUS_MIGRATION v5: CREATE TABLE scheduled_signals")
                 await conn.execute("""
                     CREATE TABLE IF NOT EXISTS scheduled_signals (
                         id TEXT PRIMARY KEY,
@@ -143,6 +151,7 @@ class CitusMigrationRunner(PostgresMigrationRunner):
                         created_at TIMESTAMPTZ NOT NULL DEFAULT NOW()
                     )
                 """)
+                logger.info("CITUS_MIGRATION v5: CREATE INDEX idx_scheduled_signals_due")
                 await conn.execute(
                     "CREATE INDEX IF NOT EXISTS idx_scheduled_signals_due "
                     "ON scheduled_signals(delivered, due_at)"
@@ -155,21 +164,30 @@ class CitusMigrationRunner(PostgresMigrationRunner):
                         WHERE c.relname = 'scheduled_signals'
                     )
                 """)
+                logger.info(
+                    f"CITUS_MIGRATION v5: scheduled_signals already distributed = {already}"
+                )
                 if not already:
+                    logger.info(
+                        "CITUS_MIGRATION v5: calling create_reference_table('scheduled_signals')"
+                    )
                     await conn.execute("SELECT create_reference_table('scheduled_signals')")
             elif migration.version == 6:
                 # V6: Add parent_run_id + parent_hook_token to stream_subscriptions.
                 # Top-level ALTERs so Citus propagates to all worker shards (see V5 note).
+                logger.info("CITUS_MIGRATION v6: ALTER stream_subscriptions ADD parent_run_id")
                 await conn.execute(
                     "ALTER TABLE stream_subscriptions "
                     "ADD COLUMN IF NOT EXISTS parent_run_id TEXT NULL"
                 )
+                logger.info("CITUS_MIGRATION v6: ALTER stream_subscriptions ADD parent_hook_token")
                 await conn.execute(
                     "ALTER TABLE stream_subscriptions "
                     "ADD COLUMN IF NOT EXISTS parent_hook_token TEXT NULL"
                 )
             elif migration.version == 7:
                 # V7: Add result column. Top-level ALTER for Citus propagation (see V5 note).
+                logger.info("CITUS_MIGRATION v7: ALTER stream_subscriptions ADD result")
                 await conn.execute(
                     "ALTER TABLE stream_subscriptions "
                     "ADD COLUMN IF NOT EXISTS result JSONB NULL"

{pyworkflow_engine-0.2.8 → pyworkflow_engine-0.2.10}/pyworkflow/storage/migrations/base.py

@@ -11,6 +11,8 @@ from dataclasses import dataclass
 from datetime import UTC, datetime
 from typing import Any
 
+from loguru import logger
+
 
 @dataclass
 class Migration:
@@ -238,11 +240,18 @@ class MigrationRunner(ABC):
         Raises:
             Exception: If any migration fails (partial migrations are rolled back)
         """
+        runner_name = type(self).__name__
+        logger.info(f"MIGRATION: run_migrations() called on {runner_name}")
+
         # Ensure we have a schema_versions table
         await self.ensure_schema_versions_table()
 
         # Get current version
         current_version = await self.get_current_version()
+        logger.info(
+            f"MIGRATION: current schema version = {current_version}",
+            runner=runner_name,
+        )
 
         # If no migrations recorded, check for existing schema
         if current_version == 0:
@@ -250,15 +259,43 @@ class MigrationRunner(ABC):
             if has_existing_schema:
                 # Database has tables but no version tracking
                 # Assume it's at V1 (original schema)
+                logger.info(
+                    "MIGRATION: existing schema detected — recording baseline v1",
+                    runner=runner_name,
+                )
                 await self.record_baseline_version(1, "Baseline: original schema")
                 current_version = 1
 
         # Get and apply pending migrations
         pending = self.registry.get_pending(current_version)
+        pending_versions = [m.version for m in pending]
+        logger.info(
+            f"MIGRATION: pending migrations = {pending_versions}",
+            runner=runner_name,
+            current_version=current_version,
+        )
         applied: list[AppliedMigration] = []
 
         for migration in pending:
-            await self.apply_migration(migration)
+            logger.info(
+                f"MIGRATION: applying v{migration.version} — {migration.description}",
+                runner=runner_name,
+                version=migration.version,
+            )
+            try:
+                await self.apply_migration(migration)
+            except Exception:
+                logger.exception(
+                    f"MIGRATION: v{migration.version} FAILED — {migration.description}",
+                    runner=runner_name,
+                    version=migration.version,
+                )
+                raise
+            logger.info(
+                f"MIGRATION: v{migration.version} applied successfully",
+                runner=runner_name,
+                version=migration.version,
+            )
             applied.append(
                 AppliedMigration(
                     version=migration.version,
@@ -267,6 +304,11 @@ class MigrationRunner(ABC):
                 )
             )
 
+        logger.info(
+            f"MIGRATION: run_migrations() completed, applied {len(applied)} migration(s)",
+            runner=runner_name,
+            applied_versions=[a.version for a in applied],
+        )
         return applied
 
 

{pyworkflow_engine-0.2.8 → pyworkflow_engine-0.2.10}/pyworkflow/storage/postgres.py

@@ -21,6 +21,7 @@ from datetime import UTC, datetime
 from typing import Any
 
 import asyncpg
+from loguru import logger
 
 from pyworkflow.engine.events import Event, EventType
 from pyworkflow.storage.base import StorageBackend
@@ -92,6 +93,10 @@ class PostgresMigrationRunner(MigrationRunner):
 
     async def apply_migration(self, migration: Migration) -> None:
         """Apply a migration with PostgreSQL-specific handling."""
+        logger.info(
+            f"PG_MIGRATION: opening transaction for v{migration.version}",
+            version=migration.version,
+        )
         async with self._pool.acquire() as conn, conn.transaction():
             if migration.version == 2:
                 # V2: Add step_id column to events table
@@ -176,14 +181,17 @@ class PostgresMigrationRunner(MigrationRunner):
                 # and bypasses the Citus ProcessUtility hook, so worker shards
                 # silently miss the column. Plain Postgres is unaffected either way,
                 # but keeping both runners symmetric avoids future foot-guns.
+                logger.info("PG_MIGRATION v5: ALTER stream_subscriptions ADD stream_run_id")
                 await conn.execute(
                     "ALTER TABLE stream_subscriptions "
                     "ADD COLUMN IF NOT EXISTS stream_run_id TEXT NULL"
                 )
+                logger.info("PG_MIGRATION v5: CREATE INDEX idx_subscriptions_stream_run")
                 await conn.execute(
                     "CREATE INDEX IF NOT EXISTS idx_subscriptions_stream_run "
                     "ON stream_subscriptions(stream_id, stream_run_id)"
                 )
+                logger.info("PG_MIGRATION v5: CREATE TABLE scheduled_signals")
                 await conn.execute("""
                     CREATE TABLE IF NOT EXISTS scheduled_signals (
                         id TEXT PRIMARY KEY,
@@ -197,6 +205,7 @@ class PostgresMigrationRunner(MigrationRunner):
                         created_at TIMESTAMPTZ NOT NULL DEFAULT NOW()
                     )
                 """)
+                logger.info("PG_MIGRATION v5: CREATE INDEX idx_scheduled_signals_due")
                 await conn.execute(
                     "CREATE INDEX IF NOT EXISTS idx_scheduled_signals_due "
                     "ON scheduled_signals(delivered, due_at)"
@@ -206,16 +215,19 @@ class PostgresMigrationRunner(MigrationRunner):
                 # stream-step dispatcher can resume the parent @workflow via
                 # resume_hook() instead of pinning a worker on an asyncio.Event.
                 # Top-level ALTERs for Citus propagation symmetry (see V5 note).
+                logger.info("PG_MIGRATION v6: ALTER stream_subscriptions ADD parent_run_id")
                 await conn.execute(
                     "ALTER TABLE stream_subscriptions "
                     "ADD COLUMN IF NOT EXISTS parent_run_id TEXT NULL"
                 )
+                logger.info("PG_MIGRATION v6: ALTER stream_subscriptions ADD parent_hook_token")
                 await conn.execute(
                     "ALTER TABLE stream_subscriptions "
                     "ADD COLUMN IF NOT EXISTS parent_hook_token TEXT NULL"
                 )
             elif migration.version == 7:
                 # V7: result column on stream_subscriptions for step output payloads.
+                logger.info("PG_MIGRATION v7: ALTER stream_subscriptions ADD result")
                 await conn.execute(
                     "ALTER TABLE stream_subscriptions "
                     "ADD COLUMN IF NOT EXISTS result JSONB NULL"