llamactl 0.3.0a19__py3-none-any.whl → 0.3.0a21__py3-none-any.whl

llama_deploy/cli/config/_config.py

@@ -1,19 +1,37 @@
 """Configuration and profile management for llamactl"""
 
+import functools
 import os
 import sqlite3
+import uuid
 from pathlib import Path
 from typing import Any
 
-from .schema import DEFAULT_ENVIRONMENT, Auth, Environment
+from ._migrations import run_migrations
+from .schema import DEFAULT_ENVIRONMENT, Auth, DeviceOIDC, Environment
+
+
+def _serialize_device_oidc(value: DeviceOIDC | None) -> str | None:
+    if value is None:
+        return None
+    return value.model_dump_json()
+
+
+def _deserialize_device_oidc(value: str | None) -> DeviceOIDC | None:
+    if not value:
+        return None
+    return DeviceOIDC.model_validate_json(value)
 
 
 def _to_auth(row: Any) -> Auth:
     return Auth(
-        name=row[0],
-        api_url=row[1],
-        project_id=row[2],
-        api_key=row[3],
+        id=row[0],
+        name=row[1],
+        api_url=row[2],
+        project_id=row[3],
+        api_key=row[4],
+        api_key_id=row[5],
+        device_oidc=_deserialize_device_oidc(row[6]),
     )
 
 
@@ -28,11 +46,12 @@ def _to_environment(row: Any) -> Environment:
 class ConfigManager:
     """Manages profiles and configuration using SQLite"""
 
-    def __init__(self):
+    def __init__(self, init_database: bool = True):
         self.config_dir = self._get_config_dir()
         self.db_path = self.config_dir / "profiles.db"
         self._ensure_config_dir()
-        self._init_database()
+        if init_database:
+            self._init_database()
 
     def _get_config_dir(self) -> Path:
         """Get the configuration directory path based on OS.
@@ -53,125 +72,58 @@ class ConfigManager:
         self.config_dir.mkdir(parents=True, exist_ok=True)
 
     def _init_database(self):
-        """Initialize SQLite database with required tables"""
-        with sqlite3.connect(self.db_path) as conn:
-            # Check if we need to migrate from old schema
-            cursor = conn.execute("PRAGMA table_info(profiles)")
-            columns = [row[1] for row in cursor.fetchall()]
+        """Initialize SQLite database and run migrations; then seed defaults."""
 
-            # Migration: handle old active_project_id -> project_id and make it required
-            if "active_project_id" in columns and "project_id" not in columns:
-                # Delete any profiles that have no active_project_id since project_id is now required
-                conn.execute(
-                    "DELETE FROM profiles WHERE active_project_id IS NULL OR active_project_id = ''"
-                )
+        with sqlite3.connect(self.db_path) as conn:
+            # Apply ad-hoc SQL migrations based on PRAGMA user_version
+            run_migrations(conn)
 
-                # Rename active_project_id to project_id
-                # Note: SQLite doesn't allow changing column constraints easily, but we enforce
-                # the NOT NULL constraint in our application code and new table creation
-                conn.execute(
-                    "ALTER TABLE profiles RENAME COLUMN active_project_id TO project_id"
-                )
+            conn.commit()
 
-                # Add api_key column if not already present
-                mig_cursor = conn.execute("PRAGMA table_info(profiles)")
-                mig_columns = [row[1] for row in mig_cursor.fetchall()]
-                if "api_key" not in mig_columns:
-                    conn.execute("ALTER TABLE profiles ADD COLUMN api_key TEXT")
-
-            # Migration: change profiles primary key from name to composite (name, api_url)
-            # Only perform if the table exists with single-column PK on name
-            pk_info_cur = conn.execute("PRAGMA table_info(profiles)")
-            pk_info_rows = pk_info_cur.fetchall()
-            pk_columns = [row[1] for row in pk_info_rows if len(row) > 5 and row[5] > 0]
-            if pk_columns == ["name"] and "api_url" in columns:
-                conn.execute("ALTER TABLE profiles RENAME TO profiles_old")
-                conn.execute(
-                    """
-                    CREATE TABLE profiles (
-                        name TEXT NOT NULL,
-                        api_url TEXT NOT NULL,
-                        project_id TEXT NOT NULL,
-                        api_key TEXT,
-                        PRIMARY KEY (name, api_url)
-                    )
-                    """
-                )
-                conn.execute(
-                    """
-                    INSERT INTO profiles (name, api_url, project_id, api_key)
-                    SELECT name, api_url, project_id, api_key FROM profiles_old
-                    """
-                )
-                conn.execute("DROP TABLE profiles_old")
-
-            # Create tables with new schema (this will only create if they don't exist)
-            conn.execute("""
-                CREATE TABLE IF NOT EXISTS profiles (
-                    name TEXT NOT NULL,
-                    api_url TEXT NOT NULL,
-                    project_id TEXT NOT NULL,
-                    api_key TEXT,
-                    PRIMARY KEY (name, api_url)
-                )
-            """)
+    def destroy_database(self):
+        """Destroy the database"""
+        self.db_path.unlink()
+        self._init_database()
 
-            conn.execute("""
-                CREATE TABLE IF NOT EXISTS settings (
-                    key TEXT PRIMARY KEY,
-                    value TEXT NOT NULL
-                )
-            """)
+    #############################################
+    ## Settings
+    #############################################
 
-            # Environments: first-class environments table
-            conn.execute(
-                """
-                CREATE TABLE IF NOT EXISTS environments (
-                    api_url TEXT PRIMARY KEY,
-                    requires_auth INTEGER NOT NULL,
-                    min_llamactl_version TEXT
-                )
-                """
-            )
+    def set_settings_current_profile(self, name: str | None):
+        """Set or clear the current active profile.
 
