db-git 0.1.0__py3-none-any.whl

git_db/backends/postgresql/backend.py

@@ -0,0 +1,176 @@
+from __future__ import annotations
+
+import os
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, ClassVar
+
+import psycopg
+
+from git_db.backends import (
+    BranchDbManager,
+    DbConnection,
+    SnapshotStrategy,
+    register_backend,
+)
+from git_db.backends.postgresql.branch_db import PostgresBranchDbManager
+from git_db.backends.postgresql.pgdump import PgDumpStrategy
+from git_db.backends.postgresql.template import TemplateStrategy
+from git_db.db import parse_database_url
+from git_db.errors import ConfigError, DatabaseError
+
+if TYPE_CHECKING:
+    from git_db.config import GitDbConfig
+
+
+@dataclass
+class PgPermissions:
+    """
+    PostgreSQL role permissions relevant to git-db operations.
+    """
+
+    can_createdb: bool
+    is_superuser: bool
+    has_pg_signal_backend: bool
+
+
+class PostgresqlBackend:
+    """
+    PostgreSQL backend implementing the DatabaseBackend protocol.
+    """
+
+    engine = "postgresql"
+    max_identifier_length = 63
+    _VALID_STRATEGIES: ClassVar[set[str]] = {"template", "pgdump"}
+
+    def apply_url_defaults(
+        self, params: dict[str, str | int | None]
+    ) -> dict[str, str | int]:
+        """
+        Apply PostgreSQL-specific defaults to parsed URL parameters.
+        """
+        dbname = params.get("dbname")
+        if not dbname:
+            raise ConfigError("Database name is required in the connection URL.")
+
+        port = params.get("port")
+        return {
+            "user": params.get("user") or "postgres",
+            "password": params.get("password") or "",
+            "host": params.get("host") or "localhost",
+            "port": port if port is not None else 5432,
+            "dbname": dbname,
+        }
+
+    def get_engine_version(self, url: str) -> int:
+        """
+        Return the major PostgreSQL version as an integer.
+        """
+        params = self.apply_url_defaults(parse_database_url(url))
+        conn = self.connect_maintenance(params)
+        try:
+            cur = conn.execute("SHOW server_version_num")
+            row = cur.fetchone()
+            if not row:
+                raise DatabaseError("SHOW server_version_num returned no rows")
+            try:
+                return int(row[0]) // 10000
+            except (TypeError, ValueError) as e:
+                raise DatabaseError(
+                    f"Could not parse server_version_num: {row[0]!r}"
+                ) from e
+        finally:
+            conn.close()
+
+    def connect_maintenance(self, params: dict[str, str | int]) -> DbConnection:
+        """
+        Connect to the 'postgres' maintenance database for admin operations.
+        """
+        conninfo = (
+            f"host={params['host']} port={params['port']} "
+            f"user={params['user']} dbname=postgres"
+        )
+        password = str(params.get("password", ""))
+        try:
+            return psycopg.connect(conninfo, autocommit=True, password=password or None)
+        except psycopg.Error as e:
+            raise DatabaseError(
+                f"Could not connect to PostgreSQL at "
+                f"{params['host']}:{params['port']}: {e}"
+            ) from e
+
+    def build_subprocess_env(self, params: dict[str, str | int]) -> dict[str, str]:
+        """
+        Build environment dict with PGPASSWORD set for subprocess calls.
+        """
+        env = os.environ.copy()
+        password = str(params.get("password", ""))
+        if password:
+            env["PGPASSWORD"] = password
+        return env
+
+    def check_permissions(self, url: str) -> PgPermissions:
+        """
+        Check the current user's PostgreSQL permissions.
+        """
+        params = self.apply_url_defaults(parse_database_url(url))
+        try:
+            conn = self.connect_maintenance(params)
+        except DatabaseError:
+            return PgPermissions(False, False, False)
+        try:
+            cur = conn.execute(
+                "SELECT rolcreatedb, rolsuper, "
+                "pg_has_role(current_user, 'pg_signal_backend', 'MEMBER') "
+                "FROM pg_roles WHERE rolname = current_user"
+            )
+            row = cur.fetchone()
+            if not row:
+                return PgPermissions(False, False, False)
+            return PgPermissions(
+                can_createdb=bool(row[0]),
+                is_superuser=bool(row[1]),
+                has_pg_signal_backend=bool(row[2]),
+            )
+        except psycopg.Error:
+            return PgPermissions(False, False, False)
+        finally:
+            conn.close()
+
+    def database_exists(self, url: str, name: str) -> bool:
+        """
+        Return whether a PostgreSQL database exists.
+        """
+        params = self.apply_url_defaults(parse_database_url(url))
+        conn = self.connect_maintenance(params)
+        try:
+            cur = conn.execute("SELECT 1 FROM pg_database WHERE datname = %s", (name,))
+            return cur.fetchone() is not None
+        except psycopg.Error as e:
+            raise DatabaseError(f"Could not inspect database '{name}': {e}") from e
+        finally:
+            conn.close()
+
+    def detect_strategy(self, config: GitDbConfig) -> SnapshotStrategy:
+        """
+        Return the configured snapshot strategy for this PostgreSQL instance.
+        """
+        if config.strategy not in self._VALID_STRATEGIES:
+            raise ConfigError(
+                f"Unknown strategy '{config.strategy}'. "
+                f"Valid strategies for PostgreSQL: "
+                f"{', '.join(sorted(self._VALID_STRATEGIES))}."
+            )
+
+        if config.strategy == "pgdump":
+            pg_version = self.get_engine_version(config.database_url)
+            return PgDumpStrategy(backend=self, pg_version=pg_version)
+        return TemplateStrategy(backend=self)
+
+    def branch_db_manager(self, config: GitDbConfig) -> BranchDbManager:
+        """
+        Return the per-branch database manager for this PostgreSQL instance.
+        """
+        return PostgresBranchDbManager(backend=self, config=config)
+
+
+register_backend("postgresql", PostgresqlBackend)

