fix(llm): reject multi-statement and non-read-only AI assistant queries

The AI Assistant's execute_sql_query tool runs LLM-generated SQL inside
a BEGIN TRANSACTION READ ONLY wrapper. However, the LLM-supplied query
was sent to psycopg as-is, so a multi-statement payload beginning with
COMMIT, END, ROLLBACK, or ABORT terminated the read-only transaction
and ran subsequent statements in autocommit mode. With ordinary write
privileges this allowed unauthorised data modification; for a superuser
or pg_execute_server_program role this chained to remote code execution
on the database host via COPY ... TO PROGRAM.

Validate the LLM-supplied query before any connection work happens:

* The input must parse to exactly one non-empty/non-comment statement.
* The leading real token (after stripping leading whitespace, comments,
  and punctuation) must be one of SELECT, WITH, EXPLAIN, SHOW, VALUES,
  TABLE. Everything else -- DML, DDL, CALL, COPY, DO, SET/RESET, the
  transaction-control verbs, and the rest -- is rejected up front.

PostgreSQL's READ ONLY mode continues to enforce the remaining cases
(data-modifying CTEs, EXPLAIN ANALYZE on writes, volatile side effects)
at runtime, so the validator is the load-bearing check for multi-
statement / top-level escapes and READ ONLY is the backstop for the
rest.

Add a 60-scenario regression suite under
web/pgadmin/llm/tests/test_database_tool_security.py covering the
original PoC payloads (COMMIT/END/ROLLBACK/ABORT/SET/BEGIN), multi-
statement masked by comments, the full allow- and deny-list of leading
keywords, dollar-quoted literals containing semicolons, and degenerate
inputs (empty, whitespace-only, comment-only, quoted identifier).

Reported by: Isaac Chen <isaac9503@gmail.com>
Reviewed by: Kundan Sable <kundan.sable@enterprisedb.com>

web/pgadmin/llm/tests/test_database_tool_security.py

