from typing import AsyncGenerator from sqlalchemy import text from sqlalchemy.ext.asyncio import AsyncSession, create_async_engine, async_sessionmaker from sqlalchemy.orm import declarative_base from app.core.config import settings from app.core.logging import get_logger logger = get_logger(__name__) DB_POOL_CONFIG = { "pool_pre_ping": True, "pool_recycle": 1800, "pool_size": 10, "max_overflow": 20, "pool_timeout": 30, } engine = create_async_engine( settings.DATABASE_URL, echo=settings.DEBUG if hasattr(settings, "DEBUG") else False, **DB_POOL_CONFIG, ) async_session_factory = async_sessionmaker(engine, class_=AsyncSession, expire_on_commit=False) Base = declarative_base() async def get_db() -> AsyncGenerator[AsyncSession, None]: async with async_session_factory() as session: try: yield session await session.commit() except Exception: await session.rollback() raise async def seed_default_datasources(session: AsyncSession): from app.core.datasource_defaults import DEFAULT_DATASOURCES from app.models.datasource import DataSource for source, info in DEFAULT_DATASOURCES.items(): existing = await session.get(DataSource, info["id"]) if existing: existing.name = info["name"] existing.source = source existing.module = info["module"] existing.priority = info["priority"] existing.frequency_minutes = info["frequency_minutes"] existing.collector_class = source if existing.config is None: existing.config = "{}" continue session.add( DataSource( id=info["id"], name=info["name"], source=source, module=info["module"], priority=info["priority"], frequency_minutes=info["frequency_minutes"], collector_class=source, config="{}", is_active=True, ) ) await session.commit() async def ensure_default_admin_user(session: AsyncSession): from app.core.security import get_password_hash from app.models.user import User result = await session.execute( text("SELECT id FROM users WHERE username = 'admin'") ) if result.fetchone(): return session.add( User( username="admin", email="admin@planet.local", password_hash=get_password_hash("admin123"), role="super_admin", is_active=True, ) ) await session.commit() async def init_db(): import app.models.user # noqa: F401 import app.models.gpu_cluster # noqa: F401 import app.models.task # noqa: F401 import app.models.data_snapshot # noqa: F401 import app.models.datasource # noqa: F401 import app.models.datasource_config # noqa: F401 import app.models.alert # noqa: F401 import app.models.bgp_anomaly # noqa: F401 import app.models.bgp_incident # noqa: F401 import app.models.bgp_observation # noqa: F401 import app.models.collected_data # noqa: F401 import app.models.system_setting # noqa: F401 import app.models.playground_session # noqa: F401 import app.models.playground_message # noqa: F401 import app.models.system_log # noqa: F401 import app.models.vessel # noqa: F401 import app.models.datasource_mapping # noqa: F401 logger.warning_event( "Database pool settings active", event="database.pool.initialized", context={ "pool_pre_ping": DB_POOL_CONFIG["pool_pre_ping"], "pool_recycle": DB_POOL_CONFIG["pool_recycle"], "pool_size": DB_POOL_CONFIG["pool_size"], "max_overflow": DB_POOL_CONFIG["max_overflow"], "pool_timeout": DB_POOL_CONFIG["pool_timeout"], }, ) async with engine.begin() as conn: await conn.run_sync(Base.metadata.create_all) await conn.execute( text( """ ALTER TABLE collected_data ADD COLUMN IF NOT EXISTS snapshot_id INTEGER, ADD COLUMN IF NOT EXISTS task_id INTEGER, ADD COLUMN IF NOT EXISTS entity_key VARCHAR(255), ADD COLUMN IF NOT EXISTS is_current BOOLEAN DEFAULT TRUE, ADD COLUMN IF NOT EXISTS previous_record_id INTEGER, ADD COLUMN IF NOT EXISTS change_type VARCHAR(20), ADD COLUMN IF NOT EXISTS change_summary JSONB DEFAULT '{}'::jsonb, ADD COLUMN IF NOT EXISTS deleted_at TIMESTAMPTZ """ ) ) await conn.execute( text( """ ALTER TABLE collection_tasks ADD COLUMN IF NOT EXISTS phase VARCHAR(30) DEFAULT 'queued', ADD COLUMN IF NOT EXISTS phase_progress DOUBLE PRECISION, ADD COLUMN IF NOT EXISTS phase_message VARCHAR(255), ADD COLUMN IF NOT EXISTS phase_current BIGINT, ADD COLUMN IF NOT EXISTS phase_total BIGINT, ADD COLUMN IF NOT EXISTS phase_unit VARCHAR(30) """ ) ) await conn.execute( text( """ CREATE INDEX IF NOT EXISTS idx_collected_data_source_source_id ON collected_data (source, source_id) """ ) ) await conn.execute( text( """ UPDATE collected_data SET entity_key = source || ':' || COALESCE(source_id, id::text) WHERE entity_key IS NULL """ ) ) await conn.execute( text( """ UPDATE collected_data SET is_current = TRUE WHERE is_current IS NULL """ ) ) async with async_session_factory() as session: await seed_default_datasources(session) await ensure_default_admin_user(session)