386 lines
13 KiB
Python
386 lines
13 KiB
Python
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": "LK12345678",
|
|
"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)
|