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
(query: str)
| 165 | |
| 166 | |
| 167 | def _validate_readonly_query(query: str) -> None: |
| 168 | """ |
| 169 | Ensure an LLM-supplied query is a single, read-only statement. |
| 170 | |
| 171 | This is the load-bearing defense against an attacker escaping the |
| 172 | read-only transaction wrapper by emitting multiple statements. The |
| 173 | PostgreSQL simple-query protocol cheerfully runs a whole |
| 174 | semicolon-separated batch in one round-trip, so a payload such as |
| 175 | ``COMMIT; <write>; SELECT 1`` would terminate the read-only |
| 176 | transaction (via the leading ``COMMIT``) and execute the remaining |
| 177 | statements in autocommit mode. The trailing ``ROLLBACK`` would then |
| 178 | be a no-op. |
| 179 | |
| 180 | Validation rules: |
| 181 | |
| 182 | * The input must contain exactly one non-empty statement. |
| 183 | * The leading keyword must be in :data:`_ALLOWED_LEADING_KEYWORDS`. |
| 184 | |
| 185 | PostgreSQL is left to enforce the rest -- ``EXPLAIN ANALYZE`` on a |
| 186 | write statement, data-modifying CTEs, and volatile function side |
| 187 | effects are all rejected at runtime by the read-only transaction. |
| 188 | |
| 189 | Args: |
| 190 | query: SQL query supplied by the LLM tool call. |
| 191 | |
| 192 | Raises: |
| 193 | DatabaseToolError: If the query is empty, contains more than |
| 194 | one statement, or is not a read-only statement. |
| 195 | """ |
| 196 | if not query or not query.strip(): |
| 197 | raise DatabaseToolError( |
| 198 | "Query is empty", |
| 199 | code="INVALID_QUERY" |
| 200 | ) |
| 201 | |
| 202 | try: |
| 203 | parsed = sqlparse.parse(query) |
| 204 | except Exception as e: |
| 205 | raise DatabaseToolError( |
| 206 | f"Failed to parse query: {e}", |
| 207 | code="INVALID_QUERY" |
| 208 | ) |
| 209 | |
| 210 | statements = [ |
| 211 | s for s in parsed |
| 212 | if s.token_first(skip_cm=True, skip_ws=True) is not None |
| 213 | ] |
| 214 | |
| 215 | if not statements: |
| 216 | raise DatabaseToolError( |
| 217 | "Query contains no SQL statement", |
| 218 | code="INVALID_QUERY" |
| 219 | ) |
| 220 | |
| 221 | if len(statements) > 1: |
| 222 | raise DatabaseToolError( |
| 223 | "Only a single SQL statement is allowed; multi-statement " |
| 224 | "queries are rejected", |