mcp-snowflake-server-no-write-system-prompt

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