@@ -0,0 +1,327 @@
+##########################################################################
+#
+# pgAdmin 4 - PostgreSQL Tools
+#
+# Copyright (C) 2013 - 2026, The pgAdmin Development Team
+# This software is released under the PostgreSQL Licence
+#
+##########################################################################
+
+"""Security regression tests for execute_sql_query input validation.
+
+These tests pin the multi-statement / leading-keyword guard added in
+``_validate_readonly_query``.  The validator is the load-bearing
+defense against an attacker escaping the ``BEGIN TRANSACTION READ
+ONLY`` wrapper by emitting a transaction-control statement (COMMIT,
+END, ROLLBACK, ABORT) followed by writes -- the original Isaac Chen
+report from 2026-06-08.
+
+The tests are pure unit tests (no DB, no Flask client) -- they exercise
+the validator directly.  Anything that reaches the connection layer is
+already too late.
+"""
+
+from pgadmin.utils.route import BaseTestGenerator
+from pgadmin.llm.tools.database import (
+    DatabaseToolError,
+    _validate_readonly_query,
+)
+
+
+class ValidateReadonlyQueryAcceptTestCase(BaseTestGenerator):
+    """Queries that MUST be accepted by the validator.
+
+    These cover the six allowlisted leading keywords plus a handful of
+    sqlparse edge cases (leading paren, leading block comment, dollar-
+    quoted literal containing ``;``, trailing comment-only "statement")
+    where over-eager rejection would break legitimate LLM usage.
+    """
+
+    scenarios = [
+        ('Plain SELECT', dict(
+            query='SELECT 1',
+        )),
+        ('SELECT with semicolon', dict(
+            query='SELECT 1;',
+        )),
+        ('Leading paren SELECT', dict(
+            query='(SELECT 1)',
+        )),
+        ('Leading block comment then SELECT', dict(
+            query='/* hint */ SELECT 1',
+        )),
+        ('SELECT followed by line comment', dict(
+            query='SELECT 1 -- trailing',
+        )),
+        ('SELECT followed by comment-only second statement', dict(
+            query='SELECT 1; -- trailing',
+        )),
+        ('WITH plain CTE', dict(
+            query='WITH x AS (SELECT 1) SELECT * FROM x',
+        )),
+        ('WITH RECURSIVE', dict(
+            query=(
+                'WITH RECURSIVE t(n) AS ('
+                '  SELECT 1 UNION ALL SELECT n+1 FROM t WHERE n < 3'
+                ') SELECT * FROM t'
+            ),
+        )),
+        ('EXPLAIN SELECT', dict(
+            query='EXPLAIN SELECT 1',
+        )),
+        ('EXPLAIN ANALYZE SELECT', dict(
+            query='EXPLAIN ANALYZE SELECT 1',
+        )),
+        ('EXPLAIN with options', dict(
+            query='EXPLAIN (ANALYZE, BUFFERS) SELECT 1',
+        )),
+        ('SHOW GUC', dict(
+            query='SHOW search_path',
+        )),
+        ('VALUES list', dict(
+            query='VALUES (1), (2), (3)',
+        )),
+        ('TABLE statement', dict(
+            query='TABLE pg_catalog.pg_class',
+        )),
+        ('Lowercase keyword', dict(
+            query='select 1',
+        )),
+        ('Mixed case keyword', dict(
+            query='SeLeCt 1',
+        )),
+        ('Dollar-quoted string containing semicolon', dict(
+            # Single SELECT whose literal contains ; -- the validator
+            # must NOT mistake this for a multi-statement payload.
+            query="SELECT $$;$$",
+        )),
+        ('Tagged dollar quote containing semicolon', dict(
+            query="SELECT $tag$;$tag$",
+        )),
+        ('Standard string with doubled quote', dict(
+            query="SELECT 'a''b'",
+        )),
+    ]
+
+    def setUp(self):
+        pass
+
+    def runTest(self):
+        # Should not raise.
+        _validate_readonly_query(self.query)
+
+
+class ValidateReadonlyQueryRejectTestCase(BaseTestGenerator):
+    """Queries that MUST be rejected by the validator.
+
+    The first block reproduces the Isaac Chen PoC payloads (multi-
+    statement with a leading transaction-control keyword to close the
+    wrapping READ ONLY transaction).  The remaining scenarios pin the
+    leading-keyword allowlist against every common write / state-
+    changing top-level statement.
+    """
+
+    scenarios = [
+        # --- Isaac Chen PoC family: transaction control + multi-stmt ---
+        ('PoC: COMMIT then COPY TO PROGRAM', dict(
+            query=(
+                "COMMIT; "
+                "COPY (SELECT 1) TO PROGRAM 'id > /tmp/id.txt 2>&1'; "
+                "SELECT 1"
+            ),
+            expected_code='INVALID_QUERY',
+        )),
+        ('PoC: COMMIT then DELETE', dict(
+            query='COMMIT; DELETE FROM t; SELECT 1',
+            expected_code='INVALID_QUERY',
+        )),
+        ('PoC: END then DELETE', dict(
+            query='END; DELETE FROM t; SELECT 1',
+            expected_code='INVALID_QUERY',
+        )),
+        ('PoC: ROLLBACK then INSERT', dict(
+            query="ROLLBACK; INSERT INTO t VALUES (1); SELECT 1",
+            expected_code='INVALID_QUERY',
+        )),
+        ('PoC: ABORT then UPDATE', dict(
+            query='ABORT; UPDATE t SET c = 1; SELECT 1',
+            expected_code='INVALID_QUERY',
+        )),
+        ('PoC: SET then DELETE', dict(
+            query=(
+                "SET role superuser; "
+                "DELETE FROM t; "
+                "SELECT 1"
+            ),
+            expected_code='INVALID_QUERY',
+        )),
+        ('PoC: BEGIN then SELECT', dict(
+            query='BEGIN; SELECT 1',
+            expected_code='INVALID_QUERY',
+        )),
+        ('PoC: trailing COMMIT after SELECT', dict(
+            query='SELECT 1; COMMIT',
+            expected_code='INVALID_QUERY',
+        )),
+
+        # --- Multi-statement: comments don't smuggle past the count ---
+        ('Multi-statement masked by block comment', dict(
+            query='SELECT 1; /* hide */ DROP TABLE t',
+            expected_code='INVALID_QUERY',
+        )),
+        ('Multi-statement masked by line comment', dict(
+            query='SELECT 1; -- ignore\nDROP TABLE t',
+            expected_code='INVALID_QUERY',
+        )),
+
+        # --- Disallowed leading keywords (single-statement) ---
+        ('Leading UPDATE', dict(
+            query='UPDATE t SET c = 1',
+            expected_code='INVALID_QUERY',
+        )),
+        ('Leading DELETE', dict(
+            query='DELETE FROM t',
+            expected_code='INVALID_QUERY',
+        )),
+        ('Leading INSERT', dict(
+            query='INSERT INTO t VALUES (1)',
+            expected_code='INVALID_QUERY',
+        )),
+        ('Leading MERGE', dict(
+            query=(
+                'MERGE INTO t USING s ON t.id = s.id '
+                'WHEN MATCHED THEN UPDATE SET c = s.c'
+            ),
+            expected_code='INVALID_QUERY',
+        )),
+        ('Leading CALL', dict(
+            query="CALL bad_proc()",
+            expected_code='INVALID_QUERY',
+        )),
+        ('Leading COPY TO PROGRAM', dict(
+            query="COPY (SELECT 1) TO PROGRAM 'id'",
+            expected_code='INVALID_QUERY',
+        )),
+        ('Leading DO block', dict(
+            query="DO $$ BEGIN PERFORM 1; END $$",
+            expected_code='INVALID_QUERY',
+        )),
+        ('Leading SET', dict(
+            query='SET role superuser',
+            expected_code='INVALID_QUERY',
+        )),
+        ('Leading RESET', dict(
+            query='RESET role',
+            expected_code='INVALID_QUERY',
+        )),
+        ('Leading CREATE', dict(
+            query='CREATE TABLE t (c int)',
+            expected_code='INVALID_QUERY',
+        )),
+        ('Leading DROP', dict(
+            query='DROP TABLE t',
+            expected_code='INVALID_QUERY',
+        )),
+        ('Leading ALTER', dict(
+            query='ALTER TABLE t ADD COLUMN c int',
+            expected_code='INVALID_QUERY',
+        )),
+        ('Leading TRUNCATE', dict(
+            query='TRUNCATE t',
+            expected_code='INVALID_QUERY',
+        )),
+        ('Leading LOCK', dict(
+            query='LOCK TABLE t',
+            expected_code='INVALID_QUERY',
+        )),
+        ('Leading GRANT', dict(
+            query='GRANT ALL ON t TO public',
+            expected_code='INVALID_QUERY',
+        )),
+        ('Leading REVOKE', dict(
+            query='REVOKE ALL ON t FROM public',
+            expected_code='INVALID_QUERY',
+        )),
+        ('Leading NOTIFY', dict(
+            query="NOTIFY chan, 'msg'",
+            expected_code='INVALID_QUERY',
+        )),
+        ('Leading LISTEN', dict(
+            query='LISTEN chan',
+            expected_code='INVALID_QUERY',
+        )),
+        ('Leading PREPARE', dict(
+            query='PREPARE p AS SELECT 1',
+            expected_code='INVALID_QUERY',
+        )),
+        ('Leading EXECUTE', dict(
+            query='EXECUTE p',
+            expected_code='INVALID_QUERY',
+        )),
+        ('Leading REFRESH MATERIALIZED VIEW', dict(
+            query='REFRESH MATERIALIZED VIEW mv',
+            expected_code='INVALID_QUERY',
+        )),
+        ('Leading VACUUM', dict(
+            query='VACUUM t',
+            expected_code='INVALID_QUERY',
+        )),
+        ('Leading ANALYZE (standalone)', dict(
+            # 'ANALYZE foo' is a maintenance command, not EXPLAIN
+            # ANALYZE -- must be rejected because the validator looks
+            # at the *first* keyword only.
+            query='ANALYZE t',
+            expected_code='INVALID_QUERY',
+        )),
+        ('Leading CHECKPOINT', dict(
+            query='CHECKPOINT',
+            expected_code='INVALID_QUERY',
+        )),
+        ('Leading CLUSTER', dict(
+            query='CLUSTER t',
+            expected_code='INVALID_QUERY',
+        )),
+        ('Leading REINDEX', dict(
+            query='REINDEX TABLE t',
+            expected_code='INVALID_QUERY',
+        )),
+
+        # --- Empty / degenerate inputs ---
+        ('Empty string', dict(
+            query='',
+            expected_code='INVALID_QUERY',
+        )),
+        ('Whitespace only', dict(
+            query='   \n\t  ',
+            expected_code='INVALID_QUERY',
+        )),
+        ('Only semicolons', dict(
+            query=';;;',
+            expected_code='INVALID_QUERY',
+        )),
+        ('Only a block comment', dict(
+            query='/* nothing here */',
+            expected_code='INVALID_QUERY',
+        )),
+        ('Quoted identifier "SELECT" is not the keyword', dict(
+            # PostgreSQL: references a column/table named SELECT.
+            # The validator must NOT treat the quoted identifier as
+            # the allowlisted keyword.
+            query='"SELECT" 1',
+            expected_code='INVALID_QUERY',
+        )),
+    ]
+
+    def setUp(self):
+        pass
+
+    def runTest(self):
+        try:
+            _validate_readonly_query(self.query)
+        except DatabaseToolError as e:
+            self.assertEqual(e.code, self.expected_code)
+            return
+        self.fail(
+            f"Validator accepted query that should have been "
+            f"rejected: {self.query!r}"
+        )

