fix(gen_ai): prevent prompt injection and SQL injection in /generate-sql

[duboc]

Fixes MiniVM finding 3053018334892130305 — Prompt Injection via unvalidated input in SQL Generation.

Problem

/generate-sql interpolated caller-controlled values (source_table, destination_table, source_schema_fields, destination_schema) straight into the Gemini prompt in InitialSQLGenerator._construct_prompt, and interpolated source_table into the sample query in TransformationPipeline:

sample_query = f"SELECT * FROM `{source_table_name}` LIMIT 3"

So the same parameter was both a prompt-injection and a SQL-injection vector, on an endpoint that is anonymously reachable. The BigQuery error from the injected identifier was caught and logged as a warning, the pipeline continued, the model's output passed the dry run, and the result was returned to the caller.

Approach

Prompt injection can't be solved by scrubbing prose — you can't enumerate "instruction-shaped" text. It can be solved when the injected values are structured, because those can be allow-listed exactly, and that covers everything in the reported taint chain. For the one genuinely free-form input (source_data_sample_json, real row data) the fix is containment instead.

Four layers:

1. New services/sql/common/input_validation.py — strict grammars for BigQuery table IDs, column names and field paths. The destination-schema validator rebuilds the schema from allow-listed keys only, so injected keys and prose are dropped rather than escaped. The data-sample sanitizer parses/re-serialises as JSON, strips control characters and code fences, and caps rows and length.

2. Validate at the boundary, then again deeper. field_validators on SQLGenerationRequest reject hostile input with a 422 before a background task or an LLM call starts. TransformationPipeline, InitialSQLGenerator._construct_prompt and SemanticEnhancer._construct_prompt re-validate as defense in depth for non-HTTP callers. Both prompts now fence parameters in a <task_parameters> block with explicit precedence rules.

3. The co-reported SQL injection is removed, not escaped. The sample fetch uses the BigQuery tabledata.list API (list_rows) with a structured TableReference — there is no SQL string to break out of, and it's cheaper than a query job.

4. Output contract. enforce_sql_contract() rejects anything that isn't a single CREATE OR REPLACE TABLE writing to the requested destination and reading only the requested source. Applied to every candidate before each dry run, so the semantic enhancer and the SQL fixer (extra LLM round-trips) are covered too. Comments and string literals are stripped before keyword scanning, so 'DROP TABLE x' AS note and -- Defaulted ... aren't false positives.

Taint chain, step by step:

Reported step	Fix
main.py:473 endpoint receives request	field_validators reject with 422
main.py:504 passed to execute_pipeline	Pipeline re-validates, fails the task
transformation_pipeline.py:102 → generate	generate() validates before building a prompt
initial_sql_generator.py:59 f-string interpolation	_construct_prompt validates + fences
client_utils.py:99 prompt reaches the model	Value can't carry a payload; response is contract-checked

Verification

48 passed in 0.02s   # src/psearch/gen_ai/services/sql/common/test_input_validation.py

Tests import only input_validation, so they run without GCP credentials. The reported payload was also replayed through the real InitialSQLGenerator with a stubbed model client:

Case	Result
Reported payload in source_table	Rejected — model never called
Model compromised, returns CREATE OR REPLACE TABLE my_table AS SELECT 'pwned'	Rejected by the contract
Model returns arbitrary prose (LLM proxying)	Rejected
Legitimate request + legitimate SQL	Unchanged, still succeeds

The committed schema.json (32 fields, nested RECORDs) validates unchanged; two checks were relaxed for the empty "mode": "" / "fields": [] placeholders it actually contains.

Still open after this PR

Warning

The endpoint remains unauthenticated. modules/gen_ai/main.tf grants roles/run.invoker to allUsers. Payloads are now validated, but every anonymous call still spends money on Gemini + BigQuery, so LLM-proxying and denial-of-wallet remain open. This PR makes the invoker list a variable (invoker_members, default unchanged so nothing breaks) rather than flipping it, because locking it down requires the UI to start sending an ID token. Recommend a follow-up setting it to the UI service account.

Note

/fix-sql has the same injection shape (caller-supplied SQL and error string into a prompt) and is untouched here. Lower impact — the caller submits and receives their own SQL — but worth a follow-up.

Residual risk, stated honestly: an attacker who controls source-table row contents can still influence the semantic-enhancer prompt, and SQL comments are a low-bandwidth text channel out. The contract check bounds this to "text inside a valid transformation script for your own tables" rather than arbitrary model output.

