team-tryouts/app/local_postgres.py

"""Manage a local PostgreSQL instance for development.

Searches for PostgreSQL binaries on the system, initialises a data directory
under ``instance/pgdata/``, starts the server on port 5433, creates the
``team_tryouts`` database, and registers cleanup on exit.

Usage from run.py::

    if os.getenv('LOCAL_PG', '').lower() == 'true':
        from app.local_postgres import start as start_local_pg
        os.environ['DATABASE_URL'] = start_local_pg()

Environment used (all optional):
    LOCAL_PG_PORT        Port to bind (default 5433)
    LOCAL_PG_USER        Database superuser name (default team_tryouts_dev)
    LOCAL_PG_DATABASE    Database name (default team_tryouts)
"""

from __future__ import annotations

import atexit
import contextlib
import logging
import os
import pathlib
import shutil
import signal
import subprocess
import sys
import tempfile

# ---------------------------------------------------------------------------
# Configuration
# ---------------------------------------------------------------------------

DATADIR = pathlib.Path(__file__).resolve().parent.parent / 'instance' / 'pgdata'
PG_LOG = DATADIR.parent / 'pg.log'

PORT = int(os.getenv('LOCAL_PG_PORT', '5433'))
_USER = os.getenv('LOCAL_PG_USER', 'team_tryouts_dev')
_DATABASE = os.getenv('LOCAL_PG_DATABASE', 'team_tryouts')

DATABASE_URL = f'postgresql://{_USER}@127.0.0.1:{PORT}/{_DATABASE}'

_started = False

#: Name of the system Windows PostgreSQL service we temporarily stopped.
_stopped_system_service_name: str | None = None

logger = logging.getLogger('app.local_postgres')
# ---------------------------------------------------------------------------
# Binary discovery
# ---------------------------------------------------------------------------

_WIN_PG_ROOTS = [
    pathlib.Path(os.environ.get('PROGRAMFILES', 'C:\\Program Files')) / 'PostgreSQL',
    pathlib.Path(os.environ.get('PROGRAMFILES(X86)', 'C:\\Program Files (x86)')) / 'PostgreSQL',
]


def _find_binary(name: str) -> pathlib.Path | None:
    """Locate a PostgreSQL binary (pg_ctl, initdb, psql)."""
    found = shutil.which(name)
    if found:
        return pathlib.Path(found)

    if sys.platform == 'win32':
        for root in _WIN_PG_ROOTS:
            if not root.is_dir():
                continue
            versions = sorted(
                (p for p in root.iterdir() if p.is_dir()),
                key=lambda p: p.name,
                reverse=True,
            )
            for version_dir in versions:
                candidate = version_dir / 'bin' / f'{name}.exe'
                if candidate.is_file():
                    return candidate

    return None


# ---------------------------------------------------------------------------
# Windows service management
# ---------------------------------------------------------------------------

#: Path to the Windows service controller.
_SC_EXE = pathlib.Path(os.environ.get('SYSTEMROOT', 'C:\\Windows')) / 'System32' / 'sc.exe'


def _detect_system_service() -> str | None:
    """Find a running Windows PostgreSQL service installed by the EDB installer.

    Returns the service name (e.g. ``postgresql-x64-18``) or None.
    """
    if sys.platform != 'win32' or not _SC_EXE.is_file():
        return None
    for root in _WIN_PG_ROOTS:
        if not root.is_dir():
            continue
        versions = sorted(
            (p for p in root.iterdir() if p.is_dir()),
            key=lambda p: p.name,
            reverse=True,
        )
        for version_dir in versions:
            srv = f'postgresql-x64-{version_dir.name}'
            try:
                result = subprocess.run(
                    [str(_SC_EXE), 'query', srv],
                    capture_output=True,
                    text=True,
                )
                if result.returncode == 0 and 'RUNNING' in result.stdout:
                    return srv
            except (FileNotFoundError, OSError):
                continue
    return None


def _stop_system_service(name: str):
    """Stop a Windows system service."""
    print(f'[INFO] Stopping system PostgreSQL service "{name}" ...')
    try:
        subprocess.run(
            [str(_SC_EXE), 'stop', name],
            capture_output=True,
            text=True,
            check=False,
        )
    except (FileNotFoundError, OSError) as exc:
        print(f'[WARNING] Could not stop service "{name}": {exc}')


def _start_system_service(name: str):
    """Start a Windows system service."""
    print(f'[INFO] Restarting system PostgreSQL service "{name}" ...')
    try:
        subprocess.run(
            [str(_SC_EXE), 'start', name],
            capture_output=True,
            text=True,
            check=False,
        )
    except (FileNotFoundError, OSError) as exc:
        print(f'[WARNING] Could not restart service "{name}": {exc}')


# ---------------------------------------------------------------------------
# Data directory management
# ---------------------------------------------------------------------------

_TRUST_HBA = """\
# Local development configuration — trust all local connections.
# Installed by app/local_postgres.py.
local   all   all                    trust
host   all   all   127.0.0.1/32     trust
host   all   all   ::1/128          trust
"""


def _init_cluster():
    """Create the data directory if it does not exist yet."""
    if DATADIR.is_dir() and (DATADIR / 'PG_VERSION').is_file():
        return

    initdb = _find_binary('initdb')
    if not initdb:
        print(
            '[ERROR] PostgreSQL not found. Install it from:',
            'https://www.postgresql.org/download/windows/',
            file=sys.stderr,
        )
        sys.exit(1)

    print(f'[INFO] Initialising PostgreSQL data directory at {DATADIR} ...')
    DATADIR.mkdir(parents=True, exist_ok=True)

    # Write a non-empty password to a temporary file. PostgreSQL >= 18
    # rejects an empty password (NUL device), and we immediately overwrite
    # pg_hba.conf with trust auth anyway, so the value is irrelevant.
    pw_file = tempfile.NamedTemporaryFile(  # noqa: SIM115 — delete=False needed for initdb --pwfile
        mode='w',
        delete=False,
        suffix='.pw',
        prefix='pg_',
    )
    pw_file.write('dev\n')
    pw_file.close()

    try:
        subprocess.run(
            [
                str(initdb),
                '--locale=C',
                '--encoding=UTF8',
                '--data-checksums',
                '--username',
                _USER,
                '--pwfile',
                pw_file.name,
                '-D',
                str(DATADIR),
            ],
            check=True,
            capture_output=True,
            text=True,
        )
    except subprocess.CalledProcessError as exc:
        print(f'[ERROR] initdb failed (exit {exc.returncode}):')
        if exc.stderr:
            print(exc.stderr)
        if exc.stdout:
            print(exc.stdout)
        sys.exit(1)
    finally:
        with contextlib.suppress(OSError):
            os.unlink(pw_file.name)

    _set_postgresql_conf()

    hba_path = DATADIR / 'pg_hba.conf'
    hba_path.write_text(_TRUST_HBA, encoding='utf-8')
    print('[INFO] pg_hba.conf set to trust local connections.')


def _set_postgresql_conf():
    """Write postgresql.conf to match our development settings."""
    conf_path = DATADIR / 'postgresql.conf'
    if not conf_path.is_file():
        return

    lines = conf_path.read_text(encoding='utf-8').splitlines()
    lines.append('\n# ---- local_postgres.py overrides ----')
    lines.append(f'port = {PORT}')
    lines.append("listen_addresses = '127.0.0.1'")
    lines.append("unix_socket_directories = ''")
    lines.append('autovacuum = off')  # avoid DLL conflicts with system PostgreSQL service
    lines.append('logging_collector = off')
    lines.append('log_destination = stderr')
    conf_path.write_text('\n'.join(lines), encoding='utf-8')


# ---------------------------------------------------------------------------
# Server start / stop
# ---------------------------------------------------------------------------


def _pg_ctl(args: list[str], cwd: str | None = None) -> subprocess.CompletedProcess:
    """Run ``pg_ctl`` with the data directory pre-filled."""
    pg_ctl = _find_binary('pg_ctl')
    if not pg_ctl:
        raise RuntimeError('pg_ctl binary not found')

    return subprocess.run(
        [str(pg_ctl), *args, '-D', str(DATADIR), '-w'],
        capture_output=True,
        text=True,
        cwd=cwd or str(DATADIR),
    )


def _server_running() -> bool:
    """Return True if the server is already up on our port."""
    try:
        result = _pg_ctl(['status'])
        return result.returncode == 0
    except (FileNotFoundError, RuntimeError):
        return False


