Files
insar-management-system-v2/backend/app/db_maintenance.py
T

634 lines
22 KiB
Python

from __future__ import annotations
import os
from typing import Any, Dict, List, Set, Tuple
from urllib.parse import urlparse
from sqlalchemy import create_engine, inspect, text
from sqlalchemy.orm import sessionmaker
from sqlalchemy.schema import CreateColumn, CreateIndex
from .config import read_bool_env, settings
POSTGIS_TABLES = {
"spatial_ref_sys",
}
POSTGIS_VIEWS = {
"geometry_columns",
"geography_columns",
"raster_columns",
"raster_overviews",
}
ALLOWED_EXTRA_TABLES = {
"alembic_version",
"spatial_query_logs",
}
ALEMBIC_HEAD_REVISION = "0004"
MIGRATION_FILES = [
"001_st_intersection_agg.sql",
"002_spatial_functions.sql",
"003_pairing_enhancement.sql",
"004_pairing_refactor.sql",
"005_pairing_task_trace.sql",
"006_result_pairing_trace.sql",
"007_timeseries_stack_plan_trace.sql",
"008_timeseries_stack_plan_edges.sql",
"009_raw_source_pairing_fields.sql",
"010_source_orbit_asset_inventory.sql",
"011_source_metadata_documents.sql",
"012_source_archive_integrity.sql",
"013_result_delivery_requests.sql",
"014_result_delivery_ortho_sources.sql",
]
def project_root() -> str:
return os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
def _database_url_to_sync_url(database_url: str) -> str:
if "postgresql+asyncpg" in database_url:
return database_url.replace("postgresql+asyncpg", "postgresql", 1)
return database_url
def _split_sql_statements(sql_text: str) -> List[str]:
statements: List[str] = []
buf: List[str] = []
in_single = False
in_double = False
dollar_tag = None
i = 0
length = len(sql_text)
while i < length:
ch = sql_text[i]
if dollar_tag:
if sql_text.startswith(dollar_tag, i):
buf.append(dollar_tag)
i += len(dollar_tag)
dollar_tag = None
continue
buf.append(ch)
i += 1
continue
if not in_single and not in_double:
if ch == "'":
in_single = True
buf.append(ch)
i += 1
continue
if ch == '"':
in_double = True
buf.append(ch)
i += 1
continue
if ch == "$":
end = sql_text.find("$", i + 1)
if end != -1:
tag = sql_text[i:end + 1]
dollar_tag = tag
buf.append(tag)
i += len(tag)
continue
if ch == ";":
stmt = "".join(buf).strip()
if stmt:
statements.append(stmt)
buf = []
i += 1
continue
else:
if in_single:
if ch == "'" and i + 1 < length and sql_text[i + 1] == "'":
buf.append(ch)
buf.append(sql_text[i + 1])
i += 2
continue
if ch == "'":
in_single = False
buf.append(ch)
i += 1
continue
if in_double and ch == '"':
in_double = False
buf.append(ch)
i += 1
continue
buf.append(ch)
i += 1
tail = "".join(buf).strip()
if tail:
statements.append(tail)
return statements
def _is_effective_sql(stmt: str) -> bool:
stripped = stmt.strip()
if not stripped:
return False
for line in stripped.splitlines():
line = line.strip()
if not line:
continue
if line.startswith("--"):
continue
if line.startswith("/*") and line.endswith("*/"):
continue
return True
return False
def _apply_sql_file(conn, path: str) -> bool:
if not os.path.exists(path):
print(f"[WARN] Migration file not found: {path}")
return False
with open(path, "r", encoding="utf-8-sig") as stream:
sql_text = stream.read()
for stmt in _split_sql_statements(sql_text):
if not _is_effective_sql(stmt):
continue
conn.exec_driver_sql(stmt)
return True
def _load_base():
from backend.app.database import Base
import backend.app.models # noqa: F401
return Base
def get_required_table_names() -> List[str]:
base = _load_base()
return sorted(base.metadata.tables.keys())
def _build_schema_diagnostics(
expected: Dict[str, Dict[str, Dict[str, Any]]],
current: Dict[str, Dict[str, Dict[str, Any]]],
) -> Dict[str, Any]:
filtered_current = {
name: cols for name, cols in current.items() if name not in ALLOWED_EXTRA_TABLES
}
expected_tables = set(expected.keys())
current_tables = set(filtered_current.keys())
missing_tables = sorted(expected_tables - current_tables)
extra_tables = sorted(current_tables - expected_tables)
missing_columns: Dict[str, List[str]] = {}
extra_columns: Dict[str, List[str]] = {}
type_mismatches: List[Dict[str, str]] = []
nullable_mismatches: List[Dict[str, Any]] = []
reasons: List[str] = []
if missing_tables:
reasons.append(f"Missing tables: {missing_tables}")
if extra_tables:
reasons.append(f"Extra tables: {extra_tables}")
bootstrap_required = not current_tables and bool(expected_tables) and set(missing_tables) == expected_tables
for table in sorted(expected_tables & current_tables):
expected_cols = expected[table]
current_cols = filtered_current[table]
expected_col_names = set(expected_cols.keys())
current_col_names = set(current_cols.keys())
table_missing_cols = sorted(expected_col_names - current_col_names)
table_extra_cols = sorted(current_col_names - expected_col_names)
if table_missing_cols:
missing_columns[table] = table_missing_cols
reasons.append(f"Table {table} missing columns: {table_missing_cols}")
if table_extra_cols:
extra_columns[table] = table_extra_cols
reasons.append(f"Table {table} extra columns: {table_extra_cols}")
for col_name in sorted(expected_col_names & current_col_names):
exp = expected_cols[col_name]
cur = current_cols[col_name]
if exp["type"] != cur["type"]:
mismatch = {
"table": table,
"column": col_name,
"expected": exp["type"],
"actual": cur["type"],
}
type_mismatches.append(mismatch)
reasons.append(
f"Table {table} column {col_name} type mismatch: expected {exp['type']} got {cur['type']}"
)
if exp["nullable"] != cur["nullable"]:
mismatch = {
"table": table,
"column": col_name,
"expected": exp["nullable"],
"actual": cur["nullable"],
}
nullable_mismatches.append(mismatch)
reasons.append(
f"Table {table} column {col_name} nullable mismatch: expected {exp['nullable']} got {cur['nullable']}"
)
return {
"mismatch": len(reasons) > 0,
"bootstrap_required": bootstrap_required,
"reasons": reasons,
"reason_count": len(reasons),
"required_tables": sorted(expected.keys()),
"required_table_count": len(expected),
"current_table_count": len(filtered_current),
"missing_tables": missing_tables,
"extra_tables": extra_tables,
"missing_columns": missing_columns,
"extra_columns": extra_columns,
"type_mismatches": type_mismatches,
"nullable_mismatches": nullable_mismatches,
}
def _normalize_type(type_obj, dialect) -> str:
if type_obj is None:
return ""
try:
compiled = type_obj.compile(dialect=dialect)
except Exception:
compiled = str(type_obj)
if compiled is None:
return ""
norm = str(compiled).lower().replace(" ", "")
float_aliases = {"doubleprecision", "float8", "float4", "real", "float"}
if norm in float_aliases:
return "float"
return norm
def _get_expected_schema(base, dialect) -> Dict[str, Dict[str, Dict[str, Any]]]:
expected: Dict[str, Dict[str, Dict[str, Any]]] = {}
for table_name, table in base.metadata.tables.items():
expected[table_name] = {}
for col in table.columns:
expected[table_name][col.name] = {
"type": _normalize_type(col.type, dialect),
"nullable": bool(col.nullable),
}
return expected
def _get_current_schema(inspector, dialect) -> Dict[str, Dict[str, Dict[str, Any]]]:
current: Dict[str, Dict[str, Dict[str, Any]]] = {}
for table_name in inspector.get_table_names():
if table_name in POSTGIS_TABLES:
continue
cols = inspector.get_columns(table_name)
current[table_name] = {}
for col in cols:
current[table_name][col["name"]] = {
"type": _normalize_type(col.get("type"), dialect),
"nullable": bool(col.get("nullable", True)),
}
return current
def _schema_mismatch(
expected: Dict[str, Dict[str, Dict[str, Any]]],
current: Dict[str, Dict[str, Dict[str, Any]]],
) -> Tuple[bool, List[str]]:
diagnostics = _build_schema_diagnostics(expected, current)
return diagnostics["mismatch"], diagnostics["reasons"]
def inspect_database_structure(bind) -> Dict[str, Any]:
base = _load_base()
expected = _get_expected_schema(base, bind.dialect)
current = _get_current_schema(inspect(bind), bind.dialect)
return _build_schema_diagnostics(expected, current)
def _drop_all_objects(conn, inspector) -> None:
for view in inspector.get_view_names():
if view in POSTGIS_VIEWS:
continue
conn.exec_driver_sql(f'DROP VIEW IF EXISTS "{view}" CASCADE')
for table in inspector.get_table_names():
if table in POSTGIS_TABLES:
continue
conn.exec_driver_sql(f'DROP TABLE IF EXISTS "{table}" CASCADE')
def _add_missing_columns(conn, inspector, base, dialect) -> List[str]:
added_columns: List[str] = []
existing_tables = set(inspector.get_table_names())
for table_name, table in base.metadata.tables.items():
if table_name not in existing_tables:
continue
current_columns = {column["name"] for column in inspector.get_columns(table_name)}
for column in table.columns:
if column.name in current_columns:
continue
try:
column_sql = str(CreateColumn(column).compile(dialect=dialect))
conn.exec_driver_sql(f'ALTER TABLE "{table_name}" ADD COLUMN {column_sql}')
added_columns.append(f"{table_name}.{column.name}")
except Exception as exc:
print(f"[WARN] Failed to add column {table_name}.{column.name}: {exc}")
return added_columns
def _create_missing_indexes(conn, inspector, base, dialect) -> List[str]:
created_indexes: List[str] = []
existing_tables = set(inspector.get_table_names())
for table_name, table in base.metadata.tables.items():
if table_name not in existing_tables:
continue
current_indexes = {index["name"] for index in inspector.get_indexes(table_name)}
for index in sorted(table.indexes, key=lambda item: item.name or ""):
if not index.name or index.name in current_indexes:
continue
try:
index_sql = str(CreateIndex(index).compile(dialect=dialect))
conn.exec_driver_sql(index_sql)
created_indexes.append(index.name)
current_indexes.add(index.name)
except Exception as exc:
print(f"[WARN] Failed to create index {table_name}.{index.name}: {exc}")
return created_indexes
def _ensure_alembic_version_marker(conn, revision: str = ALEMBIC_HEAD_REVISION) -> Dict[str, Any]:
status: Dict[str, Any] = {"revision": revision, "created": False, "updated": False}
conn.exec_driver_sql(
"""
CREATE TABLE IF NOT EXISTS alembic_version (
version_num VARCHAR(32) NOT NULL
)
"""
)
rows = conn.execute(text("SELECT version_num FROM alembic_version")).fetchall()
previous = [row[0] for row in rows]
status["previous"] = previous
if not rows:
conn.execute(text("INSERT INTO alembic_version (version_num) VALUES (:revision)"), {"revision": revision})
status["created"] = True
return status
if previous != [revision] or len(rows) != 1:
conn.execute(text("DELETE FROM alembic_version"))
conn.execute(text("INSERT INTO alembic_version (version_num) VALUES (:revision)"), {"revision": revision})
status["updated"] = True
return status
def _resolve_hazard_shapefile() -> str:
hazard_dir = settings.HAZARD_POINTS_DIR
hazard_filename = settings.HAZARD_POINTS_FILENAME or "Point.shp"
if hazard_dir:
return os.path.join(hazard_dir, hazard_filename)
return os.path.join(project_root(), "backend", "Point", hazard_filename)
def bootstrap_admin_user(session) -> Dict[str, Any]:
from backend.app.auth_utils import hash_password, normalize_username
from backend.app.models import AuthUserORM
status = {"created": False, "updated": False, "username": None, "message": ""}
admin_online = session.query(AuthUserORM).filter(
AuthUserORM.role == "admin",
AuthUserORM.is_active == True,
).first()
if admin_online:
status["username"] = admin_online.username
status["message"] = f"Admin account exists: {admin_online.username}"
return status
admin_username = normalize_username(settings.INIT_ADMIN_USERNAME or "admin") or "admin"
admin_password = settings.INIT_ADMIN_PASSWORD or ""
reset_existing_password = bool(settings.INIT_ADMIN_RESET_PASSWORD)
existing_user = session.query(AuthUserORM).filter(
AuthUserORM.username == admin_username
).one_or_none()
password_required = existing_user is None or reset_existing_password or not existing_user.password_hash
if password_required and not admin_password:
raise RuntimeError(
"INIT_ADMIN_PASSWORD is required to create/reset admin account. "
"Please set INIT_ADMIN_PASSWORD before startup."
)
if existing_user:
existing_user.role = "admin"
existing_user.is_active = True
if reset_existing_password or not existing_user.password_hash:
existing_user.password_hash = hash_password(admin_password)
existing_user.created_by = existing_user.created_by or "system:db_maintenance"
session.commit()
status["updated"] = True
status["username"] = admin_username
status["message"] = f"Promoted existing user to admin: {admin_username}"
return status
admin_user = AuthUserORM(
username=admin_username,
password_hash=hash_password(admin_password),
role="admin",
is_active=True,
created_by="system:db_maintenance",
)
session.add(admin_user)
session.commit()
status["created"] = True
status["username"] = admin_username
status["message"] = f"Created initial admin user: {admin_username}"
return status
def seed_hazard_points(session) -> Dict[str, Any]:
from backend.app.models import HazardPointORM
status = {"seeded": False, "count": 0, "message": ""}
count = session.query(HazardPointORM).count()
if count > 0:
status["count"] = count
status["message"] = f"Hazard points already exist ({count} items)."
return status
shp_path = _resolve_hazard_shapefile()
if not os.path.exists(shp_path):
status["message"] = f"Hazard points file not found: {shp_path}"
return status
import geopandas as gpd
from geoalchemy2.shape import from_shape
from shapely.geometry import Point
gdf = gpd.read_file(shp_path, engine="pyogrio")
if gdf.crs and gdf.crs.to_epsg() != 4326:
gdf = gdf.to_crs(epsg=4326)
field_tybh = ["TYBH", "tybh", "统一编", "统一编号", "UNIFIED_ID"]
field_hazard_type = ["灾害类型", "灾害类", "ZHLX", "hazard_type", "TYPE"]
field_hazard_name = ["灾害名", "ZHMC", "hazard_name", "NAME"]
field_city = ["市", "CITY", "city"]
field_county = ["县", "COUNTY", "county"]
field_township = ["乡", "TOWNSHIP", "township", "乡镇"]
field_lon = ["经度", "LON", "longitude"]
field_lat = ["纬度", "维度", "LAT", "latitude"]
def pick_value(row_obj, candidates):
for key in candidates:
if key in row_obj and row_obj[key] not in (None, ""):
return row_obj[key]
return None
def to_float(value):
try:
if value is None or value == "":
return None
return float(value)
except (TypeError, ValueError):
return None
added_count = 0
for _, row in gdf.iterrows():
tybh_value = pick_value(row, field_tybh)
if tybh_value is None:
continue
geom_x = getattr(row.geometry, "x", None)
geom_y = getattr(row.geometry, "y", None)
lon_value = to_float(geom_x)
lat_value = to_float(geom_y)
if lon_value is None:
lon_value = to_float(pick_value(row, field_lon))
if lat_value is None:
lat_value = to_float(pick_value(row, field_lat))
if lon_value is None or lat_value is None:
continue
point = HazardPointORM(
tybh=str(tybh_value).strip(),
hazard_type=pick_value(row, field_hazard_type),
hazard_name=pick_value(row, field_hazard_name),
city=pick_value(row, field_city),
county=pick_value(row, field_county),
township=pick_value(row, field_township),
longitude=lon_value,
latitude=lat_value,
geom=from_shape(Point(lon_value, lat_value), srid=4326),
)
session.add(point)
added_count += 1
session.commit()
status["seeded"] = True
status["count"] = added_count
status["message"] = f"Imported {added_count} hazard points."
return status
def ensure_database_ready(
database_url: str | None = None,
*,
bootstrap_admin: bool = True,
seed_hazard: bool = True,
) -> Dict[str, Any]:
database_url = database_url or settings.DATABASE_URL
if not database_url:
raise RuntimeError("DATABASE_URL is not configured.")
reset_on_mismatch = read_bool_env("DB_SCHEMA_RESET_ON_MISMATCH", False)
reset_confirm = read_bool_env("DB_SCHEMA_RESET_CONFIRM", False)
allow_schema_reset = reset_on_mismatch and reset_confirm
sync_url = _database_url_to_sync_url(database_url)
parsed = urlparse(sync_url)
engine = create_engine(sync_url)
Session = sessionmaker(bind=engine)
result: Dict[str, Any] = {
"database": parsed.path.strip("/"),
"host": parsed.hostname,
"schema_reset": False,
"mismatch_detected": False,
"mismatch_reasons": [],
"added_columns": [],
"created_indexes": [],
"alembic_version": None,
"applied_sql_files": [],
"admin": None,
"hazard_seed": None,
}
try:
with engine.begin() as conn:
conn.exec_driver_sql("CREATE EXTENSION IF NOT EXISTS postgis")
base = _load_base()
inspector = inspect(conn)
schema_diagnostics = inspect_database_structure(conn)
bootstrap_required = bool(schema_diagnostics.get("bootstrap_required"))
result["mismatch_detected"] = schema_diagnostics["mismatch"] and not bootstrap_required
result["mismatch_reasons"] = [] if bootstrap_required else schema_diagnostics["reasons"]
result["bootstrap_initialized"] = False
if schema_diagnostics["mismatch"]:
if bootstrap_required:
print("[INFO] Empty application schema detected. Bootstrapping database tables...")
else:
print("[WARN] Database schema mismatch detected.")
for reason in schema_diagnostics["reasons"]:
print(f" - {reason}")
if allow_schema_reset:
_drop_all_objects(conn, inspector)
base.metadata.create_all(bind=conn)
result["schema_reset"] = True
else:
base.metadata.create_all(bind=conn)
inspector = inspect(conn)
result["added_columns"] = _add_missing_columns(conn, inspector, base, engine.dialect)
inspector = inspect(conn)
result["created_indexes"] = _create_missing_indexes(conn, inspector, base, engine.dialect)
if bootstrap_required:
result["bootstrap_initialized"] = True
else:
base.metadata.create_all(bind=conn)
inspector = inspect(conn)
result["created_indexes"] = _create_missing_indexes(conn, inspector, base, engine.dialect)
migrations_dir = os.path.join(project_root(), "backend", "migrations")
for migration_file in MIGRATION_FILES:
migration_path = os.path.join(migrations_dir, migration_file)
if _apply_sql_file(conn, migration_path):
result["applied_sql_files"].append(migration_file)
final_diagnostics = inspect_database_structure(conn)
if not final_diagnostics.get("mismatch"):
result["alembic_version"] = _ensure_alembic_version_marker(conn)
else:
result["post_maintenance_mismatch_reasons"] = final_diagnostics.get("reasons", [])
session = Session()
try:
if bootstrap_admin:
result["admin"] = bootstrap_admin_user(session)
if seed_hazard:
result["hazard_seed"] = seed_hazard_points(session)
finally:
session.close()
return result
finally:
engine.dispose()