[gemini-code-assist]

Code Review

This pull request introduces a robust input validation and output contract enforcement layer to secure the SQL generation pipeline against prompt and SQL injection vulnerabilities. It updates the API request schema, the initial SQL generator, the semantic enhancer, and the transformation pipeline to validate all caller-controlled inputs, and transitions BigQuery sample fetching to the safer tabledata API. The review feedback highlights critical improvements to the validation logic, including updating regex patterns to support multi-part backticked table paths, using a single-pass regex in _strip_sql_noise to prevent incorrect stripping when comments contain quotes, handling falsy schema types correctly, and adding UPDATE to the forbidden SQL constructs for enhanced defense-in-depth.

[gemini-code-assist]

The _FROM_JOIN_RE regex only matches table references where the entire path is either unquoted or enclosed in a single set of backticks (e.g., `project.dataset.table`). However, BigQuery also allows individual parts of a table path to be backticked separately (e.g., `project`.`dataset`.`table`).

If an attacker steers the model to output separately backticked table paths, _FROM_JOIN_RE will only match the first part (e.g., `project`), which will be skipped because it lacks a dot, completely bypassing the allowed tables check.

We should update _FROM_JOIN_RE to match multi-part backticked table references and normalize them by removing backticks and spaces before checking against allowed_tables.

_FROM_JOIN_RE = re.compile(
    r"\b(?:FROM|JOIN)\s+("
    r"(?:`[^`]+`|[A-Za-z0-9_\-]+)"
    r"(?:\s*\.\s*(?:`[^`]+`|[A-Za-z0-9_\-]+))*"
    r")",
    re.IGNORECASE
)

[gemini-code-assist]

To support multi-part backticked table references matched by the updated _FROM_JOIN_RE regex, we should normalize the matched reference by removing all backticks and spaces before performing the dot and allowed tables checks.

Suggested change

	for raw_ref in _FROM_JOIN_RE.findall(stripped):
	ref = raw_ref.strip("`")
	if "." not in ref:
	# A CTE name, a table alias or an UNNEST/subquery target: harmless.
	continue
	if ref.lower() not in allowed_tables:
	raise UnsafeSQLError(
	f"Generated SQL reads from unexpected table {ref!r}. Only "
	f"{sorted(allowed_tables)} are allowed."
	)
	for raw_ref in _FROM_JOIN_RE.findall(stripped):
	ref = raw_ref.replace("`", "").replace(" ", "")
	if "." not in ref:
	# A CTE name, a table alias or an UNNEST/subquery target: harmless.
	continue
	if ref.lower() not in allowed_tables:
	raise UnsafeSQLError(
	f"Generated SQL reads from unexpected table {ref!r}. Only "
	f"{sorted(allowed_tables)} are allowed."
	)

[gemini-code-assist]

