pyworkflow-engine 0.1.7__py3-none-any.whl → 0.1.9__py3-none-any.whl

pyworkflow/storage/config.py

@@ -73,6 +73,32 @@ def storage_to_config(storage: StorageBackend | None) -> dict[str, Any] | None:
             "region": getattr(storage, "region", "us-east-1"),
             "endpoint_url": getattr(storage, "endpoint_url", None),
         }
+    elif class_name == "CassandraStorageBackend":
+        return {
+            "type": "cassandra",
+            "contact_points": getattr(storage, "contact_points", ["localhost"]),
+            "port": getattr(storage, "port", 9042),
+            "keyspace": getattr(storage, "keyspace", "pyworkflow"),
+            "username": getattr(storage, "username", None),
+            "password": getattr(storage, "password", None),
+            "read_consistency": getattr(storage, "read_consistency", "LOCAL_QUORUM"),
+            "write_consistency": getattr(storage, "write_consistency", "LOCAL_QUORUM"),
+            "replication_strategy": getattr(storage, "replication_strategy", "SimpleStrategy"),
+            "replication_factor": getattr(storage, "replication_factor", 3),
+            "datacenter": getattr(storage, "datacenter", None),
+        }
+    elif class_name == "MySQLStorageBackend":
+        config = {"type": "mysql"}
+        dsn = getattr(storage, "dsn", None)
+        if dsn:
+            config["dsn"] = dsn
+        else:
+            config["host"] = getattr(storage, "host", "localhost")
+            config["port"] = getattr(storage, "port", 3306)
+            config["user"] = getattr(storage, "user", "pyworkflow")
+            config["password"] = getattr(storage, "password", "")
+            config["database"] = getattr(storage, "database", "pyworkflow")
+        return config
     else:
         # Unknown backend - return minimal config
         return {"type": "unknown"}
@@ -181,5 +207,48 @@ def config_to_storage(config: dict[str, Any] | None = None) -> StorageBackend:
             endpoint_url=config.get("endpoint_url"),
         )
 
+    elif storage_type == "cassandra":
+        try:
+            from pyworkflow.storage.cassandra import CassandraStorageBackend
+        except ImportError:
+            raise ValueError(
+                "Cassandra storage backend is not available. "
+                "Please install the required dependencies with: pip install 'pyworkflow[cassandra]'"
+            )
+
+        return CassandraStorageBackend(
+            contact_points=config.get("contact_points", ["localhost"]),
+            port=config.get("port", 9042),
+            keyspace=config.get("keyspace", "pyworkflow"),
+            username=config.get("username"),
+            password=config.get("password"),
+            read_consistency=config.get("read_consistency", "LOCAL_QUORUM"),
+            write_consistency=config.get("write_consistency", "LOCAL_QUORUM"),
+            replication_strategy=config.get("replication_strategy", "SimpleStrategy"),
+            replication_factor=config.get("replication_factor", 3),
+            datacenter=config.get("datacenter"),
+        )
+
+    elif storage_type == "mysql":
+        try:
+            from pyworkflow.storage.mysql import MySQLStorageBackend
+        except ImportError:
+            raise ValueError(
+                "MySQL storage backend is not available. "
+                "Please install the required dependencies with: pip install 'pyworkflow[mysql]'"
+            )
+
+        # Support both DSN and individual parameters
+        if "dsn" in config:
+            return MySQLStorageBackend(dsn=config["dsn"])
+        else:
+            return MySQLStorageBackend(
+                host=config.get("host", "localhost"),
+                port=config.get("port", 3306),
+                user=config.get("user", "pyworkflow"),
+                password=config.get("password", ""),
+                database=config.get("database", "pyworkflow"),
+            )
+
     else:
         raise ValueError(f"Unknown storage type: {storage_type}")

pyworkflow/storage/dynamodb.py

@@ -234,7 +234,7 @@ class DynamoDBStorageBackend(StorageBackend):
                 "error": run.error,
                 "idempotency_key": run.idempotency_key,
                 "max_duration": run.max_duration,
-                "metadata": json.dumps(run.metadata),
+                "context": json.dumps(run.context),
                 "recovery_attempts": run.recovery_attempts,
                 "max_recovery_attempts": run.max_recovery_attempts,
                 "recover_on_worker_loss": run.recover_on_worker_loss,
@@ -362,6 +362,34 @@ class DynamoDBStorageBackend(StorageBackend):
                 },
             )
 
+    async def update_run_context(
+        self,
+        run_id: str,
+        context: dict,
+    ) -> None:
+        """Update the step context for a workflow run."""
+        async with self._get_client() as client:
+            await client.update_item(
+                TableName=self.table_name,
+                Key={
+                    "PK": {"S": f"RUN#{run_id}"},
+                    "SK": {"S": "#METADATA"},
+                },
+                UpdateExpression="SET #ctx = :ctx, updated_at = :now",
+                ExpressionAttributeNames={
+                    "#ctx": "context",
+                },
+                ExpressionAttributeValues={
+                    ":ctx": {"S": json.dumps(context)},
+                    ":now": {"S": datetime.now(UTC).isoformat()},
+                },
+            )
+
+    async def get_run_context(self, run_id: str) -> dict:
+        """Get the current step context for a workflow run."""
+        run = await self.get_run(run_id)
+        return run.context if run else {}
+
     async def list_runs(
         self,
         query: str | None = None,
@@ -1221,7 +1249,8 @@ class DynamoDBStorageBackend(StorageBackend):
             error=item.get("error"),
             idempotency_key=item.get("idempotency_key"),
             max_duration=item.get("max_duration"),
-            metadata=json.loads(item.get("metadata", "{}")),
+            # Support both 'context' (new) and 'metadata' (legacy) keys
+            context=json.loads(item.get("context", item.get("metadata", "{}"))),
             recovery_attempts=item.get("recovery_attempts", 0),
             max_recovery_attempts=item.get("max_recovery_attempts", 3),
             recover_on_worker_loss=item.get("recover_on_worker_loss", True),

pyworkflow/storage/file.py

@@ -190,6 +190,34 @@ class FileStorageBackend(StorageBackend):
 
         await asyncio.to_thread(_update)
 
+    async def update_run_context(
+        self,
+        run_id: str,
+        context: dict,
+    ) -> None:
+        """Update the step context for a workflow run."""
+        run_file = self.runs_dir / f"{run_id}.json"
+
+        if not run_file.exists():
+            raise ValueError(f"Workflow run {run_id} not found")
+
+        lock_file = self.locks_dir / f"{run_id}.lock"
+        lock = FileLock(str(lock_file))
+
+        def _update() -> None:
+            with lock:
+                data = json.loads(run_file.read_text())
+                data["context"] = context
+                data["updated_at"] = datetime.now(UTC).isoformat()
+                run_file.write_text(json.dumps(data, indent=2))
+
+        await asyncio.to_thread(_update)
+
+    async def get_run_context(self, run_id: str) -> dict:
+        """Get the current step context for a workflow run."""
+        run = await self.get_run(run_id)
+        return run.context if run else {}
+
     async def list_runs(
         self,
         query: str | None = None,

pyworkflow/storage/memory.py

@@ -109,6 +109,24 @@ class InMemoryStorageBackend(StorageBackend):
                 run.recovery_attempts = recovery_attempts
                 run.updated_at = datetime.now(UTC)
 
+    async def update_run_context(
+        self,
+        run_id: str,
+        context: dict,
+    ) -> None:
+        """Update the step context for a workflow run."""
+        with self._lock:
+            run = self._runs.get(run_id)
+            if run:
+                run.context = context
+                run.updated_at = datetime.now(UTC)
+
+    async def get_run_context(self, run_id: str) -> dict:
+        """Get the current step context for a workflow run."""
+        with self._lock:
+            run = self._runs.get(run_id)
+            return run.context if run else {}
+
     async def list_runs(
         self,
         query: str | None = None,