teleo-infrastructure/tests/test_kb_apply_prereqs.py

from __future__ import annotations

import shutil
import subprocess
import time
import uuid
from collections.abc import Iterator
from pathlib import Path

import pytest

REPO_ROOT = Path(__file__).resolve().parents[1]
MIGRATION = REPO_ROOT / "scripts" / "kb_apply_prereqs.sql"
POSTGRES_IMAGE = "postgres:16-alpine"
DATABASE = "teleo_clone"

BOOTSTRAP_SQL = r"""
create role kb_gate_owner nologin inherit;
create role kb_review login inherit;
create role kb_apply login inherit;
create role broad_writer nologin;
create role rogue_owner nologin;
create role rogue_grantee nologin;

create schema kb_stage;

create table public.agents (
  id uuid primary key,
  handle text not null unique,
  kind text not null
);
create table public.strategies (
  id uuid primary key,
  agent_id uuid not null,
  active boolean not null default true
);
create table public.strategy_nodes (id uuid primary key);
create table public.claim_evidence (
  id uuid primary key,
  claim_id uuid not null,
  source_id uuid not null,
  role text not null
);
create table public.claim_edges (id uuid primary key);
create table public.claims (id uuid primary key);
create table public.sources (
  id uuid primary key,
  hash text not null
);
create table public.reasoning_tools (id uuid primary key);

insert into public.agents (id, handle, kind)
values ('11111111-1111-1111-1111-111111111111', 'm3ta', 'human');

create table kb_stage.kb_proposals (
  id uuid primary key,
  proposal_type text not null,
  status text not null,
  payload jsonb not null,
  reviewed_by_handle text,
  reviewed_by_agent_id uuid,
  reviewed_at timestamptz,
  review_note text,
  applied_by_handle text,
  applied_by_agent_id uuid,
  applied_at timestamptz,
  updated_at timestamptz not null default now()
);
create table kb_stage.kb_review_principals (
  db_role name primary key,
  reviewed_by_handle text not null,
  reviewed_by_agent_id uuid not null references public.agents(id),
  active boolean not null default true,
  created_at timestamptz not null default now()
);
create table kb_stage.kb_proposal_approvals (
  proposal_id uuid primary key references kb_stage.kb_proposals(id),
  proposal_type text not null,
  payload jsonb not null,
  reviewed_by_handle text not null,
  reviewed_by_agent_id uuid not null references public.agents(id),
  reviewed_by_db_role name not null,
  reviewed_at timestamptz not null,
  review_note text not null
);

create function kb_stage.approve_strict_proposal(uuid, text, text)
returns jsonb language sql security definer
as $function$ select '{}'::jsonb $function$;
create function kb_stage.approve_strict_proposal(uuid, text, jsonb, text, text)
returns jsonb language sql security definer
as $function$ select '{}'::jsonb $function$;
create function kb_stage.assert_approved_proposal(
  uuid, text, jsonb, text, uuid, timestamptz, text
) returns void language sql security definer
as $function$ select $function$;
create function kb_stage.finish_approved_proposal(
  uuid, text, jsonb, text, uuid, timestamptz, text, text
) returns jsonb language sql security definer
as $function$ select '{}'::jsonb $function$;

grant all privileges on all tables in schema public
  to kb_gate_owner, kb_apply, kb_review;
grant all privileges on all tables in schema kb_stage
  to public, kb_gate_owner, kb_apply, kb_review;
grant all privileges on all functions in schema kb_stage
  to public, kb_gate_owner, kb_apply, kb_review;
grant update (review_note) on kb_stage.kb_proposals to kb_apply;
grant select (review_note) on kb_stage.kb_proposals to public;
"""


def _psql(container: str, sql: str, *, check: bool = True) -> subprocess.CompletedProcess[str]:
    completed = subprocess.run(
        [
            "docker",
            "exec",
            "-i",
            container,
            "psql",
            "-X",
            "-v",
            "ON_ERROR_STOP=1",
            "-U",
            "postgres",
            "-d",
            DATABASE,
            "-At",
            "-F",
            "|",
        ],
        input=sql,
        text=True,
        capture_output=True,
        check=False,
    )
    if check and completed.returncode != 0:
        raise AssertionError(
            f"psql failed with exit {completed.returncode}\nstdout:\n{completed.stdout}\nstderr:\n{completed.stderr}"
        )
    return completed


