dbos 0.26.0a7__py3-none-any.whl → 0.26.0a9__py3-none-any.whl

dbos/_cloudutils/authentication.py

@@ -1,163 +0,0 @@
-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

@@ -1,254 +0,0 @@
-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
-    Organization: str
-
-    def __init__(self, **kwargs: Any) -> None:
-        self.Name = kwargs.get("Name", "")
-        self.Organization = kwargs.get("Organization", "")
-
-
-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

@@ -1,241 +0,0 @@
-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
-
-    def __init__(self, **kwargs: Any) -> None:
-        self.RoleName = kwargs.get("RoleName", "")
-        self.Password = kwargs.get("Password", "")
-
-
-@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