Spaces:
Sleeping
Sleeping
Download auth_storage.py from automl-team/AutoML: direct link, hf CLI and curl.
- Browser
- Download file 57.5 kB
-
https://huggingface.co/spaces/automl-team/AutoML/resolve/main/auth_storage.py
- Command line
-
hf download hf://spaces/automl-team/AutoML/auth_storage.py
-
curl -L -o auth_storage.py https://huggingface.co/spaces/automl-team/AutoML/resolve/main/auth_storage.py
57.5 kB
| from __future__ import annotations | |
| import base64 | |
| import binascii | |
| import hashlib | |
| import hmac | |
| import json | |
| import os | |
| import re | |
| import secrets | |
| import shutil | |
| import sqlite3 | |
| import tempfile | |
| import threading | |
| import time | |
| from contextlib import contextmanager | |
| from datetime import datetime, timezone | |
| from pathlib import Path | |
| from typing import Any, Dict, Iterator, List, Tuple | |
| AUTH_DB_FILENAME = "automl_accounts.sqlite3" | |
| AUTH_DB_SNAPSHOT_FILENAME = "automl_accounts.snapshot.sqlite3" | |
| AUTH_DB_PREVIOUS_SNAPSHOT_FILENAME = "automl_accounts.snapshot.previous.sqlite3" | |
| DEFAULT_MODEL_CHUNK_SIZE = 4 * 1024 * 1024 | |
| ENCRYPTED_MODEL_MAGIC = b"AUTOMLENC1\n" | |
| USERNAME_RE = re.compile(r"^[A-Za-z0-9_.@-]{3,64}$") | |
| ACCOUNT_STATUS_PENDING = "pending" | |
| ACCOUNT_STATUS_APPROVED = "approved" | |
| ACCOUNT_STATUS_DENIED = "denied" | |
| ACCOUNT_STATUSES = frozenset( | |
| {ACCOUNT_STATUS_PENDING, ACCOUNT_STATUS_APPROVED, ACCOUNT_STATUS_DENIED} | |
| ) | |
| ADMIN_USERNAMES = frozenset({"zukhriddin", "alex", "vishnu", "snigdha", "flex"}) | |
| _DB_LOCK = threading.RLock() | |
| _BOOTSTRAPPED = False | |
| _RUNTIME_DB_PREPARED = False | |
| _PROCESS_SESSION_SECRET = secrets.token_bytes(32) | |
| _BROWSER_STATE_CONTEXT = b"automl-browser-state-v1" | |
| class StorageConfigurationError(RuntimeError): | |
| """Raised when secure account/model storage is not configured correctly.""" | |
| def env_truthy(name: str, default: bool = False) -> bool: | |
| value = os.getenv(name) | |
| if value is None: | |
| return default | |
| return value.strip().lower() in {"1", "true", "yes", "on"} | |
| def storage_root() -> Path: | |
| configured = os.getenv("AUTOML_STORAGE_DIR", "").strip() | |
| if configured: | |
| return Path(configured).expanduser() | |
| data_root = Path("/data") | |
| if data_root.is_dir() and os.access(data_root, os.W_OK): | |
| return data_root / "automl" | |
| return Path.cwd() / ".automl_storage" | |
| def auth_db_path() -> Path: | |
| return storage_root() / AUTH_DB_FILENAME | |
| def _snapshot_database_enabled() -> bool: | |
| configured = os.getenv("AUTOML_DB_SNAPSHOT_MODE") | |
| if configured is not None: | |
| return env_truthy("AUTOML_DB_SNAPSHOT_MODE") | |
| root = storage_root() | |
| return root == Path("/data") or Path("/data") in root.parents | |
| def _persistent_storage_requires_existing_database(root: Path) -> bool: | |
| """Fail closed when authenticated snapshot storage appears temporarily empty.""" | |
| return bool( | |
| _snapshot_database_enabled() | |
| and env_truthy("AUTOML_REQUIRE_AUTH", False) | |
| and not env_truthy("AUTOML_ALLOW_EMPTY_PERSISTENT_STORAGE_INIT", False) | |
| ) | |
| def _storage_mount_wait_seconds() -> float: | |
| raw = os.getenv("AUTOML_STORAGE_MOUNT_WAIT_SECONDS", "30").strip() | |
| try: | |
| return min(120.0, max(0.0, float(raw))) | |
| except ValueError: | |
| return 30.0 | |
| def _wait_for_persistent_database(candidates: Tuple[Path, ...]) -> None: | |
| deadline = time.monotonic() + _storage_mount_wait_seconds() | |
| while not any(candidate.is_file() for candidate in candidates): | |
| remaining = deadline - time.monotonic() | |
| if remaining <= 0: | |
| break | |
| time.sleep(min(0.25, remaining)) | |
| def _runtime_auth_db_path() -> Path: | |
| if not _snapshot_database_enabled(): | |
| return auth_db_path() | |
| configured = os.getenv("AUTOML_DB_RUNTIME_DIR", "").strip() | |
| runtime_root = ( | |
| Path(configured).expanduser() | |
| if configured | |
| else Path(tempfile.gettempdir()) / "automl_gateway_state" | |
| ) | |
| root_key = hashlib.sha256(str(storage_root()).encode("utf-8")).hexdigest()[:12] | |
| return runtime_root / root_key / AUTH_DB_FILENAME | |
| def _database_integrity_ok(path: Path) -> bool: | |
| if not path.is_file(): | |
| return False | |
| conn: sqlite3.Connection | None = None | |
| try: | |
| uri = f"file:{path.resolve()}?mode=ro&immutable=1" | |
| conn = sqlite3.connect(uri, uri=True, timeout=5) | |
| row = conn.execute("PRAGMA quick_check").fetchone() | |
| return bool(row and str(row[0]).lower() == "ok") | |
| except sqlite3.DatabaseError: | |
| return False | |
| finally: | |
| if conn is not None: | |
| conn.close() | |
| def _copy_valid_database(source: Path, destination: Path) -> bool: | |
| if not source.is_file(): | |
| return False | |
| destination.parent.mkdir(parents=True, exist_ok=True) | |
| staged = destination.with_suffix(destination.suffix + ".restore") | |
| staged.unlink(missing_ok=True) | |
| try: | |
| shutil.copy2(source, staged) | |
| if not _database_integrity_ok(staged): | |
| return False | |
| os.replace(staged, destination) | |
| return True | |
| finally: | |
| staged.unlink(missing_ok=True) | |
| def _salvage_database(source: Path, destination: Path) -> bool: | |
| """Rebuild readable application tables without trusting damaged indexes.""" | |
| if not source.is_file(): | |
| return False | |
| staged = destination.with_suffix(destination.suffix + ".salvage") | |
| staged.unlink(missing_ok=True) | |
| source_conn = None | |
| destination_conn = None | |
| tables = ( | |
| ( | |
| "users", | |
| frozenset( | |
| { | |
| "user_id", | |
| "username", | |
| "username_normalized", | |
| "password_hash", | |
| "password_salt", | |
| } | |
| ), | |
| ), | |
| ( | |
| "saved_models", | |
| frozenset( | |
| { | |
| "model_id", | |
| "user_id", | |
| "manifest_json", | |
| "storage_path", | |
| "sha256", | |
| } | |
| ), | |
| ), | |
| ( | |
| "scheduled_jobs", | |
| frozenset( | |
| { | |
| "request_id", | |
| "user_id", | |
| "idempotency_key", | |
| "provider", | |
| "status", | |
| } | |
| ), | |
| ), | |
| ) | |
| try: | |
| uri = f"file:{source.resolve()}?mode=ro&immutable=1" | |
| source_conn = sqlite3.connect(uri, uri=True, timeout=5) | |
| source_conn.row_factory = sqlite3.Row | |
| source_conn.execute("PRAGMA writable_schema=ON") | |
| destination.parent.mkdir(parents=True, exist_ok=True) | |
| destination_conn = sqlite3.connect(str(staged), timeout=30) | |
| destination_conn.row_factory = sqlite3.Row | |
| destination_conn.execute("PRAGMA journal_mode=DELETE") | |
| destination_conn.execute("PRAGMA foreign_keys=ON") | |
| _create_schema(destination_conn) | |
| for table, required_columns in tables: | |
| columns = [ | |
| str(row["name"]) | |
| for row in source_conn.execute( | |
| f"PRAGMA table_info({table})" | |
| ).fetchall() | |
| ] | |
| if not required_columns.issubset(columns) or any( | |
| not re.fullmatch(r"[A-Za-z_][A-Za-z0-9_]*", column) | |
| for column in columns | |
| ): | |
| return False | |
| rows = source_conn.execute(f"SELECT * FROM {table}").fetchall() | |
| source_count = int( | |
| source_conn.execute(f"SELECT COUNT(*) FROM {table}").fetchone()[0] | |
| ) | |
| if len(rows) != source_count: | |
| return False | |
| if rows: | |
| quoted_columns = ", ".join(f'"{column}"' for column in columns) | |
| placeholders = ", ".join("?" for _ in columns) | |
| destination_conn.executemany( | |
| f"INSERT INTO {table} ({quoted_columns}) VALUES ({placeholders})", | |
| [tuple(row[column] for column in columns) for row in rows], | |
| ) | |
| if destination_conn.execute("PRAGMA foreign_key_check").fetchone(): | |
| return False | |
| destination_conn.commit() | |
| destination_conn.close() | |
| destination_conn = None | |
| if not _database_integrity_ok(staged): | |
| return False | |
| os.replace(staged, destination) | |
| return True | |
| except (OSError, sqlite3.DatabaseError, TypeError, ValueError): | |
| return False | |
| finally: | |
| if source_conn is not None: | |
| source_conn.close() | |
| if destination_conn is not None: | |
| destination_conn.close() | |
| staged.unlink(missing_ok=True) | |
| def _prepare_runtime_database() -> None: | |
| global _RUNTIME_DB_PREPARED | |
| if _RUNTIME_DB_PREPARED or not _snapshot_database_enabled(): | |
| return | |
| runtime_path = _runtime_auth_db_path() | |
| runtime_path.parent.mkdir(parents=True, exist_ok=True) | |
| if _database_integrity_ok(runtime_path): | |
| _RUNTIME_DB_PREPARED = True | |
| return | |
| runtime_path.unlink(missing_ok=True) | |
| root = storage_root() | |
| candidates = ( | |
| root / AUTH_DB_SNAPSHOT_FILENAME, | |
| root / AUTH_DB_PREVIOUS_SNAPSHOT_FILENAME, | |
| auth_db_path(), | |
| ) | |
| if ( | |
| not any(candidate.is_file() for candidate in candidates) | |
| and _persistent_storage_requires_existing_database(root) | |
| ): | |
| _wait_for_persistent_database(candidates) | |
| if not any(candidate.is_file() for candidate in candidates): | |
| raise StorageConfigurationError( | |
| "Persistent account storage is unavailable or unexpectedly empty. " | |
| "The gateway stopped without creating a replacement database." | |
| ) | |
| if any(candidate.is_file() for candidate in candidates): | |
| for candidate in candidates: | |
| if _copy_valid_database(candidate, runtime_path) or _salvage_database( | |
| candidate, | |
| runtime_path, | |
| ): | |
| _RUNTIME_DB_PREPARED = True | |
| return | |
| raise StorageConfigurationError( | |
| "The persisted account database and its recovery snapshots failed " | |
| "integrity checks. The gateway stopped without replacing account data." | |
| ) | |
| _RUNTIME_DB_PREPARED = True | |
| def _persist_runtime_database(conn: sqlite3.Connection) -> None: | |
| if not _snapshot_database_enabled(): | |
| return | |
| root = storage_root() | |
| root.mkdir(parents=True, exist_ok=True) | |
| runtime_path = _runtime_auth_db_path() | |
| local_snapshot = runtime_path.with_suffix(".snapshot.tmp") | |
| durable_snapshot = root / AUTH_DB_SNAPSHOT_FILENAME | |
| previous_snapshot = root / AUTH_DB_PREVIOUS_SNAPSHOT_FILENAME | |
| staged_snapshot = root / f"{AUTH_DB_SNAPSHOT_FILENAME}.next" | |
| staged_previous = root / f"{AUTH_DB_PREVIOUS_SNAPSHOT_FILENAME}.next" | |
| local_snapshot.unlink(missing_ok=True) | |
| staged_snapshot.unlink(missing_ok=True) | |
| staged_previous.unlink(missing_ok=True) | |
| try: | |
| destination = sqlite3.connect(str(local_snapshot)) | |
| try: | |
| conn.backup(destination) | |
| finally: | |
| destination.close() | |
| if not _database_integrity_ok(local_snapshot): | |
| raise StorageConfigurationError( | |
| "The account database snapshot failed its integrity check." | |
| ) | |
| if _database_integrity_ok(durable_snapshot): | |
| shutil.copy2(durable_snapshot, staged_previous) | |
| os.replace(staged_previous, previous_snapshot) | |
| shutil.copy2(local_snapshot, staged_snapshot) | |
| if not _database_integrity_ok(staged_snapshot): | |
| raise StorageConfigurationError( | |
| "The durable account database snapshot failed verification." | |
| ) | |
| os.replace(staged_snapshot, durable_snapshot) | |
| finally: | |
| local_snapshot.unlink(missing_ok=True) | |
| staged_snapshot.unlink(missing_ok=True) | |
| staged_previous.unlink(missing_ok=True) | |
| def _utcnow() -> str: | |
| return datetime.now(timezone.utc).replace(microsecond=0).isoformat() | |
| def _connect() -> Iterator[sqlite3.Connection]: | |
| root = storage_root() | |
| _prepare_runtime_database() | |
| root.mkdir(parents=True, exist_ok=True) | |
| database_path = _runtime_auth_db_path() | |
| database_path.parent.mkdir(parents=True, exist_ok=True) | |
| conn = sqlite3.connect(str(database_path), timeout=30) | |
| conn.row_factory = sqlite3.Row | |
| conn.execute("PRAGMA journal_mode=DELETE") | |
| conn.execute("PRAGMA foreign_keys=ON") | |
| conn.execute("PRAGMA busy_timeout=30000") | |
| initial_changes = conn.total_changes | |
| snapshot_needs_refresh = ( | |
| _snapshot_database_enabled() | |
| and not _database_integrity_ok( | |
| storage_root() / AUTH_DB_SNAPSHOT_FILENAME | |
| ) | |
| ) | |
| try: | |
| yield conn | |
| except Exception: | |
| conn.rollback() | |
| raise | |
| else: | |
| conn.commit() | |
| if conn.total_changes > initial_changes or snapshot_needs_refresh: | |
| _persist_runtime_database(conn) | |
| finally: | |
| conn.close() | |
| def _create_schema(conn: sqlite3.Connection) -> None: | |
| conn.execute( | |
| """ | |
| CREATE TABLE IF NOT EXISTS users ( | |
| user_id TEXT PRIMARY KEY, | |
| username TEXT NOT NULL, | |
| username_normalized TEXT NOT NULL UNIQUE, | |
| password_alg TEXT NOT NULL, | |
| password_hash TEXT NOT NULL, | |
| password_salt TEXT NOT NULL, | |
| password_iterations INTEGER NOT NULL, | |
| session_version INTEGER NOT NULL DEFAULT 1, | |
| account_status TEXT NOT NULL DEFAULT 'approved', | |
| requested_at TEXT, | |
| reviewed_at TEXT, | |
| reviewed_by TEXT, | |
| created_at TEXT NOT NULL, | |
| last_login_at TEXT | |
| ) | |
| """ | |
| ) | |
| user_columns = {str(row[1]) for row in conn.execute("PRAGMA table_info(users)")} | |
| migrations = { | |
| "session_version": ( | |
| "ALTER TABLE users ADD COLUMN session_version INTEGER NOT NULL DEFAULT 1" | |
| ), | |
| "account_status": ( | |
| "ALTER TABLE users ADD COLUMN account_status TEXT NOT NULL DEFAULT 'approved'" | |
| ), | |
| "requested_at": "ALTER TABLE users ADD COLUMN requested_at TEXT", | |
| "reviewed_at": "ALTER TABLE users ADD COLUMN reviewed_at TEXT", | |
| "reviewed_by": "ALTER TABLE users ADD COLUMN reviewed_by TEXT", | |
| } | |
| for column_name, statement in migrations.items(): | |
| if column_name not in user_columns: | |
| conn.execute(statement) | |
| conn.execute( | |
| """ | |
| CREATE TABLE IF NOT EXISTS saved_models ( | |
| model_id TEXT PRIMARY KEY, | |
| user_id TEXT NOT NULL, | |
| display_name TEXT NOT NULL, | |
| artifact_type TEXT, | |
| predictor TEXT, | |
| task_type TEXT, | |
| target_column TEXT, | |
| manifest_json TEXT NOT NULL, | |
| storage_path TEXT NOT NULL, | |
| storage_encrypted INTEGER NOT NULL, | |
| storage_key_id TEXT, | |
| source_package_name TEXT, | |
| byte_size INTEGER NOT NULL, | |
| sha256 TEXT NOT NULL, | |
| created_at TEXT NOT NULL, | |
| last_used_at TEXT, | |
| FOREIGN KEY(user_id) REFERENCES users(user_id) ON DELETE CASCADE | |
| ) | |
| """ | |
| ) | |
| conn.execute("CREATE INDEX IF NOT EXISTS idx_saved_models_user_created ON saved_models(user_id, created_at)") | |
| conn.execute( | |
| "CREATE INDEX IF NOT EXISTS idx_users_account_status " | |
| "ON users(account_status, requested_at)" | |
| ) | |
| conn.execute( | |
| """ | |
| CREATE TABLE IF NOT EXISTS scheduled_jobs ( | |
| request_id TEXT PRIMARY KEY, | |
| user_id TEXT NOT NULL, | |
| idempotency_key TEXT NOT NULL, | |
| profile TEXT NOT NULL, | |
| provider TEXT NOT NULL, | |
| external_job_id TEXT, | |
| status TEXT NOT NULL, | |
| request_json TEXT NOT NULL, | |
| result_json TEXT, | |
| display_name TEXT NOT NULL, | |
| result_location TEXT, | |
| remote_result_location TEXT, | |
| package_sha256 TEXT, | |
| routing_reason TEXT, | |
| cancellation_requested INTEGER NOT NULL DEFAULT 0, | |
| cleanup_pending INTEGER NOT NULL DEFAULT 0, | |
| queued_at TEXT NOT NULL, | |
| submitted_at TEXT, | |
| started_at TEXT, | |
| finished_at TEXT, | |
| cancelled_at TEXT, | |
| updated_at TEXT NOT NULL, | |
| FOREIGN KEY(user_id) REFERENCES users(user_id) ON DELETE CASCADE, | |
| UNIQUE(user_id, idempotency_key) | |
| ) | |
| """ | |
| ) | |
| conn.execute( | |
| "CREATE INDEX IF NOT EXISTS idx_scheduled_jobs_active " | |
| "ON scheduled_jobs(status, provider, updated_at)" | |
| ) | |
| conn.execute( | |
| "CREATE INDEX IF NOT EXISTS idx_scheduled_jobs_user_created " | |
| "ON scheduled_jobs(user_id, queued_at)" | |
| ) | |
| scheduled_columns = { | |
| row["name"] | |
| for row in conn.execute("PRAGMA table_info(scheduled_jobs)").fetchall() | |
| } | |
| scheduled_migrations = { | |
| "remote_result_location": ( | |
| "ALTER TABLE scheduled_jobs ADD COLUMN remote_result_location TEXT" | |
| ), | |
| "cleanup_pending": ( | |
| "ALTER TABLE scheduled_jobs ADD COLUMN cleanup_pending INTEGER NOT NULL DEFAULT 0" | |
| ), | |
| "routing_reason": ( | |
| "ALTER TABLE scheduled_jobs ADD COLUMN routing_reason TEXT" | |
| ), | |
| } | |
| for column_name, statement in scheduled_migrations.items(): | |
| if column_name not in scheduled_columns: | |
| conn.execute(statement) | |
| conn.commit() | |
| def normalize_username(username: str | None) -> str: | |
| return str(username or "").strip() | |
| def _username_key(username: str | None) -> str: | |
| return normalize_username(username).lower() | |
| def is_admin(username: str | None) -> bool: | |
| return _username_key(username) in ADMIN_USERNAMES | |
| def _is_approved_user(user: Dict[str, Any] | sqlite3.Row | None) -> bool: | |
| return bool( | |
| user | |
| and str(user["account_status"] or ACCOUNT_STATUS_APPROVED).lower() | |
| == ACCOUNT_STATUS_APPROVED | |
| ) | |
| def validate_username(username: str) -> None: | |
| if not USERNAME_RE.fullmatch(username): | |
| raise ValueError("Usernames must be 3-64 characters using letters, numbers, dots, underscores, hyphens, or @.") | |
| def min_password_length() -> int: | |
| raw = os.getenv("AUTOML_MIN_PASSWORD_LENGTH", "10").strip() | |
| try: | |
| return max(8, int(raw)) | |
| except ValueError: | |
| return 10 | |
| def validate_password(password: str) -> None: | |
| if len(password or "") < min_password_length(): | |
| raise ValueError(f"Passwords must be at least {min_password_length()} characters.") | |
| def _password_iterations() -> int: | |
| raw = os.getenv("AUTOML_PASSWORD_HASH_ITERATIONS", "390000").strip() | |
| try: | |
| return max(200000, int(raw)) | |
| except ValueError: | |
| return 390000 | |
| def _hash_password(password: str, salt: bytes | None = None, iterations: int | None = None) -> Tuple[str, str, int]: | |
| salt = salt or secrets.token_bytes(16) | |
| iterations = iterations or _password_iterations() | |
| digest = hashlib.pbkdf2_hmac("sha256", password.encode("utf-8"), salt, iterations) | |
| return ( | |
| base64.b64encode(digest).decode("ascii"), | |
| base64.b64encode(salt).decode("ascii"), | |
| iterations, | |
| ) | |
| def _verify_password(password: str, stored_hash: str, stored_salt: str, iterations: int) -> bool: | |
| try: | |
| salt = base64.b64decode(stored_salt.encode("ascii"), validate=True) | |
| expected = base64.b64decode(stored_hash.encode("ascii"), validate=True) | |
| except (binascii.Error, ValueError): | |
| return False | |
| digest = hashlib.pbkdf2_hmac("sha256", password.encode("utf-8"), salt, int(iterations)) | |
| return hmac.compare_digest(digest, expected) | |
| def initialize_storage() -> None: | |
| with _DB_LOCK: | |
| with _connect() as conn: | |
| _create_schema(conn) | |
| _bootstrap_users_once() | |
| def _bootstrap_users_once() -> None: | |
| global _BOOTSTRAPPED | |
| if _BOOTSTRAPPED: | |
| return | |
| payload = os.getenv("AUTOML_BOOTSTRAP_USERS_JSON", "").strip() | |
| if not payload: | |
| _BOOTSTRAPPED = True | |
| return | |
| try: | |
| loaded = json.loads(payload) | |
| except json.JSONDecodeError as exc: | |
| raise StorageConfigurationError("AUTOML_BOOTSTRAP_USERS_JSON must be valid JSON.") from exc | |
| if isinstance(loaded, dict): | |
| entries = [{"username": username, "password": password} for username, password in loaded.items()] | |
| elif isinstance(loaded, list): | |
| entries = loaded | |
| else: | |
| raise StorageConfigurationError("AUTOML_BOOTSTRAP_USERS_JSON must be a JSON object or list.") | |
| for entry in entries: | |
| if not isinstance(entry, dict): | |
| raise StorageConfigurationError("Each bootstrap user must be an object with username and password.") | |
| username = normalize_username(str(entry.get("username", ""))) | |
| password = str(entry.get("password", "")) | |
| create_user(username, password, if_missing=True) | |
| _BOOTSTRAPPED = True | |
| def get_user(username: str | None) -> Dict[str, Any] | None: | |
| initialize_storage() | |
| key = _username_key(username) | |
| if not key: | |
| return None | |
| with _DB_LOCK: | |
| with _connect() as conn: | |
| row = conn.execute( | |
| "SELECT * FROM users WHERE username_normalized = ?", | |
| (key,), | |
| ).fetchone() | |
| return dict(row) if row else None | |
| def get_user_by_id(user_id: str | None) -> Dict[str, Any] | None: | |
| initialize_storage() | |
| normalized = str(user_id or "").strip() | |
| if not normalized: | |
| return None | |
| with _DB_LOCK: | |
| with _connect() as conn: | |
| row = conn.execute( | |
| "SELECT * FROM users WHERE user_id = ?", | |
| (normalized,), | |
| ).fetchone() | |
| return dict(row) if row else None | |
| def user_count() -> int: | |
| initialize_storage() | |
| with _DB_LOCK: | |
| with _connect() as conn: | |
| return int(conn.execute("SELECT COUNT(*) FROM users").fetchone()[0]) | |
| def create_user(username: str, password: str, *, if_missing: bool = False) -> Dict[str, Any]: | |
| """Create an approved account for trusted bootstrap/internal provisioning.""" | |
| username = normalize_username(username) | |
| validate_username(username) | |
| validate_password(password) | |
| key = _username_key(username) | |
| with _DB_LOCK: | |
| with _connect() as conn: | |
| _create_schema(conn) | |
| existing = conn.execute( | |
| "SELECT * FROM users WHERE username_normalized = ?", | |
| (key,), | |
| ).fetchone() | |
| if existing: | |
| if if_missing: | |
| return dict(existing) | |
| raise ValueError("That username already exists.") | |
| password_hash, password_salt, iterations = _hash_password(password) | |
| now = _utcnow() | |
| user_id = secrets.token_hex(16) | |
| conn.execute( | |
| """ | |
| INSERT INTO users ( | |
| user_id, username, username_normalized, password_alg, password_hash, | |
| password_salt, password_iterations, account_status, created_at | |
| ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?) | |
| """, | |
| ( | |
| user_id, | |
| username, | |
| key, | |
| "pbkdf2_sha256", | |
| password_hash, | |
| password_salt, | |
| iterations, | |
| ACCOUNT_STATUS_APPROVED, | |
| now, | |
| ), | |
| ) | |
| conn.commit() | |
| return dict( | |
| conn.execute( | |
| "SELECT * FROM users WHERE username_normalized = ?", | |
| (key,), | |
| ).fetchone() | |
| ) | |
| def request_account_creation(username: str, password: str) -> Dict[str, Any]: | |
| """Store a self-service signup as pending administrator review.""" | |
| initialize_storage() | |
| username = normalize_username(username) | |
| validate_username(username) | |
| validate_password(password) | |
| if is_admin(username): | |
| raise ValueError( | |
| "That username is reserved for an administrator and must be provisioned by the deployment owner." | |
| ) | |
| key = _username_key(username) | |
| with _DB_LOCK: | |
| with _connect() as conn: | |
| existing = conn.execute( | |
| "SELECT account_status FROM users WHERE username_normalized = ?", | |
| (key,), | |
| ).fetchone() | |
| if existing: | |
| status = str(existing["account_status"] or ACCOUNT_STATUS_APPROVED) | |
| if status == ACCOUNT_STATUS_PENDING: | |
| raise ValueError( | |
| "An account request for that username is already awaiting administrator approval." | |
| ) | |
| if status == ACCOUNT_STATUS_DENIED: | |
| raise ValueError( | |
| "That request was denied. Log in with those credentials to resend or remove it." | |
| ) | |
| raise ValueError("That username already exists.") | |
| password_hash, password_salt, iterations = _hash_password(password) | |
| now = _utcnow() | |
| user_id = secrets.token_hex(16) | |
| conn.execute( | |
| """ | |
| INSERT INTO users ( | |
| user_id, username, username_normalized, password_alg, password_hash, | |
| password_salt, password_iterations, session_version, account_status, | |
| requested_at, created_at | |
| ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) | |
| """, | |
| ( | |
| user_id, | |
| username, | |
| key, | |
| "pbkdf2_sha256", | |
| password_hash, | |
| password_salt, | |
| iterations, | |
| 1, | |
| ACCOUNT_STATUS_PENDING, | |
| now, | |
| now, | |
| ), | |
| ) | |
| conn.commit() | |
| return dict( | |
| conn.execute( | |
| "SELECT * FROM users WHERE username_normalized = ?", | |
| (key,), | |
| ).fetchone() | |
| ) | |
| def account_authentication_status(username: str, password: str) -> str: | |
| """Return an account status only after the supplied password verifies.""" | |
| initialize_storage() | |
| username = normalize_username(username) | |
| if not username or not password: | |
| return "invalid" | |
| key = _username_key(username) | |
| with _DB_LOCK: | |
| with _connect() as conn: | |
| row = conn.execute( | |
| "SELECT * FROM users WHERE username_normalized = ?", | |
| (key,), | |
| ).fetchone() | |
| if row is None or not _verify_password( | |
| password, | |
| str(row["password_hash"]), | |
| str(row["password_salt"]), | |
| int(row["password_iterations"]), | |
| ): | |
| return "invalid" | |
| status = str(row["account_status"] or ACCOUNT_STATUS_APPROVED).lower() | |
| if status not in ACCOUNT_STATUSES: | |
| return "invalid" | |
| if status == ACCOUNT_STATUS_APPROVED: | |
| conn.execute( | |
| "UPDATE users SET last_login_at = ? WHERE user_id = ?", | |
| (_utcnow(), row["user_id"]), | |
| ) | |
| conn.commit() | |
| return status | |
| def list_pending_account_requests() -> List[Dict[str, Any]]: | |
| initialize_storage() | |
| with _DB_LOCK: | |
| with _connect() as conn: | |
| rows = conn.execute( | |
| """ | |
| SELECT username, requested_at | |
| FROM users | |
| WHERE account_status = ? | |
| ORDER BY requested_at ASC, username_normalized ASC | |
| """, | |
| (ACCOUNT_STATUS_PENDING,), | |
| ).fetchall() | |
| return [dict(row) for row in rows] | |
| def review_account_request( | |
| admin_username: str, | |
| username: str, | |
| decision: str, | |
| ) -> Dict[str, Any]: | |
| initialize_storage() | |
| admin = get_user(admin_username) | |
| if admin is None or not is_admin(admin_username) or not _is_approved_user(admin): | |
| raise PermissionError( | |
| "Only an approved administrator can review account requests." | |
| ) | |
| decision = str(decision or "").strip().lower() | |
| if decision not in {ACCOUNT_STATUS_APPROVED, ACCOUNT_STATUS_DENIED}: | |
| raise ValueError("Decision must be approved or denied.") | |
| key = _username_key(username) | |
| if not key: | |
| raise ValueError("Select a pending account request.") | |
| if is_admin(username): | |
| raise ValueError( | |
| "Administrator accounts must be provisioned through the bootstrap secret." | |
| ) | |
| with _DB_LOCK: | |
| with _connect() as conn: | |
| row = conn.execute( | |
| "SELECT * FROM users WHERE username_normalized = ?", | |
| (key,), | |
| ).fetchone() | |
| if row is None or str(row["account_status"]) != ACCOUNT_STATUS_PENDING: | |
| raise ValueError("That account request is no longer pending.") | |
| cursor = conn.execute( | |
| """ | |
| UPDATE users | |
| SET account_status = ?, reviewed_at = ?, reviewed_by = ? | |
| WHERE user_id = ? AND account_status = ? | |
| """, | |
| ( | |
| decision, | |
| _utcnow(), | |
| normalize_username(admin["username"]), | |
| row["user_id"], | |
| ACCOUNT_STATUS_PENDING, | |
| ), | |
| ) | |
| if cursor.rowcount != 1: | |
| raise ValueError("That account request is no longer pending.") | |
| conn.commit() | |
| return dict( | |
| conn.execute( | |
| "SELECT * FROM users WHERE user_id = ?", | |
| (row["user_id"],), | |
| ).fetchone() | |
| ) | |
| def resubmit_denied_account(username: str, password: str) -> None: | |
| if account_authentication_status(username, password) != ACCOUNT_STATUS_DENIED: | |
| raise PermissionError("The denied account credentials could not be verified.") | |
| key = _username_key(username) | |
| with _DB_LOCK: | |
| with _connect() as conn: | |
| cursor = conn.execute( | |
| """ | |
| UPDATE users | |
| SET account_status = ?, requested_at = ?, reviewed_at = NULL, | |
| reviewed_by = NULL | |
| WHERE username_normalized = ? AND account_status = ? | |
| """, | |
| (ACCOUNT_STATUS_PENDING, _utcnow(), key, ACCOUNT_STATUS_DENIED), | |
| ) | |
| if cursor.rowcount != 1: | |
| raise ValueError("That account is no longer denied.") | |
| conn.commit() | |
| def delete_denied_account(username: str, password: str) -> None: | |
| if account_authentication_status(username, password) != ACCOUNT_STATUS_DENIED: | |
| raise PermissionError("The denied account credentials could not be verified.") | |
| key = _username_key(username) | |
| with _DB_LOCK: | |
| with _connect() as conn: | |
| row = conn.execute( | |
| "SELECT user_id FROM users " | |
| "WHERE username_normalized = ? AND account_status = ?", | |
| (key, ACCOUNT_STATUS_DENIED), | |
| ).fetchone() | |
| if row is None: | |
| raise ValueError("That account is no longer denied.") | |
| user_id = str(row["user_id"]) | |
| conn.execute("DELETE FROM users WHERE user_id = ?", (user_id,)) | |
| conn.commit() | |
| model_dir = _model_dir_for_user(user_id) | |
| if model_dir.is_dir(): | |
| shutil.rmtree(model_dir) | |
| def signups_allowed() -> bool: | |
| return env_truthy("AUTOML_ALLOW_SIGNUPS", False) | |
| def app_auth_required() -> bool: | |
| if env_truthy("AUTOML_AUTH_DISABLED", False): | |
| return False | |
| initialize_storage() | |
| return env_truthy("AUTOML_REQUIRE_AUTH", True) or signups_allowed() or user_count() > 0 | |
| def authenticate_user(username: str, password: str, *, create_if_missing: bool = False) -> bool: | |
| # Retained for API compatibility. Self-service accounts must be submitted | |
| # explicitly and approved before authentication can succeed. | |
| del create_if_missing | |
| return ( | |
| account_authentication_status(username, password) | |
| == ACCOUNT_STATUS_APPROVED | |
| ) | |
| def update_password(username: str, old_password: str, new_password: str) -> None: | |
| initialize_storage() | |
| username = normalize_username(username) | |
| if not username: | |
| raise PermissionError("Sign in before changing your password.") | |
| validate_password(new_password) | |
| key = _username_key(username) | |
| with _DB_LOCK: | |
| with _connect() as conn: | |
| row = conn.execute( | |
| "SELECT * FROM users WHERE username_normalized = ?", | |
| (key,), | |
| ).fetchone() | |
| if row is None: | |
| raise PermissionError("The signed-in user was not found.") | |
| if not _is_approved_user(row): | |
| raise PermissionError("Only approved accounts can change passwords.") | |
| if not _verify_password( | |
| old_password, | |
| str(row["password_hash"]), | |
| str(row["password_salt"]), | |
| int(row["password_iterations"]), | |
| ): | |
| raise PermissionError("The current password is incorrect.") | |
| password_hash, password_salt, iterations = _hash_password(new_password) | |
| conn.execute( | |
| """ | |
| UPDATE users | |
| SET password_hash = ?, password_salt = ?, password_iterations = ?, password_alg = ?, | |
| session_version = session_version + 1 | |
| WHERE user_id = ? | |
| """, | |
| ( | |
| password_hash, | |
| password_salt, | |
| iterations, | |
| "pbkdf2_sha256", | |
| row["user_id"], | |
| ), | |
| ) | |
| conn.commit() | |
| def get_gradio_auth(): | |
| if env_truthy("AUTOML_AUTH_DISABLED", False): | |
| return None | |
| initialize_storage() | |
| if user_count() == 0 and not signups_allowed(): | |
| if env_truthy("AUTOML_REQUIRE_AUTH", False): | |
| raise StorageConfigurationError( | |
| "No AutoML users are configured. Set AUTOML_BOOTSTRAP_USERS_JSON as a Hugging Face Space Secret." | |
| ) | |
| return None | |
| return lambda username, password: authenticate_user( | |
| username, | |
| password, | |
| create_if_missing=False, | |
| ) | |
| def _session_secret() -> bytes: | |
| configured = os.getenv("AUTOML_SESSION_SECRET", "").strip() | |
| if configured: | |
| padded = configured + ("=" * ((4 - len(configured) % 4) % 4)) | |
| try: | |
| secret = base64.urlsafe_b64decode(padded.encode("ascii")) | |
| except (binascii.Error, ValueError) as exc: | |
| raise StorageConfigurationError( | |
| "AUTOML_SESSION_SECRET must be URL-safe base64 for at least 32 random bytes." | |
| ) from exc | |
| if len(secret) < 32: | |
| raise StorageConfigurationError("AUTOML_SESSION_SECRET must decode to at least 32 bytes.") | |
| return secret | |
| model_key = _decode_model_storage_key() | |
| if model_key is not None: | |
| return hmac.new(model_key, b"automl-session-v1", hashlib.sha256).digest() | |
| return _PROCESS_SESSION_SECRET | |
| def _b64url_encode(data: bytes) -> str: | |
| return base64.urlsafe_b64encode(data).decode("ascii").rstrip("=") | |
| def browser_state_secret() -> str: | |
| """Return a stable purpose-separated key for encrypted browser state.""" | |
| derived = hmac.new( | |
| _session_secret(), | |
| _BROWSER_STATE_CONTEXT, | |
| hashlib.sha256, | |
| ).digest() | |
| return _b64url_encode(derived) | |
| def _b64url_decode(value: str) -> bytes: | |
| padded = value + ("=" * ((4 - len(value) % 4) % 4)) | |
| return base64.urlsafe_b64decode(padded.encode("ascii")) | |
| def issue_session_token(username: str) -> str: | |
| user = get_user(username) | |
| if user is None or not _is_approved_user(user): | |
| raise PermissionError("Cannot create a session for an unapproved user.") | |
| payload = { | |
| "user_id": user["user_id"], | |
| "username": user["username"], | |
| "session_version": int(user.get("session_version") or 1), | |
| "iat": int(datetime.now(timezone.utc).timestamp()), | |
| "nonce": secrets.token_urlsafe(16), | |
| } | |
| payload_bytes = json.dumps(payload, separators=(",", ":"), sort_keys=True).encode("utf-8") | |
| encoded_payload = _b64url_encode(payload_bytes) | |
| signature = hmac.new(_session_secret(), encoded_payload.encode("ascii"), hashlib.sha256).digest() | |
| return f"{encoded_payload}.{_b64url_encode(signature)}" | |
| def username_from_session_token(session_token: str | None) -> str: | |
| token = str(session_token or "").strip() | |
| if not token or "." not in token: | |
| return "" | |
| encoded_payload, encoded_signature = token.rsplit(".", 1) | |
| try: | |
| expected = hmac.new(_session_secret(), encoded_payload.encode("ascii"), hashlib.sha256).digest() | |
| provided = _b64url_decode(encoded_signature) | |
| except (binascii.Error, ValueError): | |
| return "" | |
| if not hmac.compare_digest(expected, provided): | |
| return "" | |
| try: | |
| payload = json.loads(_b64url_decode(encoded_payload).decode("utf-8")) | |
| except (json.JSONDecodeError, UnicodeDecodeError, binascii.Error, ValueError): | |
| return "" | |
| issued_at = int(payload.get("iat") or 0) | |
| max_age_raw = os.getenv("AUTOML_SESSION_MAX_AGE_SECONDS", str(30 * 24 * 60 * 60)).strip() | |
| try: | |
| max_age = max(3600, int(max_age_raw)) | |
| except ValueError: | |
| max_age = 30 * 24 * 60 * 60 | |
| now = int(datetime.now(timezone.utc).timestamp()) | |
| if issued_at <= 0 or issued_at > now + 300 or now - issued_at > max_age: | |
| return "" | |
| user = get_user_by_id(payload.get("user_id")) | |
| if user is None or not _is_approved_user(user): | |
| return "" | |
| try: | |
| token_version = int(payload.get("session_version")) | |
| stored_version = int(user.get("session_version") or 1) | |
| except (TypeError, ValueError): | |
| return "" | |
| if token_version != stored_version: | |
| return "" | |
| token_username = normalize_username(payload.get("username")) | |
| if _username_key(token_username) != _username_key(user.get("username")): | |
| return "" | |
| return str(user["username"]) | |
| def current_username(request: Any | None = None, session_token: str | None = None) -> str: | |
| request_username = normalize_username(getattr(request, "username", None)) | |
| if request_username: | |
| return request_username | |
| return username_from_session_token(session_token) | |
| def require_authenticated_username(request: Any | None = None, session_token: str | None = None) -> str: | |
| username = current_username(request, session_token) | |
| if not username: | |
| raise PermissionError("Sign in before using saved cloud models.") | |
| return username | |
| def _decode_model_storage_key() -> bytes | None: | |
| raw = os.getenv("AUTOML_MODEL_STORAGE_KEY", "").strip() | |
| if not raw: | |
| return None | |
| padded = raw + ("=" * ((4 - len(raw) % 4) % 4)) | |
| try: | |
| key = base64.urlsafe_b64decode(padded.encode("ascii")) | |
| except (binascii.Error, ValueError) as exc: | |
| raise StorageConfigurationError( | |
| "AUTOML_MODEL_STORAGE_KEY must be URL-safe base64 for 32 random bytes." | |
| ) from exc | |
| if len(key) != 32: | |
| raise StorageConfigurationError("AUTOML_MODEL_STORAGE_KEY must decode to exactly 32 bytes.") | |
| return key | |
| def encryption_enabled() -> bool: | |
| return _decode_model_storage_key() is not None | |
| def _model_storage_key_id() -> str: | |
| key = _decode_model_storage_key() | |
| if key is None: | |
| return "" | |
| return hashlib.sha256(key).hexdigest()[:16] | |
| def _chunk_aad(header: bytes, index: int) -> bytes: | |
| return ENCRYPTED_MODEL_MAGIC + header + index.to_bytes(8, "big") | |
| def _encrypt_file(source: Path, destination: Path, *, model_id: str) -> Tuple[int, str]: | |
| key = _decode_model_storage_key() | |
| if key is None: | |
| raise StorageConfigurationError("Set AUTOML_MODEL_STORAGE_KEY before saving encrypted models.") | |
| try: | |
| from cryptography.hazmat.primitives.ciphers.aead import AESGCM | |
| except ImportError as exc: | |
| raise StorageConfigurationError("Install the cryptography package to save encrypted models.") from exc | |
| destination.parent.mkdir(parents=True, exist_ok=True) | |
| nonce_prefix = secrets.token_bytes(8) | |
| header_obj = { | |
| "alg": "AES-256-GCM-CHUNKED", | |
| "chunk_size": DEFAULT_MODEL_CHUNK_SIZE, | |
| "nonce_prefix": base64.b64encode(nonce_prefix).decode("ascii"), | |
| "key_id": _model_storage_key_id(), | |
| "model_id": model_id, | |
| } | |
| header = json.dumps(header_obj, separators=(",", ":")).encode("utf-8") | |
| aesgcm = AESGCM(key) | |
| sha = hashlib.sha256() | |
| total = 0 | |
| index = 0 | |
| with source.open("rb") as src, destination.open("wb") as dst: | |
| dst.write(ENCRYPTED_MODEL_MAGIC) | |
| dst.write(len(header).to_bytes(4, "big")) | |
| dst.write(header) | |
| while True: | |
| chunk = src.read(DEFAULT_MODEL_CHUNK_SIZE) | |
| if not chunk: | |
| break | |
| sha.update(chunk) | |
| total += len(chunk) | |
| nonce = nonce_prefix + index.to_bytes(4, "big") | |
| encrypted = aesgcm.encrypt(nonce, chunk, _chunk_aad(header, index)) | |
| dst.write(len(encrypted).to_bytes(4, "big")) | |
| dst.write(encrypted) | |
| index += 1 | |
| return total, sha.hexdigest() | |
| def _decrypt_file(source: Path, destination: Path, *, expected_size: int, expected_sha256: str) -> None: | |
| key = _decode_model_storage_key() | |
| if key is None: | |
| raise StorageConfigurationError("Set AUTOML_MODEL_STORAGE_KEY before loading encrypted saved models.") | |
| try: | |
| from cryptography.hazmat.primitives.ciphers.aead import AESGCM | |
| except ImportError as exc: | |
| raise StorageConfigurationError("Install the cryptography package to load encrypted models.") from exc | |
| aesgcm = AESGCM(key) | |
| sha = hashlib.sha256() | |
| total = 0 | |
| index = 0 | |
| destination.parent.mkdir(parents=True, exist_ok=True) | |
| try: | |
| with source.open("rb") as src, destination.open("wb") as dst: | |
| magic = src.read(len(ENCRYPTED_MODEL_MAGIC)) | |
| if magic != ENCRYPTED_MODEL_MAGIC: | |
| raise StorageConfigurationError("Saved model file is not in the expected encrypted format.") | |
| header_len_raw = src.read(4) | |
| if len(header_len_raw) != 4: | |
| raise StorageConfigurationError("Saved model encryption header is incomplete.") | |
| header_len = int.from_bytes(header_len_raw, "big") | |
| header = src.read(header_len) | |
| if len(header) != header_len: | |
| raise StorageConfigurationError("Saved model encryption header is truncated.") | |
| header_obj = json.loads(header.decode("utf-8")) | |
| nonce_prefix = base64.b64decode(header_obj["nonce_prefix"].encode("ascii"), validate=True) | |
| while True: | |
| length_raw = src.read(4) | |
| if not length_raw: | |
| break | |
| if len(length_raw) != 4: | |
| raise StorageConfigurationError("Saved model ciphertext length is incomplete.") | |
| encrypted_len = int.from_bytes(length_raw, "big") | |
| encrypted = src.read(encrypted_len) | |
| if len(encrypted) != encrypted_len: | |
| raise StorageConfigurationError("Saved model ciphertext is truncated.") | |
| nonce = nonce_prefix + index.to_bytes(4, "big") | |
| chunk = aesgcm.decrypt(nonce, encrypted, _chunk_aad(header, index)) | |
| dst.write(chunk) | |
| sha.update(chunk) | |
| total += len(chunk) | |
| index += 1 | |
| except Exception: | |
| if destination.exists(): | |
| destination.unlink(missing_ok=True) | |
| raise | |
| if total != int(expected_size) or sha.hexdigest() != expected_sha256: | |
| destination.unlink(missing_ok=True) | |
| raise StorageConfigurationError("Saved model integrity check failed after decryption.") | |
| def _copy_plaintext_file(source: Path, destination: Path) -> Tuple[int, str]: | |
| destination.parent.mkdir(parents=True, exist_ok=True) | |
| shutil.copy2(source, destination) | |
| sha = hashlib.sha256() | |
| total = 0 | |
| with source.open("rb") as src: | |
| while True: | |
| chunk = src.read(DEFAULT_MODEL_CHUNK_SIZE) | |
| if not chunk: | |
| break | |
| sha.update(chunk) | |
| total += len(chunk) | |
| return total, sha.hexdigest() | |
| def _model_dir_for_user(user_id: str) -> Path: | |
| return storage_root() / "models" / user_id | |
| def save_model_for_user( | |
| username: str, | |
| package_path: str, | |
| *, | |
| display_name: str, | |
| manifest: Dict[str, Any], | |
| artifact_type: str, | |
| ) -> Dict[str, Any]: | |
| initialize_storage() | |
| user = get_user(username) | |
| if user is None or not _is_approved_user(user): | |
| raise PermissionError("The signed-in user was not found in the account database.") | |
| source = Path(package_path) | |
| if not source.is_file(): | |
| raise FileNotFoundError(f"Model package was not found: {package_path}") | |
| model_id = secrets.token_urlsafe(12) | |
| encrypted = encryption_enabled() | |
| if encrypted: | |
| destination = _model_dir_for_user(user["user_id"]) / f"{model_id}.zip.enc" | |
| byte_size, digest = _encrypt_file(source, destination, model_id=model_id) | |
| key_id = _model_storage_key_id() | |
| elif env_truthy("AUTOML_ALLOW_UNENCRYPTED_MODEL_STORAGE", False): | |
| destination = _model_dir_for_user(user["user_id"]) / f"{model_id}.zip" | |
| byte_size, digest = _copy_plaintext_file(source, destination) | |
| key_id = "" | |
| else: | |
| raise StorageConfigurationError( | |
| "Encrypted model storage is not configured. Set AUTOML_MODEL_STORAGE_KEY as a Space Secret." | |
| ) | |
| now = _utcnow() | |
| manifest_json = json.dumps(manifest or {}, sort_keys=True) | |
| with _DB_LOCK: | |
| with _connect() as conn: | |
| conn.execute( | |
| """ | |
| INSERT INTO saved_models ( | |
| model_id, user_id, display_name, artifact_type, predictor, task_type, target_column, | |
| manifest_json, storage_path, storage_encrypted, storage_key_id, source_package_name, | |
| byte_size, sha256, created_at | |
| ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) | |
| """, | |
| ( | |
| model_id, | |
| user["user_id"], | |
| display_name[:120] or "AutoML model", | |
| artifact_type or "", | |
| str(manifest.get("predictor") or ""), | |
| str(manifest.get("task_type") or ""), | |
| str(manifest.get("target_column") or manifest.get("label_column") or ""), | |
| manifest_json, | |
| str(destination), | |
| 1 if encrypted else 0, | |
| key_id, | |
| source.name, | |
| int(byte_size), | |
| digest, | |
| now, | |
| ), | |
| ) | |
| conn.commit() | |
| return get_saved_model(username, model_id) or {"model_id": model_id, "display_name": display_name} | |
| def list_saved_models(username: str | None) -> List[Dict[str, Any]]: | |
| initialize_storage() | |
| user = get_user(username) | |
| if user is None or not _is_approved_user(user): | |
| return [] | |
| with _DB_LOCK: | |
| with _connect() as conn: | |
| rows = conn.execute( | |
| """ | |
| SELECT * FROM saved_models | |
| WHERE user_id = ? | |
| ORDER BY created_at DESC, display_name ASC | |
| """, | |
| (user["user_id"],), | |
| ).fetchall() | |
| return [dict(row) for row in rows] | |
| def get_saved_model(username: str | None, model_id: str | None) -> Dict[str, Any] | None: | |
| initialize_storage() | |
| user = get_user(username) | |
| if user is None or not _is_approved_user(user) or not model_id: | |
| return None | |
| with _DB_LOCK: | |
| with _connect() as conn: | |
| row = conn.execute( | |
| """ | |
| SELECT * FROM saved_models | |
| WHERE user_id = ? AND model_id = ? | |
| """, | |
| (user["user_id"], str(model_id)), | |
| ).fetchone() | |
| return dict(row) if row else None | |
| def manifest_from_record(record: Dict[str, Any]) -> Dict[str, Any]: | |
| try: | |
| loaded = json.loads(record.get("manifest_json") or "{}") | |
| return loaded if isinstance(loaded, dict) else {} | |
| except json.JSONDecodeError: | |
| return {} | |
| def materialize_saved_model(username: str, model_id: str) -> Tuple[str, str]: | |
| record = get_saved_model(username, model_id) | |
| if record is None: | |
| raise PermissionError("Saved model not found for this account.") | |
| source = Path(record["storage_path"]) | |
| if not source.is_file(): | |
| raise FileNotFoundError("The saved model file is missing from storage.") | |
| temp_dir = Path(tempfile.mkdtemp(prefix="automl_saved_model_")) | |
| destination = temp_dir / f"{model_id}.zip" | |
| if int(record.get("storage_encrypted") or 0): | |
| _decrypt_file( | |
| source, | |
| destination, | |
| expected_size=int(record["byte_size"]), | |
| expected_sha256=str(record["sha256"]), | |
| ) | |
| else: | |
| if not env_truthy("AUTOML_ALLOW_UNENCRYPTED_MODEL_STORAGE", False): | |
| raise StorageConfigurationError("Plaintext saved model storage is disabled.") | |
| shutil.copy2(source, destination) | |
| with _DB_LOCK: | |
| with _connect() as conn: | |
| conn.execute( | |
| "UPDATE saved_models SET last_used_at = ? WHERE model_id = ?", | |
| (_utcnow(), model_id), | |
| ) | |
| conn.commit() | |
| return str(destination), str(temp_dir) | |
| _SCHEDULED_JOB_UPDATE_FIELDS = frozenset( | |
| { | |
| "provider", | |
| "external_job_id", | |
| "status", | |
| "result_json", | |
| "result_location", | |
| "remote_result_location", | |
| "package_sha256", | |
| "routing_reason", | |
| "cancellation_requested", | |
| "cleanup_pending", | |
| "submitted_at", | |
| "started_at", | |
| "finished_at", | |
| "cancelled_at", | |
| } | |
| ) | |
| _ACTIVE_SCHEDULED_STATUSES = ( | |
| "queued", | |
| "submitting", | |
| "pending", | |
| "running", | |
| "importing", | |
| "cancelling", | |
| "cleanup_pending", | |
| ) | |
| def _scheduled_job_dict(row: sqlite3.Row | None) -> Dict[str, Any] | None: | |
| if row is None: | |
| return None | |
| value = dict(row) | |
| value["cancellation_requested"] = bool( | |
| value.get("cancellation_requested", 0) | |
| ) | |
| value["cleanup_pending"] = bool(value.get("cleanup_pending", 0)) | |
| for key in ("request_json", "result_json"): | |
| raw = value.get(key) | |
| if not raw: | |
| value[key.removesuffix("_json") + "_payload"] = None | |
| continue | |
| try: | |
| parsed = json.loads(raw) | |
| except json.JSONDecodeError: | |
| parsed = None | |
| value[key.removesuffix("_json") + "_payload"] = parsed | |
| return value | |
| def create_or_get_scheduled_job( | |
| username: str, | |
| *, | |
| request_id: str, | |
| idempotency_key: str, | |
| profile: str, | |
| request_payload: Dict[str, Any], | |
| display_name: str, | |
| ) -> Dict[str, Any]: | |
| """Persist one scheduler request, returning the prior idempotent row.""" | |
| initialize_storage() | |
| user = get_user(username) | |
| if user is None or not _is_approved_user(user): | |
| raise PermissionError("The signed-in user was not found in the account database.") | |
| normalized_request = str(request_id or "").strip() | |
| normalized_key = str(idempotency_key or "").strip() | |
| normalized_profile = str(profile or "").strip().lower() | |
| if not normalized_request or not normalized_key: | |
| raise ValueError("Scheduler request and idempotency IDs are required.") | |
| if normalized_profile not in {"cpu", "gpu"}: | |
| raise ValueError("Scheduler profile must be cpu or gpu.") | |
| encoded_request = json.dumps( | |
| dict(request_payload), separators=(",", ":"), sort_keys=True | |
| ) | |
| now = _utcnow() | |
| with _DB_LOCK: | |
| with _connect() as conn: | |
| existing = conn.execute( | |
| """ | |
| SELECT jobs.*, users.username | |
| FROM scheduled_jobs AS jobs | |
| JOIN users ON users.user_id = jobs.user_id | |
| WHERE jobs.user_id = ? AND jobs.idempotency_key = ? | |
| """, | |
| (user["user_id"], normalized_key), | |
| ).fetchone() | |
| if existing is not None: | |
| if ( | |
| str(existing["profile"]) != normalized_profile | |
| or str(existing["request_json"]) != encoded_request | |
| ): | |
| raise ValueError( | |
| "The idempotency key is already bound to a different scheduler request." | |
| ) | |
| return _scheduled_job_dict(existing) or {} | |
| conn.execute( | |
| """ | |
| INSERT INTO scheduled_jobs ( | |
| request_id, user_id, idempotency_key, profile, provider, | |
| status, request_json, display_name, queued_at, updated_at | |
| ) VALUES (?, ?, ?, ?, 'local', 'queued', ?, ?, ?, ?) | |
| """, | |
| ( | |
| normalized_request, | |
| user["user_id"], | |
| normalized_key, | |
| normalized_profile, | |
| encoded_request, | |
| str(display_name or "AutoML model")[:120], | |
| now, | |
| now, | |
| ), | |
| ) | |
| conn.commit() | |
| return get_scheduled_job(username, normalized_request) or {} | |
| def get_scheduled_job( | |
| username: str | None, request_id: str | None | |
| ) -> Dict[str, Any] | None: | |
| initialize_storage() | |
| user = get_user(username) | |
| if user is None or not _is_approved_user(user) or not request_id: | |
| return None | |
| with _DB_LOCK: | |
| with _connect() as conn: | |
| row = conn.execute( | |
| """ | |
| SELECT jobs.*, users.username | |
| FROM scheduled_jobs AS jobs | |
| JOIN users ON users.user_id = jobs.user_id | |
| WHERE jobs.user_id = ? AND jobs.request_id = ? | |
| """, | |
| (user["user_id"], str(request_id)), | |
| ).fetchone() | |
| return _scheduled_job_dict(row) | |
| def update_scheduled_job(request_id: str, **changes: Any) -> Dict[str, Any]: | |
| """Update only the persisted provider lifecycle fields.""" | |
| initialize_storage() | |
| unknown = set(changes) - _SCHEDULED_JOB_UPDATE_FIELDS | |
| if unknown: | |
| raise ValueError(f"Unsupported scheduled job fields: {sorted(unknown)}") | |
| if not changes: | |
| raise ValueError("At least one scheduled job field must be updated.") | |
| assignments: List[str] = [] | |
| values: List[Any] = [] | |
| for key, value in changes.items(): | |
| column = key | |
| if key == "result_json" and isinstance(value, (dict, list)): | |
| value = json.dumps(value, separators=(",", ":"), sort_keys=True) | |
| if key == "cancellation_requested": | |
| value = 1 if value else 0 | |
| assignments.append(f"{column} = ?") | |
| values.append(value) | |
| assignments.append("updated_at = ?") | |
| values.extend([_utcnow(), str(request_id)]) | |
| with _DB_LOCK: | |
| with _connect() as conn: | |
| cursor = conn.execute( | |
| f"UPDATE scheduled_jobs SET {', '.join(assignments)} WHERE request_id = ?", | |
| values, | |
| ) | |
| if cursor.rowcount != 1: | |
| raise KeyError("Scheduled job was not found.") | |
| row = conn.execute( | |
| """ | |
| SELECT jobs.*, users.username | |
| FROM scheduled_jobs AS jobs | |
| JOIN users ON users.user_id = jobs.user_id | |
| WHERE jobs.request_id = ? | |
| """, | |
| (str(request_id),), | |
| ).fetchone() | |
| conn.commit() | |
| return _scheduled_job_dict(row) or {} | |
| def list_active_scheduled_jobs() -> List[Dict[str, Any]]: | |
| initialize_storage() | |
| placeholders = ",".join("?" for _ in _ACTIVE_SCHEDULED_STATUSES) | |
| with _DB_LOCK: | |
| with _connect() as conn: | |
| rows = conn.execute( | |
| f""" | |
| SELECT jobs.*, users.username | |
| FROM scheduled_jobs AS jobs | |
| JOIN users ON users.user_id = jobs.user_id | |
| WHERE jobs.status IN ({placeholders}) | |
| ORDER BY jobs.queued_at, jobs.request_id | |
| """, | |
| _ACTIVE_SCHEDULED_STATUSES, | |
| ).fetchall() | |
| return [value for row in rows if (value := _scheduled_job_dict(row))] | |
| def request_scheduled_job_cancellation( | |
| username: str, request_id: str | |
| ) -> Dict[str, Any]: | |
| record = get_scheduled_job(username, request_id) | |
| if record is None: | |
| raise PermissionError("Scheduled job was not found for this account.") | |
| return update_scheduled_job( | |
| request_id, | |
| cancellation_requested=True, | |
| status=( | |
| "cancelling" | |
| if record.get("status") in {"pending", "running", "submitting"} | |
| else record.get("status") | |
| ), | |
| ) | |
| def delete_saved_model(username: str, model_id: str) -> None: | |
| record = get_saved_model(username, model_id) | |
| if record is None: | |
| raise PermissionError("Saved model not found for this account.") | |
| storage_path = Path(record["storage_path"]) | |
| with _DB_LOCK: | |
| with _connect() as conn: | |
| conn.execute("DELETE FROM saved_models WHERE model_id = ?", (model_id,)) | |
| conn.commit() | |
| if storage_path.is_file(): | |
| storage_path.unlink() | |