dbos 0.17.0a4__py3-none-any.whl → 0.18.0a1__py3-none-any.whl

dbos/_cloudutils/authentication.py

@@ -0,0 +1,163 @@
+import os
+import time
+from dataclasses import dataclass
+from typing import Any, Dict, Optional
+
+import jwt
+import requests
+from cryptography.hazmat.primitives.asymmetric.rsa import RSAPublicKey
+from cryptography.x509 import load_pem_x509_certificate
+from rich import print
+
+from .._logger import dbos_logger
+
+# Constants
+DBOS_CLOUD_HOST = os.getenv("DBOS_DOMAIN", "cloud.dbos.dev")
+PRODUCTION_ENVIRONMENT = DBOS_CLOUD_HOST == "cloud.dbos.dev"
+AUTH0_DOMAIN = "login.dbos.dev" if PRODUCTION_ENVIRONMENT else "dbos-inc.us.auth0.com"
+DBOS_CLIENT_ID = (
+    "6p7Sjxf13cyLMkdwn14MxlH7JdhILled"
+    if PRODUCTION_ENVIRONMENT
+    else "G38fLmVErczEo9ioCFjVIHea6yd0qMZu"
+)
+DBOS_CLOUD_IDENTIFIER = "dbos-cloud-api"
+
+
+@dataclass
+class DeviceCodeResponse:
+    device_code: str
+    user_code: str
+    verification_uri: str
+    verification_uri_complete: str
+    expires_in: int
+    interval: int
+
+    @classmethod
+    def from_dict(cls, data: Dict[str, Any]) -> "DeviceCodeResponse":
+        return cls(
+            device_code=data["device_code"],
+            user_code=data["user_code"],
+            verification_uri=data["verification_uri"],
+            verification_uri_complete=data["verification_uri_complete"],
+            expires_in=data["expires_in"],
+            interval=data["interval"],
+        )
+
+
+@dataclass
+class TokenResponse:
+    access_token: str
+    token_type: str
+    expires_in: int
+    refresh_token: Optional[str] = None
+
+    @classmethod
+    def from_dict(cls, data: Dict[str, Any]) -> "TokenResponse":
+        return cls(
+            access_token=data["access_token"],
+            token_type=data["token_type"],
+            expires_in=data["expires_in"],
+            refresh_token=data.get("refresh_token"),
+        )
+
+
+@dataclass
+class AuthenticationResponse:
+    token: str
+    refresh_token: Optional[str] = None
+
+
+class JWKSClient:
+    def __init__(self, jwks_uri: str):
+        self.jwks_uri = jwks_uri
+
+    def get_signing_key(self, kid: str) -> RSAPublicKey:
+        response = requests.get(self.jwks_uri)
+        jwks = response.json()
+        for key in jwks["keys"]:
+            if key["kid"] == kid:
+                cert_text = f"-----BEGIN CERTIFICATE-----\n{key['x5c'][0]}\n-----END CERTIFICATE-----"
+                cert = load_pem_x509_certificate(cert_text.encode())
+                return cert.public_key()  # type: ignore
+        raise Exception(f"Unable to find signing key with kid: {kid}")
+
+
+def verify_token(token: str) -> None:
+    header = jwt.get_unverified_header(token)
+
+    if not header.get("kid"):
+        raise ValueError("Invalid token: No 'kid' in header")
+
+    client = JWKSClient(f"https://{AUTH0_DOMAIN}/.well-known/jwks.json")
+    signing_key = client.get_signing_key(header["kid"])
+    jwt.decode(
+        token,
+        signing_key,
+        algorithms=["RS256"],
+        audience=DBOS_CLOUD_IDENTIFIER,
+        options={
+            "verify_iat": False,
+            "clock_tolerance": 60,
+        },
+    )
+
+
+def authenticate(get_refresh_token: bool = False) -> Optional[AuthenticationResponse]:
+    print(
+        "[bold blue]Please authenticate with DBOS Cloud to access a Postgres database[/bold blue]"
+    )
+
+    # Get device code
+    device_code_data = {
+        "client_id": DBOS_CLIENT_ID,
+        "scope": "offline_access" if get_refresh_token else "sub",
+        "audience": DBOS_CLOUD_IDENTIFIER,
+    }
+
+    try:
+        response = requests.post(
+            f"https://{AUTH0_DOMAIN}/oauth/device/code",
+            data=device_code_data,
+            headers={"content-type": "application/x-www-form-urlencoded"},
+        )
+        device_code_response = DeviceCodeResponse.from_dict(response.json())
+    except Exception as e:
+        dbos_logger.error(f"Failed to log in: {str(e)}")
+        return None
+
+    login_url = device_code_response.verification_uri_complete
+    print(f"[bold blue]Login URL:[/bold blue] {login_url}")
+
+    # Poll for token
+    token_data = {
+        "grant_type": "urn:ietf:params:oauth:grant-type:device_code",
+        "device_code": device_code_response.device_code,
+        "client_id": DBOS_CLIENT_ID,
+    }
+
+    elapsed_time_sec = 0
+    token_response = None
+
+    while elapsed_time_sec < device_code_response.expires_in:
+        try:
+            time.sleep(device_code_response.interval)
+            elapsed_time_sec += device_code_response.interval
+
+            response = requests.post(
+                f"https://{AUTH0_DOMAIN}/oauth/token",
+                data=token_data,
+                headers={"content-type": "application/x-www-form-urlencoded"},
+            )
+            if response.status_code == 200:
+                token_response = TokenResponse.from_dict(response.json())
+                break
+        except Exception:
+            dbos_logger.info("Waiting for login...")
+
+    if not token_response:
+        return None
+
+    verify_token(token_response.access_token)
+    return AuthenticationResponse(
+        token=token_response.access_token, refresh_token=token_response.refresh_token
+    )

