92 lines
2.5 KiB
Python
92 lines
2.5 KiB
Python
"""Utility to reset development Postgres schema artifacts.
|
|
|
|
This script drops managed tables and enum types created by `scripts.init_db`.
|
|
It is intended for local development only; it refuses to run if CALMINER_ENV
|
|
indicates production or staging. The operation is idempotent: missing objects
|
|
are ignored. Use with caution.
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import logging
|
|
import os
|
|
from dataclasses import dataclass
|
|
from typing import Iterable
|
|
|
|
from sqlalchemy import text
|
|
from sqlalchemy.engine import Engine
|
|
|
|
from config.database import DATABASE_URL
|
|
from scripts.init_db import ENUM_DEFINITIONS, _create_engine
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
@dataclass(slots=True)
|
|
class ResetOptions:
|
|
drop_tables: bool = True
|
|
drop_enums: bool = True
|
|
|
|
|
|
MANAGED_TABLES: tuple[str, ...] = (
|
|
"simulation_parameters",
|
|
"financial_inputs",
|
|
"scenarios",
|
|
"projects",
|
|
"pricing_impurity_settings",
|
|
"pricing_metal_settings",
|
|
"pricing_settings",
|
|
"user_roles",
|
|
"users",
|
|
"roles",
|
|
)
|
|
|
|
|
|
FORBIDDEN_ENVIRONMENTS: set[str] = {"production", "staging", "prod", "stage"}
|
|
|
|
|
|
def _ensure_safe_environment() -> None:
|
|
env = os.getenv("CALMINER_ENV", "development").lower()
|
|
if env in FORBIDDEN_ENVIRONMENTS:
|
|
raise RuntimeError(
|
|
f"Refusing to reset database in environment '{env}'. "
|
|
"Set CALMINER_ENV to 'development' to proceed."
|
|
)
|
|
|
|
|
|
def _drop_tables(engine: Engine, tables: Iterable[str]) -> None:
|
|
if not tables:
|
|
return
|
|
with engine.begin() as conn:
|
|
for table in tables:
|
|
logger.info("Dropping table if exists: %s", table)
|
|
conn.execute(text(f"DROP TABLE IF EXISTS {table} CASCADE"))
|
|
|
|
|
|
def _drop_enums(engine: Engine, enum_names: Iterable[str]) -> None:
|
|
if not enum_names:
|
|
return
|
|
with engine.begin() as conn:
|
|
for enum_name in enum_names:
|
|
logger.info("Dropping enum type if exists: %s", enum_name)
|
|
conn.execute(text(f"DROP TYPE IF EXISTS {enum_name} CASCADE"))
|
|
|
|
|
|
def reset_database(*, options: ResetOptions | None = None, database_url: str | None = None) -> None:
|
|
"""Drop managed tables and enums for a clean slate."""
|
|
_ensure_safe_environment()
|
|
opts = options or ResetOptions()
|
|
engine = _create_engine(database_url or DATABASE_URL)
|
|
|
|
if opts.drop_tables:
|
|
_drop_tables(engine, MANAGED_TABLES)
|
|
|
|
if opts.drop_enums:
|
|
_drop_enums(engine, ENUM_DEFINITIONS.keys())
|
|
|
|
logger.info("Database reset complete")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
logging.basicConfig(level=logging.INFO)
|
|
reset_database()
|