git_db/backends/postgresql/branch_db.py

@@ -0,0 +1,203 @@
+from __future__ import annotations
+
+import shutil
+import subprocess
+from pathlib import Path
+from typing import TYPE_CHECKING
+
+import psycopg
+from psycopg import sql
+
+from git_db.backends import DatabaseBackend
+from git_db.backends.postgresql.connections import handle_active_connections
+from git_db.backends.postgresql.template import _create_from_template
+from git_db.db import parse_database_url
+from git_db.errors import SnapshotError, ToolNotFoundError
+from git_db.state import BranchDbEntry, load_state, record_branch_db, remove_branch_db
+
+if TYPE_CHECKING:
+    from git_db.config import GitDbConfig
+
+
+class PostgresBranchDbManager:
+    """
+    Per-branch database operations for PostgreSQL.
+    """
+
+    def __init__(self, backend: DatabaseBackend, config: GitDbConfig) -> None:
+        self._backend = backend
+        self._config = config
+        self._params = backend.apply_url_defaults(
+            parse_database_url(config.database_url)
+        )
+
+    def exists(self, name: str) -> bool:
+        conn = self._backend.connect_maintenance(self._params)
+        try:
+            cur = conn.execute(
+                "SELECT 1 FROM pg_database WHERE datname = %s",
+                (name,),
+            )
+            return cur.fetchone() is not None
+        finally:
+            conn.close()
+
+    def create(
+        self,
+        target: str,
+        source: str,
+        branch: str,
+        created_from: str,
+        git_dir: Path,
+    ) -> None:
+        strategy = self._backend.detect_strategy(self._config)
+
+        if strategy.name == "template":
+            _create_via_template(
+                self._backend, self._params, target, source, self._config
+            )
+        else:
+            _create_via_pgdump(
+                self._backend,
+                self._params,
+                self._backend.build_subprocess_env(self._params),
+                target,
+                source,
+                self._config,
+            )
+
+        record_branch_db(git_dir, branch, target, created_from)
+
+    def drop(self, name: str, branch: str, git_dir: Path) -> None:
+        conn = self._backend.connect_maintenance(self._params)
+        try:
+            handle_active_connections(conn, name, self._config)
+            conn.execute(
+                sql.SQL("DROP DATABASE IF EXISTS {} WITH (FORCE)").format(
+                    sql.Identifier(name)
+                )
+            )
+        finally:
+            conn.close()
+
+        remove_branch_db(git_dir, branch)
+
+    def list(self, git_dir: Path) -> list[tuple[str, BranchDbEntry, bool]]:
+        state = load_state(git_dir)
+        return [
+            (branch, entry, self.exists(entry.db_name))
+            for branch, entry in state.databases.items()
+        ]
+
+
+def _create_via_template(
+    backend: DatabaseBackend,
+    params: dict[str, str | int],
+    target: str,
+    source: str,
+    config: GitDbConfig,
+) -> None:
+    """
+    Create a database using CREATE DATABASE ... TEMPLATE.
+    """
+    conn = backend.connect_maintenance(params)
+    try:
+        handle_active_connections(conn, source, config)
+        handle_active_connections(conn, target, config)
+        conn.execute(
+            sql.SQL("DROP DATABASE IF EXISTS {} WITH (FORCE)").format(
+                sql.Identifier(target)
+            )
+        )
+
+        _create_from_template(conn, target, source)
+    except psycopg.Error as e:
+        raise SnapshotError(f"Template clone failed: {e}") from e
+    finally:
+        conn.close()
+
+
+def _create_via_pgdump(
+    backend: DatabaseBackend,
+    params: dict[str, str | int],
+    env: dict[str, str],
+    target: str,
+    source: str,
+    config: GitDbConfig,
+) -> None:
+    """
+    Create a database by piping pg_dump to pg_restore.
+    """
+    pg_dump = shutil.which("pg_dump")
+    pg_restore = shutil.which("pg_restore")
+
+    if not pg_dump:
+        raise ToolNotFoundError("pg_dump not found in PATH.")
+    if not pg_restore:
+        raise ToolNotFoundError("pg_restore not found in PATH.")
+
+    common_args = [
+        "-h",
+        str(params["host"]),
+        "-p",
+        str(params["port"]),
+        "-U",
+        str(params["user"]),
+    ]
+
+    conn = backend.connect_maintenance(params)
+    try:
+        handle_active_connections(conn, target, config)
+        target_ident = sql.Identifier(target)
+        conn.execute(
+            sql.SQL("DROP DATABASE IF EXISTS {} WITH (FORCE)").format(target_ident)
+        )
+        conn.execute(
+            sql.SQL("CREATE DATABASE {} TEMPLATE template0").format(target_ident)
+        )
+    except psycopg.Error as e:
+        raise SnapshotError(f"Create branch database failed: {e}") from e
+    finally:
+        conn.close()
+
+    dump_cmd = [
+        pg_dump,
+        "-Fc",
+        "--no-owner",
+        "--no-privileges",
+        *common_args,
+        source,
+    ]
+    restore_cmd = [
+        pg_restore,
+        "--no-owner",
+        "--no-privileges",
+        *common_args,
+        "-d",
+        target,
+    ]
+
+    dump_proc = subprocess.Popen(
+        dump_cmd,
+        stdout=subprocess.PIPE,
+        stderr=subprocess.PIPE,
+        env=env,
+    )
+    restore_result = subprocess.run(
+        restore_cmd,
+        stdin=dump_proc.stdout,
+        capture_output=True,
+        text=True,
+        env=env,
+        timeout=300,
+    )
+    if dump_proc.stdout is not None:
+        dump_proc.stdout.close()
+    dump_proc.wait(timeout=300)
+
+    if dump_proc.returncode != 0:
+        stderr = dump_proc.stderr.read().decode() if dump_proc.stderr else ""
+        raise SnapshotError(f"pg_dump failed: {stderr.strip()}")
+
+    if restore_result.returncode != 0 and "ERROR" in restore_result.stderr:
+        raise SnapshotError(f"pg_restore failed: {restore_result.stderr.strip()}")

