bootnode/api/tests/test_datastore.py

"""Datastore client: basic auth, and a schema that applies idempotently."""

import base64
import json
import re
import threading
from http.server import BaseHTTPRequestHandler, HTTPServer
from importlib import resources

import pytest
from datastore_connect.driver.exceptions import DatabaseError

from bootnode.core.datastore import client as ds
from bootnode.core.datastore.client import DataStoreClient, statements

SCHEMA = resources.files("bootnode.core.datastore").joinpath("schema")
FILES = sorted(p.name for p in SCHEMA.iterdir() if p.name.endswith(".sql"))


def schema_statements() -> list[str]:
    return [s for name in FILES for s in statements(SCHEMA.joinpath(name).read_text())]


def test_statements_drop_comments_and_split():
    sql = "-- a; comment\nCREATE TABLE IF NOT EXISTS t (x UInt8) -- trailing; too\nENGINE = Memory;\n\n;"
    assert statements(sql) == ["CREATE TABLE IF NOT EXISTS t (x UInt8) \nENGINE = Memory"]


def test_every_schema_statement_is_idempotent():
    stmts = schema_statements()
    assert stmts
    for s in stmts:
        assert "IF NOT EXISTS" in s, s


def columns() -> dict[str, set[str]]:
    tables: dict[str, set[str]] = {}
    for s in schema_statements():
        m = re.match(r"CREATE TABLE IF NOT EXISTS (\w+) \((.*)\)\s*ENGINE", s, re.S)
        if m:
            assert m.group(1) not in tables, f"{m.group(1)} is defined twice"
            tables[m.group(1)] = {
                c.group(1)
                for c in re.finditer(r"^\s*(\w+)\s+[A-Z]", m.group(2), re.M)
                if c.group(1) != "INDEX"
            }
    return tables


def test_views_group_by_columns_their_source_defines():
    tables = columns()
    views = 0
    for s in schema_statements():
        m = re.match(
            r"CREATE MATERIALIZED VIEW IF NOT EXISTS (\w+).*?AS SELECT(.*?)FROM (\w+)\s+GROUP BY (.*)",
            s,
            re.S,
        )
        if not m:
            continue
        views += 1
        name, select, source, keys = m.groups()
        aliases = set(re.findall(r"\bas (\w+)", select))
        for key in (k.strip() for k in keys.split(",")):
            assert key in tables[source] or key in aliases, f"{name}: {source} has no {key}"
    assert views >= 8


class _Result:
    def __init__(self, rows):
        self._rows = rows
        self.result_rows = [tuple(r.values()) for r in rows]

    def named_results(self):
        return iter(self._rows)


class _Datastore:
    """Records statements; keeps schema_migrations rows as the server would."""

    def __init__(self):
        self.commands: list[str] = []
        self.migrations: list[dict] = []

    async def command(self, query, **_):
        self.commands.append(query)

    async def query(self, query, **_):
        assert query == "SELECT version, checksum FROM schema_migrations"
        return _Result(
            [{"version": r["version"], "checksum": r["checksum"]} for r in self.migrations]
        )

    async def raw_insert(self, table, insert_block, fmt, **_):
        assert (table, fmt) == ("schema_migrations", "JSONEachRow")
        self.migrations += [json.loads(line) for line in insert_block.splitlines()]

    async def close(self):
        pass


async def test_migrate_applies_each_file_once():
    store = _Datastore()
    client = DataStoreClient()
    client._client = store

    assert await client.migrate() == FILES
    assert [r["version"] for r in store.migrations] == FILES
    assert len(store.commands) == 1 + len(schema_statements())

    store.commands.clear()
    assert await client.migrate() == []
    assert store.commands == [ds.MIGRATIONS]


async def test_migrate_reapplies_a_file_whose_content_changed():
    store = _Datastore()
    client = DataStoreClient()
    client._client = store
    await client.migrate()
    store.migrations[0]["checksum"] = "0" * 64

    assert await client.migrate() == [FILES[0]]


class _Recorder(BaseHTTPRequestHandler):
    seen: list[dict[str, str]] = []

    def do_POST(self):
        _Recorder.seen.append(dict(self.headers))
        self.send_response(401)
        self.end_headers()

    do_GET = do_POST

    def log_message(self, *args):
        pass


async def test_credentials_travel_as_basic_auth(monkeypatch):
    server = HTTPServer(("127.0.0.1", 0), _Recorder)
    threading.Thread(target=server.serve_forever, daemon=True).start()
    _Recorder.seen.clear()
    port = server.server_address[1]
    monkeypatch.setattr(
        ds.settings, "datastore_url", f"clickhouse://ds%2Duser:p%[email protected]:{port}/bootnode"
    )

    client = DataStoreClient()
    with pytest.raises(DatabaseError):
        await client.initialize()
    server.shutdown()

    assert not client.is_connected
    assert _Recorder.seen
    expected = "Basic " + base64.b64encode(b"ds-user:p@ss").decode()
    for headers in _Recorder.seen:
        assert headers["Authorization"] == expected
        assert not [h for h in headers if h.lower().startswith("x-clickhouse")]