pyworkflow-engine 0.2.3__tar.gz → 0.2.6__tar.gz

{pyworkflow_engine-0.2.3 → pyworkflow_engine-0.2.6}/PKG-INFO

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

{pyworkflow_engine-0.2.3 → pyworkflow_engine-0.2.6}/pyproject.toml

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

{pyworkflow_engine-0.2.3 → pyworkflow_engine-0.2.6}/pyworkflow/__init__.py

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

{pyworkflow_engine-0.2.3 → pyworkflow_engine-0.2.6}/pyworkflow/celery/tasks.py

@@ -2462,6 +2462,9 @@ async def _resume_workflow_on_worker(
                     run_id=run_id,
                     step_id=step_id,
                 )
+                resume_workflow_task.release_lock(
+                    task_args=[run_id], task_kwargs={"storage_config": storage_config}
+                )
                 schedule_workflow_resumption(
                     run_id,
                     datetime.now(UTC),
@@ -2481,6 +2484,9 @@ async def _resume_workflow_on_worker(
                     run_id=run_id,
                     hook_id=hook_id,
                 )
+                resume_workflow_task.release_lock(
+                    task_args=[run_id], task_kwargs={"storage_config": storage_config}
+                )
                 schedule_workflow_resumption(
                     run_id,
                     datetime.now(UTC),
@@ -2492,6 +2498,13 @@ async def _resume_workflow_on_worker(
         # Schedule automatic resumption if we have a resume_at time
         resume_at = e.data.get("resume_at") if e.data else None
         if resume_at:
+            # Release the singleton lock BEFORE scheduling the next resume.
+            # Without this, the new resume_workflow_task.apply_async() sees
+            # the lock still held by the current (calling) task and silently
+            # drops the schedule as a duplicate.
+            resume_workflow_task.release_lock(
+                task_args=[run_id], task_kwargs={"storage_config": storage_config}
+            )
             schedule_workflow_resumption(
                 run_id, resume_at, storage_config=storage_config, triggered_by="start_sleep_hook"
             )

{pyworkflow_engine-0.2.3 → pyworkflow_engine-0.2.6}/pyworkflow/context/local.py

@@ -121,31 +121,6 @@ class LocalContext(WorkflowContext):
             self._replay_events(event_log)
             self._is_replaying = False
 
-    def _extract_counter_from_id(self, id_string: str) -> int:
-        """Extract counter value from hook_id or sleep_id.
-
-        Formats:
-        - hook_{name}_{counter}
-        - sleep_{counter}_{duration}s
-
-        Args:
-            id_string: The hook_id or sleep_id string
-
-        Returns:
-            The counter value, or 0 if parsing fails
-        """
-        try:
-            parts = id_string.split("_")
-            if id_string.startswith("hook_"):
-                # hook_{name}_{counter} - counter is last part
-                return int(parts[-1])
-            elif id_string.startswith("sleep_"):
-                # sleep_{counter}_{duration}s - counter is second part
-                return int(parts[1])
-        except (ValueError, IndexError):
-            pass
-        return 0
-
     def _replay_events(self, events: list[Any]) -> None:
         """Replay events to restore state."""
         from pyworkflow.engine.events import EventType

{pyworkflow_engine-0.2.3 → pyworkflow_engine-0.2.6}/pyworkflow/storage/base.py

@@ -755,6 +755,7 @@ class StorageBackend(ABC):
         signal_type: str,
         payload: dict,
         source_run_id: str | None = None,
+        stream_run_id: str | None = None,
         metadata: dict | None = None,
     ) -> int:
         """
@@ -766,6 +767,7 @@ class StorageBackend(ABC):
             signal_type: Signal type (e.g., "task.created")
             payload: Signal payload data
             source_run_id: Optional source workflow run ID
+            stream_run_id: Optional stream run identifier for grouping signals
             metadata: Optional signal metadata
 
         Returns:
@@ -798,6 +800,40 @@ class StorageBackend(ABC):
         """
         raise NotImplementedError("This storage backend does not support streams")
 
+    async def query_stream_signals(
+        self,
+        stream_id: str,
+        stream_run_id: str,
+        *,
+        source_run_id: str | None = None,
+        signal_type: str | None = None,
+        after_sequence: int | None = None,
+        before_sequence: int | None = None,
+        last_n: int | None = None,
+        limit: int = 50,
+    ) -> list[dict]:
+        """
+        Query signals from a stream with rich filtering.
+
+        Args:
+            stream_id: Stream identifier.
+            source_run_id: Filter to signals from a specific workflow run.
+            stream_run_id: Filter to signals from a specific stream run.
+            signal_type: Filter by signal type (e.g. ``"task.completed"``).
+            after_sequence: Only include signals with ``sequence >= value``.
+            before_sequence: Only include signals with ``sequence <= value``.
+            last_n: Return the N most recent signals (overrides
+                ``after_sequence`` and ``limit``).
+            limit: Maximum number of signals to return.
+
+        Returns:
+            List of signal dicts ordered by sequence ASC.
+
+        Raises:
+            NotImplementedError: If backend doesn't support streams.
+        """
+        raise NotImplementedError("This storage backend does not support streams")
+
     # Subscription Operations
 
     async def register_stream_subscription(
@@ -839,6 +875,28 @@ class StorageBackend(ABC):
         """
         raise NotImplementedError("This storage backend does not support streams")
 
+    async def get_subscriptions_for_stream(
+        self,
+        stream_id: str,
+        signal_type: str,
+    ) -> list[dict]:
+        """
+        Get ALL subscriptions for a signal type on a stream, regardless of status.
+
+        Unlike ``get_waiting_steps`` which only returns ``"waiting"``
+        subscriptions, this returns subscriptions in any status (waiting,
+        running, etc.).  Used to avoid creating duplicate subscriptions
+        when a step is currently executing.
+
+        Args:
+            stream_id: Stream identifier
+            signal_type: Signal type to match
+
+        Returns:
+            List of dicts with step subscription info
+        """
+        raise NotImplementedError("This storage backend does not support streams")
+
     async def update_subscription_status(
         self,
         stream_id: str,

{pyworkflow_engine-0.2.3 → pyworkflow_engine-0.2.6}/pyworkflow/storage/cassandra.py

@@ -179,6 +179,7 @@ class CassandraStorageBackend(StorageBackend):
         await self._create_hooks_tables(session)
         await self._create_cancellation_flags_table(session)
         await self._create_checkpoints_table(session)
+        await self._create_streams_tables(session)
         await self._create_schedules_tables(session)
 
         # Create query pattern tables
@@ -325,6 +326,52 @@ class CassandraStorageBackend(StorageBackend):
             )
         """)
 
+    async def _create_streams_tables(self, session: Session) -> None:
+        """Create stream-related tables."""
+        session.execute("""
+            CREATE TABLE IF NOT EXISTS streams (
+                stream_id TEXT PRIMARY KEY,
+                status TEXT,
+                created_at TIMESTAMP,
+                metadata TEXT
+            )
+        """)
+
+        session.execute("""
+            CREATE TABLE IF NOT EXISTS signals (
+                stream_run_id TEXT,
+                sequence INT,
+                signal_id TEXT,
+                stream_id TEXT,
+                signal_type TEXT,
+                payload TEXT,
+                published_at TIMESTAMP,
+                source_run_id TEXT,
+                metadata TEXT,
+                PRIMARY KEY ((stream_id, stream_run_id), sequence)
+            ) WITH CLUSTERING ORDER BY (sequence ASC)
+        """)
+
+        session.execute("""
+            CREATE TABLE IF NOT EXISTS stream_subscriptions (
+                stream_id TEXT,
+                step_run_id TEXT,
+                signal_types TEXT,
+                status TEXT,
+                created_at TIMESTAMP,
+                PRIMARY KEY (stream_id, step_run_id)
+            )
+        """)
+
+        session.execute("""
+            CREATE TABLE IF NOT EXISTS signal_acknowledgments (
+                signal_id TEXT,
+                step_run_id TEXT,
+                acknowledged_at TIMESTAMP,
+                PRIMARY KEY (signal_id, step_run_id)
+            )
+        """)
+
     async def _create_schedules_tables(self, session: Session) -> None:
         """Create schedules tables."""
         # Main schedules table
@@ -1931,3 +1978,276 @@ class CassandraStorageBackend(StorageBackend):
             running_run_ids=json.loads(row.running_run_ids) if row.running_run_ids else [],
             buffered_count=row.buffered_count or 0,
         )
+
+    # Stream Operations
+
+    async def create_stream(self, stream_id: str, metadata: dict | None = None) -> None:
+        """Create a new stream."""
+        session = self._ensure_connected()
+        now = datetime.now(UTC)
+
+        session.execute(
+            SimpleStatement(
+                "INSERT INTO streams (stream_id, status, created_at, metadata) "
+                "VALUES (%s, %s, %s, %s)",
+                consistency_level=self.write_consistency,
+            ),
+            (stream_id, "active", now, json.dumps(metadata or {})),
+        )
+
+    async def get_stream(self, stream_id: str) -> dict | None:
+        """Get a stream by ID."""
+        session = self._ensure_connected()
+
+        row = session.execute(
+            SimpleStatement(
+                "SELECT stream_id, status, created_at, metadata FROM streams WHERE stream_id = %s",
+                consistency_level=self.read_consistency,
+            ),
+            (stream_id,),
+        ).one()
+
+        if row is None:
+            return None
+        return {
+            "stream_id": row.stream_id,
+            "status": row.status,
+            "created_at": row.created_at.isoformat() if row.created_at else None,
+            "metadata": json.loads(row.metadata) if row.metadata else {},
+        }
+
+    async def publish_signal(
+        self,
+        signal_id: str,
+        stream_id: str,
+        signal_type: str,
+        payload: dict,
+        source_run_id: str | None = None,
+        stream_run_id: str | None = None,
+        metadata: dict | None = None,
+    ) -> int:
+        """Publish a signal to a stream."""
+        session = self._ensure_connected()
+        now = datetime.now(UTC)
+
+        # Get next sequence number (scoped per stream_run_id)
+        run_key = stream_run_id
+        row = session.execute(
+            SimpleStatement(
+                "SELECT MAX(sequence) AS max_seq FROM signals WHERE stream_run_id = %s",
+                consistency_level=self.read_consistency,
+            ),
+            (run_key,),
+        ).one()
+
+        seq = (row.max_seq + 1) if row and row.max_seq is not None else 0
+
+        session.execute(
+            SimpleStatement(
+                "INSERT INTO signals (stream_run_id, sequence, signal_id, stream_id, signal_type, payload, "
+                "published_at, source_run_id, metadata) VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s)",
+                consistency_level=self.write_consistency,
+            ),
+            (
+                run_key,
+                seq,
+                signal_id,
+                stream_id,
+                signal_type,
+                json.dumps(payload),
+                now,
+                source_run_id,
+                json.dumps(metadata or {}),
+            ),
+        )
+        return seq
+
+    async def get_signals(
+        self,
+        stream_id: str,
+        after_sequence: int = 0,
+        limit: int = 100,
+    ) -> list[dict]:
+        """Get signals from a stream after a given sequence number."""
+        session = self._ensure_connected()
+
+        rows = session.execute(
+            SimpleStatement(
+                "SELECT signal_id, stream_id, signal_type, payload, published_at, "
+                "sequence, source_run_id, stream_run_id, metadata FROM signals "
+                "WHERE stream_id = %s AND sequence >= %s LIMIT %s ALLOW FILTERING",
+                consistency_level=self.read_consistency,
+            ),
+            (stream_id, after_sequence, limit),
+        )
+
+        return [
+            {
+                "signal_id": row.signal_id,
+                "stream_id": row.stream_id,
+                "signal_type": row.signal_type,
+                "payload": json.loads(row.payload) if row.payload else {},
+                "published_at": row.published_at.isoformat() if row.published_at else None,
+                "sequence": row.sequence,
+                "source_run_id": row.source_run_id,
+                "stream_run_id": row.stream_run_id,
+                "metadata": json.loads(row.metadata) if row.metadata else {},
+            }
+            for row in rows
+        ]
+
+    async def register_stream_subscription(
+        self,
+        stream_id: str,
+        step_run_id: str,
+        signal_types: list[str],
+    ) -> None:
+        """Register a stream step's subscription to signal types."""
+        session = self._ensure_connected()
+        now = datetime.now(UTC)
+
+        session.execute(
+            SimpleStatement(
+                "INSERT INTO stream_subscriptions (stream_id, step_run_id, signal_types, "
+                "status, created_at) VALUES (%s, %s, %s, %s, %s)",
+                consistency_level=self.write_consistency,
+            ),
+            (stream_id, step_run_id, json.dumps(signal_types), "waiting", now),
+        )
+
+    async def get_waiting_steps(
+        self,
+        stream_id: str,
+        signal_type: str,
+    ) -> list[dict]:
+        """Get step_run_ids waiting for a specific signal type on a stream."""
+        session = self._ensure_connected()
+
+        rows = session.execute(
+            SimpleStatement(
+                "SELECT stream_id, step_run_id, signal_types, status, created_at "
+                "FROM stream_subscriptions WHERE stream_id = %s",
+                consistency_level=self.read_consistency,
+            ),
+            (stream_id,),
+        )
+
+        result = []
+        for row in rows:
+            if row.status != "waiting":
+                continue
+            types = json.loads(row.signal_types) if row.signal_types else []
+            if signal_type in types:
+                result.append(
+                    {
+                        "stream_id": row.stream_id,
+                        "step_run_id": row.step_run_id,
+                        "signal_types": types,
+                        "status": row.status,
+                        "created_at": row.created_at.isoformat() if row.created_at else None,
+                    }
+                )
+        return result
+
+    async def update_subscription_status(
+        self,
+        stream_id: str,
+        step_run_id: str,
+        status: str,
+    ) -> None:
+        """Update a subscription's status."""
+        session = self._ensure_connected()
+
+        session.execute(
+            SimpleStatement(
+                "UPDATE stream_subscriptions SET status = %s "
+                "WHERE stream_id = %s AND step_run_id = %s",
+                consistency_level=self.write_consistency,
+            ),
+            (status, stream_id, step_run_id),
+        )
+
+    async def acknowledge_signal(
+        self,
+        signal_id: str,
+        step_run_id: str,
+    ) -> None:
+        """Acknowledge that a signal has been processed by a step."""
+        session = self._ensure_connected()
+        now = datetime.now(UTC)
+
+        session.execute(
+            SimpleStatement(
+                "INSERT INTO signal_acknowledgments (signal_id, step_run_id, acknowledged_at) "
+                "VALUES (%s, %s, %s)",
+                consistency_level=self.write_consistency,
+            ),
+            (signal_id, step_run_id, now),
+        )
+
+    async def get_pending_signals(
+        self,
+        stream_id: str,
+        step_run_id: str,
+    ) -> list[dict]:
+        """Get signals that arrived for a step but haven't been acknowledged."""
+        session = self._ensure_connected()
+
+        # Get subscription
+        sub_row = session.execute(
+            SimpleStatement(
+                "SELECT signal_types FROM stream_subscriptions "
+                "WHERE stream_id = %s AND step_run_id = %s",
+                consistency_level=self.read_consistency,
+            ),
+            (stream_id, step_run_id),
+        ).one()
+
+        if not sub_row:
+            return []
+
+        signal_types = json.loads(sub_row.signal_types) if sub_row.signal_types else []
+        if not signal_types:
+            return []
+
+        # Get all signals for the stream
+        signals = session.execute(
+            SimpleStatement(
+                "SELECT signal_id, stream_id, signal_type, payload, published_at, "
+                "sequence, source_run_id, metadata FROM signals "
+                "WHERE stream_id = %s ORDER BY sequence ASC",
+                consistency_level=self.read_consistency,
+            ),
+            (stream_id,),
+        )
+
+        pending = []
+        for sig in signals:
+            if sig.signal_type not in signal_types:
+                continue
+
+            # Check if acknowledged
+            ack = session.execute(
+                SimpleStatement(
+                    "SELECT signal_id FROM signal_acknowledgments "
+                    "WHERE signal_id = %s AND step_run_id = %s",
+                    consistency_level=self.read_consistency,
+                ),
+                (sig.signal_id, step_run_id),
+            ).one()
+
+            if ack is None:
+                pending.append(
+                    {
+                        "signal_id": sig.signal_id,
+                        "stream_id": sig.stream_id,
+                        "signal_type": sig.signal_type,
+                        "payload": json.loads(sig.payload) if sig.payload else {},
+                        "published_at": sig.published_at.isoformat() if sig.published_at else None,
+                        "sequence": sig.sequence,
+                        "source_run_id": sig.source_run_id,
+                        "metadata": json.loads(sig.metadata) if sig.metadata else {},
+                    }
+                )
+
+        return pending

{pyworkflow_engine-0.2.3 → pyworkflow_engine-0.2.6}/pyworkflow/storage/citus.py

@@ -79,6 +79,37 @@ class CitusMigrationRunner(PostgresMigrationRunner):
                     """)
                 # NOTE: No UNIQUE constraints added here — Citus requires distribution
                 # column in every unique constraint, which V2 doesn't add.
+            elif migration.version == 3:
+                # V3: Restructure signals table — PK changes to (stream_id, stream_run_id, sequence).
+                # Must undistribute before dropping on Citus.
+                await conn.execute("DROP TABLE IF EXISTS signal_acknowledgments")
+                await conn.execute("DROP TABLE IF EXISTS signals")
+                await conn.execute("""
+                    CREATE TABLE signals (
+                        stream_run_id TEXT NOT NULL,
+                        sequence INTEGER NOT NULL,
+                        signal_id TEXT NOT NULL,
+                        stream_id TEXT NOT NULL,
+                        signal_type TEXT NOT NULL,
+                        payload JSONB NOT NULL DEFAULT '{}',
+                        published_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),
+                        source_run_id TEXT,
+                        metadata JSONB DEFAULT '{}',
+                        PRIMARY KEY (stream_id, stream_run_id, sequence)
+                    )
+                """)
+                await conn.execute("CREATE INDEX idx_signals_signal_id ON signals(signal_id)")
+                await conn.execute("""
+                    CREATE TABLE signal_acknowledgments (
+                        signal_id TEXT NOT NULL,
+                        step_run_id TEXT NOT NULL,
+                        acknowledged_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),
+                        PRIMARY KEY (signal_id, step_run_id)
+                    )
+                """)
+                # Re-distribute the new tables
+                await conn.execute("SELECT create_distributed_table('signals', 'stream_run_id')")
+                await conn.execute("SELECT create_reference_table('signal_acknowledgments')")
             elif migration.up_func:
                 await migration.up_func(conn)
             elif migration.up_sql and migration.up_sql != "SELECT 1":
@@ -337,6 +368,68 @@ class CitusStorageBackend(PostgresStorageBackend):
                 )
             """)
 
