ArchiveBox/archivebox/tests/test_postgres_backend.py

#!/usr/bin/env python3
"""
End-to-end tests for the PostgreSQL database backend (DATABASE_ENGINE=postgres).

Spins up a real throwaway PostgreSQL cluster (initdb + pg_ctl) for the module,
then exercises real archivebox CLI flows against it: init, status, re-init,
add --index-only, list, remove, and a full schema-vs-models parity check.

Requires PostgreSQL server binaries (initdb/pg_ctl) to be installed.
"""

import os
import shlex
import shutil
import socket
import subprocess
import tempfile
import uuid
from pathlib import Path

import pytest

from .conftest import cli_env, run_archivebox_cmd, run_queued_crawls


def _find_pg_bindir() -> Path:
    """Locate PostgreSQL server binaries (initdb) on this machine."""
    initdb_on_path = shutil.which("initdb")
    if initdb_on_path:
        return Path(initdb_on_path).resolve().parent
    candidates = []
    for base in (Path("/usr/lib/postgresql"), Path("/opt/homebrew/opt"), Path("/usr/local/opt"), Path("/opt/homebrew/Cellar/postgresql")):
        if base.is_dir():
            for sub in sorted(base.iterdir(), reverse=True):
                initdb = sub / "bin" / "initdb"
                if initdb.is_file():
                    candidates.append(initdb.parent)
    assert candidates, "PostgreSQL server binaries (initdb) are required for test_postgres_backend tests"
    return candidates[0]


def _free_port() -> int:
    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
        sock.bind(("127.0.0.1", 0))
        return sock.getsockname()[1]


PG_TEST_USER = "abxtest"


@pytest.fixture(scope="module")
def pg_server():
    """A real throwaway PostgreSQL server for the whole test module."""
    bindir = _find_pg_bindir()
    workdir = Path(tempfile.mkdtemp(prefix="abx-pg-test-"))
    datadir = workdir / "data"
    logfile = workdir / "server.log"
    port = _free_port()

    # postgres refuses to run as root; when the suite runs as root (e.g. in a
    # dev container), run the server as the system postgres user instead.
    run_as_postgres_user = os.geteuid() == 0
    if run_as_postgres_user:
        workdir.chmod(0o755)
        shutil.chown(workdir, user="postgres")

    def pg_cmd(args: list[str]) -> subprocess.CompletedProcess:
        if run_as_postgres_user:
            full = ["su", "postgres", "-s", "/bin/sh", "-c", shlex.join(args)]
        else:
            full = args
        return subprocess.run(full, capture_output=True, text=True, check=False, timeout=120)

    result = pg_cmd([str(bindir / "initdb"), "-D", str(datadir), "-E", "UTF8", "-A", "trust", "-U", PG_TEST_USER])
    assert result.returncode == 0, f"initdb failed: {result.stderr}"

    server_opts = f"-p {port} -k {workdir} -c listen_addresses=127.0.0.1 -c fsync=off -c synchronous_commit=off"
    result = pg_cmd([str(bindir / "pg_ctl"), "-D", str(datadir), "-l", str(logfile), "-o", server_opts, "-w", "start"])
    assert result.returncode == 0, f"pg_ctl start failed: {result.stderr}\n{logfile.read_text() if logfile.exists() else ''}"

    try:
        yield {"host": "127.0.0.1", "port": port, "user": PG_TEST_USER}
    finally:
        pg_cmd([str(bindir / "pg_ctl"), "-D", str(datadir), "-m", "immediate", "stop"])
        shutil.rmtree(workdir, ignore_errors=True)


def pg_cli_env(pg_server: dict, dbname: str, **extra) -> dict:
    env = cli_env(**extra)
    env.update(
        {
            "ARCHIVEBOX_DATABASE_ENGINE": "postgres",
            "ARCHIVEBOX_DATABASE_HOST": pg_server["host"],
            "ARCHIVEBOX_DATABASE_PORT": str(pg_server["port"]),
            "ARCHIVEBOX_DATABASE_USER": pg_server["user"],
            "ARCHIVEBOX_DATABASE_NAME": dbname,
        },
    )
    return env


def unique_dbname() -> str:
    return f"abx_test_{uuid.uuid4().hex[:12]}"


def pg_query(pg_server: dict, dbname: str, query: str) -> list[tuple]:
    import psycopg

    with psycopg.connect(host=pg_server["host"], port=pg_server["port"], user=pg_server["user"], dbname=dbname) as conn:
        return conn.execute(query).fetchall()


def test_init_creates_postgres_schema(pg_server, tmp_path):
    """Fresh init against postgres should create the database + full schema, and no sqlite file."""
    dbname = unique_dbname()
    env = pg_cli_env(pg_server, dbname)

    result = run_archivebox_cmd(["init", "--quick"], cwd=tmp_path, env=env, timeout=300)
    assert result.returncode == 0, f"Init failed: {result.stderr}\n{result.stdout}"

    assert not (tmp_path / "index.sqlite3").exists(), "sqlite file should not be created when using postgres"
    assert (tmp_path / "archive").is_dir(), "Archive dir not created"

    (applied_migrations,) = pg_query(pg_server, dbname, "SELECT COUNT(*) FROM django_migrations")[0]
    assert applied_migrations > 50, f"Expected all migrations applied, got {applied_migrations}"

    for table in ("core_snapshot", "core_archiveresult", "core_tag", "crawls_crawl", "machine_machine", "api_apitoken", "personas_persona"):
        (regclass,) = pg_query(pg_server, dbname, f"SELECT to_regclass('{table}')")[0]
        assert regclass == table, f"Table {table} missing from postgres schema"


