insights/tools/infra-scripts/datastore-multinode/verify-multinode-datastore.py

#!/usr/bin/env python
# ruff: noqa: T201
"""
Verify that the multinode Datastore smoke-test stack ended up in the expected
shape after `manage.py migrate_datastore` ran.

Connects to each Datastore node directly (via its published port) and asserts:

1. `getMacro('hostClusterRole')` returns the expected role per node — confirms
   the per-node config files mounted correctly.
2. `system.clusters` is consistent across nodes and lists the expected logical
   clusters with the expected member hosts — confirms `<remote_servers>` is
   wired correctly.
3. Tables that production routes to a satellite cluster (ai_events / aux / ops
   / sessions) live on the matching node and **not** on the data node. Catches
   migrations that forgot `node_roles=[NodeRole.X]` and accidentally created a
   table everywhere.

Exits 0 on success, non-zero on the first failed assertion (with context).
"""

from __future__ import annotations

import sys
from dataclasses import dataclass

from datastore_driver import Client


@dataclass(frozen=True)
class Node:
    name: str
    host: str
    port: int
    expected_role: str


# Match the published ports in docker-compose.multinode-datastore.yml.
NODES: list[Node] = [
    Node("datastore-data", "localhost", 9000, "data"),
    Node("datastore-ai-events", "localhost", 9100, "ai_events"),
    Node("datastore-aux", "localhost", 9200, "aux"),
    Node("datastore-ops", "localhost", 9300, "ops"),
    Node("datastore-sessions", "localhost", 9400, "sessions"),
    Node("datastore-logs", "localhost", 9500, "logs"),
]

# Logical cluster -> expected member hosts (sorted). Each multinode XML file
# declares the same `<remote_servers>` block, so every node must agree.
EXPECTED_CLUSTERS: dict[str, list[str]] = {
    "insights": ["datastore-data"],
    "insights_single_shard": ["datastore-data"],
    "insights_writable": ["datastore-data"],
    "insights_primary_replica": ["datastore-data"],
    "insights_migrations": [
        "datastore-ai-events",
        "datastore-aux",
        "datastore-data",
        "datastore-logs",
        "datastore-ops",
        "datastore-sessions",
    ],
    "ai_events": ["datastore-ai-events"],
    "aux": ["datastore-aux"],
    "ops": ["datastore-ops"],
    "sessions": ["datastore-sessions"],
    "logs": ["datastore-logs"],
}


@dataclass(frozen=True)
class SatelliteTables:
    # Tables that must exist on the satellite node.
    on_satellite: list[str]
    # Tables that must NOT exist on the data node. Distributed wrappers are
    # routinely created on data too (so the data node can query the satellite
    # shard cross-cluster) — only the underlying sharded data table belongs
    # exclusively to the satellite.
    forbidden_on_data: list[str]


# Manifest of tables (in DATASTORE_DATABASE) per satellite role. Update when
# adding new per-cluster tables.
SATELLITE_TABLE_MANIFEST: dict[str, SatelliteTables] = {
    "ai_events": SatelliteTables(
        on_satellite=["ai_events", "sharded_ai_events"],
        forbidden_on_data=["sharded_ai_events"],
    ),
    "aux": SatelliteTables(on_satellite=[], forbidden_on_data=[]),
    "ops": SatelliteTables(on_satellite=[], forbidden_on_data=[]),
    "sessions": SatelliteTables(on_satellite=[], forbidden_on_data=[]),
    "logs": SatelliteTables(
        on_satellite=["logs34"],
        forbidden_on_data=["logs34"],
    ),
}

DATABASE = "insights"


def client(node: Node) -> Client:
    return Client(host=node.host, port=node.port, database="default")


def fail(msg: str) -> None:
    print(f"FAIL: {msg}", file=sys.stderr)
    sys.exit(1)


def check_macros() -> None:
    print("==> Checking <macros> per node")
    for node in NODES:
        with client(node) as c:
            ((role,),) = c.execute("SELECT getMacro('hostClusterRole')")
        if role != node.expected_role:
            fail(f"{node.name}: hostClusterRole={role!r}, expected {node.expected_role!r}")
        print(f"  {node.name}: hostClusterRole={role}")


def fetch_clusters(node: Node) -> dict[str, list[str]]:
    with client(node) as c:
        rows = c.execute(
            "SELECT cluster, host_name FROM system.clusters WHERE cluster IN %(clusters)s ORDER BY cluster, host_name",
            {"clusters": tuple(EXPECTED_CLUSTERS.keys())},
        )
    layout: dict[str, list[str]] = {}
    for cluster, host in rows:
        layout.setdefault(cluster, []).append(host)
    return layout


def check_cluster_topology() -> None:
    print("==> Checking system.clusters consistency across nodes")
    expected = {k: sorted(set(v)) for k, v in EXPECTED_CLUSTERS.items()}
    for node in NODES:
        layout = {k: sorted(set(v)) for k, v in fetch_clusters(node).items()}
        for cluster, hosts in expected.items():
            actual = layout.get(cluster)
            if actual != hosts:
                fail(f"{node.name}: cluster {cluster!r} hosts={actual}, expected {hosts}")
        print(f"  {node.name}: {len(layout)} clusters match")


def fetch_local_tables(node: Node) -> set[str]:
    with client(node) as c:
        rows = c.execute(
            "SELECT name FROM system.tables WHERE database = %(db)s",
            {"db": DATABASE},
        )
    return {name for (name,) in rows}


def find_node(role: str) -> Node:
    for node in NODES:
        if node.expected_role == role:
            return node
    raise KeyError(role)


def check_satellite_tables() -> None:
    print("==> Checking per-cluster table placement (manifest)")
    data_node = find_node("data")
    data_tables = fetch_local_tables(data_node)
    print(f"  {data_node.name}: {len(data_tables)} tables in {DATABASE!r}")

    any_checked = False
    for role, manifest in SATELLITE_TABLE_MANIFEST.items():
        if not manifest.on_satellite and not manifest.forbidden_on_data:
            continue
        any_checked = True
        node = find_node(role)
        node_tables = fetch_local_tables(node)
        print(f"  {node.name}: {len(node_tables)} tables in {DATABASE!r}")

        missing = [t for t in manifest.on_satellite if t not in node_tables]
        if missing:
            fail(f"{node.name}: missing expected tables {missing}")

        leaked = [t for t in manifest.forbidden_on_data if t in data_tables]
        if leaked:
            fail(
                f"{data_node.name} should not have satellite-only tables {leaked} "
                f"(they belong on {node.name}). Did a migration forget node_roles=?"
            )

    if not any_checked:
        print("  (no satellite manifests populated yet — skipping presence checks)")


def main() -> int:
    check_macros()
    check_cluster_topology()
    check_satellite_tables()
    print("OK: multinode topology and table layout look correct.")
    return 0


if __name__ == "__main__":
    sys.exit(main())