Spaces:
Running
Running
File size: 7,438 Bytes
d646dd7 95f5290 d646dd7 95f5290 d646dd7 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 | """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")
|