-            # Ensure there is a current environment setting. If missing, set to default.
-            setting_cursor = conn.execute(
-                "SELECT value FROM settings WHERE key = 'current_environment_api_url'"
-            )
-            setting_row = setting_cursor.fetchone()
-            if not setting_row:
+        If name is None, the setting is removed.
+        """
+        with sqlite3.connect(self.db_path) as conn:
+            if name is None:
+                conn.execute("DELETE FROM settings WHERE key = 'current_profile'")
+            else:
                 conn.execute(
-                    "INSERT OR REPLACE INTO settings (key, value) VALUES ('current_environment_api_url', ?)",
-                    (DEFAULT_ENVIRONMENT.api_url,),
+                    "INSERT OR REPLACE INTO settings (key, value) VALUES ('current_profile', ?)",
+                    (name,),
                 )
+            conn.commit()
 
-            # Seed environments from existing profiles if environments is empty
-            env_count_cur = conn.execute("SELECT COUNT(*) FROM environments")
-            env_count = env_count_cur.fetchone()[0]
-            if env_count == 0:
-                # Insert distinct api_url values from profiles with requires_auth = 0 (False)
-                conn.execute(
-                    """
-                    INSERT OR IGNORE INTO environments (api_url, requires_auth)
-                    SELECT DISTINCT api_url, 0 FROM profiles
-                    """
-                )
+    def get_settings_current_profile_name(self) -> str | None:
+        """Get the name of the current active profile"""
+        with sqlite3.connect(self.db_path) as conn:
+            cursor = conn.execute(
+                "SELECT value FROM settings WHERE key = 'current_profile'"
+            )
+            row = cursor.fetchone()
+            return row[0] if row else None
 
-                # Also ensure the current environment exists as a row
-                cur_env_cursor = conn.execute(
-                    "SELECT value FROM settings WHERE key = 'current_environment_api_url'"
-                )
-                cur_env_row = cur_env_cursor.fetchone()
-                if cur_env_row:
-                    conn.execute(
-                        "INSERT OR IGNORE INTO environments (api_url, requires_auth, min_llamactl_version) VALUES (?, ?, ?)",
-                        (
-                            cur_env_row[0],
-                            1 if DEFAULT_ENVIRONMENT.requires_auth else 0,
-                            DEFAULT_ENVIRONMENT.min_llamactl_version,
-                        ),
-                    )
+    def set_settings_current_environment(self, api_url: str) -> None:
+        """Set the current environment by URL.
 