def _apply_migration(container: str, *, check: bool = True) -> subprocess.CompletedProcess[str]:
    return _psql(container, MIGRATION.read_text(), check=check)


@pytest.fixture(scope="module")
def migrated_postgres() -> Iterator[str]:
    if shutil.which("docker") is None:
        pytest.skip("Docker is required for the disposable PostgreSQL migration canary")

    container = f"kb-apply-prereqs-{uuid.uuid4().hex[:12]}"
    started = subprocess.run(
        [
            "docker",
            "run",
            "--rm",
            "-d",
            "--name",
            container,
            "-e",
            "POSTGRES_PASSWORD=postgres",
            "-e",
            f"POSTGRES_DB={DATABASE}",
            POSTGRES_IMAGE,
        ],
        text=True,
        capture_output=True,
        check=False,
    )
    if started.returncode != 0:
        pytest.fail(f"could not start disposable PostgreSQL: {started.stderr}")

    try:
        for _ in range(120):
            logs = subprocess.run(
                ["docker", "logs", container],
                text=True,
                capture_output=True,
                check=False,
            )
            query = _psql(container, "select 1;", check=False)
            ready_transitions = (logs.stdout + logs.stderr).count("database system is ready to accept connections")
            if ready_transitions >= 2 and query.returncode == 0:
                break
            time.sleep(0.25)
        else:
            pytest.fail("disposable PostgreSQL did not become ready")

        _psql(container, BOOTSTRAP_SQL)
        _apply_migration(container)
        yield container
    finally:
        subprocess.run(
            ["docker", "rm", "-f", container],
            capture_output=True,
            check=False,
        )


def _lines(completed: subprocess.CompletedProcess[str]) -> set[str]:
    return {line for line in completed.stdout.splitlines() if line}


def _catalog_snapshot(container: str) -> str:
    return _psql(
        container,
        r"""
select 'role|' || rolname || '|' || rolcanlogin || '|' || rolinherit || '|' || rolsuper
  from pg_catalog.pg_roles
 where rolname in ('kb_gate_owner', 'kb_review', 'kb_apply')
 order by rolname;
select 'owner|' || namespace.nspname || '|' || relation.relname || '|' || owner_role.rolname
  from pg_catalog.pg_class relation
  join pg_catalog.pg_namespace namespace on namespace.oid = relation.relnamespace
  join pg_catalog.pg_roles owner_role on owner_role.oid = relation.relowner
 where (namespace.nspname, relation.relname) in (
   ('public', 'agents'), ('public', 'strategies'), ('public', 'strategy_nodes'),
   ('public', 'claim_evidence'), ('public', 'claim_edges'), ('public', 'claims'),
   ('public', 'sources'), ('public', 'reasoning_tools'),
   ('kb_stage', 'kb_proposals'), ('kb_stage', 'kb_review_principals'),
   ('kb_stage', 'kb_proposal_approvals')
 )
 order by namespace.nspname, relation.relname;
select 'function|' || procedure.proname || '(' || pg_catalog.oidvectortypes(procedure.proargtypes) ||
       ')|' || owner_role.rolname || '|' || procedure.prosecdef || '|' || procedure.proconfig::text
  from pg_catalog.pg_proc procedure
  join pg_catalog.pg_namespace namespace on namespace.oid = procedure.pronamespace
  join pg_catalog.pg_roles owner_role on owner_role.oid = procedure.proowner
 where namespace.nspname = 'kb_stage'
   and procedure.proname in (
     'approve_strict_proposal', 'assert_approved_proposal', 'finish_approved_proposal'
   )
 order by procedure.proname, procedure.proargtypes::text;
select 'table_acl|' || namespace.nspname || '|' || relation.relname || '|' ||
       case when acl.grantee = 0 then 'PUBLIC' else grantee_role.rolname end || '|' ||
       pg_catalog.upper(acl.privilege_type) || '|' || acl.is_grantable
  from pg_catalog.pg_class relation
  join pg_catalog.pg_namespace namespace on namespace.oid = relation.relnamespace
  cross join lateral pg_catalog.aclexplode(
    coalesce(relation.relacl, pg_catalog.acldefault('r', relation.relowner))
  ) acl
  left join pg_catalog.pg_roles grantee_role on grantee_role.oid = acl.grantee
 where (namespace.nspname, relation.relname) in (
   ('public', 'agents'), ('public', 'strategies'), ('public', 'strategy_nodes'),
   ('public', 'claim_evidence'), ('public', 'claim_edges'), ('public', 'claims'),
   ('public', 'sources'), ('public', 'reasoning_tools'),
   ('kb_stage', 'kb_proposals'), ('kb_stage', 'kb_review_principals'),
   ('kb_stage', 'kb_proposal_approvals')
 )
   and acl.grantee <> relation.relowner
   and (acl.grantee = 0 or grantee_role.rolname in ('kb_gate_owner', 'kb_review', 'kb_apply'))
 order by namespace.nspname, relation.relname, grantee_role.rolname, acl.privilege_type;
select 'function_acl|' || procedure.proname || '(' ||
       pg_catalog.oidvectortypes(procedure.proargtypes) || ')|' ||
       case when acl.grantee = 0 then 'PUBLIC' else pg_catalog.pg_get_userbyid(acl.grantee) end ||
       '|' || pg_catalog.upper(acl.privilege_type) || '|' || acl.is_grantable
  from pg_catalog.pg_proc procedure
  join pg_catalog.pg_namespace namespace on namespace.oid = procedure.pronamespace
  cross join lateral pg_catalog.aclexplode(
    coalesce(procedure.proacl, pg_catalog.acldefault('f', procedure.proowner))
  ) acl
 where namespace.nspname = 'kb_stage'
   and procedure.proname in (
     'approve_strict_proposal', 'assert_approved_proposal', 'finish_approved_proposal'
   )
   and acl.grantee <> procedure.proowner
 order by procedure.proname, procedure.proargtypes::text, acl.grantee;
""",
    ).stdout