web/pgadmin/llm/tools/database.py

@@ -19,6 +19,7 @@ Uses pgAdmin's SQL template infrastructure for version-aware queries.
 import secrets
 from typing import Optional
 
+import sqlparse
 from flask import render_template
 
 from pgadmin.utils.driver import get_driver
@@ -38,6 +39,23 @@ INDEXES_TEMPLATE_PATH = 'indexes/sql'
 LLM_APP_NAME_PREFIX = 'pgAdmin 4 - LLM'
 
 
+# Statement keywords permitted at the start of an LLM-supplied query.
+# The BEGIN TRANSACTION READ ONLY wrapper around the query is only
+# effective if the LLM cannot exit that transaction. A multi-statement
+# payload such as `COMMIT; <write>; SELECT 1` would otherwise terminate
+# the read-only transaction (via COMMIT/END/ROLLBACK/ABORT) and run
+# subsequent statements in autocommit mode. The allowlist below, in
+# combination with the single-statement check, ensures the LLM cannot
+# emit transaction-control statements, SET/RESET, DML/DDL, CALL, COPY,
+# or anything else that could escape or weaken the read-only sandbox.
+# EXPLAIN ANALYZE on a SELECT remains supported; EXPLAIN ANALYZE on a
+# write statement is blocked by PostgreSQL itself inside the read-only
+# transaction.
+_ALLOWED_LEADING_KEYWORDS = frozenset({
+    'SELECT', 'WITH', 'EXPLAIN', 'SHOW', 'VALUES', 'TABLE',
+})
+
+
 class DatabaseToolError(Exception):
     """Exception raised when a database tool operation fails."""
 