+        Requires the environment row to already exist (validated elsewhere, e.g. via
+        a probe before creation). Raises ValueError if not found.
+        """
+        with sqlite3.connect(self.db_path) as conn:
+            conn.execute(
+                "INSERT OR REPLACE INTO settings (key, value) VALUES ('current_environment_api_url', ?)",
+                (api_url,),
+            )
             conn.commit()
 
     def create_profile(
@@ -180,26 +132,34 @@ class ConfigManager:
         api_url: str,
         project_id: str,
         api_key: str | None = None,
+        api_key_id: str | None = None,
+        device_oidc: DeviceOIDC | None = None,
     ) -> Auth:
         """Create a new auth profile"""
         if not project_id.strip():
             raise ValueError("Project ID is required")
         profile = Auth(
+            id=str(uuid.uuid4()),
             name=name,
             api_url=api_url,
             project_id=project_id,
             api_key=api_key,
+            api_key_id=api_key_id,
+            device_oidc=device_oidc,
         )
 
         with sqlite3.connect(self.db_path) as conn:
             try:
                 conn.execute(
-                    "INSERT INTO profiles (name, api_url, project_id, api_key) VALUES (?, ?, ?, ?)",
+                    "INSERT INTO profiles (id, name, api_url, project_id, api_key, api_key_id, device_oidc) VALUES (?, ?, ?, ?, ?, ?, ?)",
                     (
+                        profile.id,
                         profile.name,
                         profile.api_url,
                         profile.project_id,
                         profile.api_key,
+                        profile.api_key_id,
+                        _serialize_device_oidc(profile.device_oidc),
                     ),
                 )
                 conn.commit()
@@ -210,24 +170,101 @@ class ConfigManager:
 
         return profile
 
+    def get_current_profile(self, env_url: str) -> Auth | None:
+        """Get the current active profile"""
+        current_name = self.get_settings_current_profile_name()
+        if current_name:
+            return self.get_profile(current_name, env_url)
+        return None
+
+    def get_current_environment(self) -> Environment:
+        """Get the current active environment"""
+        with sqlite3.connect(self.db_path) as conn:
+            cursor = conn.execute(
+                "SELECT value FROM settings WHERE key = 'current_environment_api_url'"
+            )
+            row = cursor.fetchone()
+            api_url = row[0] if row else DEFAULT_ENVIRONMENT.api_url
+
+            env_row = conn.execute(
+                "SELECT api_url, requires_auth, min_llamactl_version FROM environments WHERE api_url = ?",
+                (api_url,),
+            ).fetchone()
+            if env_row:
+                return _to_environment(env_row)
+
+        # Fallback: return an in-memory default without writing to DB
+        if api_url == DEFAULT_ENVIRONMENT.api_url:
+            return DEFAULT_ENVIRONMENT
+        return Environment(
+            api_url=api_url, requires_auth=False, min_llamactl_version=None
+        )
+
+    ##################################
+    ## Profiles
+    ##################################
+
     def get_profile(self, name: str, env_url: str) -> Auth | None:
         """Get a profile by name"""
         with sqlite3.connect(self.db_path) as conn:
             row = conn.execute(
-                "SELECT name, api_url, project_id, api_key FROM profiles WHERE name = ? AND api_url = ?",
+                "SELECT id, name, api_url, project_id, api_key, api_key_id, device_oidc FROM profiles WHERE name = ? AND api_url = ?",
                 (name, env_url),
             ).fetchone()
             if row:
                 return _to_auth(row)
         return None
 
+    def get_profile_by_id(self, id: str) -> Auth | None:
+        """Get a profile by ID"""
+        with sqlite3.connect(self.db_path) as conn:
+            row = conn.execute(
+                "SELECT id, name, api_url, project_id, api_key, api_key_id, device_oidc FROM profiles WHERE id = ?",
+                (id,),
+            ).fetchone()
+            if row:
+                return _to_auth(row)
+        return None
+
+    def get_profile_by_api_key(self, env_url: str, api_key: str) -> Auth | None:
+        """Get a profile by api_key within an environment URL."""
+        with sqlite3.connect(self.db_path) as conn:
+            row = conn.execute(
+                """
+                SELECT id, name, api_url, project_id, api_key, api_key_id, device_oidc
+                FROM profiles
+                WHERE api_url = ? AND api_key = ?
+                LIMIT 1
+                """,
+                (env_url, api_key),
+            ).fetchone()
+            if row:
+                return _to_auth(row)
+        return None
+
+    def get_profile_by_device_user_id(self, env_url: str, user_id: str) -> Auth | None:
+        """Get a profile by device OIDC user_id within an environment URL."""
+        with sqlite3.connect(self.db_path) as conn:
+            row = conn.execute(
+                """
+                SELECT id, name, api_url, project_id, api_key, api_key_id, device_oidc
+                FROM profiles
+                WHERE api_url = ? AND JSON_EXTRACT(device_oidc, '$.user_id') = ?
+                LIMIT 1
+                """,
+                (env_url, user_id),
+            ).fetchone()
+            if row:
+                return _to_auth(row)
+        return None
+
     def list_profiles(self, env_url: str) -> list[Auth]:
         """List all profiles"""
         with sqlite3.connect(self.db_path) as conn:
             return [
                 _to_auth(row)
                 for row in conn.execute(
-                    "SELECT name, api_url, project_id, api_key FROM profiles WHERE api_url = ? ORDER BY name",
+                    "SELECT id, name, api_url, project_id, api_key, api_key_id, device_oidc FROM profiles WHERE api_url = ? ORDER BY name",
                     (env_url,),
                 ).fetchall()
             ]
@@ -246,37 +283,6 @@ class ConfigManager:
 
             return cursor.rowcount > 0
 
-    def set_settings_current_profile(self, name: str | None):
-        """Set or clear the current active profile.
-
-        If name is None, the setting is removed.
-        """
-        with sqlite3.connect(self.db_path) as conn:
-            if name is None:
-                conn.execute("DELETE FROM settings WHERE key = 'current_profile'")
-            else:
-                conn.execute(
-                    "INSERT OR REPLACE INTO settings (key, value) VALUES ('current_profile', ?)",
-                    (name,),
-                )
-            conn.commit()
-
-    def get_settings_current_profile_name(self) -> str | None:
-        """Get the name of the current active profile"""
-        with sqlite3.connect(self.db_path) as conn:
-            cursor = conn.execute(
-                "SELECT value FROM settings WHERE key = 'current_profile'"
-            )
-            row = cursor.fetchone()
-            return row[0] if row else None
-
-    def get_current_profile(self, env_url: str) -> Auth | None:
-        """Get the current active profile"""
-        current_name = self.get_settings_current_profile_name()
-        if current_name:
-            return self.get_profile(current_name, env_url)
-        return None
-
     def set_project(self, profile_name: str, env_url: str, project_id: str) -> bool:
         """Set the project for a profile. Returns True if profile exists."""
         with sqlite3.connect(self.db_path) as conn:
@@ -287,7 +293,26 @@ class ConfigManager:
             conn.commit()
             return cursor.rowcount > 0
 
-    # Environment management APIs
+    def update_profile(self, profile: Auth) -> None:
+        """Update a profile"""
+        with sqlite3.connect(self.db_path) as conn:
+            conn.execute(
+                "UPDATE profiles SET name = ?, api_url = ?, project_id = ?, api_key = ?, api_key_id = ?, device_oidc = ? WHERE id = ?",
+                (
+                    profile.name,
+                    profile.api_url,
+                    profile.project_id,
+                    profile.api_key,
+                    profile.api_key_id,
+                    _serialize_device_oidc(profile.device_oidc),
+                    profile.id,
+                ),
+            )
+            conn.commit()
+
+    ##################################
+    ## Environments
+    ##################################
     def create_or_update_environment(
         self, api_url: str, requires_auth: bool, min_llamactl_version: str | None = None
     ) -> None:
@@ -323,57 +348,6 @@ class ConfigManager:
                 envs = [DEFAULT_ENVIRONMENT]
             return envs
 
-    def set_settings_current_environment(self, api_url: str) -> None:
-        """Set the current environment by URL.
-
-        Requires the environment row to already exist (validated elsewhere, e.g. via
-        a probe before creation). Raises ValueError if not found.
-        """
-        with sqlite3.connect(self.db_path) as conn:
-            conn.execute(
-                "INSERT OR REPLACE INTO settings (key, value) VALUES ('current_environment_api_url', ?)",
-                (api_url,),
-            )
-            conn.commit()
-
-    def get_current_environment(self) -> Environment:
-        """Get the current environment.
-
-        Ensures there is a settings entry and a corresponding row in `environments`.
-        """
-        with sqlite3.connect(self.db_path) as conn:
-            cursor = conn.execute(
-                "SELECT value FROM settings WHERE key = 'current_environment_api_url'"
-            )
-            row = cursor.fetchone()
-            api_url = row[0] if row else DEFAULT_ENVIRONMENT.api_url
-            if not row:
-                conn.execute(
-                    "INSERT OR REPLACE INTO settings (key, value) VALUES ('current_environment_api_url', ?)",
-                    (api_url,),
-                )
-
-            # Ensure environment exists
-            conn.execute(
-                "INSERT OR IGNORE INTO environments (api_url, requires_auth, min_llamactl_version) VALUES (?, ?, ?)",
-                (
-                    api_url,
-                    1 if DEFAULT_ENVIRONMENT.requires_auth else 0,
-                    DEFAULT_ENVIRONMENT.min_llamactl_version,
-                ),
-            )
-
-            # Read the environment
-            env_row = conn.execute(
-                "SELECT api_url, requires_auth, min_llamactl_version FROM environments WHERE api_url = ?",
-                (api_url,),
-            ).fetchone()
-            if env_row:
-                return _to_environment(env_row)
-
-        # If we somehow got here without returning, raise an error
-        raise RuntimeError("Failed to load current environment")
-
     def delete_environment(self, api_url: str) -> bool:
         """Delete an environment and all associated profiles.
 