In _strip_sql_noise, comments and string literals are stripped in separate, sequential regex passes. This creates a correctness issue if a comment contains a single quote (e.g., -- don't do this or -- Defaulted name to NULL as no direct source match found.). In such cases, _STRING_LITERAL_RE will match from that single quote to the next single quote in the actual SQL query, stripping out valid SQL code in between and causing the contract check to fail or behave unexpectedly.

To resolve this, use a single-pass regex with alternation to match both comments and string literals simultaneously, replacing them appropriately based on what was matched.

def _strip_sql_noise(sql: str) -> str:
    """Remove comments and string literals so keyword scanning cannot be fooled."""
    pattern = re.compile(
        r"'''.*?'''|"
        r'"""[^"]*"""|'
        r"'(?:\\.|[^'\\])*'|"
        r'"(?:\\.|[^"\\])*"|'
        r"/\*.*?\*/|"
        r"(?:--|#)[^\n]*",
        re.DOTALL
    )
    def replace(match):
        text = match.group(0)
        if text.startswith(('--', '#', '/*')):
            return ' '
        return "''"
    return pattern.sub(replace, sql)

[gemini-code-assist]

In _validate_schema_fields, the check raw_type = field.get("type") or "STRING" will evaluate to "STRING" if type is any falsy value (such as an empty list [] or empty dict {}). This bypasses the isinstance(raw_type, str) check on the next line, silently accepting invalid types as "STRING".

Instead, check if raw_type is None explicitly before validating that it is a string.

        raw_type = field.get("type")
        if raw_type is None:
            field_type = "STRING"
        elif not isinstance(raw_type, str):
            raise InputValidationError(f"{field_path}.type must be a string.")
        else:
            field_type = raw_type.strip().upper() or "STRING"

[gemini-code-assist]

The _FORBIDDEN_CONSTRUCTS list contains DML statements like DELETE, INSERT, and MERGE, but is missing UPDATE. Adding UPDATE provides additional defense-in-depth to prevent any unauthorized data modification.

Suggested change

	_FORBIDDEN_CONSTRUCTS = [
	(re.compile(r"\bEXECUTE\s+IMMEDIATE\b", re.IGNORECASE), "EXECUTE IMMEDIATE (dynamic SQL)"),
	(re.compile(r"\bEXPORT\s+DATA\b", re.IGNORECASE), "EXPORT DATA"),
	(re.compile(r"\bLOAD\s+DATA\b", re.IGNORECASE), "LOAD DATA"),
	(re.compile(r"\bDROP\s+\w", re.IGNORECASE), "DROP"),
	(re.compile(r"\bTRUNCATE\s+TABLE\b", re.IGNORECASE), "TRUNCATE TABLE"),
	(re.compile(r"\bDELETE\s+FROM\b", re.IGNORECASE), "DELETE"),
	(re.compile(r"\bINSERT\s+INTO\b", re.IGNORECASE), "INSERT"),
	(re.compile(r"\bMERGE\s+INTO\b", re.IGNORECASE), "MERGE"),
	(re.compile(r"\bALTER\s+(?:TABLE|SCHEMA|VIEW|MODEL|ORGANIZATION|PROJECT)\b", re.IGNORECASE), "ALTER"),
	(re.compile(r"\b(?:GRANT|REVOKE)\s+", re.IGNORECASE), "GRANT/REVOKE"),
	(re.compile(r"\bCALL\s+\w", re.IGNORECASE), "CALL"),
	# Any CREATE after the single expected header is unexpected.
	(re.compile(r"\bCREATE\s+(?:OR\s+REPLACE\s+)?\w", re.IGNORECASE), "additional CREATE statement"),
	(re.compile(r"\bEXTERNAL_QUERY\s*\(", re.IGNORECASE), "EXTERNAL_QUERY"),
	(re.compile(r"\bSET\s+@@", re.IGNORECASE), "system variable assignment"),
	(re.compile(r"\bBEGIN\b|\bDECLARE\b", re.IGNORECASE), "scripting block"),
	]
	_FORBIDDEN_CONSTRUCTS = [
	(re.compile(r"\bEXECUTE\s+IMMEDIATE\b", re.IGNORECASE), "EXECUTE IMMEDIATE (dynamic SQL)"),
	(re.compile(r"\bEXPORT\s+DATA\b", re.IGNORECASE), "EXPORT DATA"),
	(re.compile(r"\bLOAD\s+DATA\b", re.IGNORECASE), "LOAD DATA"),
	(re.compile(r"\bDROP\s+\w", re.IGNORECASE), "DROP"),
	(re.compile(r"\bTRUNCATE\s+TABLE\b", re.IGNORECASE), "TRUNCATE TABLE"),
	(re.compile(r"\bDELETE\s+FROM\b", re.IGNORECASE), "DELETE"),
	(re.compile(r"\bINSERT\s+INTO\b", re.IGNORECASE), "INSERT"),
	(re.compile(r"\bUPDATE\s+\w", re.IGNORECASE), "UPDATE"),
	(re.compile(r"\bMERGE\s+INTO\b", re.IGNORECASE), "MERGE"),
	(re.compile(r"\bALTER\s+(?:TABLE|SCHEMA|VIEW|MODEL|ORGANIZATION|PROJECT)\b", re.IGNORECASE), "ALTER"),
	(re.compile(r"\b(?:GRANT|REVOKE)\s+", re.IGNORECASE), "GRANT/REVOKE"),
	(re.compile(r"\bCALL\s+\w", re.IGNORECASE), "CALL"),
	# Any CREATE after the single expected header is unexpected.
	(re.compile(r"\bCREATE\s+(?:OR\s+REPLACE\s+)?\w", re.IGNORECASE), "additional CREATE statement"),
	(re.compile(r"\bEXTERNAL_QUERY\s*\(", re.IGNORECASE), "EXTERNAL_QUERY"),
	(re.compile(r"\bSET\s+@@", re.IGNORECASE), "system variable assignment"),
	(re.compile(r"\bBEGIN\b|\bDECLARE\b", re.IGNORECASE), "scripting block"),
	]