def test_postgres_schema_matches_models(pg_server, tmp_path):
    """Every model column must exist in postgres and vice versa (no state/schema drift)."""
    dbname = unique_dbname()
    env = pg_cli_env(pg_server, dbname)

    result = run_archivebox_cmd(["init", "--quick"], cwd=tmp_path, env=env, timeout=300)
    assert result.returncode == 0, f"Init failed: {result.stderr}"

    parity_script = (
        "from django.apps import apps\n"
        "from django.db import connection\n"
        "problems = []\n"
        "with connection.cursor() as cursor:\n"
        "    tables = set(connection.introspection.table_names(cursor))\n"
        "    for model in apps.get_models(include_auto_created=True):\n"
        "        meta = model._meta\n"
        "        if not meta.managed or meta.proxy:\n"
        "            continue\n"
        "        if meta.db_table not in tables:\n"
        "            problems.append(f'missing table {meta.db_table}')\n"
        "            continue\n"
        "        db_cols = {col.name for col in connection.introspection.get_table_description(cursor, meta.db_table)}\n"
        "        model_cols = {field.column for field in meta.local_concrete_fields}\n"
        "        for col in sorted(model_cols - db_cols):\n"
        "            problems.append(f'{meta.db_table}: missing column {col}')\n"
        "        for col in sorted(db_cols - model_cols):\n"
        "            problems.append(f'{meta.db_table}: extra column {col}')\n"
        "print('SCHEMA_PROBLEMS=' + repr(sorted(problems)))\n"
    )
    result = run_archivebox_cmd(["manage", "shell", "-c", parity_script], cwd=tmp_path, env=env, timeout=120)
    assert result.returncode == 0, f"manage shell failed: {result.stderr}"
    assert "SCHEMA_PROBLEMS=[]" in result.stdout, f"Postgres schema diverges from models:\n{result.stdout}\n{result.stderr}"

    result = run_archivebox_cmd(["manage", "makemigrations", "--check", "--dry-run"], cwd=tmp_path, env=env, timeout=120)
    assert result.returncode == 0, f"Model state does not match migrations: {result.stdout}\n{result.stderr}"


def test_status_and_reinit_on_postgres(pg_server, tmp_path):
    """status works against postgres, and a second init takes the 'verify existing' path."""
    dbname = unique_dbname()
    env = pg_cli_env(pg_server, dbname)

    result = run_archivebox_cmd(["init", "--quick"], cwd=tmp_path, env=env, timeout=300)
    assert result.returncode == 0, f"Init failed: {result.stderr}"
    assert "Initializing a new ArchiveBox" in result.stdout

    result = run_archivebox_cmd(["status"], cwd=tmp_path, env=env, timeout=120)
    assert result.returncode == 0, f"Status failed: {result.stderr}"
    assert f"postgresql://{PG_TEST_USER}@" in result.stdout, f"status should show the postgres DSN:\n{result.stdout}"

    result = run_archivebox_cmd(["init", "--quick"], cwd=tmp_path, env=env, timeout=300)
    assert result.returncode == 0, f"Re-init failed: {result.stderr}"
    assert "Verifying and updating existing ArchiveBox collection" in result.stdout


def test_add_list_remove_on_postgres(pg_server, tmp_path):
    """Real add/list/remove CLI flows store and retrieve rows from postgres."""
    dbname = unique_dbname()
    env = pg_cli_env(pg_server, dbname, disable_extractors=True)
    test_url = "https://example.com/abx-postgres-test"

    result = run_archivebox_cmd(["init", "--quick"], cwd=tmp_path, env=env, timeout=300)
    assert result.returncode == 0, f"Init failed: {result.stderr}"

    result = run_archivebox_cmd(["add", "--index-only", test_url], cwd=tmp_path, env=env, timeout=300)
    assert result.returncode == 0, f"Add failed: {result.stderr}\n{result.stdout}"
    run_queued_crawls(tmp_path, env)

    (snapshot_count,) = pg_query(pg_server, dbname, "SELECT COUNT(*) FROM core_snapshot")[0]
    assert snapshot_count >= 1, "Snapshot row not written to postgres"
    (crawl_count,) = pg_query(pg_server, dbname, "SELECT COUNT(*) FROM crawls_crawl")[0]
    assert crawl_count >= 1, "Crawl row not written to postgres"

    result = run_archivebox_cmd(["list"], cwd=tmp_path, env=env, timeout=120)
    assert result.returncode == 0, f"List failed: {result.stderr}"
    assert "example.com/abx-postgres-test" in result.stdout, f"Added URL missing from list output:\n{result.stdout}"

    result = run_archivebox_cmd(["search", "abx-postgres-test"], cwd=tmp_path, env=env, timeout=120)
    assert result.returncode == 0, f"Search failed: {result.stderr}"
    assert "example.com/abx-postgres-test" in result.stdout, f"Added URL missing from search output:\n{result.stdout}"

    result = run_archivebox_cmd(
        ["remove", "--yes", "--delete", "--filter-type=exact", test_url],
        cwd=tmp_path,
        env=env,
        timeout=120,
    )
    assert result.returncode == 0, f"Remove failed: {result.stderr}\n{result.stdout}"
    (snapshot_count,) = pg_query(pg_server, dbname, f"SELECT COUNT(*) FROM core_snapshot WHERE url = '{test_url}'")[0]
    assert snapshot_count == 0, "Snapshot row should be deleted from postgres"


def test_sqlite_remains_the_default(tmp_path):
    """Without DATABASE_ENGINE config, init keeps using the sqlite file backend."""
    result = run_archivebox_cmd(["init", "--quick"], cwd=tmp_path, env=cli_env(), timeout=300)
    assert result.returncode == 0, f"Init failed: {result.stderr}"
    assert (tmp_path / "index.sqlite3").exists(), "sqlite file should be created by default"