fix(security): Add table blocklist and fix MCP SQL validation bypass (#37411)

(cherry picked from commit 15b3c96f8e)
This commit is contained in:
Amin Ghadersohi
2026-02-09 08:12:06 -05:00
committed by Joe Li
parent a219f6cbe6
commit b5cbbe86b7
6 changed files with 1276 additions and 0 deletions

View File

@@ -452,6 +452,15 @@ class BaseSQLStatement(Generic[InternalRepresentation]):
"""
raise NotImplementedError()
def check_tables_present(self, tables: set[str]) -> bool:
"""
Check if any of the given tables are present in the statement.
:param tables: Set of table names to check for (case-insensitive)
:return: True if any of the tables are present
"""
raise NotImplementedError()
def get_limit_value(self) -> int | None:
"""
Get the limit value of the statement.
@@ -766,6 +775,16 @@ class SQLStatement(BaseSQLStatement[exp.Expression]):
}
return any(function.upper() in present for function in functions)
def check_tables_present(self, tables: set[str]) -> bool:
"""
Check if any of the given tables are present in the statement.
:param tables: Set of table names to check for (case-insensitive)
:return: True if any of the tables are present
"""
present = {table.table.lower() for table in self.tables}
return any(table.lower() in present for table in tables)
def get_limit_value(self) -> int | None:
"""
Parse a SQL query and return the `LIMIT` or `TOP` value, if present.
@@ -1172,6 +1191,16 @@ class KustoKQLStatement(BaseSQLStatement[str]):
logger.warning("Kusto KQL doesn't support checking for functions present.")
return False
def check_tables_present(self, tables: set[str]) -> bool:
"""
Check if any of the given tables are present in the statement.
:param tables: Set of table names to check for (case-insensitive)
:return: True if any of the tables are present
"""
logger.warning("Kusto KQL doesn't support checking for tables present.")
return False
def get_limit_value(self) -> int | None:
"""
Get the limit value of the statement.
@@ -1313,6 +1342,17 @@ class SQLScript:
for statement in self.statements
)
def check_tables_present(self, tables: set[str]) -> bool:
"""
Check if any of the given tables are present in the script.
:param tables: Set of table names to check for (case-insensitive)
:return: True if any of the tables are present
"""
return any(
statement.check_tables_present(tables) for statement in self.statements
)
def is_valid_ctas(self) -> bool:
"""
Check if the script contains a valid CTAS statement.