tsapi-seriesdb-shared-sql/run_setup.py

"""
Provision the TSAPI series database schema from the canonical SQL files.

This is the SINGLE SOURCE OF TRUTH provisioner for the TSAPI timeseries
PostgreSQL schema. Both the Python library (TSAPI_SERIESDB_PG_PY) and the C++
library (TSAPI_POSTGRES_CPP) target the schema it creates. Do not fork the SQL
into either project -- edit it here.

It executes the numbered DDL files in ./sql/ in order:
    00_extensions, 01_schemas, 02_sequences, 03_tables, 04_indexes,
    05_procedures, 06_triggers, 07_reference_data

Connection settings come from (highest precedence first):
    1. command-line flags  (--host --port --dbname --user --password)
    2. db_config.json next to this script
    3. built-in defaults (localhost:5432 / timeseries_db__dev / tsapi)

Usage:
    python run_setup.py                       # use db_config.json / defaults
    python run_setup.py --host 192.168.1.7 --user claude --password claudepwd
    python run_setup.py --drop-existing       # drop meta+series schemas first
    python run_setup.py --clean-slate         # drop ALL user objects first
"""

import sys
import json
import argparse
from pathlib import Path
from typing import Dict

import psycopg2

SQL_DIR = Path(__file__).parent / "sql"
CONFIG_FILE = Path(__file__).parent / "db_config.json"

# Executed in order. 08_script_log.sql and sql/migrate/*.sql are auxiliary and
# are intentionally NOT part of a fresh install.
SQL_FILES = [
    "00_extensions.sql",
    "01_schemas.sql",
    "02_sequences.sql",
    "03_tables.sql",
    "04_indexes.sql",
    "05_procedures.sql",
    "06_triggers.sql",
    "07_reference_data.sql",
]

DEFAULTS = {
    "host": "localhost",
    "port": 5432,
    "database": "timeseries_db__dev",
    "user": "tsapi",
    "password": "tsapipwd",
}


def load_config(args) -> Dict:
    config = dict(DEFAULTS)
    if CONFIG_FILE.exists():
        try:
            config.update(json.loads(CONFIG_FILE.read_text()))
        except json.JSONDecodeError as e:
            print(f"[FAIL] Invalid JSON in {CONFIG_FILE}: {e}")
            sys.exit(1)
    # CLI overrides
    for key, val in (
        ("host", args.host), ("port", args.port), ("database", args.dbname),
        ("user", args.user), ("password", args.password),
    ):
        if val is not None:
            config[key] = val
    return config


def execute_sql_file(conn, sql_file: Path):
    try:
        sql = sql_file.read_text(encoding="utf-8")
        with conn.cursor() as cur:
            cur.execute(sql)
            conn.commit()
        return True, f"[OK] Executed {sql_file.name}"
    except Exception as e:
        conn.rollback()
        return False, f"[FAIL] {sql_file.name}: {e}"


def verify_setup(conn) -> None:
    print("\n" + "=" * 70)
    print("Verifying setup")
    print("=" * 70)
    with conn.cursor() as cur:
        cur.execute(
            "SELECT schema_name FROM information_schema.schemata "
            "WHERE schema_name IN ('meta','series','log') ORDER BY 1"
        )
        print("[OK] Schemas:", ", ".join(r[0] for r in cur.fetchall()))

        cur.execute("SELECT count(*) FROM information_schema.tables WHERE table_schema='meta'")
        print("[OK] meta tables:", cur.fetchone()[0])

        cur.execute("SELECT count(*) FROM information_schema.routines WHERE routine_schema IN ('meta','log')")
        print("[OK] procedures:", cur.fetchone()[0])

        cur.execute("SELECT count(*) FROM pg_event_trigger WHERE evtname='prevent_series_drops'")
        print("[OK] event triggers:", cur.fetchone()[0])

        cur.execute("SELECT count(*) FROM meta.timezones")
        tz = cur.fetchone()[0]
        cur.execute("SELECT count(*) FROM meta.currencies")
        cu = cur.fetchone()[0]
        print(f"[OK] reference data: {tz} timezones, {cu} currencies")

        cur.execute("SELECT version, description FROM meta.schema_version ORDER BY installed_at DESC LIMIT 1")
        row = cur.fetchone()
        if row:
            print(f"[OK] schema_version: {row[0]} ({row[1]})")