def test_upgrade_normalizes_exact_role_ownership_and_acl_matrix(migrated_postgres: str) -> None:
    role_rows = _lines(
        _psql(
            migrated_postgres,
            """
select rolname || '|' || rolcanlogin || '|' || rolinherit || '|' || rolsuper || '|' ||
       rolcreatedb || '|' || rolcreaterole || '|' || rolreplication || '|' || rolbypassrls
  from pg_catalog.pg_roles
 where rolname in ('kb_gate_owner', 'kb_review', 'kb_apply')
 order by rolname;
""",
        )
    )
    assert role_rows == {
        "kb_apply|true|false|false|false|false|false|false",
        "kb_gate_owner|false|false|false|false|false|false|false",
        "kb_review|true|false|false|false|false|false|false",
    }

    assert (
        _psql(
            migrated_postgres,
            """
select count(*)
  from pg_catalog.pg_auth_members membership
  join pg_catalog.pg_roles member_role on member_role.oid = membership.member
  join pg_catalog.pg_roles granted_role on granted_role.oid = membership.roleid
 where member_role.rolname in ('kb_gate_owner', 'kb_review', 'kb_apply')
    or granted_role.rolname in ('kb_gate_owner', 'kb_review', 'kb_apply');
""",
        ).stdout.strip()
        == "0"
    )

    schema_rows = _lines(
        _psql(
            migrated_postgres,
            """
select protected_role.rolname || '|' || protected_schema.nspname || '|' ||
       pg_catalog.has_schema_privilege(protected_role.rolname, protected_schema.nspname, 'USAGE') || '|' ||
       pg_catalog.has_schema_privilege(protected_role.rolname, protected_schema.nspname, 'CREATE')
  from (values ('kb_gate_owner'), ('kb_review'), ('kb_apply')) protected_role(rolname)
  cross join (values ('public'), ('kb_stage')) protected_schema(nspname);
""",
        )
    )
    assert schema_rows == {
        f"{role}|{schema}|true|false"
        for role in ("kb_gate_owner", "kb_review", "kb_apply")
        for schema in ("public", "kb_stage")
    }

    function_rows = _lines(
        _psql(
            migrated_postgres,
            """
select procedure.proname || '(' || pg_catalog.oidvectortypes(procedure.proargtypes) || ')|' ||
       owner_role.rolname || '|' || procedure.prosecdef || '|' || procedure.proconfig::text
  from pg_catalog.pg_proc procedure
  join pg_catalog.pg_namespace namespace on namespace.oid = procedure.pronamespace
  join pg_catalog.pg_roles owner_role on owner_role.oid = procedure.proowner
 where namespace.nspname = 'kb_stage'
   and procedure.proname in (
     'approve_strict_proposal', 'assert_approved_proposal', 'finish_approved_proposal'
   );
""",
        )
    )
    assert function_rows == {
        'approve_strict_proposal(uuid, text, jsonb, text, text)|kb_gate_owner|true|{"search_path=pg_catalog, pg_temp"}',
        "assert_approved_proposal(uuid, text, jsonb, text, uuid, timestamp with time zone, text)|"
        'kb_gate_owner|true|{"search_path=pg_catalog, pg_temp"}',
        "finish_approved_proposal(uuid, text, jsonb, text, uuid, timestamp with time zone, text, text)|"
        'kb_gate_owner|true|{"search_path=pg_catalog, pg_temp"}',
    }
    assert (
        _psql(
            migrated_postgres,
            "select pg_catalog.to_regprocedure('kb_stage.approve_strict_proposal(uuid,text,text)') is null;",
        ).stdout.strip()
        == "t"
    )
    assert (
        _psql(
            migrated_postgres,
            """
select count(*)
  from pg_catalog.pg_attribute attribute
  join pg_catalog.pg_class relation on relation.oid = attribute.attrelid
  join pg_catalog.pg_namespace namespace on namespace.oid = relation.relnamespace
  cross join lateral pg_catalog.aclexplode(attribute.attacl) acl
  left join pg_catalog.pg_roles grantee_role on grantee_role.oid = acl.grantee
 where (namespace.nspname, relation.relname) in (
   ('public', 'agents'), ('public', 'strategies'), ('public', 'strategy_nodes'),
   ('public', 'claim_evidence'), ('public', 'claim_edges'), ('public', 'claims'),
   ('public', 'sources'), ('public', 'reasoning_tools'),
   ('kb_stage', 'kb_proposals'), ('kb_stage', 'kb_review_principals'),
   ('kb_stage', 'kb_proposal_approvals')
 )
   and (acl.grantee = 0 or grantee_role.rolname in ('kb_gate_owner', 'kb_review', 'kb_apply'));
""",
        ).stdout.strip()
        == "0"
    )

    for runtime_role, target_role in (
        ("kb_apply", "broad_writer"),
        ("kb_apply", "kb_gate_owner"),
        ("kb_review", "kb_gate_owner"),
    ):
        denied = _psql(
            migrated_postgres,
            f"set session authorization {runtime_role}; set role {target_role};",
            check=False,
        )
        assert denied.returncode != 0
        assert "permission denied to set role" in denied.stderr

    snapshot_lines = set(_catalog_snapshot(migrated_postgres).splitlines())
    expected_table_acls = {
        "table_acl|public|agents|kb_gate_owner|SELECT|false",
        "table_acl|public|agents|kb_apply|SELECT|false",
        "table_acl|public|agents|kb_review|SELECT|false",
        "table_acl|public|strategies|kb_apply|SELECT|false",
        "table_acl|public|strategies|kb_apply|INSERT|false",
        "table_acl|public|strategies|kb_apply|UPDATE|false",
        "table_acl|public|strategy_nodes|kb_apply|SELECT|false",
        "table_acl|public|strategy_nodes|kb_apply|INSERT|false",
        "table_acl|public|strategy_nodes|kb_apply|UPDATE|false",
        "table_acl|public|claim_evidence|kb_apply|SELECT|false",
        "table_acl|public|claim_evidence|kb_apply|INSERT|false",
        "table_acl|public|claim_edges|kb_apply|SELECT|false",
        "table_acl|public|claim_edges|kb_apply|INSERT|false",
        "table_acl|public|claims|kb_apply|SELECT|false",
        "table_acl|public|claims|kb_apply|INSERT|false",
        "table_acl|public|sources|kb_apply|SELECT|false",
        "table_acl|public|sources|kb_apply|INSERT|false",
        "table_acl|public|reasoning_tools|kb_apply|SELECT|false",
        "table_acl|public|reasoning_tools|kb_apply|INSERT|false",
        "table_acl|kb_stage|kb_proposals|kb_gate_owner|SELECT|false",
        "table_acl|kb_stage|kb_proposals|kb_gate_owner|UPDATE|false",
        "table_acl|kb_stage|kb_proposals|kb_apply|SELECT|false",
        "table_acl|kb_stage|kb_proposals|kb_review|SELECT|false",
    }
    assert {line for line in snapshot_lines if line.startswith("table_acl|")} == expected_table_acls
    assert {line for line in snapshot_lines if line.startswith("function_acl|")} == {
        "function_acl|approve_strict_proposal(uuid, text, jsonb, text, text)|kb_review|EXECUTE|false",
        "function_acl|assert_approved_proposal(uuid, text, jsonb, text, uuid, timestamp with time zone, text)|"
        "kb_apply|EXECUTE|false",
        "function_acl|finish_approved_proposal(uuid, text, jsonb, text, uuid, timestamp with time zone, text, text)|"
        "kb_apply|EXECUTE|false",
    }