dbos/_cloudutils/cloudutils.py

@@ -0,0 +1,252 @@
+import json
+import os
+import re
+import time
+from dataclasses import dataclass
+from enum import Enum
+from typing import Any, Optional, Union
+
+import jwt
+import requests
+import typer
+from rich import print
+
+from dbos._error import DBOSInitializationError
+
+from .._logger import dbos_logger
+from .authentication import authenticate
+
+
+# Must be the same as in TS
+@dataclass
+class DBOSCloudCredentials:
+    token: str
+    userName: str
+    organization: str
+    refreshToken: Optional[str] = None
+
+
+@dataclass
+class UserProfile:
+    Name: str
+    Email: str
+    Organization: str
+    SubscriptionPlan: str
+
+
+class AppLanguages(Enum):
+    Node = "node"
+    Python = "python"
+
+
+dbos_config_file_path = "dbos-config.yaml"
+DBOS_CLOUD_HOST = os.getenv("DBOS_DOMAIN", "cloud.dbos.dev")
+dbos_env_path = ".dbos"
+
+
+def is_token_expired(token: str) -> bool:
+    try:
+        decoded = jwt.decode(token, options={"verify_signature": False})
+        exp: int = decoded.get("exp")
+        if not exp:
+            return False
+        return time.time() >= exp
+    except Exception:
+        return True
+
+
+def credentials_exist() -> bool:
+    return os.path.exists(os.path.join(dbos_env_path, "credentials"))
+
+
+def delete_credentials() -> None:
+    credentials_path = os.path.join(dbos_env_path, "credentials")
+    if os.path.exists(credentials_path):
+        os.unlink(credentials_path)
+
+
+def write_credentials(credentials: DBOSCloudCredentials) -> None:
+    os.makedirs(dbos_env_path, exist_ok=True)
+    with open(os.path.join(dbos_env_path, "credentials"), "w", encoding="utf-8") as f:
+        json.dump(credentials.__dict__, f)
+
+
+def check_read_file(path: str, encoding: str = "utf-8") -> Union[str, bytes]:
+    # Check if file exists and is a file
+    if not os.path.exists(path):
+        raise FileNotFoundError(f"File {path} does not exist")
+    if not os.path.isfile(path):
+        raise IsADirectoryError(f"Path {path} is not a file")
+
+    # Read file content
+    with open(path, encoding=encoding) as f:
+        return f.read()
+
+
+@dataclass
+class CloudAPIErrorResponse:
+    message: str
+    status_code: int
+    request_id: str
+    detailed_error: Optional[str] = None
+
+
+def is_cloud_api_error_response(obj: Any) -> bool:
+    return (
+        isinstance(obj, dict)
+        and "message" in obj
+        and isinstance(obj["message"], str)
+        and "statusCode" in obj
+        and isinstance(obj["statusCode"], int)
+        and "requestID" in obj
+        and isinstance(obj["requestID"], str)
+    )
+
+
+def handle_api_errors(label: str, e: requests.exceptions.RequestException) -> None:
+    if hasattr(e, "response") and e.response is not None:
+        resp = e.response.json()
+        if is_cloud_api_error_response(resp):
+            message = f"[{resp['requestID']}] {label}: {resp['message']}."
+            dbos_logger.error(message)
+            raise DBOSInitializationError(message)
+
+
+def is_valid_username(value: str) -> Union[bool, str]:
+    if len(value) < 3 or len(value) > 30:
+        return "Username must be 3~30 characters long"
+    if not re.match("^[a-z0-9_]+$", value):
+        return "Username must contain only lowercase letters, numbers, and underscores."
+    return True
+
+
+def check_user_profile(credentials: DBOSCloudCredentials) -> bool:
+    bearer_token = f"Bearer {credentials.token}"
+    try:
+        response = requests.get(
+            f"https://{DBOS_CLOUD_HOST}/v1alpha1/user/profile",
+            headers={
+                "Content-Type": "application/json",
+                "Authorization": bearer_token,
+            },
+        )
+        response.raise_for_status()
+        profile = UserProfile(**response.json())
+        credentials.userName = profile.Name
+        credentials.organization = profile.Organization
+        return True
+    except requests.exceptions.RequestException as e:
+        error_label = "Failed to login"
+        if hasattr(e, "response") and e.response is not None:
+            resp = e.response.json()
+            if is_cloud_api_error_response(resp):
+                if "user not found in DBOS Cloud" not in resp["message"]:
+                    handle_api_errors(error_label, e)
+                    exit(1)
+        else:
+            dbos_logger.error(f"{error_label}: {str(e)}")
+            exit(1)
+    return False
+
+
+def register_user(credentials: DBOSCloudCredentials) -> None:
+    print("Please register for DBOS Cloud")
+
+    user_name = None
+    while not user_name:
+        user_name = typer.prompt("Choose your username")
+        validation_result = is_valid_username(user_name)
+        if validation_result is not True:
+            print(f"[red]Invalid username: {validation_result}[/red]")
+            user_name = None
+            continue
+
+    bearer_token = f"Bearer {credentials.token}"
+    try:
+        # Register user
+        response = requests.put(
+            f"https://{DBOS_CLOUD_HOST}/v1alpha1/user",
+            json={
+                "name": user_name,
+            },
+            headers={
+                "Content-Type": "application/json",
+                "Authorization": bearer_token,
+            },
+        )
+        response.raise_for_status()
+
+        # Get user profile
+        response = requests.get(
+            f"https://{DBOS_CLOUD_HOST}/v1alpha1/user/profile",
+            headers={
+                "Content-Type": "application/json",
+                "Authorization": bearer_token,
+            },
+        )
+        response.raise_for_status()
+        profile = UserProfile(**response.json())
+        credentials.userName = profile.Name
+        credentials.organization = profile.Organization
+        print(f"[green]Successfully registered as {credentials.userName}[/green]")
+
+    except requests.exceptions.RequestException as e:
+        error_label = f"Failed to register user {user_name}"
+        if hasattr(e, "response") and e.response is not None:
+            handle_api_errors(error_label, e)
+        else:
+            dbos_logger.error(f"{error_label}: {str(e)}")
+        exit(1)
+
+
+def check_credentials() -> DBOSCloudCredentials:
+    empty_credentials = DBOSCloudCredentials(token="", userName="", organization="")
+
+    if not credentials_exist():
+        return empty_credentials
+
+    try:
+        with open(os.path.join(dbos_env_path, "credentials"), "r") as f:
+            cred_data = json.load(f)
+            credentials = DBOSCloudCredentials(**cred_data)
+
+        # Trim trailing /r /n
+        credentials.token = credentials.token.strip()
+
+        if is_token_expired(credentials.token):
+            print("Credentials expired. Logging in again...")
+            delete_credentials()
+            return empty_credentials
+
+        return credentials
+    except Exception as e:
+        dbos_logger.error(f"Error loading credentials: {str(e)}")
+        return empty_credentials
+
+
+def get_cloud_credentials() -> DBOSCloudCredentials:
+    # Check if credentials exist and are not expired
+    credentials = check_credentials()
+
+    # Log in the user
+    if not credentials.token:
+        auth_response = authenticate()
+        if auth_response is None:
+            dbos_logger.error("Failed to login. Exiting...")
+            exit(1)
+        credentials.token = auth_response.token
+        credentials.refreshToken = auth_response.refresh_token
+        write_credentials(credentials)
+
+    # Check if the user exists in DBOS Cloud
+    user_exists = check_user_profile(credentials)
+    if user_exists:
+        write_credentials(credentials)
+        print(f"[green]Successfully logged in as {credentials.userName}[/green]")
+        return credentials
+
+    # User doesn't exist, register the user in DBOS Cloud
+    register_user(credentials)
+    write_credentials(credentials)
+
+    return credentials