def clean_slate(conn) -> None:
    print("Performing clean slate...")
    with conn.cursor() as cur:
        cur.execute(
            "SELECT schema_name FROM information_schema.schemata "
            "WHERE schema_name NOT IN ('pg_catalog','information_schema','pg_toast','public') "
            "AND schema_name NOT LIKE 'pg_temp_%' AND schema_name NOT LIKE 'pg_toast_temp_%'"
        )
        for (schema,) in cur.fetchall():
            print(f"  drop schema {schema}")
            cur.execute(f"DROP SCHEMA IF EXISTS {schema} CASCADE")
        cur.execute("SELECT evtname FROM pg_event_trigger")
        for (trig,) in cur.fetchall():
            cur.execute(f"DROP EVENT TRIGGER IF EXISTS {trig}")
        cur.execute("CREATE SCHEMA IF NOT EXISTS public")
        conn.commit()
    print("[OK] clean slate done\n")


def main() -> int:
    p = argparse.ArgumentParser(description="Provision the TSAPI series DB schema")
    p.add_argument("--host")
    p.add_argument("--port", type=int)
    p.add_argument("--dbname")
    p.add_argument("--user")
    p.add_argument("--password")
    p.add_argument("--drop-existing", action="store_true",
                   help="Drop meta and series schemas before creating (destructive)")
    p.add_argument("--clean-slate", action="store_true",
                   help="Drop ALL user objects before creating (very destructive)")
    p.add_argument("--yes", action="store_true", help="Skip confirmation prompts")
    args = p.parse_args()

    cfg = load_config(args)

    print("=" * 70)
    print("TSAPI series DB schema setup")
    print("=" * 70)
    print(f"Target : {cfg['host']}:{cfg['port']}/{cfg['database']} (user {cfg['user']})")
    print(f"SQL dir: {SQL_DIR}")

    if (args.clean_slate or args.drop_existing) and not args.yes:
        what = "ALL user objects" if args.clean_slate else "meta and series schemas"
        if input(f"This will DROP {what}. Continue? (yes/no): ").lower() != "yes":
            print("Aborted.")
            return 1

    try:
        conn = psycopg2.connect(
            host=cfg["host"], port=cfg["port"], dbname=cfg["database"],
            user=cfg["user"], password=cfg["password"],
        )
    except Exception as e:
        print(f"[FAIL] connect: {e}")
        return 1

    try:
        if args.clean_slate:
            clean_slate(conn)
        elif args.drop_existing:
            with conn.cursor() as cur:
                cur.execute("DROP SCHEMA IF EXISTS meta CASCADE")
                cur.execute("DROP SCHEMA IF EXISTS series CASCADE")
                cur.execute("DROP SCHEMA IF EXISTS log CASCADE")
                cur.execute("DROP EVENT TRIGGER IF EXISTS prevent_series_drops")
                conn.commit()
            print("[OK] existing schemas dropped\n")

        failed = []
        for name in SQL_FILES:
            f = SQL_DIR / name
            if not f.exists():
                print(f"[FAIL] missing {name}")
                failed.append(name)
                continue
            ok, msg = execute_sql_file(conn, f)
            print(msg)
            if not ok:
                failed.append(name)

        if failed:
            print(f"\n[FAIL] {len(failed)} file(s) failed: {', '.join(failed)}")
            return 1
        verify_setup(conn)
        print("\n[SUCCESS] schema setup complete")
        return 0
    finally:
        conn.close()


if __name__ == "__main__":
    sys.exit(main())