write_detector.py•4.36 kB
import sqlparse
from sqlparse.sql import Token, TokenList
from sqlparse.tokens import Keyword, DML, DDL
from typing import Dict, List, Set, Tuple
class SQLWriteDetector:
def __init__(self):
# Data-changing DML (row-level or bulk)
self.dml_write_keywords: Set[str] = {
"INSERT", "UPDATE", "DELETE", "MERGE",
"UPSERT", "REPLACE", # vendor variants
"COPY", "LOAD", "BULK", # bulk-ingest / export family
"IMPORT", "EXPORT",
"PUT", "REMOVE", # Snowflake stage operations
}
# DDL & maintenance (schema / storage)
self.ddl_keywords: Set[str] = {
"CREATE", "ALTER", "DROP", "TRUNCATE", "RENAME",
"COMMENT", "REINDEX", "CLUSTER",
"OPTIMIZE", "VACUUM", "ANALYZE", "REFRESH",
}
# DCL (privilege & role management)
self.dcl_keywords: Set[str] = {"GRANT", "REVOKE", "DENY"}
# TCL / session state that can persist writes (optional)
self.tcl_keywords: Set[str] = {
"COMMIT", "ROLLBACK", "SAVEPOINT",
"SET", "PRAGMA",
"LOCK", "UNLOCK",
}
# Combine everything you want to forbid
self.write_keywords: Set[str] = (
self.dml_write_keywords
| self.ddl_keywords
| self.dcl_keywords
| self.tcl_keywords
)
def analyze_query(self, sql_query: str) -> Dict:
"""
Analyze a SQL query to determine if it contains write operations.
Args:
sql_query: The SQL query string to analyze
Returns:
Dictionary containing analysis results
"""
# Parse the SQL query
parsed = sqlparse.parse(sql_query)
if not parsed:
return {"contains_write": False, "write_operations": set(), "has_cte_write": False}
# Initialize result tracking
found_operations = set()
has_cte_write = False
# Analyze each statement in the query
for statement in parsed:
# Check for write operations in CTEs (WITH clauses)
if self._has_cte(statement):
cte_write = self._analyze_cte(statement)
if cte_write:
has_cte_write = True
found_operations.add("CTE_WRITE")
# Analyze the main query
operations = self._find_write_operations(statement)
found_operations.update(operations)
return {
"contains_write": bool(found_operations) or has_cte_write,
"write_operations": found_operations,
"has_cte_write": has_cte_write,
}
def _has_cte(self, statement: TokenList) -> bool:
"""Check if the statement has a WITH clause."""
return any(token.is_keyword and token.normalized == "WITH" for token in statement.tokens)
def _analyze_cte(self, statement: TokenList) -> bool:
"""
Analyze CTEs (WITH clauses) for write operations.
Returns True if any CTE contains a write operation.
"""
in_cte = False
for token in statement.tokens:
if token.is_keyword and token.normalized == "WITH":
in_cte = True
elif in_cte:
if any(write_kw in token.normalized for write_kw in self.write_keywords):
return True
return False
def _find_write_operations(self, statement: TokenList) -> Set[str]:
"""
Find all write operations in a statement.
Returns a set of found write operation keywords.
"""
operations = set()
for token in statement.tokens:
# Skip comments and whitespace
if token.is_whitespace or token.ttype in (sqlparse.tokens.Comment,):
continue
# Check if token is a keyword
if token.ttype in (Keyword, DML, DDL):
normalized = token.normalized.upper()
if normalized in self.write_keywords:
operations.add(normalized)
# Recursively check child tokens
if isinstance(token, TokenList):
child_ops = self._find_write_operations(token)
operations.update(child_ops)
return operations