dbos/_cloudutils/databases.py

@@ -0,0 +1,237 @@
+import base64
+import random
+import time
+from dataclasses import dataclass
+from typing import Any, List, Optional
+
+import requests
+from rich import print
+
+from dbos._cloudutils.cloudutils import (
+    DBOS_CLOUD_HOST,
+    DBOSCloudCredentials,
+    handle_api_errors,
+    is_cloud_api_error_response,
+)
+from dbos._error import DBOSInitializationError
+
+from .._logger import dbos_logger
+
+
+@dataclass
+class UserDBCredentials:
+    RoleName: str
+    Password: str
+
+
+@dataclass
+class UserDBInstance:
+    PostgresInstanceName: str = ""
+    Status: str = ""
+    HostName: str = ""
+    Port: int = 0
+    DatabaseUsername: str = ""
+    IsLinked: bool = False
+    SupabaseReference: Optional[str] = None
+
+    def __init__(self, **kwargs: Any) -> None:
+        self.PostgresInstanceName = kwargs.get("PostgresInstanceName", "")
+        self.Status = kwargs.get("Status", "")
+        self.HostName = kwargs.get("HostName", "")
+        self.Port = kwargs.get("Port", 0)
+        self.DatabaseUsername = kwargs.get("DatabaseUsername", "")
+        self.IsLinked = kwargs.get("IsLinked", False)
+        self.SupabaseReference = kwargs.get("SupabaseReference", None)
+
+
+def get_user_db_info(credentials: DBOSCloudCredentials, db_name: str) -> UserDBInstance:
+    bearer_token = f"Bearer {credentials.token}"
+
+    try:
+        response = requests.get(
+            f"https://{DBOS_CLOUD_HOST}/v1alpha1/{credentials.organization}/databases/userdb/info/{db_name}",
+            headers={
+                "Content-Type": "application/json",
+                "Authorization": bearer_token,
+            },
+        )
+        response.raise_for_status()
+        data = response.json()
+        return UserDBInstance(**data)
+    except requests.exceptions.RequestException as e:
+        error_label = f"Failed to get status of database {db_name}"
+        if hasattr(e, "response") and e.response is not None:
+            resp = e.response.json()
+            if is_cloud_api_error_response(resp):
+                handle_api_errors(error_label, e)
+        else:
+            dbos_logger.error(f"{error_label}: {str(e)}")
+        raise DBOSInitializationError(f"{error_label}: {str(e)}")
+
+
+def get_user_db_credentials(
+    credentials: DBOSCloudCredentials, db_name: str
+) -> UserDBCredentials:
+    bearer_token = f"Bearer {credentials.token}"
+    try:
+        response = requests.get(
+            f"https://{DBOS_CLOUD_HOST}/v1alpha1/{credentials.organization}/databases/userdb/{db_name}/credentials",
+            headers={
+                "Content-Type": "application/json",
+                "Authorization": bearer_token,
+            },
+        )
+        response.raise_for_status()
+        data = response.json()
+        return UserDBCredentials(**data)
+    except requests.exceptions.RequestException as e:
+        error_label = f"Failed to get credentials for database {db_name}"
+        if hasattr(e, "response") and e.response is not None:
+            resp = e.response.json()
+            if is_cloud_api_error_response(resp):
+                handle_api_errors(error_label, e)
+        else:
+            dbos_logger.error(f"{error_label}: {str(e)}")
+        raise DBOSInitializationError(f"{error_label}: {str(e)}")
+
+
+def create_user_role(credentials: DBOSCloudCredentials, db_name: str) -> None:
+    bearer_token = f"Bearer {credentials.token}"
+    try:
+        response = requests.post(
+            f"https://{DBOS_CLOUD_HOST}/v1alpha1/{credentials.organization}/databases/userdb/{db_name}/createuserdbrole",
+            headers={
+                "Content-Type": "application/json",
+                "Authorization": bearer_token,
+            },
+        )
+        response.raise_for_status()
+    except requests.exceptions.RequestException as e:
+        error_label = f"Failed to create a user role for database {db_name}"
+        if hasattr(e, "response") and e.response is not None:
+            resp = e.response.json()
+            if is_cloud_api_error_response(resp):
+                handle_api_errors(error_label, e)
+        else:
+            dbos_logger.error(f"{error_label}: {str(e)}")
+        raise DBOSInitializationError(f"{error_label}: {str(e)}")
+
+
+def create_user_db(
+    credentials: DBOSCloudCredentials,
+    db_name: str,
+    app_db_username: str,
+    app_db_password: str,
+) -> int:
+    bearer_token = f"Bearer {credentials.token}"
+
+    try:
+        response = requests.post(
+            f"https://{DBOS_CLOUD_HOST}/v1alpha1/{credentials.organization}/databases/userdb",
+            json={
+                "Name": db_name,
+                "AdminName": app_db_username,
+                "AdminPassword": app_db_password,
+            },
+            headers={
+                "Content-Type": "application/json",
+                "Authorization": bearer_token,
+            },
+        )
+        response.raise_for_status()
+
+        print(f"Successfully started provisioning {db_name}")
+
+        status = ""
+        while status not in ["available", "backing-up"]:
+            if status == "":
+                time.sleep(5)  # First time sleep 5 sec
+            else:
+                time.sleep(30)  # Otherwise, sleep 30 sec
+
+            user_db_info = get_user_db_info(credentials, db_name)
+            status = user_db_info.Status
+            print(
+                f"[bold blue]Waiting for cloud database to finish provisioning. Status:[/bold blue] [yellow]{status}[/yellow]"
+            )
+
+        print("[green]Database successfully provisioned![/green]")
+        return 0
+
+    except requests.exceptions.RequestException as e:
+        error_label = f"Failed to create database {db_name}"
+        if hasattr(e, "response") and e.response is not None:
+            resp = e.response.json()
+            if is_cloud_api_error_response(resp):
+                handle_api_errors(error_label, e)
+        else:
+            dbos_logger.error(f"{error_label}: {str(e)}")
+        return 1
+
+
+def choose_database(credentials: DBOSCloudCredentials) -> Optional[UserDBInstance]:
+    # List existing database instances
+    user_dbs: List[UserDBInstance] = []
+    bearer_token = f"Bearer {credentials.token}"
+
+    try:
+        response = requests.get(
+            f"https://{DBOS_CLOUD_HOST}/v1alpha1/{credentials.organization}/databases",
+            headers={
+                "Content-Type": "application/json",
+                "Authorization": bearer_token,
+            },
+        )
+        response.raise_for_status()
+        data = response.json()
+        user_dbs = [UserDBInstance(**db) for db in data]
+
+    except requests.exceptions.RequestException as e:
+        error_label = "Failed to list databases"
+        if hasattr(e, "response") and e.response is not None:
+            resp = e.response.json()
+            if is_cloud_api_error_response(resp):
+                handle_api_errors(error_label, e)
+        else:
+            dbos_logger.error(f"{error_label}: {str(e)}")
+        return None
+
+    if not user_dbs:
+        # If not, prompt the user to provision one
+        print("Provisioning a cloud Postgres database server")
+        user_db_name = f"{credentials.userName}-db-server"
+
+        # Use a default user name and auto generated password
+        app_db_username = "dbos_user"
+        app_db_password = base64.b64encode(str(random.random()).encode()).decode()
+        res = create_user_db(
+            credentials, user_db_name, app_db_username, app_db_password
+        )
+        if res != 0:
+            return None
+    elif len(user_dbs) > 1:
+        # If there is more than one database instance, prompt the user to select one
+        choices = [db.PostgresInstanceName for db in user_dbs]
+        print("Choose a database instance for this app:")
+        for i, entry in enumerate(choices, 1):
+            print(f"{i}. {entry}")
+        while True:
+            try:
+                choice = int(input("Enter number: ")) - 1
+                if 0 <= choice < len(choices):
+                    user_db_name = choices[choice]
+                    break
+            except ValueError:
+                continue
+            print("Invalid choice, please try again")
+    else:
+        # Use the only available database server
+        user_db_name = user_dbs[0].PostgresInstanceName
+        print(f"[green]Using database instance:[/green] {user_db_name}")
+
+    info = get_user_db_info(credentials, user_db_name)
+
+    if not info.IsLinked:
+        create_user_role(credentials, user_db_name)
+
+    return info

