teleo-infrastructure/tests/test_observatory_read_role_postgres.py

from __future__ import annotations

import shutil
import subprocess
import time
import uuid
from pathlib import Path

import pytest

ROOT = Path(__file__).resolve().parents[1]
ROLE_SQL = ROOT / "ops" / "observatory_read_role.sql"
POSTGRES_IMAGE = "postgres:16-alpine"
DATABASE = "teleo_canonical"
IAM_USER = "sa-observatory-read-adapter@teleo-501523.iam"
TEST_PASSWORD = "generated-test-only-password"

BOOTSTRAP_SQL = f"""
begin;
create schema kb_stage;
create role "{IAM_USER}" login inherit password '{TEST_PASSWORD}';
create table public.claims (id uuid primary key);
create table public.sources (id uuid primary key);
create table public.claim_evidence (id uuid primary key);
create table public.claim_edges (id uuid primary key);
create table public.private_notes (id uuid primary key);
create table kb_stage.kb_proposals (id uuid primary key);
insert into public.claims values ('11111111-1111-1111-1111-111111111111');
commit;
"""


def run(command: list[str], *, input_text: str | None = None, check: bool = True) -> subprocess.CompletedProcess[str]:
    return subprocess.run(
        command,
        input=input_text,
        text=True,
        capture_output=True,
        check=check,
        timeout=30,
    )


@pytest.mark.skipif(shutil.which("docker") is None, reason="Docker is required")
def test_observatory_role_is_read_only_and_exactly_allowlisted() -> None:
    container = f"teleo-observatory-role-{uuid.uuid4().hex[:10]}"
    try:
        run(
            [
                "docker",
                "run",
                "--detach",
                "--rm",
                "--name",
                container,
                "--env",
                f"POSTGRES_PASSWORD={TEST_PASSWORD}",
                "--env",
                f"POSTGRES_DB={DATABASE}",
                POSTGRES_IMAGE,
            ]
        )
        deadline = time.monotonic() + 25
        while time.monotonic() < deadline:
            ready = run(
                ["docker", "exec", container, "pg_isready", "-U", "postgres", "-d", DATABASE],
                check=False,
            )
            if ready.returncode == 0:
                break
            time.sleep(0.25)
        else:
            raise AssertionError("ephemeral Postgres did not become ready")

        bootstrap: subprocess.CompletedProcess[str] | None = None
        deadline = time.monotonic() + 15
        while time.monotonic() < deadline:
            bootstrap = run(
                [
                    "docker",
                    "exec",
                    "-i",
                    container,
                    "psql",
                    "--set",
                    "ON_ERROR_STOP=1",
                    "-U",
                    "postgres",
                    "-d",
                    DATABASE,
                ],
                input_text=BOOTSTRAP_SQL,
                check=False,
            )
            if bootstrap.returncode == 0:
                break
            time.sleep(0.25)
        else:
            assert bootstrap is not None
            raise AssertionError(f"ephemeral Postgres bootstrap failed: {bootstrap.stderr[-1000:]}")
        migration = run(
            [
                "docker",
                "exec",
                "-e",
                f"OBSERVATORY_DB_IAM_USER={IAM_USER}",
                "-i",
                container,
                "psql",
                "-U",
                "postgres",
                "-d",
                DATABASE,
            ],
            input_text=ROLE_SQL.read_text(encoding="utf-8"),
        )
        assert '"effective_table_writes_denied": true' in migration.stdout

        readback = run(
            [
                "docker",
                "exec",
                "-e",
                f"PGPASSWORD={TEST_PASSWORD}",
                container,
                "psql",
                "-h",
                "127.0.0.1",
                "-U",
                IAM_USER,
                "-d",
                DATABASE,
                "-At",
                "-c",
                "show default_transaction_read_only; select count(*) from public.claims;",
            ]
        )
        assert readback.stdout.splitlines() == ["on", "1"]

        write_attempt = run(
            [
                "docker",
                "exec",
                "-e",
                f"PGPASSWORD={TEST_PASSWORD}",
                container,
                "psql",
                "-h",
                "127.0.0.1",
                "-U",
                IAM_USER,
                "-d",
                DATABASE,
                "-c",
                "insert into public.claims values ('22222222-2222-2222-2222-222222222222');",
            ],
            check=False,
        )
        assert write_attempt.returncode != 0

        outside_allowlist = run(
            [
                "docker",
                "exec",
                "-e",
                f"PGPASSWORD={TEST_PASSWORD}",
                container,
                "psql",
                "-h",
                "127.0.0.1",
                "-U",
                IAM_USER,
                "-d",
                DATABASE,
                "-c",
                "select * from public.private_notes;",
            ],
            check=False,
        )
        assert outside_allowlist.returncode != 0
    finally:
        run(["docker", "rm", "--force", container], check=False)