git_db/backends/postgresql/connections.py

@@ -0,0 +1,82 @@
+from __future__ import annotations
+
+import contextlib
+import time
+from typing import TYPE_CHECKING, NamedTuple
+
+from psycopg import sql
+
+from git_db.backends import DbConnection
+from git_db.errors import ActiveConnectionsError, TerminationTimeout
+
+if TYPE_CHECKING:
+    from git_db.config import GitDbConfig
+
+
+class ConnectionInfo(NamedTuple):
+    pid: int
+    application_name: str
+    state: str
+    query_start: str | None
+
+
+def check_connections(conn: DbConnection, dbname: str) -> list[ConnectionInfo]:
+    """
+    Return list of active connections to the given database,
+    excluding our own connection.
+    """
+    cur = conn.execute(
+        "SELECT pid, application_name, state, query_start::text "
+        "FROM pg_stat_activity "
+        "WHERE datname = %s AND pid <> pg_backend_pid()",
+        (dbname,),
+    )
+    return [ConnectionInfo(*row) for row in cur.fetchall()]
+
+
+def handle_active_connections(
+    conn: DbConnection,
+    dbname: str,
+    config: GitDbConfig,
+) -> bool:
+    """
+    Handle active connections based on config policy.
+    """
+    active = check_connections(conn, dbname)
+    if not active:
+        return True
+
+    if config.on_active_connections == "fail":
+        raise ActiveConnectionsError(
+            f"{len(active)} active connection(s) to '{dbname}'. "
+            "Stop your dev server or set on_active_connections = 'terminate'."
+        )
+
+    _terminate_all(conn, dbname, config.force_terminate_timeout_ms)
+    return True
+
+
+def _terminate_all(conn: DbConnection, dbname: str, timeout_ms: int = 5000) -> None:
+    """
+    Block new connections, terminate existing ones, then restore access.
+    """
+    db_ident = sql.Identifier(dbname)
+    conn.execute(sql.SQL("ALTER DATABASE {} ALLOW_CONNECTIONS false").format(db_ident))
+    try:
+        deadline = time.monotonic() + timeout_ms / 1000
+        while time.monotonic() < deadline:
+            remaining = check_connections(conn, dbname)
+            if not remaining:
+                break
+            for c in remaining:
+                conn.execute("SELECT pg_terminate_backend(%s)", (c.pid,))
+            time.sleep(0.1)
+        else:
+            raise TerminationTimeout(
+                f"Could not terminate all connections within {timeout_ms}ms"
+            )
+    finally:
+        with contextlib.suppress(Exception):
+            conn.execute(
+                sql.SQL("ALTER DATABASE {} ALLOW_CONNECTIONS true").format(db_ident)
+            )