dbos/_db_wizard.py

@@ -0,0 +1,167 @@
+import time
+from typing import TYPE_CHECKING, Optional
+
+import docker  # type: ignore
+import typer
+import yaml
+from rich import print
+from sqlalchemy import URL, create_engine, text
+
+if TYPE_CHECKING:
+    from ._dbos_config import ConfigFile
+
+from ._cloudutils.cloudutils import get_cloud_credentials
+from ._cloudutils.databases import choose_database, get_user_db_credentials
+from ._error import DBOSInitializationError
+from ._logger import dbos_logger
+
+
+def db_connect(config: "ConfigFile", config_file_path: str) -> "ConfigFile":
+    # 1. Check the connectivity to the database. Return if successful. If cannot connect, continue to the following steps.
+    db_connection_error = _check_db_connectivity(config)
+    if db_connection_error is None:
+        return config
+
+    # 2. If the error is due to password authentication or the configuration is non-default, surface the error and exit.
+    if "password authentication failed" in str(db_connection_error) or "28P01" in str(
+        db_connection_error
+    ):
+        raise DBOSInitializationError(
+            f"Could not connect to Postgres: password authentication failed: {db_connection_error}"
+        )
+    db_config = config["database"]
+    if (
+        db_config["hostname"] != "localhost"
+        or db_config["port"] != 5432
+        or db_config["username"] != "postgres"
+    ):
+        raise DBOSInitializationError(
+            f"Could not connect to the database. Exception: {db_connection_error}"
+        )
+    print("[yellow]Postgres not detected locally[/yellow]")
+
+    # 3. If the database config is the default one, check if the user has Docker properly installed.
+    print("Attempting to start Postgres via Docker")
+    has_docker = _check_docker_installed()
+
+    # 4. If Docker is installed, prompt the user to start a local Docker based Postgres, and then set the PGPASSWORD to 'dbos' and try to connect to the database.
+    docker_started = False
+    if has_docker:
+        docker_started = _start_docker_postgres(config)
+    else:
+        print("[yellow]Docker not detected locally[/yellow]")
+
+    # 5. If no Docker, then prompt the user to log in to DBOS Cloud and provision a DB there. Wait for the remote DB to be ready, and then create a copy of the original config file, and then load the remote connection string to the local config file.
+    if not docker_started:
+        print("Attempting to connect to Postgres via DBOS Cloud")
+        cred = get_cloud_credentials()
+        db = choose_database(cred)
+        if db is None:
+            raise DBOSInitializationError("Error connecting to cloud database")
+        config["database"]["hostname"] = db.HostName
+        config["database"]["port"] = db.Port
+        if db.SupabaseReference is not None:
+            config["database"]["username"] = f"postgres.{db.SupabaseReference}"
+            supabase_password = typer.prompt(
+                "Enter your Supabase database password", hide_input=True
+            )
+            config["database"]["password"] = supabase_password
+        else:
+            config["database"]["username"] = db.DatabaseUsername
+            db_credentials = get_user_db_credentials(cred, db.PostgresInstanceName)
+            config["database"]["password"] = db_credentials.Password
+        config["database"]["local_suffix"] = True
+
+        # Verify these new credentials work
+        db_connection_error = _check_db_connectivity(config)
+        if db_connection_error is not None:
+            raise DBOSInitializationError(
+                f"Could not connect to the database. Exception: {db_connection_error}"
+            )
+
+    # 6. Save the config to the config file and return the updated config.
+    # TODO: make the config file prettier
+    with open(config_file_path, "w") as file:
+        file.write(yaml.dump(config))
+
+    return config
+
+
+def _start_docker_postgres(config: "ConfigFile") -> bool:
+    print("Starting a Postgres Docker container...")
+    config["database"]["password"] = "dbos"
+    client = docker.from_env()
+    pg_data = "/var/lib/postgresql/data"
+    container_name = "dbos-db"
+    client.containers.run(
+        image="pgvector/pgvector:pg16",
+        detach=True,
+        environment={
+            "POSTGRES_PASSWORD": config["database"]["password"],
+            "PGDATA": pg_data,
+        },
+        volumes={pg_data: {"bind": pg_data, "mode": "rw"}},
+        ports={"5432/tcp": config["database"]["port"]},
+        name=container_name,
+        remove=True,
+    )
+
+    container = client.containers.get(container_name)
+    attempts = 30
+    while attempts > 0:
+        if attempts % 5 == 0:
+            print("Waiting for Postgres Docker container to start...")
+        try:
+            res = container.exec_run("psql -U postgres -c 'SELECT 1;'")
+            if res.exit_code != 0:
+                attempts -= 1
+                time.sleep(1)
+                continue
+            print("[green]Postgres Docker container started successfully![/green]")
+            break
+        except Exception as e:
+            attempts -= 1
+            time.sleep(1)
+
+    if attempts == 0:
+        print("[yellow]Failed to start Postgres Docker container.[/yellow]")
+        return False
+
+    return True
+
+
+def _check_docker_installed() -> bool:
+    # Check if Docker is installed
+    try:
+        client = docker.from_env()
+        client.ping()
+    except Exception:
+        return False
+    return True
+
+
+def _check_db_connectivity(config: "ConfigFile") -> Optional[Exception]:
+    postgres_db_url = URL.create(
+        "postgresql+psycopg",
+        username=config["database"]["username"],
+        password=config["database"]["password"],
+        host=config["database"]["hostname"],
+        port=config["database"]["port"],
+        database="postgres",
+        query={"connect_timeout": "2"},
+    )
+    postgres_db_engine = create_engine(postgres_db_url)
+    try:
+        with postgres_db_engine.connect() as conn:
+            val = conn.execute(text("SELECT 1")).scalar()
+            if val != 1:
+                dbos_logger.error(
+                    f"Unexpected value returned from database: expected 1, received {val}"
+                )
+                return Exception()
+    except Exception as e:
+        return e
+    finally:
+        postgres_db_engine.dispose()
+
+    return None

dbos/_dbos.py

@@ -76,15 +76,14 @@ else:
 from ._admin_server import AdminServer
 from ._app_db import ApplicationDatabase
 from ._context import (
-    DBOSContext,
     EnterDBOSStep,
     TracedAttributes,
     assert_current_dbos_context,
     get_local_dbos_context,
 )
-from ._dbos_config import ConfigFile, _set_env_vars, load_config
+from ._dbos_config import ConfigFile, load_config, set_env_vars
 from ._error import DBOSException, DBOSNonExistentWorkflowError
-from ._logger import add_otlp_to_all_loggers, config_logger, dbos_logger, init_logger
+from ._logger import add_otlp_to_all_loggers, dbos_logger, init_logger
 from ._sys_db import SystemDatabase
 
 # Most DBOS functions are just any callable F, so decorators / wrappers work on F
@@ -264,11 +263,9 @@ class DBOS:
             return
 
         self._initialized: bool = True
-        init_logger()
         if config is None:
             config = load_config()
-        config_logger(config)
-        _set_env_vars(config)
+        set_env_vars(config)
         dbos_tracer.config(config)
         dbos_logger.info("Initializing DBOS")
         self.config: ConfigFile = config

dbos/_dbos_config.py

@@ -8,8 +8,9 @@ import yaml
 from jsonschema import ValidationError, validate
 from sqlalchemy import URL
 
