Spaces:
Running
Running
| """Async database engine — PostgreSQL via asyncpg.""" | |
| from sqlalchemy import text | |
| from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker, create_async_engine | |
| from sqlalchemy.orm import DeclarativeBase | |
| from app.config import settings | |
| class Base(DeclarativeBase): | |
| pass | |
| engine = create_async_engine( | |
| settings.DATABASE_URL, | |
| echo=False, | |
| pool_size=10, | |
| max_overflow=20, | |
| pool_timeout=settings.DB_POOL_TIMEOUT, | |
| connect_args={ | |
| "timeout": settings.DB_CONNECT_TIMEOUT, | |
| "command_timeout": settings.DB_COMMAND_TIMEOUT, | |
| }, | |
| ) | |
| async_session = async_sessionmaker(engine, class_=AsyncSession, expire_on_commit=False) | |
| async def create_database_if_not_exists(): | |
| """Create data_acq database if missing.""" | |
| import asyncpg | |
| from urllib.parse import urlparse | |
| u = urlparse(settings.DATABASE_URL.replace("+asyncpg", "")) | |
| admin_url = f"postgresql://{u.username}:{u.password}@{u.hostname}:{u.port or 5432}/postgres" | |
| try: | |
| conn = await asyncpg.connect(admin_url) | |
| except Exception as exc: | |
| # Some deployments expose only the application database through the | |
| # SSH tunnel and reject connections to the postgres maintenance DB. | |
| # The application database is already provisioned in that setup, so | |
| # let init_db() continue and verify it through the main engine. | |
| print(f"[DB] Skipping maintenance-database check: {exc}") | |
| return | |
| try: | |
| exists = await conn.fetchval("SELECT 1 FROM pg_database WHERE datname=$1", u.path.lstrip("/")) | |
| if not exists: | |
| await conn.execute(f'CREATE DATABASE "{u.path.lstrip("/")}"') | |
| print(f"[DB] Created database: {u.path.lstrip('/')}") | |
| finally: | |
| await conn.close() | |
| async def init_db(): | |
| # Import model classes before create_all(); otherwise SQLAlchemy's metadata | |
| # is empty and a fresh EGOData database receives no tables. | |
| from app import models # noqa: F401 | |
| await create_database_if_not_exists() | |
| async with engine.begin() as conn: | |
| await conn.run_sync(Base.metadata.create_all) | |
| # Ensure new columns exist on existing tables (create_all won't add columns) | |
| await _migrate_episode_columns() | |
| await _migrate_task_definition_columns() | |
| await _migrate_device_columns() | |
| await _migrate_workflow_run_columns() | |
| await _migrate_session_columns() | |
| await _migrate_user_columns() | |
| await _seed_default_users() | |
| print("[DB] Tables created/verified") | |
| async def _migrate_episode_columns(): | |
| """Add new pipeline columns to episodes table if they don't exist yet.""" | |
| new_columns = [ | |
| ("received_at", "TIMESTAMP WITH TIME ZONE"), | |
| ("processing_started_at", "TIMESTAMP WITH TIME ZONE"), | |
| ("review_ready_at", "TIMESTAMP WITH TIME ZONE"), | |
| ("approved_at", "TIMESTAMP WITH TIME ZONE"), | |
| ("rejected_at", "TIMESTAMP WITH TIME ZONE"), | |
| ("cleaning_report", "JSONB"), | |
| ] | |
| async with engine.begin() as conn: | |
| for col_name, col_type in new_columns: | |
| try: | |
| await conn.execute( | |
| text(f"ALTER TABLE episodes ADD COLUMN IF NOT EXISTS {col_name} {col_type}") | |
| ) | |
| except Exception: | |
| pass # column already exists or DB doesn't support IF NOT EXISTS | |
| async def _migrate_task_definition_columns(): | |
| """Add new columns to task_definitions table.""" | |
| new_columns = [ | |
| ("claimer", "VARCHAR(128)"), | |
| ("params", "JSONB"), | |
| ] | |
| async with engine.begin() as conn: | |
| for col_name, col_type in new_columns: | |
| try: | |
| await conn.execute( | |
| text(f"ALTER TABLE task_definitions ADD COLUMN IF NOT EXISTS {col_name} {col_type}") | |
| ) | |
| except Exception: | |
| pass | |
| async def _migrate_device_columns(): | |
| """Add columns to devices table if they don't exist yet.""" | |
| new_columns = [ | |
| ("meta", "JSONB"), | |
| ("first_seen_at", "TIMESTAMP WITH TIME ZONE"), | |
| ] | |
| async with engine.begin() as conn: | |
| for col_name, col_type in new_columns: | |
| try: | |
| await conn.execute( | |
| text(f"ALTER TABLE devices ADD COLUMN IF NOT EXISTS {col_name} {col_type}") | |
| ) | |
| except Exception: | |
| pass | |
| async def _migrate_workflow_run_columns(): | |
| """Add worker queue columns to existing workflow_runs tables.""" | |
| new_columns = [ | |
| ("episode_id", "UUID"), | |
| ("worker_id", "VARCHAR(128)"), | |
| ("lease_until", "TIMESTAMP WITH TIME ZONE"), | |
| ("heartbeat_at", "TIMESTAMP WITH TIME ZONE"), | |
| ("attempt", "INTEGER NOT NULL DEFAULT 0"), | |
| ("progress", "DOUBLE PRECISION NOT NULL DEFAULT 0"), | |
| ("outputs", "JSONB NOT NULL DEFAULT '{}'::jsonb"), | |
| ] | |
| async with engine.begin() as conn: | |
| for col_name, col_type in new_columns: | |
| try: | |
| await conn.execute(text( | |
| f"ALTER TABLE workflow_runs ADD COLUMN IF NOT EXISTS {col_name} {col_type}" | |
| )) | |
| except Exception: | |
| pass | |
| async def _migrate_session_columns(): | |
| """Add project binding column to sessions table.""" | |
| async with engine.begin() as conn: | |
| try: | |
| await conn.execute(text( | |
| "ALTER TABLE sessions ADD COLUMN IF NOT EXISTS project_id UUID" | |
| )) | |
| except Exception: | |
| pass | |
| async def _migrate_user_columns(): | |
| """Add account lifecycle fields to existing EGOData installations.""" | |
| new_columns = [ | |
| ("expires_at", "TIMESTAMP WITH TIME ZONE"), | |
| ("last_login_at", "TIMESTAMP WITH TIME ZONE"), | |
| ] | |
| async with engine.begin() as conn: | |
| for col_name, col_type in new_columns: | |
| try: | |
| await conn.execute(text( | |
| f"ALTER TABLE users ADD COLUMN IF NOT EXISTS {col_name} {col_type}" | |
| )) | |
| except Exception: | |
| pass | |
| async def get_session() -> AsyncSession: | |
| async with async_session() as session: | |
| yield session | |
| async def _seed_default_users(): | |
| """Insert the configurable bootstrap administrator if the table is empty. | |
| The public template uses Root/Root for a first local demo only. Existing | |
| databases are never modified because seeding stops when any user exists. | |
| Production deployments should set both bootstrap environment variables to | |
| private values before the first startup. | |
| """ | |
| import os | |
| from app.models import User | |
| from app.auth import hash_password | |
| from sqlalchemy import select, func | |
| async with async_session() as db: | |
| cnt = (await db.execute(select(func.count(User.id)))).scalar() or 0 | |
| if cnt > 0: | |
| return # Already seeded | |
| username = (os.environ.get("EGODATA_BOOTSTRAP_USERNAME", "Root").strip() | |
| or "Root") | |
| password = os.environ.get("EGODATA_BOOTSTRAP_PASSWORD", "Root") | |
| email = (os.environ.get("EGODATA_BOOTSTRAP_EMAIL", "").strip() | |
| or f"{username.lower()}@egodata.local") | |
| defaults = [User( | |
| username=username, | |
| password_hash=hash_password(password), | |
| email=email, | |
| role="admin", | |
| status="active", | |
| )] | |
| db.add_all(defaults) | |
| await db.commit() | |
| print(f"[DB] Seeded {len(defaults)} default users") | |