254 lines
9.9 KiB
Python
254 lines
9.9 KiB
Python
"""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())
|