"""Read-only readiness check for workbench query indexes and test resources.""" from __future__ import annotations import json import sys from pathlib import Path import psycopg from psycopg.rows import dict_row ROOT = Path(__file__).resolve().parents[1] sys.path.insert(0, str(ROOT)) from app.core.config import settings from samples.peripheral.config import SampleSettings RESOURCES = { "analysis_runtime_queries": ( "analysis.runs", "analysis.runs_created_at_run_id_idx", "analysis.results", "analysis.results_run_created_at_result_id_idx", ), "task_alerts": ( "workbench.alert_events", "workbench.alert_events_alert_idx", "workbench.alerts", "workbench.alerts_active_recent_idx", ), "peripheral_work_orders": ( "peripheral_sample.requests", "peripheral_sample.scenarios", "peripheral_sample.work_orders", "peripheral_sample.work_order_revisions", ), "peripheral_test_data": ( "peripheral_sample.test_data_records", "peripheral_sample.test_data_record_revisions", ), } COLUMN_REQUIREMENTS = { "analysis_runtime_queries": { "analysis.runs": {"run_id", "created_at"}, "analysis.results": {"result_id", "run_id", "created_at"}, }, "task_alerts": { "workbench.alert_events": { "project_id", "producer_id", "event_id", "alert_id", "event_type", "occurred_at", "ingested_at", "document", }, "workbench.alerts": { "project_id", "alert_id", "revision", "status", "occurred_at", "updated_at", "document", }, }, } INDEX_REQUIREMENTS = { "analysis.runs_created_at_run_id_idx": ( "on analysis.runs", "(created_at desc, run_id)", ), "analysis.results_run_created_at_result_id_idx": ( "on analysis.results", "(run_id, created_at desc, result_id desc)", ), "workbench.alert_events_alert_idx": ( "(project_id, alert_id, occurred_at, event_id)", ), "workbench.alerts_active_recent_idx": ( "(project_id, occurred_at desc, alert_id)", "status = 'active'", ), } RUNTIME_PRIVILEGES = { "analysis.runs": ("SELECT",), "analysis.results": ("SELECT",), "workbench.alert_events": ("SELECT", "INSERT"), "workbench.alerts": ("SELECT", "INSERT", "UPDATE"), "peripheral_sample.requests": ("SELECT", "INSERT"), "peripheral_sample.scenarios": ("SELECT", "INSERT", "UPDATE"), "peripheral_sample.tasks": ("SELECT", "INSERT", "UPDATE"), "peripheral_sample.revisions": ("SELECT", "INSERT"), "peripheral_sample.feedback": ("SELECT", "INSERT"), "peripheral_sample.demo_datasets": ("SELECT",), "peripheral_sample.work_orders": ("SELECT", "INSERT", "UPDATE"), "peripheral_sample.work_order_revisions": ("SELECT", "INSERT"), "peripheral_sample.test_data_records": ("SELECT", "INSERT", "UPDATE"), "peripheral_sample.test_data_record_revisions": ("SELECT", "INSERT"), } PROHIBITED_PRIVILEGES = { relation: ("DELETE", "TRUNCATE") for relation in RUNTIME_PRIVILEGES } | { "workbench.alert_events": ("UPDATE", "DELETE", "TRUNCATE"), } def main() -> int: sample_settings = SampleSettings() with psycopg.connect( host=settings.DB_HOST, port=settings.DB_PORT, user=settings.DB_USER, password=settings.DB_PASSWORD, dbname=sample_settings.db_name, connect_timeout=5, options="-c statement_timeout=5000", row_factory=dict_row, ) as connection: connection.execute("SET TRANSACTION READ ONLY") database = connection.execute("SELECT current_database() AS name").fetchone()["name"] availability: dict[str, dict[str, object]] = {} for resource, relations in RESOURCES.items(): rows = connection.execute( "SELECT value,to_regclass(value)::text AS relation FROM unnest(%s::text[]) value", (list(relations),), ).fetchall() missing = [row["value"] for row in rows if row["relation"] is None] missing_columns: dict[str, list[str]] = {} for relation, required in COLUMN_REQUIREMENTS.get(resource, {}).items(): schema_name, table_name = relation.split(".", 1) columns = connection.execute( """SELECT column_name FROM information_schema.columns WHERE table_schema=%s AND table_name=%s""", (schema_name, table_name), ).fetchall() present = {row["column_name"] for row in columns} absent = sorted(required - present) if absent: missing_columns[relation] = absent invalid_indexes: dict[str, list[str]] = {} for relation in relations: required_fragments = INDEX_REQUIREMENTS.get(relation) if required_fragments is None or relation in missing: continue definition_row = connection.execute( "SELECT pg_get_indexdef(to_regclass(%s)) AS definition", (relation,), ).fetchone() definition = " ".join( str(definition_row["definition"] or "").lower().split() ) absent_fragments = [ fragment for fragment in required_fragments if fragment not in definition ] if absent_fragments: invalid_indexes[relation] = absent_fragments availability[resource] = { "status": ( "available" if not missing and not missing_columns and not invalid_indexes else "not_configured" ), "missing": missing, "missing_columns": missing_columns, "invalid_indexes": invalid_indexes, } role = connection.execute( """SELECT current_user AS name,rolsuper,rolcreatedb,rolcreaterole FROM pg_roles WHERE rolname=current_user""" ).fetchone() database_owner = connection.execute( """SELECT pg_get_userbyid(datdba) AS name FROM pg_database WHERE datname=current_database()""" ).fetchone()["name"] owned_relations = [ f"{row['schema_name']}.{row['relation_name']}" for row in connection.execute( """SELECT namespace.nspname AS schema_name, relation.relname AS relation_name FROM pg_class relation JOIN pg_namespace namespace ON namespace.oid=relation.relnamespace WHERE pg_get_userbyid(relation.relowner)=current_user AND namespace.nspname=ANY(%s::text[]) AND relation.relkind IN ('r','p') ORDER BY namespace.nspname,relation.relname""", (["analysis", "peripheral_sample", "workbench"],), ).fetchall() ] missing_runtime_privileges: list[str] = [] prohibited_runtime_privileges: list[str] = [] for relation, privileges in RUNTIME_PRIVILEGES.items(): exists = connection.execute( "SELECT to_regclass(%s) IS NOT NULL AS available", (relation,) ).fetchone()["available"] if not exists: missing_runtime_privileges.extend( f"{relation}:{privilege}" for privilege in privileges ) continue for privilege in privileges: allowed = connection.execute( "SELECT has_table_privilege(current_user,%s,%s) AS allowed", (relation, privilege), ).fetchone()["allowed"] if not allowed: missing_runtime_privileges.append(f"{relation}:{privilege}") for relation, privileges in PROHIBITED_PRIVILEGES.items(): exists = connection.execute( "SELECT to_regclass(%s) IS NOT NULL AS available", (relation,) ).fetchone()["available"] if not exists: continue for privilege in privileges: allowed = connection.execute( "SELECT has_table_privilege(current_user,%s,%s) AS allowed", (relation, privilege), ).fetchone()["allowed"] if allowed: prohibited_runtime_privileges.append( f"{relation}:{privilege}" ) production_safe_role = bool( role and not role["rolsuper"] and not role["rolcreatedb"] and not role["rolcreaterole"] and role["name"] != database_owner and not owned_relations and not missing_runtime_privileges and not prohibited_runtime_privileges ) ready = all(item["status"] == "available" for item in availability.values()) print(json.dumps({ "database": database, "ready": ready, "production_ready": ready and production_safe_role, "connection_role": { "name": role["name"] if role else None, "superuser": role["rolsuper"] if role else None, "create_database": role["rolcreatedb"] if role else None, "create_role": role["rolcreaterole"] if role else None, "owns_database": role["name"] == database_owner if role else None, "owned_relations": owned_relations, "missing_runtime_privileges": missing_runtime_privileges, "prohibited_runtime_privileges": prohibited_runtime_privileges, "production_safe": production_safe_role, }, "resources": availability, }, ensure_ascii=False, indent=2, sort_keys=True)) return 0 if ready else 2 if __name__ == "__main__": raise SystemExit(main())