+            # streams — reference table (independent, keyed by stream_id)
+            await conn.execute("""
+                CREATE TABLE IF NOT EXISTS streams (
+                    stream_id TEXT PRIMARY KEY,
+                    status TEXT NOT NULL DEFAULT 'active',
+                    created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),
+                    metadata JSONB DEFAULT '{}'
+                )
+            """)
+
+            # signals — distributed on stream_run_id
+            # Differences vs postgres:
+            #   - PK: (stream_id, stream_run_id, sequence) — distribution column is stream_run_id
+            #   - No FK constraints (cross-shard FKs unsupported)
+            await conn.execute("""
+                CREATE TABLE IF NOT EXISTS signals (
+                    stream_run_id TEXT NOT NULL,
+                    sequence INTEGER NOT NULL,
+                    signal_id TEXT NOT NULL,
+                    stream_id TEXT NOT NULL,
+                    signal_type TEXT NOT NULL,
+                    payload JSONB NOT NULL DEFAULT '{}',
+                    published_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),
+                    source_run_id TEXT,
+                    metadata JSONB DEFAULT '{}',
+                    PRIMARY KEY (stream_id, stream_run_id, sequence)
+                )
+            """)
+            await conn.execute(
+                "CREATE INDEX IF NOT EXISTS idx_signals_signal_id ON signals(signal_id)"
+            )
+
+            # stream_subscriptions — distributed on stream_id, co-located with signals
+            # Differences vs postgres:
+            #   - PK: (stream_id, step_run_id) to include distribution column
+            #   - No FK constraints (cross-shard FKs unsupported)
+            await conn.execute("""
+                CREATE TABLE IF NOT EXISTS stream_subscriptions (
+                    stream_id TEXT NOT NULL,
+                    step_run_id TEXT NOT NULL,
+                    signal_types JSONB NOT NULL DEFAULT '[]',
+                    status TEXT NOT NULL DEFAULT 'waiting',
+                    created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),
+                    PRIMARY KEY (stream_id, step_run_id)
+                )
+            """)
+            await conn.execute(
+                "CREATE INDEX IF NOT EXISTS idx_subscriptions_status ON stream_subscriptions(stream_id, status)"
+            )
+
+            # signal_acknowledgments — reference table (needs cross-shard joins)
+            # Differences vs postgres:
+            #   - No FK constraints (cross-shard FKs unsupported)
+            await conn.execute("""
+                CREATE TABLE IF NOT EXISTS signal_acknowledgments (
+                    signal_id TEXT NOT NULL,
+                    step_run_id TEXT NOT NULL,
+                    acknowledged_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),
+                    PRIMARY KEY (signal_id, step_run_id)
+                )
+            """)
+
         # Step 4: Distribute tables (idempotent — already-distributed tables are skipped)
         await self._distribute_tables(pool)
 
@@ -385,3 +478,21 @@ class CitusStorageBackend(PostgresStorageBackend):
             # checkpoints: reference table (keyed by step_run_id, can't co-locate with run_id)
             if "checkpoints" not in distributed:
                 await conn.execute("SELECT create_reference_table('checkpoints')")
+
+            # streams: reference table (independent, small)
+            if "streams" not in distributed:
+                await conn.execute("SELECT create_reference_table('streams')")
+
+            # signals: distributed on stream_run_id
+            if "signals" not in distributed:
+                await conn.execute("SELECT create_distributed_table('signals', 'stream_run_id')")
+
+            # stream_subscriptions: distributed on stream_id (independent from signals)
+            if "stream_subscriptions" not in distributed:
+                await conn.execute(
+                    "SELECT create_distributed_table('stream_subscriptions', 'stream_id')"
+                )
+
+            # signal_acknowledgments: reference table (needs cross-shard joins)
+            if "signal_acknowledgments" not in distributed:
+                await conn.execute("SELECT create_reference_table('signal_acknowledgments')")