Spaces:
Running
Running
| from __future__ import annotations | |
| import re | |
| MAX_READ_QUERY_ROWS = 100 | |
| READ_QUERY_TIMEOUT_MS = 10000 | |
| _READ_QUERY_START = re.compile( | |
| r"^\s*(SELECT|WITH|EXPLAIN|TABLE)\b", | |
| re.IGNORECASE | re.DOTALL, | |
| ) | |
| _FORBIDDEN_SQL = re.compile( | |
| r"\b(" | |
| r"INSERT|UPDATE|DELETE|DROP|ALTER|CREATE|TRUNCATE|GRANT|REVOKE|" | |
| r"COPY|MERGE|CALL|DO|VACUUM|REINDEX|CLUSTER|" | |
| r"LISTEN|NOTIFY|LOAD|LOCK|DISCARD|RESET" | |
| r")\b", | |
| re.IGNORECASE, | |
| ) | |
| _CYPHER_MARKERS = re.compile( | |
| r"\b(MATCH|RETURN|OPTIONAL\s+MATCH|CREATE\s*\(|DETACH\s+DELETE|SET\s+n\.|REMOVE\s+n\.)\b", | |
| re.IGNORECASE, | |
| ) | |
| _LEGACY_TABLE_NAMES = re.compile( | |
| r"\b(FROM|JOIN|TABLE|INTO|UPDATE)\s+(nodes|relationships)\b", | |
| re.IGNORECASE, | |
| ) | |
| _WITH_RECURSIVE = re.compile(r"\bWITH\s+RECURSIVE\b", re.IGNORECASE) | |
| class ReadQueryValidationError(Exception): | |
| pass | |
| def _paren_balance(sql: str) -> int: | |
| """Return open-paren count minus close-paren count, ignoring quoted strings.""" | |
| depth = 0 | |
| in_single = False | |
| in_double = False | |
| i = 0 | |
| while i < len(sql): | |
| ch = sql[i] | |
| if in_single: | |
| if ch == "'" and i + 1 < len(sql) and sql[i + 1] == "'": | |
| i += 2 | |
| continue | |
| if ch == "'": | |
| in_single = False | |
| i += 1 | |
| continue | |
| if in_double: | |
| if ch == '"': | |
| in_double = False | |
| i += 1 | |
| continue | |
| if ch == "'": | |
| in_single = True | |
| elif ch == '"': | |
| in_double = True | |
| elif ch == "(": | |
| depth += 1 | |
| elif ch == ")": | |
| depth -= 1 | |
| i += 1 | |
| return depth | |
| def validate_read_only_sql(query: str) -> str: | |
| normalized = query.strip() | |
| if not normalized: | |
| raise ReadQueryValidationError("Empty query") | |
| parts = [part.strip() for part in normalized.split(";") if part.strip()] | |
| if len(parts) != 1: | |
| raise ReadQueryValidationError("Multiple statements are not allowed") | |
| normalized = parts[0] | |
| if _CYPHER_MARKERS.search(normalized): | |
| raise ReadQueryValidationError( | |
| "Cypher is not supported on PostgreSQL. Use SQL against kg_nodes " | |
| "and kg_relationships (for example: SELECT uuid, data->>'name' AS name " | |
| "FROM kg_nodes WHERE data->>'name' ILIKE '%alice%' LIMIT 50)." | |
| ) | |
| if _LEGACY_TABLE_NAMES.search(normalized): | |
| raise ReadQueryValidationError( | |
| "Unknown table. Use kg_nodes and kg_relationships " | |
| "(not nodes/relationships)." | |
| ) | |
| if _FORBIDDEN_SQL.search(normalized): | |
| raise ReadQueryValidationError("Only read-only SELECT queries are allowed") | |
| if not _READ_QUERY_START.match(normalized): | |
| raise ReadQueryValidationError( | |
| "Query must start with SELECT, WITH, EXPLAIN, or TABLE" | |
| ) | |
| # EXPLAIN of Cypher still starts with EXPLAIN; catch residual Cypher body. | |
| if normalized.upper().startswith("EXPLAIN") and _CYPHER_MARKERS.search( | |
| normalized[7:] | |
| ): | |
| raise ReadQueryValidationError( | |
| "Cypher is not supported on PostgreSQL. Use SQL EXPLAIN SELECT ..." | |
| ) | |
| balance = _paren_balance(normalized) | |
| if balance != 0: | |
| raise ReadQueryValidationError( | |
| "Incomplete SQL (unbalanced parentheses). Send one complete statement " | |
| "in a single tool call — do not split WITH RECURSIVE CTEs across calls. " | |
| "Example: WITH RECURSIVE walk AS (" | |
| "SELECT source_uuid, target_uuid, rel_type, 1 AS depth " | |
| "FROM kg_relationships WHERE source_uuid = '<uuid>' " | |
| "UNION ALL " | |
| "SELECT r.source_uuid, r.target_uuid, r.rel_type, w.depth + 1 " | |
| "FROM walk w JOIN kg_relationships r ON r.source_uuid = w.target_uuid " | |
| "WHERE w.depth < 3" | |
| ") SELECT * FROM walk LIMIT 50" | |
| ) | |
| if _WITH_RECURSIVE.search(normalized): | |
| # A recursive CTE must both define the CTE and select from it. | |
| if not re.search(r"\)\s*SELECT\b", normalized, re.IGNORECASE): | |
| raise ReadQueryValidationError( | |
| "Incomplete WITH RECURSIVE query. Include the outer SELECT after " | |
| "the CTE, e.g. ') SELECT * FROM walk LIMIT 50'. Do not send only " | |
| "the CTE body or a fragment of the UNION ALL branch." | |
| ) | |
| return normalized | |