home-automation/scripts/location_db_adopt.py

from __future__ import annotations

import sqlite3
import sys
from pathlib import Path

from alembic import command
from alembic.config import Config

PROJECT_ROOT = Path(__file__).resolve().parents[1]
if str(PROJECT_ROOT) not in sys.path:
    sys.path.insert(0, str(PROJECT_ROOT))

from app.config import get_settings

LOCATION_BASELINE_REVISION = "20260419_01_location_baseline"
EXPECTED_USER_VERSION = 2
EXPECTED_LOCATION_TABLE_INFO = [
    (0, "person", "TEXT", 1, None, 1),
    (1, "datetime", "TEXT", 1, None, 2),
    (2, "latitude", "REAL", 1, None, 0),
    (3, "longitude", "REAL", 1, None, 0),
    (4, "altitude", "REAL", 0, None, 0),
]


class LocationDatabaseAdoptionError(RuntimeError):
    """Raised when a legacy location database does not match the expected baseline."""


def _database_path_from_url(database_url: str) -> Path:
    prefix = "sqlite:///"
    if not database_url.startswith(prefix):
        raise LocationDatabaseAdoptionError(
            f"Only sqlite URLs are supported for location DB adoption, got: {database_url}"
        )
    return Path(database_url[len(prefix) :])


def _make_alembic_config(database_url: str) -> Config:
    config = Config("alembic.ini")
    config.set_main_option("sqlalchemy.url", database_url)
    return config


def _location_table_exists(database_path: Path) -> bool:
    conn = sqlite3.connect(database_path)
    try:
        row = conn.execute(
            "SELECT 1 FROM sqlite_master WHERE type = 'table' AND name = 'location'"
        ).fetchone()
        return row is not None
    finally:
        conn.close()


def _alembic_version_table_exists(database_path: Path) -> bool:
    conn = sqlite3.connect(database_path)
    try:
        row = conn.execute(
            "SELECT 1 FROM sqlite_master WHERE type = 'table' AND name = 'alembic_version'"
        ).fetchone()
        return row is not None
    finally:
        conn.close()


def _fetch_alembic_revision(database_path: Path) -> str:
    conn = sqlite3.connect(database_path)
    try:
        row = conn.execute("SELECT version_num FROM alembic_version").fetchone()
        if row is None:
            raise LocationDatabaseAdoptionError(
                "Alembic version table exists but contains no revision"
            )
        return row[0]
    finally:
        conn.close()


def _fetch_location_table_info(database_path: Path) -> list[tuple]:
    conn = sqlite3.connect(database_path)
    try:
        return list(conn.execute("PRAGMA table_info(location)"))
    finally:
        conn.close()


def _fetch_user_version(database_path: Path) -> int:
    conn = sqlite3.connect(database_path)
    try:
        return conn.execute("PRAGMA user_version").fetchone()[0]
    finally:
        conn.close()


def validate_legacy_location_db(database_url: str) -> None:
    database_path = _database_path_from_url(database_url)
    if not database_path.exists():
        raise LocationDatabaseAdoptionError(f"Location DB file does not exist: {database_path}")

    if not _location_table_exists(database_path):
        raise LocationDatabaseAdoptionError("Expected table 'location' was not found in the DB")

    table_info = _fetch_location_table_info(database_path)
    if table_info != EXPECTED_LOCATION_TABLE_INFO:
        raise LocationDatabaseAdoptionError(
            "Location table schema does not match the expected baseline schema"
        )

    user_version = _fetch_user_version(database_path)
    if user_version != EXPECTED_USER_VERSION:
        raise LocationDatabaseAdoptionError(
            f"Expected PRAGMA user_version = {EXPECTED_USER_VERSION}, got {user_version}"
        )


def validate_location_runtime_db(database_url: str) -> None:
    database_path = _database_path_from_url(database_url)
    if not database_path.exists():
        raise LocationDatabaseAdoptionError(
            "Location DB file was not found. Run 'python scripts/location_db_adopt.py' "
            "first to initialize or adopt the location DB before starting the app."
        )

    if not _alembic_version_table_exists(database_path):
        raise LocationDatabaseAdoptionError(
            "Location DB exists but is not yet Alembic-managed. Run "
            "'python scripts/location_db_adopt.py' first to adopt the legacy DB "
            "before starting the app."
        )

    current_revision = _fetch_alembic_revision(database_path)
    if current_revision != LOCATION_BASELINE_REVISION:
        raise LocationDatabaseAdoptionError(
            "Location DB revision mismatch. Refusing to start the app: "
            f"expected {LOCATION_BASELINE_REVISION}, got {current_revision}"
        )


def adopt_or_initialize_location_db(database_url: str) -> str:
    database_path = _database_path_from_url(database_url)
    alembic_config = _make_alembic_config(database_url)

    if database_path.exists():
        if _alembic_version_table_exists(database_path):
            current_revision = _fetch_alembic_revision(database_path)
            if current_revision != LOCATION_BASELINE_REVISION:
                raise LocationDatabaseAdoptionError(
                    "Location DB is already Alembic-managed but revision does not match "
                    f"the expected baseline: expected {LOCATION_BASELINE_REVISION}, "
                    f"got {current_revision}"
                )
            return "already_managed"

        validate_legacy_location_db(database_url)
        command.stamp(alembic_config, LOCATION_BASELINE_REVISION)
        return "adopted"

    database_path.parent.mkdir(parents=True, exist_ok=True)
    command.upgrade(alembic_config, "head")
    return "initialized"


def main() -> None:
    settings = get_settings()
    result = adopt_or_initialize_location_db(settings.location_database_url)
    if result == "initialized":
        print("Initialized a new location DB via Alembic upgrade head.")
    elif result == "already_managed":
        print("Location DB is already Alembic-managed at the expected baseline revision.")
    else:
        print("Validated legacy location DB and stamped Alembic baseline successfully.")


if __name__ == "__main__":
    main()