def test_migration_rerun_preserves_catalog_snapshot(migrated_postgres: str) -> None:
    before = _catalog_snapshot(migrated_postgres)
    _apply_migration(migrated_postgres)
    assert _catalog_snapshot(migrated_postgres) == before


def _assert_migration_fails(container: str, expected_error: str) -> None:
    completed = _apply_migration(container, check=False)
    assert completed.returncode != 0
    assert expected_error in completed.stderr


def test_migration_fails_closed_on_protected_role_membership(migrated_postgres: str) -> None:
    _psql(migrated_postgres, "grant broad_writer to kb_apply;")
    try:
        _assert_migration_fails(migrated_postgres, "unexpected protected role membership")
    finally:
        _psql(migrated_postgres, "revoke broad_writer from kb_apply;")


def test_migration_fails_closed_on_ownership_drift(migrated_postgres: str) -> None:
    _psql(migrated_postgres, "alter table kb_stage.kb_proposal_approvals owner to rogue_owner;")
    try:
        _assert_migration_fails(migrated_postgres, "unexpected protected table owner")
    finally:
        _psql(migrated_postgres, "alter table kb_stage.kb_proposal_approvals owner to kb_gate_owner;")


def test_migration_fails_closed_on_unexpected_table_grantee(migrated_postgres: str) -> None:
    _psql(migrated_postgres, "grant select on kb_stage.kb_proposal_approvals to rogue_grantee;")
    try:
        _assert_migration_fails(migrated_postgres, "unexpected protected table grantee")
    finally:
        _psql(migrated_postgres, "revoke all on kb_stage.kb_proposal_approvals from rogue_grantee;")