git_db/backends/postgresql/pgdump.py

@@ -0,0 +1,212 @@
+from __future__ import annotations
+
+import shutil
+import subprocess
+from pathlib import Path
+from typing import TYPE_CHECKING
+
+import psycopg
+from psycopg import sql
+
+from git_db.backends import DatabaseBackend
+from git_db.backends.postgresql.connections import handle_active_connections
+from git_db.db import parse_database_url
+from git_db.errors import DatabaseError, SnapshotError, ToolNotFoundError
+from git_db.storage import (
+    ensure_snapshot_dir,
+    make_metadata,
+    metadata_path,
+    snapshot_dump_path,
+    write_metadata,
+)
+
+if TYPE_CHECKING:
+    from git_db.config import GitDbConfig
+
+
+class PgDumpStrategy:
+    """
+    Snapshot strategy using pg_dump/pg_restore.
+    """
+
+    name = "pgdump"
+
+    def __init__(self, backend: DatabaseBackend, pg_version: int) -> None:
+        self._backend = backend
+        self._pg_version = pg_version
+
+    def save(
+        self,
+        db_url: str,
+        branch: str,
+        snapshot_dir: Path,
+        config: GitDbConfig,
+    ) -> None:
+        pg_dump = shutil.which("pg_dump")
+        if not pg_dump:
+            raise ToolNotFoundError(
+                "pg_dump not found in PATH. Install PostgreSQL client tools."
+            )
+
+        params = self._backend.apply_url_defaults(parse_database_url(db_url))
+        env = self._backend.build_subprocess_env(params)
+
+        ensure_snapshot_dir(snapshot_dir)
+        dump_file = snapshot_dump_path(snapshot_dir, branch)
+
+        cmd = _build_pg_dump_cmd(
+            pg_dump,
+            params,
+            str(dump_file),
+            self._pg_version,
+        )
+        result = subprocess.run(
+            cmd,
+            capture_output=True,
+            text=True,
+            env=env,
+            timeout=300,
+        )
+        if result.returncode != 0:
+            raise SnapshotError(f"pg_dump failed: {result.stderr.strip()}")
+
+        write_metadata(
+            snapshot_dir,
+            make_metadata(
+                branch=branch,
+                database=str(params["dbname"]),
+                strategy=self.name,
+                engine=self._backend.engine,
+                engine_version=str(self._pg_version),
+                file_size_bytes=dump_file.stat().st_size,
+            ),
+        )
+
+    def restore(
+        self,
+        db_url: str,
+        branch: str,
+        snapshot_dir: Path,
+        config: GitDbConfig,
+    ) -> None:
+        pg_restore = shutil.which("pg_restore")
+        if not pg_restore:
+            raise ToolNotFoundError(
+                "pg_restore not found in PATH. Install PostgreSQL client tools."
+            )
+
+        params = self._backend.apply_url_defaults(parse_database_url(db_url))
+        env = self._backend.build_subprocess_env(params)
+        dump_file = snapshot_dump_path(snapshot_dir, branch)
+
+        if not dump_file.exists():
+            raise SnapshotError(
+                f"No dump file found for branch '{branch}' at {dump_file}"
+            )
+
+        self._drop_and_create_db(params, config)
+
+        cmd = _build_pg_restore_cmd(pg_restore, params, str(dump_file))
+        result = subprocess.run(
+            cmd,
+            capture_output=True,
+            text=True,
+            env=env,
+            timeout=300,
+        )
+        if result.returncode != 0 and "ERROR" in result.stderr:
+            raise SnapshotError(f"pg_restore failed: {result.stderr.strip()}")
+
+    def cleanup(
+        self,
+        branch: str,
+        snapshot_dir: Path,
+        config: GitDbConfig,
+    ) -> None:
+        dump = snapshot_dump_path(snapshot_dir, branch)
+        meta = metadata_path(snapshot_dir, branch)
+        if dump.exists():
+            dump.unlink()
+        if meta.exists():
+            meta.unlink()
+
+    def _drop_and_create_db(
+        self,
+        params: dict[str, str | int],
+        config: GitDbConfig,
+    ) -> None:
+        """
+        Drop and recreate the target database via the maintenance connection.
+        """
+        dbname = str(params["dbname"])
+        conn = self._backend.connect_maintenance(params)
+        try:
+            handle_active_connections(conn, dbname, config)
+            db_ident = sql.Identifier(dbname)
+            conn.execute(
+                sql.SQL("DROP DATABASE IF EXISTS {} WITH (FORCE)").format(db_ident)
+            )
+            conn.execute(
+                sql.SQL("CREATE DATABASE {} TEMPLATE template0").format(db_ident)
+            )
+        except (psycopg.Error, DatabaseError) as e:
+            raise SnapshotError(f"Drop/create database failed: {e}") from e
+        finally:
+            conn.close()
+
+
+def _build_pg_dump_cmd(
+    pg_dump: str,
+    params: dict[str, str | int],
+    dump_path: str,
+    pg_version: int,
+) -> list[str]:
+    """
+    Build the pg_dump command list.
+    """
+    cmd = [
+        pg_dump,
+        "-Fc",
+        "--no-owner",
+        "--no-privileges",
+        "-h",
+        str(params["host"]),
+        "-p",
+        str(params["port"]),
+        "-U",
+        str(params["user"]),
+        "-f",
+        dump_path,
+    ]
+
+    if pg_version >= 16:
+        cmd.extend(["--compress", "zstd:3"])
+    else:
+        cmd.extend(["-Z", "1"])
+
+    cmd.append(str(params["dbname"]))
+    return cmd
+
+
+def _build_pg_restore_cmd(
+    pg_restore: str,
+    params: dict[str, str | int],
+    dump_path: str,
+) -> list[str]:
+    """
+    Build the pg_restore command list.
+    """
+    return [
+        pg_restore,
+        "--no-owner",
+        "--no-privileges",
+        "-h",
+        str(params["host"]),
+        "-p",
+        str(params["port"]),
+        "-U",
+        str(params["user"]),
+        "-d",
+        str(params["dbname"]),
+        dump_path,
+    ]