Don't Give Your LLM the Database: Safer Text-to-SQL on AWS

import re

from sqlglot import parse, exp
from sqlglot.errors import ParseError

ALLOWED_TABLE = "sales"

ALLOWED_COLUMNS = {
"order_id",
"order_date",
"region",
"category",
"revenue",
"cost",
"quantity"
}

EXTERNAL_FUNCTION_RE = re.compile(
r"\bUSING\s+EXTERNAL\s+FUNCTION\b",
re.IGNORECASE
)


def validate_sql(sql: str) -> str:
if EXTERNAL_FUNCTION_RE.search(sql):
raise ValueError(
"Athena external functions are outside this prototype's scope."
)

try:
statements = parse(sql, read="athena")
except ParseError as exc:
raise ValueError("SQL could not be parsed.") from exc

if len(statements) != 1:
raise ValueError(
"Only one SQL statement is allowed."
)

tree = statements[0]

if not isinstance(tree, exp.Select):
raise ValueError(
"Only SELECT queries are allowed."
)

if tree.find(exp.Star):
raise ValueError(
"Wildcard references are disabled."
)

if tree.find(exp.Join):
raise ValueError(
"JOINs are outside this prototype's scope."
)

if tree.find(exp.Subquery):
raise ValueError(
"Subqueries are outside this prototype's scope."
)

if tree.find(exp.CTE):
raise ValueError(
"CTEs are outside this prototype's scope."
)

tables = list(tree.find_all(exp.Table))

if len(tables) != 1:
raise ValueError(
"Exactly one approved table is required."
)

table = tables[0]

if table.name.lower() != ALLOWED_TABLE:
raise ValueError(
"The referenced table is not allowed."
)

if table.db or table.catalog:
raise ValueError(
"Qualified table references are disabled."
)

table_alias = table.alias_or_name.lower()

projection_aliases = {
projection.alias.lower()
for projection in tree.expressions
if isinstance(projection, exp.Alias)
}

order = tree.args.get("order")

order_alias_columns = set()

if order:
order_alias_columns = {
id(column)
for column in order.find_all(exp.Column)
if (
not column.table
and not column.db
and not column.catalog
and column.name.lower() in projection_aliases
)
}

for column in tree.find_all(exp.Column):
if id(column) in order_alias_columns:
continue

if column.catalog or column.db:
raise ValueError(
"Qualified column references are disabled."
)

if column.name.lower() not in ALLOWED_COLUMNS:
raise ValueError(
f"Column is not allowed: {column.name}"
)

if column.table and column.table.lower() != table_alias:
raise ValueError(
f"Column qualifier is not allowed: {column.table}"
)

return sql.strip()