@@ -127,6 +145,96 @@ def _connect_readonly(
         return False, str(e)
 
 
+def _first_real_keyword(statement) -> str:
+    """
+    Return the first non-trivial leaf token of a parsed statement.
+
+    Whitespace, comments, and punctuation are skipped so that a leading
+    open-parenthesis (e.g. ``(SELECT 1) UNION (SELECT 2)``) or a leading
+    comment block does not mask the real leading keyword. The result is
+    upper-cased; an empty string is returned for an empty statement.
+    """
+    for tok in statement.flatten():
+        if tok.is_whitespace:
+            continue
+        ttype = str(tok.ttype) if tok.ttype is not None else ''
+        if 'Comment' in ttype or 'Punctuation' in ttype:
+            continue
+        return (tok.normalized or '').upper()
+    return ''
+
+
+def _validate_readonly_query(query: str) -> None:
+    """
+    Ensure an LLM-supplied query is a single, read-only statement.
+
+    This is the load-bearing defense against an attacker escaping the
+    read-only transaction wrapper by emitting multiple statements. The
+    PostgreSQL simple-query protocol cheerfully runs a whole
+    semicolon-separated batch in one round-trip, so a payload such as
+    ``COMMIT; <write>; SELECT 1`` would terminate the read-only
+    transaction (via the leading ``COMMIT``) and execute the remaining
+    statements in autocommit mode. The trailing ``ROLLBACK`` would then
+    be a no-op.
+
+    Validation rules:
+
+    * The input must contain exactly one non-empty statement.
+    * The leading keyword must be in :data:`_ALLOWED_LEADING_KEYWORDS`.
+
+    PostgreSQL is left to enforce the rest -- ``EXPLAIN ANALYZE`` on a
+    write statement, data-modifying CTEs, and volatile function side
+    effects are all rejected at runtime by the read-only transaction.
+
+    Args:
+        query: SQL query supplied by the LLM tool call.
+
+    Raises:
+        DatabaseToolError: If the query is empty, contains more than
+            one statement, or is not a read-only statement.
+    """
+    if not query or not query.strip():
+        raise DatabaseToolError(
+            "Query is empty",
+            code="INVALID_QUERY"
+        )
+
+    try:
+        parsed = sqlparse.parse(query)
+    except Exception as e:
+        raise DatabaseToolError(
+            f"Failed to parse query: {e}",
+            code="INVALID_QUERY"
+        )
+
+    statements = [
+        s for s in parsed
+        if s.token_first(skip_cm=True, skip_ws=True) is not None
+    ]
+
+    if not statements:
+        raise DatabaseToolError(
+            "Query contains no SQL statement",
+            code="INVALID_QUERY"
+        )
+
+    if len(statements) > 1:
+        raise DatabaseToolError(
+            "Only a single SQL statement is allowed; multi-statement "
+            "queries are rejected",
+            code="INVALID_QUERY"
+        )
+
+    keyword = _first_real_keyword(statements[0])
+    if keyword not in _ALLOWED_LEADING_KEYWORDS:
+        raise DatabaseToolError(
+            f"Statement type '{keyword or 'UNKNOWN'}' is not permitted; "
+            "only read-only statements (SELECT, WITH, EXPLAIN, SHOW, "
+            "VALUES, TABLE) are allowed",
+            code="INVALID_QUERY"
+        )
+
+
 def _execute_readonly_query(conn, query: str) -> dict:
     """
     Execute a query in a read-only transaction.
@@ -231,9 +339,16 @@ def execute_readonly_query(
         - truncated: True if results were limited
 
     Raises:
-        DatabaseToolError: If the query fails or connection
+        DatabaseToolError: If the query is rejected by validation, the
-            cannot be established
+            query fails at runtime, or a connection cannot be
+            established.
     """
+    # Validate the LLM-supplied query before allocating a connection.
+    # The BEGIN TRANSACTION READ ONLY wrapper below is only effective
+    # if the query is a single read-only statement; see
+    # _validate_readonly_query for the threat model and rules.
+    _validate_readonly_query(query)
+
     # Generate unique connection ID for this LLM query
     conn_id = f"llm_{secrets.choice(range(1, 9999999))}"