dbos 1.10.0a1__tar.gz → 1.11.0a1__tar.gz

{dbos-1.10.0a1 → dbos-1.11.0a1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: dbos
-Version: 1.10.0a1
+Version: 1.11.0a1
 Summary: Ultra-lightweight durable execution in Python
 Author-Email: "DBOS, Inc." <contact@dbos.dev>
 License: MIT

{dbos-1.10.0a1 → dbos-1.11.0a1}/dbos/_client.py

@@ -1,7 +1,18 @@
 import asyncio
 import sys
+import time
 import uuid
-from typing import Any, Generic, List, Optional, TypedDict, TypeVar, Union
+from typing import (
+    Any,
+    AsyncGenerator,
+    Generator,
+    Generic,
+    List,
+    Optional,
+    TypedDict,
+    TypeVar,
+    Union,
+)
 
 from dbos._app_db import ApplicationDatabase
 from dbos._context import MaxPriority, MinPriority
@@ -24,6 +35,8 @@ from dbos._sys_db import (
     WorkflowStatus,
     WorkflowStatusInternal,
     WorkflowStatusString,
+    _dbos_stream_closed_sentinel,
+    workflow_is_active,
 )
 from dbos._workflow_commands import (
     fork_workflow,
@@ -449,3 +462,71 @@ class DBOSClient:
             application_version=application_version,
         )
         return WorkflowHandleClientAsyncPolling[Any](forked_workflow_id, self._sys_db)
+
+    def read_stream(self, workflow_id: str, key: str) -> Generator[Any, Any, None]:
+        """
+        Read values from a stream as a generator.
+        This function reads values from a stream identified by the workflow_id and key,
+        yielding each value in order until the stream is closed or the workflow terminates.
+
+        Args:
+            workflow_id: The ID of the workflow that wrote to the stream
+            key: The stream key to read from
+
+        Yields:
+            The values written to the stream in order
+        """
+        offset = 0
+        while True:
+            try:
+                value = self._sys_db.read_stream(workflow_id, key, offset)
+                if value == _dbos_stream_closed_sentinel:
+                    break
+                yield value
+                offset += 1
+            except ValueError:
+                # Poll the offset until a value arrives or the workflow terminates
+                status = get_workflow(self._sys_db, workflow_id)
+                if status is None:
+                    break
+                if not workflow_is_active(status.status):
+                    break
+                time.sleep(1.0)
+                continue
+
+    async def read_stream_async(
+        self, workflow_id: str, key: str
+    ) -> AsyncGenerator[Any, None]:
+        """
+        Read values from a stream as an async generator.
+        This function reads values from a stream identified by the workflow_id and key,
+        yielding each value in order until the stream is closed or the workflow terminates.
+
+        Args:
+            workflow_id: The ID of the workflow that wrote to the stream
+            key: The stream key to read from
+
+        Yields:
+            The values written to the stream in order
+        """
+        offset = 0
+        while True:
+            try:
+                value = await asyncio.to_thread(
+                    self._sys_db.read_stream, workflow_id, key, offset
+                )
+                if value == _dbos_stream_closed_sentinel:
+                    break
+                yield value
+                offset += 1
+            except ValueError:
+                # Poll the offset until a value arrives or the workflow terminates
+                status = await asyncio.to_thread(
+                    get_workflow, self._sys_db, workflow_id
+                )
+                if status is None:
+                    break
+                if not workflow_is_active(status.status):
+                    break
+                await asyncio.sleep(1.0)
+                continue

{dbos-1.10.0a1 → dbos-1.11.0a1}/dbos/_dbos.py

@@ -6,14 +6,17 @@ import inspect
 import os
 import sys
 import threading
+import time
 import uuid
 from concurrent.futures import ThreadPoolExecutor
 from logging import Logger
 from typing import (
     TYPE_CHECKING,
     Any,
+    AsyncGenerator,
     Callable,
     Coroutine,
+    Generator,
     Generic,
     List,
     Literal,
@@ -62,7 +65,14 @@ from ._registrations import (
 )
 from ._roles import default_required_roles, required_roles
 from ._scheduler import ScheduledWorkflow, scheduled
-from ._sys_db import StepInfo, SystemDatabase, WorkflowStatus, reset_system_database
+from ._sys_db import (
+    StepInfo,
+    SystemDatabase,
+    WorkflowStatus,
+    _dbos_stream_closed_sentinel,
+    reset_system_database,
+    workflow_is_active,
+)
 from ._tracer import DBOSTracer, dbos_tracer
 
 if TYPE_CHECKING:
@@ -1304,6 +1314,171 @@ class DBOS:
         ctx.authenticated_user = authenticated_user
         ctx.authenticated_roles = authenticated_roles
 
+    @classmethod
+    def write_stream(cls, key: str, value: Any) -> None:
+        """
+        Write a value to a stream.
+
+        Args:
+            key(str): The stream key / name within the workflow
+            value(Any): A serializable value to write to the stream
+
+        """
+        ctx = get_local_dbos_context()
+        if ctx is not None:
+            # Must call it within a workflow
+            if ctx.is_workflow():
+                attributes: TracedAttributes = {
+                    "name": "write_stream",
+                }
+                with EnterDBOSStep(attributes):
+                    ctx = assert_current_dbos_context()
+                    _get_dbos_instance()._sys_db.write_stream_from_workflow(
+                        ctx.workflow_id, ctx.function_id, key, value
+                    )
+            elif ctx.is_step():
+                _get_dbos_instance()._sys_db.write_stream_from_step(
+                    ctx.workflow_id, key, value
+                )
+            else:
+                raise DBOSException(
+                    "write_stream() must be called from within a workflow or step"
+                )
+        else:
+            # Cannot call it from outside of a workflow
+            raise DBOSException(
+                "write_stream() must be called from within a workflow or step"
+            )
+
+    @classmethod
+    def close_stream(cls, key: str) -> None:
+        """
+        Close a stream.
+
+        Args:
+            key(str): The stream key / name within the workflow
+
+        """
+        ctx = get_local_dbos_context()
+        if ctx is not None:
+            # Must call it within a workflow
+            if ctx.is_workflow():
+                attributes: TracedAttributes = {
+                    "name": "close_stream",
+                }
+                with EnterDBOSStep(attributes):
+                    ctx = assert_current_dbos_context()
+                    _get_dbos_instance()._sys_db.close_stream(
+                        ctx.workflow_id, ctx.function_id, key
+                    )
+            else:
+                raise DBOSException(
+                    "close_stream() must be called from within a workflow"
+                )
+        else:
+            # Cannot call it from outside of a workflow
+            raise DBOSException("close_stream() must be called from within a workflow")
+
+    @classmethod
+    def read_stream(cls, workflow_id: str, key: str) -> Generator[Any, Any, None]:
+        """
+        Read values from a stream as a generator.
+
+        This function reads values from a stream identified by the workflow_id and key,
+        yielding each value in order until the stream is closed or the workflow terminates.
+
+        Args:
+            workflow_id(str): The workflow instance ID that owns the stream
+            key(str): The stream key / name within the workflow
+
+        Yields:
+            Any: Each value in the stream until the stream is closed
+
+        """
+        offset = 0
+        sys_db = _get_dbos_instance()._sys_db
+
+        while True:
+            try:
+                value = sys_db.read_stream(workflow_id, key, offset)
+                if value == _dbos_stream_closed_sentinel:
+                    break
+                yield value
+                offset += 1
+            except ValueError:
+                # Poll the offset until a value arrives or the workflow terminates
+                status = cls.retrieve_workflow(workflow_id).get_status().status
+                if not workflow_is_active(status):
+                    break
+                time.sleep(1.0)
+                continue
+
+    @classmethod
+    async def write_stream_async(cls, key: str, value: Any) -> None:
+        """
+        Write a value to a stream asynchronously.
+
+        Args:
+            key(str): The stream key / name within the workflow
+            value(Any): A serializable value to write to the stream
+
+        """
+        await cls._configure_asyncio_thread_pool()
+        await asyncio.to_thread(lambda: DBOS.write_stream(key, value))
+
+    @classmethod
+    async def close_stream_async(cls, key: str) -> None:
+        """
+        Close a stream asynchronously.
+
+        Args:
+            key(str): The stream key / name within the workflow
+
+        """
+        await cls._configure_asyncio_thread_pool()
+        await asyncio.to_thread(lambda: DBOS.close_stream(key))
+
+    @classmethod
+    async def read_stream_async(
+        cls, workflow_id: str, key: str
+    ) -> AsyncGenerator[Any, None]:
+        """
+        Read values from a stream as an async generator.
+
+        This function reads values from a stream identified by the workflow_id and key,
+        yielding each value in order until the stream is closed or the workflow terminates.
+
+        Args:
+            workflow_id(str): The workflow instance ID that owns the stream
+            key(str): The stream key / name within the workflow
+
+        Yields:
+            Any: Each value in the stream until the stream is closed
+
+        """
+        await cls._configure_asyncio_thread_pool()
+        offset = 0
+        sys_db = _get_dbos_instance()._sys_db
+
+        while True:
+            try:
+                value = await asyncio.to_thread(
+                    sys_db.read_stream, workflow_id, key, offset
+                )
+                if value == _dbos_stream_closed_sentinel:
+                    break
+                yield value
+                offset += 1
+            except ValueError:
+                # Poll the offset until a value arrives or the workflow terminates
+                status = (
+                    await (await cls.retrieve_workflow_async(workflow_id)).get_status()
+                ).status
+                if not workflow_is_active(status):
+                    break
+                await asyncio.sleep(1.0)
+                continue
+
     @classproperty
     def tracer(self) -> DBOSTracer:
         """Return the DBOS OpenTelemetry tracer."""

dbos-1.11.0a1/dbos/_migrations/versions/01ce9f07bd10_streaming.py

@@ -0,0 +1,42 @@
+"""streaming
+
+Revision ID: 01ce9f07bd10
+Revises: d994145b47b6
+Create Date: 2025-08-05 10:20:46.424975
+
+"""
+
+from typing import Sequence, Union
+
+import sqlalchemy as sa
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision: str = "01ce9f07bd10"
+down_revision: Union[str, None] = "d994145b47b6"
+branch_labels: Union[str, Sequence[str], None] = None
+depends_on: Union[str, Sequence[str], None] = None
+
+
+def upgrade() -> None:
+    # Create streams table
+    op.create_table(
+        "streams",
+        sa.Column("workflow_uuid", sa.Text(), nullable=False),
+        sa.Column("key", sa.Text(), nullable=False),
+        sa.Column("value", sa.Text(), nullable=False),
+        sa.Column("offset", sa.Integer(), nullable=False),
+        sa.ForeignKeyConstraint(
+            ["workflow_uuid"],
+            ["dbos.workflow_status.workflow_uuid"],
+            onupdate="CASCADE",
+            ondelete="CASCADE",
+        ),
+        sa.PrimaryKeyConstraint("workflow_uuid", "key", "offset"),
+        schema="dbos",
+    )
+
+
+def downgrade() -> None:
+    # Drop streams table
+    op.drop_table("streams", schema="dbos")

{dbos-1.10.0a1 → dbos-1.11.0a1}/dbos/_schemas/system_database.py

@@ -132,3 +132,20 @@ class SystemSchema:
         Column("value", Text, nullable=False),
         PrimaryKeyConstraint("workflow_uuid", "key"),
     )
+
+    streams = Table(
+        "streams",
+        metadata_obj,
+        Column(
+            "workflow_uuid",
+            Text,
+            ForeignKey(
+                "workflow_status.workflow_uuid", onupdate="CASCADE", ondelete="CASCADE"
+            ),
+            nullable=False,
+        ),
+        Column("key", Text, nullable=False),
+        Column("value", Text, nullable=False),
+        Column("offset", Integer, nullable=False),
+        PrimaryKeyConstraint("workflow_uuid", "key", "offset"),
+    )

{dbos-1.10.0a1 → dbos-1.11.0a1}/dbos/_sys_db.py

@@ -62,6 +62,13 @@ class WorkflowStatusString(Enum):
     ENQUEUED = "ENQUEUED"
 
 
+def workflow_is_active(status: str) -> bool:
+    return (
+        status == WorkflowStatusString.ENQUEUED.value
+        or status == WorkflowStatusString.PENDING.value
+    )
+
+
 WorkflowStatuses = Literal[
     "PENDING",
     "SUCCESS",
@@ -243,6 +250,7 @@ class StepInfo(TypedDict):
 
 
 _dbos_null_topic = "__null__topic__"
+_dbos_stream_closed_sentinel = "__DBOS_STREAM_CLOSED__"
 
 
 class ConditionCount(TypedDict):
@@ -1882,6 +1890,127 @@ class SystemDatabase:
             dbos_logger.error(f"Error connecting to the DBOS system database: {e}")
             raise
 
+    def write_stream_from_step(self, workflow_uuid: str, key: str, value: Any) -> None:
+        """
+        Write a key-value pair to the stream at the first unused offset.
+        """
+        if self._debug_mode:
+            raise Exception("called write_stream in debug mode")
+
+        with self.engine.begin() as c:
+            # Find the maximum offset for this workflow_uuid and key combination
+            max_offset_result = c.execute(
+                sa.select(sa.func.max(SystemSchema.streams.c.offset)).where(
+                    SystemSchema.streams.c.workflow_uuid == workflow_uuid,
+                    SystemSchema.streams.c.key == key,
+                )
+            ).fetchone()
+
+            # Next offset is max + 1, or 0 if no records exist
+            next_offset = (
+                (max_offset_result[0] + 1)
+                if max_offset_result is not None and max_offset_result[0] is not None
+                else 0
+            )
+
+            # Serialize the value before storing
+            serialized_value = _serialization.serialize(value)
+
+            # Insert the new stream entry
+            c.execute(
+                sa.insert(SystemSchema.streams).values(
+                    workflow_uuid=workflow_uuid,
+                    key=key,
+                    value=serialized_value,
+                    offset=next_offset,
+                )
+            )
+
+    @db_retry()
+    def write_stream_from_workflow(
+        self, workflow_uuid: str, function_id: int, key: str, value: Any
+    ) -> None:
+        """
+        Write a key-value pair to the stream at the first unused offset.
+        """
+        function_name = (
+            "DBOS.closeStream"
+            if value == _dbos_stream_closed_sentinel
+            else "DBOS.writeStream"
+        )
+
+        with self.engine.begin() as c:
+
+            recorded_output = self._check_operation_execution_txn(
+                workflow_uuid, function_id, function_name, conn=c
+            )
+            if self._debug_mode and recorded_output is None:
+                raise Exception(
+                    "called set_event in debug mode without a previous execution"
+                )
+            # Find the maximum offset for this workflow_uuid and key combination
+            max_offset_result = c.execute(
+                sa.select(sa.func.max(SystemSchema.streams.c.offset)).where(
+                    SystemSchema.streams.c.workflow_uuid == workflow_uuid,
+                    SystemSchema.streams.c.key == key,
+                )
+            ).fetchone()
+
+            # Next offset is max + 1, or 0 if no records exist
+            next_offset = (
+                (max_offset_result[0] + 1)
+                if max_offset_result is not None and max_offset_result[0] is not None
+                else 0
+            )
+
+            # Serialize the value before storing
+            serialized_value = _serialization.serialize(value)
+
+            # Insert the new stream entry
+            c.execute(
+                sa.insert(SystemSchema.streams).values(
+                    workflow_uuid=workflow_uuid,
+                    key=key,
+                    value=serialized_value,
+                    offset=next_offset,
+                )
+            )
+            output: OperationResultInternal = {
+                "workflow_uuid": workflow_uuid,
+                "function_id": function_id,
+                "function_name": function_name,
+                "output": None,
+                "error": None,
+            }
+            self._record_operation_result_txn(output, conn=c)
+
+    def close_stream(self, workflow_uuid: str, function_id: int, key: str) -> None:
+        """Write a sentinel value to the stream at the first unused offset to mark it as closed."""
+        self.write_stream_from_workflow(
+            workflow_uuid, function_id, key, _dbos_stream_closed_sentinel
+        )
+
+    @db_retry()
+    def read_stream(self, workflow_uuid: str, key: str, offset: int) -> Any:
+        """Read the value at the specified offset for the given workflow_uuid and key."""
+
+        with self.engine.begin() as c:
+            result = c.execute(
+                sa.select(SystemSchema.streams.c.value).where(
+                    SystemSchema.streams.c.workflow_uuid == workflow_uuid,
+                    SystemSchema.streams.c.key == key,
+                    SystemSchema.streams.c.offset == offset,
+                )
+            ).fetchone()
+
+            if result is None:
+                raise ValueError(
+                    f"No value found for workflow_uuid={workflow_uuid}, key={key}, offset={offset}"
+                )
+
+            # Deserialize the value before returning
+            return _serialization.deserialize(result[0])
+
     def garbage_collect(
         self, cutoff_epoch_timestamp_ms: Optional[int], rows_threshold: Optional[int]
     ) -> Optional[tuple[int, list[str]]]:

{dbos-1.10.0a1 → dbos-1.11.0a1}/dbos/cli/cli.py

@@ -260,11 +260,9 @@ def _resolve_project_name_and_template(
     return project_name, template
 
 
-@app.command(
-    help="Run your database schema migrations using the migration commands in 'dbos-config.yaml'"
-)
+@app.command(help="Create DBOS system tables.")
 def migrate(
-    db_url: Annotated[
+    app_database_url: Annotated[
         typing.Optional[str],
         typer.Option(
             "--db-url",
@@ -272,30 +270,34 @@ def migrate(
             help="Your DBOS application database URL",
         ),
     ] = None,
-    sys_db_name: Annotated[
+    system_database_url: Annotated[
         typing.Optional[str],
         typer.Option(
-            "--sys-db-name",
+            "--sys-db-url",
             "-s",
-            help="Specify the name of the system database to reset",
+            help="Your DBOS system database URL",
         ),
     ] = None,
 ) -> None:
-    config = load_config(run_process_config=False, silent=True)
-    connection_string = _get_db_url(db_url)
-    app_db_name = sa.make_url(connection_string).database
-    assert app_db_name is not None, "Database name is required in URL"
-    if sys_db_name is None:
-        sys_db_name = app_db_name + SystemSchema.sysdb_suffix
+    app_database_url = _get_db_url(app_database_url)
+    system_database_url = get_system_database_url(
+        {
+            "system_database_url": system_database_url,
+            "database_url": app_database_url,
+            "database": {},
+        }
+    )
 
-    typer.echo(f"Starting schema migration for database {app_db_name}")
+    typer.echo(f"Starting DBOS migrations")
+    typer.echo(f"Application database: {sa.make_url(app_database_url)}")
+    typer.echo(f"System database: {sa.make_url(system_database_url)}")
 
     # First, run DBOS migrations on the system database and the application database
     app_db = None
     sys_db = None
     try:
         sys_db = SystemDatabase(
-            system_database_url=get_system_database_url(config),
+            system_database_url=system_database_url,
             engine_kwargs={
                 "pool_timeout": 30,
                 "max_overflow": 0,
@@ -303,7 +305,7 @@ def migrate(
             },
         )
         app_db = ApplicationDatabase(
-            database_url=connection_string,
+            database_url=app_database_url,
             engine_kwargs={
                 "pool_timeout": 30,
                 "max_overflow": 0,
@@ -313,17 +315,19 @@ def migrate(
         sys_db.run_migrations()
         app_db.run_migrations()
     except Exception as e:
-        typer.echo(f"DBOS system schema migration failed: {e}")
+        typer.echo(f"DBOS migrations failed: {e}")
+        raise typer.Exit(code=1)
     finally:
         if sys_db:
             sys_db.destroy()
         if app_db:
             app_db.destroy()
 
+    typer.echo(f"DBOS migrations successful")
+
     # Next, run any custom migration commands specified in the configuration
-    typer.echo("Executing migration commands from 'dbos-config.yaml'")
-    try:
-        # handle the case where the user has not specified migrations commands
+    if os.path.exists("dbos-config.yaml"):
+        config = load_config(run_process_config=False, silent=True)
         if "database" not in config:
             config["database"] = {}
         migrate_commands = (
@@ -331,20 +335,21 @@ def migrate(
             if "migrate" in config["database"] and config["database"]["migrate"]
             else []
         )
-        for command in migrate_commands:
-            typer.echo(f"Executing migration command: {command}")
-            result = subprocess.run(command, shell=True, text=True)
-            if result.returncode != 0:
-                typer.echo(f"Migration command failed: {command}")
-                typer.echo(result.stderr)
-                raise typer.Exit(1)
-            if result.stdout:
-                typer.echo(result.stdout.rstrip())
-    except Exception as e:
-        typer.echo(f"An error occurred during schema migration: {e}")
-        raise typer.Exit(code=1)
-
-    typer.echo(f"Completed schema migration for database {app_db_name}")
+        if migrate_commands:
+            typer.echo("Executing migration commands from 'dbos-config.yaml'")
+            try:
+                for command in migrate_commands:
+                    typer.echo(f"Executing migration command: {command}")
+                    result = subprocess.run(command, shell=True, text=True)
+                    if result.returncode != 0:
+                        typer.echo(f"Migration command failed: {command}")
+                        typer.echo(result.stderr)
+                        raise typer.Exit(1)
+                    if result.stdout:
+                        typer.echo(result.stdout.rstrip())
+            except Exception as e:
+                typer.echo(f"An error occurred during schema migration: {e}")
+                raise typer.Exit(code=1)
 
 
 @app.command(help="Reset the DBOS system database")

{dbos-1.10.0a1 → dbos-1.11.0a1}/pyproject.toml

@@ -27,7 +27,7 @@ dependencies = [
 ]
 requires-python = ">=3.9"
 readme = "README.md"
-version = "1.10.0a1"
+version = "1.11.0a1"
 
 [project.license]
 text = "MIT"