@@ -411,4 +385,6 @@ class ConfigManager:
 
 
 # Global config manager instance
-config_manager = ConfigManager()
+@functools.cache
+def config_manager() -> ConfigManager:
+    return ConfigManager()

llama_deploy/cli/config/_migrations.py

@@ -0,0 +1,65 @@
+"""Ad-hoc SQLite schema migrations using PRAGMA user_version.
+
+Inspired by https://eskerda.com/sqlite-schema-migrations-python/
+"""
+
+from __future__ import annotations
+
+import logging
+import re
+import sqlite3
+from importlib import import_module, resources
+from importlib.resources.abc import Traversable
+
+logger = logging.getLogger(__name__)
+
+
+_MIGRATIONS_PKG = "llama_deploy.cli.config.migrations"
+_USER_VERSION_PATTERN = re.compile(r"pragma\s+user_version\s*=\s*(\d+)", re.IGNORECASE)
+
+
+def _iter_migration_files() -> list[Traversable]:
+    """Yield packaged SQL migration files in lexicographic order."""
+    pkg = import_module(_MIGRATIONS_PKG)
+    root = resources.files(pkg)
+    files = (p for p in root.iterdir() if p.name.endswith(".sql"))
+    if not files:
+        raise ValueError("No migration files found")
+    return sorted(files, key=lambda p: p.name)
+
+
+def _parse_target_version(sql_text: str) -> int | None:
+    """Return target schema version declared in the first PRAGMA line, if any."""
+    first_line = sql_text.splitlines()[0] if sql_text else ""
+    match = _USER_VERSION_PATTERN.search(first_line)
+    return int(match.group(1)) if match else None
+
+
+def run_migrations(conn: sqlite3.Connection) -> None:
+    """Apply pending migrations found under the migrations package.
+
+    Each migration file should start with a `PRAGMA user_version=N;` line.
+    Files are applied in lexicographic order and only when N > current_version.
+    """
+    cur = conn.cursor()
+    current_version_row = cur.execute("PRAGMA user_version").fetchone()
+    current_version = int(current_version_row[0]) if current_version_row else 0
+
+    for path in _iter_migration_files():
+        sql_text = path.read_text()
+        target_version = _parse_target_version(sql_text) or 0
+        if target_version <= current_version:
+            continue
+
+        try:
+            logger.debug(
+                "Applying migration %s → target version %s", path.name, target_version
+            )
+            cur.executescript("BEGIN;\n" + sql_text)
+        except Exception as exc:  # noqa: BLE001 – we surface the exact error
+            logger.error("Failed migration %s: %s", path.name, exc)
+            cur.execute("ROLLBACK")
+            raise
+        else:
+            cur.execute("COMMIT")
+            current_version = target_version