def _create_database():
    """Create the database if it does not exist."""
    psql = _find_binary('psql')
    if not psql:
        raise RuntimeError('psql binary not found')

    check = subprocess.run(
        [
            str(psql),
            '-h',
            '127.0.0.1',
            '-p',
            str(PORT),
            '-U',
            _USER,
            '-d',
            'postgres',
            '-tAc',
            f"SELECT 1 FROM pg_database WHERE datname = '{_DATABASE}'",
        ],
        capture_output=True,
        text=True,
        env={**os.environ, 'PGPASSWORD': ''},
    )

    if check.stdout.strip() == '1':
        print(f'[INFO] Database "{_DATABASE}" already exists.')
        return

    subprocess.run(
        [
            str(psql),
            '-h',
            '127.0.0.1',
            '-p',
            str(PORT),
            '-U',
            _USER,
            '-d',
            'postgres',
            '-c',
            f'CREATE DATABASE "{_DATABASE}"',
        ],
        check=True,
        capture_output=True,
        text=True,
        env={**os.environ, 'PGPASSWORD': ''},
    )
    print(f'[INFO] Created database "{_DATABASE}".')


def _stop(*_args):
    """Stop the PostgreSQL server with a fast shutdown.

    Registered as ``atexit`` and signal handler so it fires on Ctrl+C,
    normal exit, and terminal close.
    """
    global _started
    global _stopped_system_service_name
    if not _started:
        return
    _started = False

    try:
        result = _pg_ctl(['stop', '-m', 'fast'])
        if result.returncode != 0:
            logger.info('pg_ctl stop returned %d: %s', result.returncode, result.stderr.strip())
        else:
            logger.info('Local PostgreSQL stopped.')
    except Exception as exc:  # noqa: BLE001 — shutdown must never raise
        logger.warning('Failed to stop local PostgreSQL: %s', exc)

    # Restart the system service if we stopped it earlier.
    if _stopped_system_service_name:
        srv = _stopped_system_service_name
        _stopped_system_service_name = None
        _start_system_service(srv)


# ---------------------------------------------------------------------------
# Public API
# ---------------------------------------------------------------------------


def start() -> str:
    """Start a local PostgreSQL server for development.

    Returns the ``DATABASE_URL`` the application should use.

    The caller is expected to set this on ``os.environ['DATABASE_URL']``
    before invoking ``create_app()``.

    This function is idempotent: if the server is already running and the
    database exists, it returns immediately.
    """
    global _started
    global _stopped_system_service_name

    # On Windows, temporarily stop the system PostgreSQL service so its child
    # processes (autovacuum, etc.) do not collide with our local instance.
    if _stopped_system_service_name is None:  # only do this once per process
        srv = _detect_system_service()
        if srv:
            _stop_system_service(srv)
            _stopped_system_service_name = srv

    # 1. Initialise if needed.
    _init_cluster()

    # 2. Start if not already running.
    if not _server_running():
        print(f'[INFO] Starting PostgreSQL on port {PORT} ...')
        PG_LOG.parent.mkdir(parents=True, exist_ok=True)
        pg_ctl = _find_binary('pg_ctl')
        try:
            subprocess.run(
                [
                    str(pg_ctl),
                    'start',
                    '-D',
                    str(DATADIR),
                    '-l',
                    str(PG_LOG),
                    '-w',  # wait for readiness
                ],
                check=True,
            )
        except subprocess.CalledProcessError:
            print('[ERROR] pg_ctl start failed.')
            # Also show the server log — it often has the real error.
            try:
                log_text = PG_LOG.read_text(encoding='utf-8')
                if log_text.strip():
                    print(f'Server log ({PG_LOG}):')
                    print(log_text)
            except Exception:  # noqa: BLE001
                pass
            sys.exit(1)
        print(f'[INFO] PostgreSQL is up (pid in {PG_LOG}).')
    else:
        print(f'[INFO] PostgreSQL is already running on port {PORT}.')

    # 3. Create the application database (idempotent).
    _create_database()

    # 4. Register cleanup (only in the *parent* process, not the reloader
    #    worker — the worker registers its own and tears the database down
    #    behind any sibling process that is still connected).
    #
    #    The Werkzeug reloader sets WERKZEUG_RUN_MAIN=true in the worker
    #    process and leaves it unset (or false) in the watcher process that
    #    acts as the parent.  We want the cleanup on the *actual process
    #    lifecycle*, i.e. when the watcher exits, not every time the
    #    reloader restarts the worker.
    if not _started:
        _started = True
        if os.environ.get('WERKZEUG_RUN_MAIN', '').lower() != 'true':
            atexit.register(_stop)
            with contextlib.suppress(ValueError, OSError):
                signal.signal(signal.SIGINT, _stop)
            with contextlib.suppress(AttributeError, ValueError, OSError):
                signal.signal(signal.SIGBREAK, _stop)  # type: ignore[attr-defined]

    logger.info('Local PostgreSQL ready — DATABASE_URL=%s', DATABASE_URL)
    return DATABASE_URL