def test_migration_fails_closed_on_unexpected_column_grantee(migrated_postgres: str) -> None:
    _psql(
        migrated_postgres,
        "grant select (review_note) on kb_stage.kb_proposal_approvals to rogue_grantee;",
    )
    try:
        _assert_migration_fails(migrated_postgres, "unexpected protected table grantee")
    finally:
        _psql(
            migrated_postgres,
            "revoke select (review_note) on kb_stage.kb_proposal_approvals from rogue_grantee;",
        )


def test_migration_fails_closed_on_unexpected_function_grantee(migrated_postgres: str) -> None:
    signature = "kb_stage.assert_approved_proposal(uuid,text,jsonb,text,uuid,timestamptz,text)"
    _psql(migrated_postgres, f"grant execute on function {signature} to rogue_grantee;")
    try:
        _assert_migration_fails(migrated_postgres, "unexpected protected function grantee")
    finally:
        _psql(migrated_postgres, f"revoke all on function {signature} from rogue_grantee;")


def test_migration_fails_closed_on_unexpected_gate_overload(migrated_postgres: str) -> None:
    _psql(
        migrated_postgres,
        """
create function kb_stage.assert_approved_proposal(text)
returns void language sql security definer as $function$ select $function$;
""",
    )
    try:
        _assert_migration_fails(migrated_postgres, "unexpected protected function overload")
    finally:
        _psql(migrated_postgres, "drop function kb_stage.assert_approved_proposal(text);")