From 08ee35e25f1416b8409ad4e7cb23615e675d45a2 Mon Sep 17 00:00:00 2001
From: nessi <sascha@nesterovic.cc>
Date: Thu, 12 Feb 2026 17:07:56 +0100
Subject: [PATCH] Add handling for pg_stat_statements compatibility checks

Introduced a new `_run_expect_failure` helper to manage cases where specific queries are expected to fail. Added smoke tests for pg_stat_statements, validating its behavior when unavailable, loaded, or enabled. Also extended connectivity checks and enhanced database discovery queries.
---
 backend/scripts/pg_compat_smoke.py | 81 ++++++++++++++++++++++++++++++
 1 file changed, 81 insertions(+)

diff --git a/backend/scripts/pg_compat_smoke.py b/backend/scripts/pg_compat_smoke.py
index 50df087..486b9bc 100644
--- a/backend/scripts/pg_compat_smoke.py
+++ b/backend/scripts/pg_compat_smoke.py
@@ -54,6 +54,24 @@ async def _run_optional(conn: asyncpg.Connection, label: str, query: str) -> Non
         print(f"[compat] SKIP optional: {label} ({exc})")
 
 
+async def _run_expect_failure(
+    conn: asyncpg.Connection,
+    label: str,
+    query: str,
+    accepted_sqlstates: set[str],
+) -> None:
+    try:
+        await conn.fetch(query)
+    except asyncpg.PostgresError as exc:
+        if exc.sqlstate in accepted_sqlstates:
+            print(f"[compat] PASS expected-failure: {label} (sqlstate={exc.sqlstate})")
+            return
+        raise RuntimeError(f"{label} failed with unexpected sqlstate={exc.sqlstate}: {exc}") from exc
+    except Exception as exc:
+        raise RuntimeError(f"{label} failed with unexpected non-Postgres error: {exc}") from exc
+    raise RuntimeError(f"{label} unexpectedly succeeded, but failure was expected")
+
+
 def _section(title: str) -> None:
     print(f"[compat] --- {title} ---")
 
@@ -94,6 +112,8 @@ async def run() -> None:
 
         _section("connectivity")
         await _run_required_fetchval(conn, "target_connection.select_1", "SELECT 1")
+        await _run_required_fetchval(conn, "connectivity.server_encoding", "SHOW server_encoding")
+        await _run_required_fetchval(conn, "connectivity.timezone", "SHOW TimeZone")
 
         _section("collector")
         # Core collector queries used in app/services/collector.py
@@ -181,6 +201,17 @@ async def run() -> None:
             LIMIT 200
             """,
         )
+        await _run_required_fetch(
+            conn,
+            "target_endpoint.discover_databases",
+            """
+            SELECT datname
+            FROM pg_database
+            WHERE datallowconn
+              AND NOT datistemplate
+            ORDER BY datname
+            """,
+        )
 
         _section("overview")
         # Overview queries used in app/services/overview_collector.py
@@ -272,6 +303,56 @@ async def run() -> None:
             LIMIT 20
             """,
         )
+        _section("pg_stat_statements modes")
+        # Validate both runtime modes NexaPG must support:
+        # 1) extension unavailable/not preloaded -> query fails with known sqlstate
+        # 2) extension available + loaded -> query succeeds
+        await conn.execute("DROP EXTENSION IF EXISTS pg_stat_statements")
+        await _run_expect_failure(
+            conn,
+            "pg_stat_statements.absent_or_not_loaded",
+            """
+            SELECT queryid::text, calls, total_exec_time, mean_exec_time, rows, left(query, 2000) AS query_text
+            FROM pg_stat_statements
+            ORDER BY total_exec_time DESC
+            LIMIT 20
+            """,
+            accepted_sqlstates={"42P01", "55000"},
+        )
+
+        available = await conn.fetchval(
+            """
+            SELECT EXISTS(
+              SELECT 1
+              FROM pg_available_extensions
+              WHERE name = 'pg_stat_statements'
+            )
+            """
+        )
+        if available:
+            try:
+                await conn.execute("CREATE EXTENSION IF NOT EXISTS pg_stat_statements")
+            except Exception as exc:
+                print(f"[compat] SKIP optional: pg_stat_statements.create_extension ({exc})")
+            else:
+                try:
+                    await conn.fetch(
+                        """
+                        SELECT queryid::text, calls, total_exec_time, mean_exec_time, rows, left(query, 2000) AS query_text
+                        FROM pg_stat_statements
+                        ORDER BY total_exec_time DESC
+                        LIMIT 20
+                        """
+                    )
+                    print("[compat] PASS optional: pg_stat_statements.enabled_query")
+                except asyncpg.PostgresError as exc:
+                    # Typical when shared_preload_libraries does not include pg_stat_statements.
+                    if exc.sqlstate == "55000":
+                        print(f"[compat] SKIP optional: pg_stat_statements.enabled_query ({exc})")
+                    else:
+                        raise RuntimeError(f"pg_stat_statements.enabled_query unexpected sqlstate={exc.sqlstate}: {exc}") from exc
+        else:
+            print("[compat] SKIP optional: pg_stat_statements.extension_unavailable")
 
         print("[compat] Smoke checks passed")
     finally: