from typing import AsyncGenerator from sqlalchemy import bindparam, 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() LEGACY_EARTH_BOUNDARY_SOURCES = ( "earth_admin0_boundaries", "earth_coastline", "earth_claim_lines", "earth_boundary_tiles", ) LEGACY_EARTH_BOUNDARY_DATATYPES = ( "earth_boundary_source", "earth_boundary_tiles", ) LEGACY_EARTH_BOUNDARY_IDS = (29, 30, 31, 32) async def purge_legacy_earth_boundary_datasources(session: AsyncSession) -> None: source_names = tuple(LEGACY_EARTH_BOUNDARY_SOURCES) source_ids = tuple(LEGACY_EARTH_BOUNDARY_IDS) data_types = tuple(LEGACY_EARTH_BOUNDARY_DATATYPES) await session.execute( text( """ DELETE FROM datasource_mapping_templates WHERE target_schema IN :data_types OR datasource_config_id IN ( SELECT id FROM datasource_configs WHERE name IN :source_names ) """ ).bindparams(bindparam("source_names", expanding=True), bindparam("data_types", expanding=True)), {"source_names": list(source_names), "data_types": list(data_types)}, ) await session.execute( text("DELETE FROM datasource_configs WHERE name IN :source_names").bindparams( bindparam("source_names", expanding=True) ), {"source_names": list(source_names)}, ) await session.execute( text( """ DELETE FROM collected_data WHERE source IN :source_names OR data_type IN :data_types """ ).bindparams(bindparam("source_names", expanding=True), bindparam("data_types", expanding=True)), {"source_names": list(source_names), "data_types": list(data_types)}, ) await session.execute( text( """ DELETE FROM data_snapshots WHERE source IN :source_names OR datasource_id IN :source_ids """ ).bindparams(bindparam("source_names", expanding=True), bindparam("source_ids", expanding=True)), {"source_names": list(source_names), "source_ids": list(source_ids)}, ) await session.execute( text("DELETE FROM collection_tasks WHERE datasource_id IN :source_ids").bindparams( bindparam("source_ids", expanding=True) ), {"source_ids": list(source_ids)}, ) await session.execute( text("DELETE FROM data_sources WHERE source IN :source_names OR id IN :source_ids").bindparams( bindparam("source_names", expanding=True), bindparam("source_ids", expanding=True) ), {"source_names": list(source_names), "source_ids": list(source_ids)}, ) await session.commit() DEFAULT_LOGIN_USERS = ( { "username": "admin", "email": "admin@planet.local", "password": "admin123", "role": "super_admin", }, { "username": "linkong", "email": "linkong@planet.local", "password": "12345678", "role": "super_admin", }, ) async def ensure_default_admin_user(session: AsyncSession): from app.core.security import get_password_hash from app.models.user import User for default_user in DEFAULT_LOGIN_USERS: result = await session.execute( text("SELECT id FROM users WHERE username = :username"), {"username": default_user["username"]}, ) if result.fetchone(): continue session.add( User( username=default_user["username"], email=default_user["email"], password_hash=get_password_hash(default_user["password"]), role=default_user["role"], is_active=True, email_verified=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_collector_location # 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.compute_center_location # 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.vessel_enrichment # noqa: F401 import app.models.datasource_mapping # noqa: F401 import app.models.earth_news # 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) users_email_verified_existed = ( await conn.execute( text( """ SELECT 1 FROM information_schema.columns WHERE table_name = 'users' AND column_name = 'email_verified' """ ) ) ).fetchone() is not None await conn.execute( text( """ ALTER TABLE users ADD COLUMN IF NOT EXISTS gatekeeper_groups JSONB DEFAULT '[]'::jsonb, ADD COLUMN IF NOT EXISTS email_verified BOOLEAN NOT NULL DEFAULT FALSE, ADD COLUMN IF NOT EXISTS pending_email VARCHAR(255) """ ) ) if not users_email_verified_existed: await conn.execute( text("UPDATE users SET email_verified = TRUE WHERE email_verified = FALSE") ) 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( """ ALTER TABLE earth_news_items ADD COLUMN IF NOT EXISTS content_language VARCHAR(32) NOT NULL DEFAULT 'en', ADD COLUMN IF NOT EXISTS localizations JSONB NOT NULL DEFAULT '{}'::jsonb, ADD COLUMN IF NOT EXISTS enrichment_status VARCHAR(80) NOT NULL DEFAULT 'pending', ADD COLUMN IF NOT EXISTS enrichment_error TEXT, ADD COLUMN IF NOT EXISTS enriched_at TIMESTAMPTZ """ ) ) 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( """ CREATE INDEX IF NOT EXISTS idx_earth_news_enrichment_status ON earth_news_items (enrichment_status) """ ) ) await conn.execute( text( """ CREATE INDEX IF NOT EXISTS idx_earth_news_enriched_at ON earth_news_items (enriched_at) """ ) ) await conn.execute( text( """ CREATE INDEX IF NOT EXISTS idx_collected_data_source_current_id ON collected_data (source, is_current, id) """ ) ) await conn.execute( text( """ CREATE INDEX IF NOT EXISTS idx_collected_data_source_task_id ON collected_data (source, task_id, id) """ ) ) await conn.execute( text( """ CREATE INDEX IF NOT EXISTS idx_ais_raw_schema_observed_entity ON ais_raw_observations (target_schema, observed_at, entity_key) """ ) ) await conn.execute( text( """ CREATE INDEX IF NOT EXISTS idx_ais_raw_schema_observed_desc ON ais_raw_observations (target_schema, observed_at DESC) """ ) ) await conn.execute( text( """ CREATE INDEX IF NOT EXISTS idx_ais_raw_payload_lon_lat ON ais_raw_observations ( ((normalized_payload->>'lon')::double precision), ((normalized_payload->>'lat')::double precision) ) WHERE target_schema = 'vessel_ais' """ ) ) 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: from app.services.bgp_collector_locations import ( seed_default_bgp_collector_locations, ) from app.services.compute_center_locations import ( seed_compute_center_locations_from_source_coords, ) await seed_default_bgp_collector_locations(session) await seed_compute_center_locations_from_source_coords(session) await seed_default_datasources(session) await purge_legacy_earth_boundary_datasources(session) await ensure_default_admin_user(session)