dbos 1.14.0a8__py3-none-any.whl → 1.15.0a1__py3-none-any.whl

dbos/_logger.py

@@ -2,14 +2,6 @@ import logging
 import os
 from typing import TYPE_CHECKING, Any
 
-from opentelemetry._logs import set_logger_provider
-from opentelemetry.exporter.otlp.proto.http._log_exporter import OTLPLogExporter
-from opentelemetry.sdk._logs import LoggerProvider, LoggingHandler
-from opentelemetry.sdk._logs.export import BatchLogRecordProcessor
-from opentelemetry.sdk.resources import Resource
-from opentelemetry.semconv.resource import ResourceAttributes
-from opentelemetry.trace.span import format_trace_id
-
 from dbos._utils import GlobalParams
 
 if TYPE_CHECKING:
@@ -24,6 +16,7 @@ class DBOSLogTransformer(logging.Filter):
         super().__init__()
         self.app_id = os.environ.get("DBOS__APPID", "")
         self.otlp_attributes: dict[str, str] = config.get("telemetry", {}).get("otlp_attributes", {})  # type: ignore
+        self.disable_otlp = config.get("telemetry", {}).get("disable_otlp", True)  # type: ignore
 
     def filter(self, record: Any) -> bool:
         record.applicationID = self.app_id
@@ -39,19 +32,15 @@ class DBOSLogTransformer(logging.Filter):
         if ctx:
             if ctx.is_within_workflow():
                 record.operationUUID = ctx.workflow_id
-            span = ctx.get_current_active_span()
-            if span:
-                trace_id = format_trace_id(span.get_span_context().trace_id)
-                record.traceId = trace_id
-
-        return True
+            if not self.disable_otlp:
+                from opentelemetry.trace.span import format_trace_id
 
+                span = ctx.get_current_active_span()
+                if span:
+                    trace_id = format_trace_id(span.get_span_context().trace_id)
+                    record.traceId = trace_id
 
-# Mitigation for https://github.com/open-telemetry/opentelemetry-python/issues/3193
-# Reduce the force flush timeout
-class PatchedOTLPLoggerProvider(LoggerProvider):
-    def force_flush(self, timeout_millis: int = 5000) -> bool:
-        return super().force_flush(timeout_millis)
+        return True
 
 
 def init_logger() -> None:
@@ -80,10 +69,18 @@ def config_logger(config: "ConfigFile") -> None:
     disable_otlp = config.get("telemetry", {}).get("disable_otlp", False)  # type: ignore
 
     if not disable_otlp and otlp_logs_endpoints:
-        log_provider = PatchedOTLPLoggerProvider(
+
+        from opentelemetry._logs import set_logger_provider
+        from opentelemetry.exporter.otlp.proto.http._log_exporter import OTLPLogExporter
+        from opentelemetry.sdk._logs import LoggerProvider, LoggingHandler
+        from opentelemetry.sdk._logs.export import BatchLogRecordProcessor
+        from opentelemetry.sdk.resources import Resource
+        from opentelemetry.semconv.attributes.service_attributes import SERVICE_NAME
+
+        log_provider = LoggerProvider(
             Resource.create(
                 attributes={
-                    ResourceAttributes.SERVICE_NAME: config["name"],
+                    SERVICE_NAME: config["name"],
                 }
             )
         )

dbos/_migration.py

@@ -1,16 +1,11 @@
-import logging
-import os
-import re
 import sys
 
 import sqlalchemy as sa
-from alembic import command
-from alembic.config import Config
 
 from ._logger import dbos_logger
 
 
-def ensure_dbos_schema(engine: sa.Engine) -> bool:
+def ensure_dbos_schema(engine: sa.Engine) -> None:
     """
     True if using DBOS migrations (DBOS schema and migrations table already exist or were created)
     False if using Alembic migrations (DBOS schema exists, but dbos_migrations table doesn't)
@@ -22,10 +17,10 @@ def ensure_dbos_schema(engine: sa.Engine) -> bool:
                 "SELECT schema_name FROM information_schema.schemata WHERE schema_name = 'dbos'"
             )
         )
-        schema_existed = schema_result.fetchone() is not None
+        schema_exists = schema_result.fetchone() is not None
 
         # Create schema if it doesn't exist
-        if not schema_existed:
+        if not schema_exists:
             conn.execute(sa.text("CREATE SCHEMA dbos"))
 
         # Check if dbos_migrations table exists
@@ -36,44 +31,12 @@ def ensure_dbos_schema(engine: sa.Engine) -> bool:
         )
         table_exists = table_result.fetchone() is not None
 
-        if table_exists:
-            return True
-        elif schema_existed:
-            return False
-        else:
+        if not table_exists:
             conn.execute(
                 sa.text(
                     "CREATE TABLE dbos.dbos_migrations (version BIGINT NOT NULL PRIMARY KEY)"
                 )
             )
-            return True
-
-
-def run_alembic_migrations(engine: sa.Engine) -> None:
-    """Run system database schema migrations with Alembic.
-    This is DEPRECATED in favor of DBOS-managed migrations.
-    It is retained only for backwards compatibility and
-    will be removed in the next major version."""
-    # Run a schema migration for the system database
-    migration_dir = os.path.join(
-        os.path.dirname(os.path.realpath(__file__)), "_alembic_migrations"
-    )
-    alembic_cfg = Config()
-    alembic_cfg.set_main_option("script_location", migration_dir)
-    logging.getLogger("alembic").setLevel(logging.WARNING)
-    # Alembic requires the % in URL-escaped parameters to itself be escaped to %%.
-    escaped_conn_string = re.sub(
-        r"%(?=[0-9A-Fa-f]{2})",
-        "%%",
-        engine.url.render_as_string(hide_password=False),
-    )
-    alembic_cfg.set_main_option("sqlalchemy.url", escaped_conn_string)
-    try:
-        command.upgrade(alembic_cfg, "head")
-    except Exception as e:
-        dbos_logger.warning(
-            f"Exception during system database construction. This is most likely because the system database was configured using a later version of DBOS: {e}"
-        )
 
 
 def run_dbos_migrations(engine: sa.Engine) -> None:

dbos/_serialization.py

@@ -1,8 +1,8 @@
+import base64
+import pickle
 import types
 from typing import Any, Dict, Optional, Tuple, TypedDict
 
-import jsonpickle  # type: ignore
-
 from ._logger import dbos_logger
 
 
@@ -11,54 +11,43 @@ class WorkflowInputs(TypedDict):
     kwargs: Dict[str, Any]
 
 
-def _validate_item(data: Any) -> None:
-    if isinstance(data, (types.MethodType)):
-        raise TypeError("Serialized data item should not be a class method")
-    if isinstance(data, (types.FunctionType)):
-        if jsonpickle.decode(jsonpickle.encode(data, unpicklable=True)) is None:
-            raise TypeError(
-                "Serialized function should be defined at the top level of a module"
-            )
-
-
 def serialize(data: Any) -> str:
-    """Serialize an object to a JSON string using jsonpickle."""
-    _validate_item(data)
-    encoded_data: str = jsonpickle.encode(data, unpicklable=True)
+    pickled_data: bytes = pickle.dumps(data)
+    encoded_data: str = base64.b64encode(pickled_data).decode("utf-8")
     return encoded_data
 
 
 def serialize_args(data: WorkflowInputs) -> str:
-    """Serialize args to a JSON string using jsonpickle."""
-    arg: Any
-    for arg in data["args"]:
-        _validate_item(arg)
-    for arg in data["kwargs"].values():
-        _validate_item(arg)
-    encoded_data: str = jsonpickle.encode(data, unpicklable=True)
+    """Serialize args to a base64-encoded string using pickle."""
+    pickled_data: bytes = pickle.dumps(data)
+    encoded_data: str = base64.b64encode(pickled_data).decode("utf-8")
     return encoded_data
 
 
 def serialize_exception(data: Exception) -> str:
-    """Serialize an Exception object to a JSON string using jsonpickle."""
-    encoded_data: str = jsonpickle.encode(data, unpicklable=True)
+    """Serialize an Exception object to a base64-encoded string using pickle."""
+    pickled_data: bytes = pickle.dumps(data)
+    encoded_data: str = base64.b64encode(pickled_data).decode("utf-8")
     return encoded_data
 
 
 def deserialize(serialized_data: str) -> Any:
-    """Deserialize a JSON string back to a Python object using jsonpickle."""
-    return jsonpickle.decode(serialized_data)
+    """Deserialize a base64-encoded string back to a Python object using pickle."""
+    pickled_data: bytes = base64.b64decode(serialized_data)
+    return pickle.loads(pickled_data)
 
 
 def deserialize_args(serialized_data: str) -> WorkflowInputs:
-    """Deserialize a JSON string back to a Python object list using jsonpickle."""
-    args: WorkflowInputs = jsonpickle.decode(serialized_data)
+    """Deserialize a base64-encoded string back to a Python object list using pickle."""
+    pickled_data: bytes = base64.b64decode(serialized_data)
+    args: WorkflowInputs = pickle.loads(pickled_data)
     return args
 
 
 def deserialize_exception(serialized_data: str) -> Exception:
-    """Deserialize JSON string back to a Python Exception using jsonpickle."""
-    exc: Exception = jsonpickle.decode(serialized_data)
+    """Deserialize a base64-encoded string back to a Python Exception using pickle."""
+    pickled_data: bytes = base64.b64decode(serialized_data)
+    exc: Exception = pickle.loads(pickled_data)
     return exc
 
 

dbos/_sys_db_postgres.py

@@ -5,11 +5,7 @@ import psycopg
 import sqlalchemy as sa
 from sqlalchemy.exc import DBAPIError
 
-from dbos._migration import (
-    ensure_dbos_schema,
-    run_alembic_migrations,
-    run_dbos_migrations,
-)
+from dbos._migration import ensure_dbos_schema, run_dbos_migrations
 from dbos._schemas.system_database import SystemSchema
 
 from ._logger import dbos_logger
@@ -66,10 +62,7 @@ class PostgresSystemDatabase(SystemDatabase):
                 conn.execute(sa.text(f"CREATE DATABASE {sysdb_name}"))
         engine.dispose()
 
-        using_dbos_migrations = ensure_dbos_schema(self.engine)
-        if not using_dbos_migrations:
-            # Complete the Alembic migrations, create the dbos_migrations table
-            run_alembic_migrations(self.engine)
+        ensure_dbos_schema(self.engine)
         run_dbos_migrations(self.engine)
 
     def _cleanup_connections(self) -> None:

dbos/_templates/dbos-db-starter/migrations/create_table.py.dbos

@@ -0,0 +1,34 @@
+"""
+Create the dbos_hello table using SQLAlchemy.
+"""
+
+import os
+from sqlalchemy import create_engine, MetaData, Table, Column, Integer, String
+
+
+def create_dbos_hello_table() -> None:
+    """
+    Create the dbos_hello table in the database.
+    
+    Args:
+        database_url: Database connection string. If not provided,
+                     uses DATABASE_URL environment variable.
+    """
+    database_url = os.environ.get("DBOS_DATABASE_URL", "postgresql+psycopg://postgres:dbos@localhost:5432/${default_db_name}?connect_timeout=5")
+    
+    engine = create_engine(database_url)
+    metadata = MetaData()
+    
+    dbos_hello = Table(
+        'dbos_hello',
+        metadata,
+        Column('greet_count', Integer, primary_key=True, autoincrement=True, nullable=False),
+        Column('name', String, nullable=False)
+    )
+    
+    metadata.create_all(engine)
+    engine.dispose()
+
+
+if __name__ == "__main__":
+    create_dbos_hello_table()

dbos/_tracer.py

@@ -1,13 +1,9 @@
 import os
 from typing import TYPE_CHECKING, Optional
 
-from opentelemetry import trace
-from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter
-from opentelemetry.sdk.resources import Resource
-from opentelemetry.sdk.trace import TracerProvider
-from opentelemetry.sdk.trace.export import BatchSpanProcessor, ConsoleSpanExporter
-from opentelemetry.semconv.resource import ResourceAttributes
-from opentelemetry.trace import Span
+if TYPE_CHECKING:
+    from opentelemetry.trace import Span
+    from opentelemetry.sdk.trace import TracerProvider
 
 from dbos._utils import GlobalParams
 
@@ -29,34 +25,47 @@ class DBOSTracer:
     def config(self, config: ConfigFile) -> None:
         self.otlp_attributes = config.get("telemetry", {}).get("otlp_attributes", {})  # type: ignore
         self.disable_otlp = config.get("telemetry", {}).get("disable_otlp", False)  # type: ignore
-        if not self.disable_otlp and not isinstance(
-            trace.get_tracer_provider(), TracerProvider
-        ):
-            resource = Resource(
-                attributes={
-                    ResourceAttributes.SERVICE_NAME: config["name"],
-                }
+        if not self.disable_otlp:
+            from opentelemetry import trace
+            from opentelemetry.exporter.otlp.proto.http.trace_exporter import (
+                OTLPSpanExporter,
             )
-
-            provider = TracerProvider(resource=resource)
-            if os.environ.get("DBOS__CONSOLE_TRACES", None) is not None:
-                processor = BatchSpanProcessor(ConsoleSpanExporter())
-                provider.add_span_processor(processor)
-            otlp_traces_endpoints = (
-                config.get("telemetry", {}).get("OTLPExporter", {}).get("tracesEndpoint")  # type: ignore
+            from opentelemetry.sdk.resources import Resource
+            from opentelemetry.sdk.trace import TracerProvider
+            from opentelemetry.sdk.trace.export import (
+                BatchSpanProcessor,
+                ConsoleSpanExporter,
             )
-            if otlp_traces_endpoints:
-                for e in otlp_traces_endpoints:
-                    processor = BatchSpanProcessor(OTLPSpanExporter(endpoint=e))
+            from opentelemetry.semconv.attributes.service_attributes import SERVICE_NAME
+
+            if not isinstance(trace.get_tracer_provider(), TracerProvider):
+                resource = Resource(
+                    attributes={
+                        SERVICE_NAME: config["name"],
+                    }
+                )
+
+                provider = TracerProvider(resource=resource)
+                if os.environ.get("DBOS__CONSOLE_TRACES", None) is not None:
+                    processor = BatchSpanProcessor(ConsoleSpanExporter())
                     provider.add_span_processor(processor)
-            trace.set_tracer_provider(provider)
-
-    def set_provider(self, provider: Optional[TracerProvider]) -> None:
+                otlp_traces_endpoints = (
+                    config.get("telemetry", {}).get("OTLPExporter", {}).get("tracesEndpoint")  # type: ignore
+                )
+                if otlp_traces_endpoints:
+                    for e in otlp_traces_endpoints:
+                        processor = BatchSpanProcessor(OTLPSpanExporter(endpoint=e))
+                        provider.add_span_processor(processor)
+                trace.set_tracer_provider(provider)
+
+    def set_provider(self, provider: "Optional[TracerProvider]") -> None:
         self.provider = provider
 
     def start_span(
-        self, attributes: "TracedAttributes", parent: Optional[Span] = None
-    ) -> Span:
+        self, attributes: "TracedAttributes", parent: "Optional[Span]" = None
+    ) -> "Span":
+        from opentelemetry import trace
+
         tracer = (
             self.provider.get_tracer("dbos-tracer")
             if self.provider is not None
@@ -74,11 +83,13 @@ class DBOSTracer:
             span.set_attribute(k, v)
         return span
 
-    def end_span(self, span: Span) -> None:
+    def end_span(self, span: "Span") -> None:
         span.end()
 
-    def get_current_span(self) -> Optional[Span]:
+    def get_current_span(self) -> "Optional[Span]":
         # Return the current active span if any. It might not be a DBOS span.
+        from opentelemetry import trace
+
         span = trace.get_current_span()
         if span.get_span_context().is_valid:
             return span

dbos/cli/_github_init.py

@@ -1,8 +1,9 @@
+import json
 import os
 from base64 import b64decode
-from typing import List, TypedDict
-
-import requests
+from typing import Any, List, TypedDict
+from urllib.error import HTTPError
+from urllib.request import Request, urlopen
 
 DEMO_REPO_API = "https://api.github.com/repos/dbos-inc/dbos-demo-apps"
 PY_DEMO_PATH = "python/"
@@ -34,43 +35,48 @@ class GitHubItem(TypedDict):
     size: int
 
 
-def _fetch_github(url: str) -> requests.Response:
+def _fetch_github(url: str) -> Any:
     headers = {}
     github_token = os.getenv("GITHUB_TOKEN")
     if github_token:
         headers["Authorization"] = f"Bearer {github_token}"
 
-    response = requests.get(url, headers=headers)
+    request = Request(url, headers=headers)
+
+    try:
+        with urlopen(request) as response:
+            data = response.read()
+            return json.loads(data.decode("utf-8"))
+    except HTTPError as e:
+        # Read response headers
+        rate_limit_remaining = e.headers.get("x-ratelimit-remaining")
 
-    if not response.ok:
-        if response.headers.get("x-ratelimit-remaining") == "0":
+        if rate_limit_remaining == "0":
             raise Exception(
                 "Error fetching from GitHub API: rate limit exceeded.\n"
                 "Please wait a few minutes and try again.\n"
                 "To increase the limit, you can create a personal access token and set it in the GITHUB_TOKEN environment variable.\n"
                 "Details: https://docs.github.com/en/rest/using-the-rest-api/rate-limits-for-the-rest-api"
             )
-        elif response.status_code == 401:
+        elif e.code == 401:
             raise Exception(
-                f"Error fetching content from GitHub {url}: {response.status_code} {response.reason}.\n"
+                f"Error fetching content from GitHub {url}: {e.code} {e.reason}.\n"
                 "Please ensure your GITHUB_TOKEN environment variable is set to a valid personal access token."
             )
         raise Exception(
-            f"Error fetching content from GitHub {url}: {response.status_code} {response.reason}"
+            f"Error fetching content from GitHub {url}: {e.code} {e.reason}"
         )
 
-    return response
-
 
 def _fetch_github_tree(tag: str) -> List[GitHubTreeItem]:
-    response = _fetch_github(f"{DEMO_REPO_API}/git/trees/{tag}?recursive=1")
-    tree_data: GitHubTree = response.json()
+    tree_data: GitHubTree = _fetch_github(
+        f"{DEMO_REPO_API}/git/trees/{tag}?recursive=1"
+    )
     return tree_data["tree"]
 
 
 def _fetch_github_item(url: str) -> str:
-    response = _fetch_github(url)
-    item: GitHubItem = response.json()
+    item: GitHubItem = _fetch_github(url)
     return b64decode(item["content"]).decode("utf-8")
 
 

dbos/cli/_template_init.py

@@ -2,10 +2,6 @@ import os
 import shutil
 import typing
 from os import path
-from typing import Any
-
-import tomlkit
-from rich import print
 
 from dbos._dbos_config import _app_name_to_db_name
 
@@ -46,7 +42,7 @@ def _copy_template_dir(src_dir: str, dst_dir: str, ctx: dict[str, str]) -> None:
 
             dst = path.join(dst_root, base if ext == ".dbos" else file)
             if path.exists(dst):
-                print(f"[yellow]File {dst} already exists, skipping[/yellow]")
+                print(f"File {dst} already exists, skipping")
                 continue
 
             if ext == ".dbos":
@@ -62,7 +58,7 @@ def copy_template(src_dir: str, project_name: str, config_mode: bool) -> None:
     package_name = project_name.replace("-", "_")
     default_migration_section = """database:
   migrate:
-    - alembic upgrade head
+    - python3 migrations/create_table.py
 """
     ctx = {
         "project_name": project_name,
@@ -89,18 +85,11 @@ def copy_template(src_dir: str, project_name: str, config_mode: bool) -> None:
 
 def get_project_name() -> typing.Union[str, None]:
     name = None
+
     try:
-        with open("pyproject.toml", "rb") as file:
-            pyproj = typing.cast(dict[str, Any], tomlkit.load(file))
-            name = typing.cast(str, pyproj["project"]["name"])
+        _, parent = path.split(path.abspath("."))
+        name = parent
     except:
         pass
 
-    if name == None:
-        try:
-            _, parent = path.split(path.abspath("."))
-            name = parent
-        except:
-            pass
-
     return name

dbos/cli/cli.py

@@ -1,3 +1,4 @@
+import json
 import os
 import platform
 import signal
@@ -5,14 +6,10 @@ import subprocess
 import time
 import typing
 from os import path
-from typing import Any, Optional, Tuple
+from typing import Annotated, Any, List, Optional, Tuple
 
-import jsonpickle  # type: ignore
 import sqlalchemy as sa
 import typer
-from rich import print as richprint
-from rich.prompt import IntPrompt
-from typing_extensions import Annotated, List
 
 from dbos._context import SetWorkflowID
 from dbos._debug import debug_workflow, parse_start_command
@@ -34,6 +31,11 @@ from ..cli._github_init import create_template_from_github
 from ._template_init import copy_template, get_project_name, get_templates_directory
 
 
+class DefaultEncoder(json.JSONEncoder):
+    def default(self, obj: Any) -> str:
+        return str(obj)
+
+
 def _get_db_url(
     *, system_database_url: Optional[str], application_database_url: Optional[str]
 ) -> Tuple[str, str]:
@@ -201,7 +203,7 @@ def init(
                 path.join(templates_dir, template), project_name, config_mode=config
             )
     except Exception as e:
-        richprint(f"[red]{e}[/red]")
+        print(e)
 
 
 def _resolve_project_name_and_template(
@@ -222,27 +224,21 @@ def _resolve_project_name_and_template(
         if template not in templates:
             raise Exception(f"Template {template} not found in {templates_dir}")
     else:
-        richprint("\n[bold]Available templates:[/bold]")
+        print("\nAvailable templates:")
         for idx, template_name in enumerate(templates, 1):
-            richprint(f"  {idx}. {template_name}")
+            print(f"  {idx}. {template_name}")
         while True:
             try:
-                choice = IntPrompt.ask(
-                    "\nSelect template number",
-                    show_choices=False,
-                    show_default=False,
-                )
+                choice = int(input("\nSelect template number: "))
                 if 1 <= choice <= len(templates):
                     template = templates[choice - 1]
                     break
                 else:
-                    richprint(
-                        "[red]Invalid selection. Please choose a number from the list.[/red]"
-                    )
+                    print("Invalid selection. Please choose a number from the list.")
             except (KeyboardInterrupt, EOFError):
                 raise typer.Abort()
             except ValueError:
-                richprint("[red]Please enter a valid number.[/red]")
+                print("Please enter a valid number.")
 
     if template in git_templates:
         if project_name is None:
@@ -499,7 +495,7 @@ def list(
         app_version=appversion,
         name=name,
     )
-    print(jsonpickle.encode(workflows, unpicklable=False))
+    print(json.dumps([w.__dict__ for w in workflows], cls=DefaultEncoder))
 
 
 @workflow.command(help="Retrieve the status of a workflow")
@@ -531,7 +527,7 @@ def get(
         system_database_url=system_database_url,
     )
     status = client.retrieve_workflow(workflow_id=workflow_id).get_status()
-    print(jsonpickle.encode(status, unpicklable=False))
+    print(json.dumps(status.__dict__, cls=DefaultEncoder))
 
 
 @workflow.command(help="List the steps of a workflow")
@@ -562,12 +558,8 @@ def steps(
         application_database_url=application_database_url,
         system_database_url=system_database_url,
     )
-    print(
-        jsonpickle.encode(
-            client.list_workflow_steps(workflow_id=workflow_id),
-            unpicklable=False,
-        )
-    )
+    steps = client.list_workflow_steps(workflow_id=workflow_id)
+    print(json.dumps(steps, cls=DefaultEncoder))
 
 
 @workflow.command(
@@ -665,7 +657,7 @@ def restart(
         system_database_url=system_database_url,
     )
     status = client.fork_workflow(workflow_id=workflow_id, start_step=1).get_status()
-    print(jsonpickle.encode(status, unpicklable=False))
+    print(json.dumps(status.__dict__, cls=DefaultEncoder))
 
 
 @workflow.command(
@@ -736,7 +728,7 @@ def fork(
             start_step=step,
             application_version=application_version,
         ).get_status()
-    print(jsonpickle.encode(status, unpicklable=False))
+    print(json.dumps(status.__dict__, cls=DefaultEncoder))
 
 
 @queue.command(name="list", help="List enqueued functions for your application")
@@ -836,7 +828,7 @@ def list_queue(
         status=status,
         name=name,
     )
-    print(jsonpickle.encode(workflows, unpicklable=False))
+    print(json.dumps([w.__dict__ for w in workflows], cls=DefaultEncoder))
 
 
 if __name__ == "__main__":