llama_deploy/cli/config/auth_service.py

@@ -1,7 +1,9 @@
 import asyncio
 
+from llama_deploy.cli.auth.client import PlatformAuthClient, RefreshMiddleware
 from llama_deploy.cli.config._config import Auth, ConfigManager, Environment
-from llama_deploy.core.client.manage_client import ControlPlaneClient
+from llama_deploy.cli.config.schema import DeviceOIDC
+from llama_deploy.core.client.manage_client import ControlPlaneClient, httpx
 from llama_deploy.core.schema import VersionResponse
 from llama_deploy.core.schema.projects import ProjectSummary
 
@@ -17,6 +19,9 @@ class AuthService:
     def get_profile(self, name: str) -> Auth | None:
         return self.config_manager.get_profile(name, self.env.api_url)
 
+    def get_profile_by_id(self, id: str) -> Auth | None:
+        return self.config_manager.get_profile_by_id(id)
+
     def set_current_profile(self, name: str) -> None:
         self.config_manager.set_settings_current_profile(name)
 
@@ -37,7 +42,35 @@ class AuthService:
         self.config_manager.set_settings_current_profile(auth.name)
         return auth
 
-    def delete_profile(self, name: str) -> bool:
+    def create_or_update_profile_from_oidc(
+        self, project_id: str, device_oidc: DeviceOIDC
+    ) -> Auth:
+        base = device_oidc.email
+        existing = self.config_manager.get_profile_by_device_user_id(
+            self.env.api_url, device_oidc.user_id
+        )
+        if existing:
+            existing.device_oidc = device_oidc
+            self.config_manager.update_profile(existing)
+            auth = existing
+        else:
+            auth = self.config_manager.create_profile(
+                base, self.env.api_url, project_id, device_oidc=device_oidc
+            )
+        self.config_manager.set_settings_current_profile(auth.name)
+        return auth
+
+    def update_profile(self, profile: Auth) -> None:
+        self.config_manager.update_profile(profile)
+
+    async def delete_profile(self, name: str) -> bool:
+        profile = self.get_profile(name)
+        if profile and profile.api_key_id:
+            async with self.profile_client(profile) as client:
+                try:
+                    await client.delete_api_key(profile.api_key_id)
+                except Exception:
+                    pass
         return self.config_manager.delete_profile(name, self.env.api_url)
 
     def set_project(self, name: str, project_id: str) -> None:
@@ -58,6 +91,35 @@ class AuthService:
 
         return asyncio.run(_run())
 
+    def auth_middleware(self, profile: Auth | None = None) -> httpx.Auth | None:
+        profile = profile or self.get_current_profile()
+        if profile and profile.device_oidc:
+            _profile = profile  # copy to assist type checker being inflexible
+
+            async def _on_refresh(updated: DeviceOIDC) -> None:
+                # Persist refreshed tokens to the database synchronously within async wrapper
+                self.refresh_to_db(_profile.id, updated)
+
+            return RefreshMiddleware(
+                profile.device_oidc,
+                _on_refresh,
+            )
+        return None
+
+    def refresh_to_db(self, profile_id: str, device_oidc: DeviceOIDC) -> None:
+        profile = self.get_profile_by_id(profile_id)
+        if profile:
+            profile.device_oidc = device_oidc
+            self.config_manager.update_profile(profile)
+
+    def profile_client(self, profile: Auth | None = None) -> PlatformAuthClient:
+        profile = profile or self.get_current_profile()
+        if not profile:
+            raise ValueError("No active profile")
+        return PlatformAuthClient(
+            profile.api_url, profile.api_key, self.auth_middleware(profile)
+        )
+
 
 def _auto_profile_name_from_token(api_key: str) -> str:
     token = api_key or "token"