+from ._db_wizard import db_connect
 from ._error import DBOSInitializationError
-from ._logger import dbos_logger
+from ._logger import config_logger, dbos_logger, init_logger
 
 
 class RuntimeConfig(TypedDict, total=False):
@@ -132,6 +133,8 @@ def load_config(config_file_path: str = "dbos-config.yaml") -> ConfigFile:
 
     """
 
+    init_logger()
+
     with open(config_file_path, "r") as file:
         content = file.read()
         substituted_content = _substitute_env_vars(content)
@@ -176,6 +179,11 @@ def load_config(config_file_path: str = "dbos-config.yaml") -> ConfigFile:
     if "app_db_name" not in data["database"]:
         data["database"]["app_db_name"] = _app_name_to_db_name(data["name"])
 
+    config_logger(data)
+
+    # Check the connectivity to the database and make sure it's properly configured
+    data = db_connect(data, config_file_path)
+
     if "local_suffix" in data["database"] and data["database"]["local_suffix"]:
         data["database"]["app_db_name"] = f"{data['database']['app_db_name']}_local"
 
@@ -196,7 +204,7 @@ def _app_name_to_db_name(app_name: str) -> str:
     return name if not name[0].isdigit() else f"_{name}"
 
 
-def _set_env_vars(config: ConfigFile) -> None:
+def set_env_vars(config: ConfigFile) -> None:
     for env, value in config.get("env", {}).items():
         if value is not None:
             os.environ[env] = str(value)

dbos/dbos-config.schema.json

@@ -36,7 +36,7 @@
             }
           },
           "password": {
-            "type": "string",
+            "type": ["string", "null"],
             "description": "The password to use when connecting to the application database. Developers are strongly encouraged to use environment variable substitution to avoid storing secrets in source."
           },
           "connectionTimeoutMillis": {

{dbos-0.17.0a4.dist-info → dbos-0.18.0a1.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: dbos
-Version: 0.17.0a4
+Version: 0.18.0a1
 Summary: Ultra-lightweight durable execution in Python
 Author-Email: "DBOS, Inc." <contact@dbos.dev>
 License: MIT
@@ -19,6 +19,10 @@ Requires-Dist: fastapi[standard]>=0.115.2
 Requires-Dist: tomlkit>=0.13.2
 Requires-Dist: psycopg[binary]>=3.1
 Requires-Dist: fastapi-cli==0.0.5
+Requires-Dist: docker>=7.1.0
+Requires-Dist: cryptography>=43.0.3
+Requires-Dist: rich>=13.9.4
+Requires-Dist: pyjwt>=2.10.1
 Description-Content-Type: text/markdown
 
 

{dbos-0.17.0a4.dist-info → dbos-0.18.0a1.dist-info}/RECORD

@@ -1,16 +1,20 @@
-dbos-0.17.0a4.dist-info/METADATA,sha256=c8YqVga1fsHCq5TQSBnCBfxjkQj1SP1z9kb7E-lSER0,5022
-dbos-0.17.0a4.dist-info/WHEEL,sha256=thaaA2w1JzcGC48WYufAs8nrYZjJm8LqNfnXFOFyCC4,90
-dbos-0.17.0a4.dist-info/entry_points.txt,sha256=z6GcVANQV7Uw_82H9Ob2axJX6V3imftyZsljdh-M1HU,54
-dbos-0.17.0a4.dist-info/licenses/LICENSE,sha256=VGZit_a5-kdw9WT6fY5jxAWVwGQzgLFyPWrcVVUhVNU,1067
+dbos-0.18.0a1.dist-info/METADATA,sha256=8i67YGwId1DCG0l6uylCoyiSjCz2881pFqnDxh-2JbI,5144
+dbos-0.18.0a1.dist-info/WHEEL,sha256=thaaA2w1JzcGC48WYufAs8nrYZjJm8LqNfnXFOFyCC4,90
+dbos-0.18.0a1.dist-info/entry_points.txt,sha256=z6GcVANQV7Uw_82H9Ob2axJX6V3imftyZsljdh-M1HU,54
+dbos-0.18.0a1.dist-info/licenses/LICENSE,sha256=VGZit_a5-kdw9WT6fY5jxAWVwGQzgLFyPWrcVVUhVNU,1067
 dbos/__init__.py,sha256=CxRHBHEthPL4PZoLbZhp3rdm44-KkRTT2-7DkK9d4QQ,724
 dbos/_admin_server.py,sha256=DOgzVp9kmwiebQqmJB1LcrZnGTxSMbZiGXdenc1wZDg,3163
 dbos/_app_db.py,sha256=_tv2vmPjjiaikwgxH3mqxgJ4nUUcG2-0uMXKWCqVu1c,5509
 dbos/_classproperty.py,sha256=f0X-_BySzn3yFDRKB2JpCbLYQ9tLwt1XftfshvY7CBs,626
+dbos/_cloudutils/authentication.py,sha256=V0fCWQN9stCkhbuuxgPTGpvuQcDqfU3KAxPAh01vKW4,5007
+dbos/_cloudutils/cloudutils.py,sha256=5e3CW1deSW-dI5G3QN0XbiVsBhyqT8wu7fuV2f8wtGU,7688
+dbos/_cloudutils/databases.py,sha256=x4187Djsyoa-QaG3Kog8JT2_GERsnqa93LIVanmVUmg,8393
 dbos/_context.py,sha256=KV3fd3-Rv6EWrYDUdHARxltSlNZGNtQtNSqeQ-gkXE8,18049
 dbos/_core.py,sha256=NWJFQX5bECBvKlYH9pVmNJgmqFGYPnkHnOGjOlOQ3Ag,33504
 dbos/_croniter.py,sha256=hbhgfsHBqclUS8VeLnJ9PSE9Z54z6mi4nnrr1aUXn0k,47561
-dbos/_dbos.py,sha256=riYx_dkYFzqeVDYpmcA5ABdAYQFhwyDi4AwxIihDNKA,34809
-dbos/_dbos_config.py,sha256=f37eccN3JpCA32kRdQ4UsERjhYGcdLWv-N21ijnDZmY,6406
+dbos/_db_wizard.py,sha256=o0OdGEsOjtL-1G6HgbXSigBYSaEN25tfSLrgBw0uh6o,6273
+dbos/_dbos.py,sha256=z12yGw2QHx7BLdxvoI2zJiwDTSes1r48E7qZ7uOn0mw,34723
+dbos/_dbos_config.py,sha256=Hs9L-PJhxeGa2R3nDX75ZFOGLRxA3AiXVf7UoemN9lM,6643
 dbos/_error.py,sha256=UETk8CoZL-TO2Utn1-E7OSWelhShWmKM-fOlODMR9PE,3893
 dbos/_fastapi.py,sha256=iyefCZq-ZDKRUjN_rgYQmFmyvWf4gPrSlC6CLbfq4a8,3419
 dbos/_flask.py,sha256=z1cijbTi5Dpq6kqikPCx1LcR2YHHv2oc41NehOWjw74,2431
@@ -49,7 +53,7 @@ dbos/_templates/hello/migrations/versions/2024_07_31_180642_init.py,sha256=U5thF
 dbos/_templates/hello/start_postgres_docker.py,sha256=lQVLlYO5YkhGPEgPqwGc7Y8uDKse9HsWv5fynJEFJHM,1681
 dbos/_tracer.py,sha256=rvBY1RQU6DO7rL7EnaJJxGcmd4tP_PpGqUEE6imZnhY,2518
 dbos/cli.py,sha256=em1uAxrp5yyg53V7ZpmHFtqD6OJp2cMJkG9vGJPoFTA,10904
-dbos/dbos-config.schema.json,sha256=tS7x-bdFbFvpobcs3pIOhwun3yr_ndvTEYOn4BJjTzs,5889
+dbos/dbos-config.schema.json,sha256=00lliu7hGT6NAIZt8UNAn8mEhQ71RGw6Q2CI3nWxULA,5899
 dbos/py.typed,sha256=QfzXT1Ktfk3Rj84akygc7_42z0lRpCq0Ilh8OXI6Zas,44
 version/__init__.py,sha256=L4sNxecRuqdtSFdpUGX3TtBi9KL3k7YsZVIvv-fv9-A,1678
-dbos-0.17.0a4.dist-info/RECORD,,
+dbos-0.18.0a1.dist-info/RECORD,,