Supabase MCP Server

from typing import Any from pglast.parser import ParseError, parse_sql from supabase_mcp.exceptions import ValidationError from supabase_mcp.logger import logger from supabase_mcp.services.database.sql.models import ( QueryValidationResults, SQLQueryCategory, SQLQueryCommand, ValidatedStatement, ) from supabase_mcp.services.safety.safety_configs import SQLSafetyConfig class SQLValidator: """SQL validator class that is based on pglast library. Responsible for: - SQL query syntax validation - SQL query categorization""" # Mapping from statement types to object types STATEMENT_TYPE_TO_OBJECT_TYPE = { "CreateFunctionStmt": "function", "ViewStmt": "view", "CreateTableAsStmt": "materialized_view", # When relkind is 'm', otherwise 'table' "CreateEnumStmt": "type", "CreateTypeStmt": "type", "CreateExtensionStmt": "extension", "CreateForeignTableStmt": "foreign_table", "CreatePolicyStmt": "policy", "CreateTrigStmt": "trigger", "IndexStmt": "index", "CreateStmt": "table", "AlterTableStmt": "table", "GrantStmt": "privilege", "RevokeStmt": "privilege", "CreateProcStmt": "procedure", # For CREATE PROCEDURE } def __init__(self, safety_config: SQLSafetyConfig | None = None) -> None: self.safety_config = safety_config or SQLSafetyConfig() def validate_schema_name(self, schema_name: str) -> str: """Validate schema name. Rules: - Must be a string - Cannot be empty - Cannot contain spaces or special characters """ if not schema_name.strip(): raise ValidationError("Schema name cannot be empty") if " " in schema_name: raise ValidationError("Schema name cannot contain spaces") return schema_name def validate_table_name(self, table: str) -> str: """Validate table name. Rules: - Must be a string - Cannot be empty - Cannot contain spaces or special characters """ if not table.strip(): raise ValidationError("Table name cannot be empty") if " " in table: raise ValidationError("Table name cannot contain spaces") return table def basic_query_validation(self, query: str) -> str: """Validate SQL query. Rules: - Must be a string - Cannot be empty """ if not query.strip(): raise ValidationError("Query cannot be empty") return query @classmethod def validate_transaction_control(cls, query: str) -> bool: """Check if the query contains transaction control statements. Args: query: SQL query string Returns: bool: True if the query contains any transaction control statements """ return any(x in query.upper() for x in ["BEGIN", "COMMIT", "ROLLBACK"]) def validate_query(self, sql_query: str) -> QueryValidationResults: """ Identify the type of SQL query using PostgreSQL's parser. Args: sql_query: A SQL query string to parse Returns: QueryValidationResults: A validation result object containing information about the SQL statements Raises: ValidationError: If the query is not valid or contains TCL statements """ try: # Validate raw input sql_query = self.basic_query_validation(sql_query) # Parse the SQL using PostgreSQL's parser parse_tree = parse_sql(sql_query) if parse_tree is None: logger.debug("No statements found in the query") # logger.debug(f"Parse tree generated with {parse_tree} statements") # Validate statements result = self.validate_statements(original_query=sql_query, parse_tree=parse_tree) # Check if the query contains transaction control statements and reject them for statement in result.statements: if statement.category == SQLQueryCategory.TCL: logger.warning(f"Transaction control statement detected: {statement.command}") raise ValidationError( "Transaction control statements (BEGIN, COMMIT, ROLLBACK) are not allowed. " "Queries will be automatically wrapped in transactions by the system." ) return result except ParseError as e: logger.exception(f"SQL syntax error: {str(e)}") raise ValidationError(f"SQL syntax error: {str(e)}") from e except ValidationError: # let it propagate raise except Exception as e: logger.exception(f"Unexpected error during SQL validation: {str(e)}") raise ValidationError(f"Unexpected error during SQL validation: {str(e)}") from e def _map_to_command(self, stmt_type: str) -> SQLQueryCommand: """Map a pglast statement type to our SQLQueryCommand enum.""" mapping = { # DQL Commands "SelectStmt": SQLQueryCommand.SELECT, # DML Commands "InsertStmt": SQLQueryCommand.INSERT, "UpdateStmt": SQLQueryCommand.UPDATE, "DeleteStmt": SQLQueryCommand.DELETE, "MergeStmt": SQLQueryCommand.MERGE, # DDL Commands "CreateStmt": SQLQueryCommand.CREATE, "CreateTableAsStmt": SQLQueryCommand.CREATE, "CreateSchemaStmt": SQLQueryCommand.CREATE, "CreateExtensionStmt": SQLQueryCommand.CREATE, "CreateFunctionStmt": SQLQueryCommand.CREATE, "CreateTrigStmt": SQLQueryCommand.CREATE, "ViewStmt": SQLQueryCommand.CREATE, "IndexStmt": SQLQueryCommand.CREATE, # Additional DDL Commands "CreateEnumStmt": SQLQueryCommand.CREATE, "CreateTypeStmt": SQLQueryCommand.CREATE, "CreateDomainStmt": SQLQueryCommand.CREATE, "CreateSeqStmt": SQLQueryCommand.CREATE, "CreateForeignTableStmt": SQLQueryCommand.CREATE, "CreatePolicyStmt": SQLQueryCommand.CREATE, "CreateCastStmt": SQLQueryCommand.CREATE, "CreateOpClassStmt": SQLQueryCommand.CREATE, "CreateOpFamilyStmt": SQLQueryCommand.CREATE, "AlterTableStmt": SQLQueryCommand.ALTER, "AlterDomainStmt": SQLQueryCommand.ALTER, "AlterEnumStmt": SQLQueryCommand.ALTER, "AlterSeqStmt": SQLQueryCommand.ALTER, "AlterOwnerStmt": SQLQueryCommand.ALTER, "AlterObjectSchemaStmt": SQLQueryCommand.ALTER, "DropStmt": SQLQueryCommand.DROP, "TruncateStmt": SQLQueryCommand.TRUNCATE, "CommentStmt": SQLQueryCommand.COMMENT, "RenameStmt": SQLQueryCommand.RENAME, # DCL Commands "GrantStmt": SQLQueryCommand.GRANT, "GrantRoleStmt": SQLQueryCommand.GRANT, "RevokeStmt": SQLQueryCommand.REVOKE, "RevokeRoleStmt": SQLQueryCommand.REVOKE, "CreateRoleStmt": SQLQueryCommand.CREATE, "AlterRoleStmt": SQLQueryCommand.ALTER, "DropRoleStmt": SQLQueryCommand.DROP, # TCL Commands "TransactionStmt": SQLQueryCommand.BEGIN, # Will need refinement for different transaction types # PostgreSQL-specific Commands "VacuumStmt": SQLQueryCommand.VACUUM, "ExplainStmt": SQLQueryCommand.EXPLAIN, "CopyStmt": SQLQueryCommand.COPY, "ListenStmt": SQLQueryCommand.LISTEN, "NotifyStmt": SQLQueryCommand.NOTIFY, "PrepareStmt": SQLQueryCommand.PREPARE, "ExecuteStmt": SQLQueryCommand.EXECUTE, "DeallocateStmt": SQLQueryCommand.DEALLOCATE, } # Try to map the statement type, default to UNKNOWN return mapping.get(stmt_type, SQLQueryCommand.UNKNOWN) def validate_statements(self, original_query: str, parse_tree: Any) -> QueryValidationResults: """Validate the statements in the parse tree. Args: parse_tree: The parse tree to validate Returns: SQLBatchValidationResult: A validation result object containing information about the SQL statements Raises: ValidationError: If the query is not valid """ result = QueryValidationResults(original_query=original_query) if parse_tree is None: return result try: for stmt in parse_tree: if not hasattr(stmt, "stmt"): continue stmt_node = stmt.stmt stmt_type = stmt_node.__class__.__name__ logger.debug(f"Processing statement node type: {stmt_type}") # logger.debug(f"DEBUGGING stmt_node: {stmt_node}") logger.debug(f"DEBUGGING stmt_node.stmt_location: {stmt.stmt_location}") # Extract the object type if available object_type = None schema_name = None if hasattr(stmt_node, "relation") and stmt_node.relation is not None: if hasattr(stmt_node.relation, "relname"): object_type = stmt_node.relation.relname if hasattr(stmt_node.relation, "schemaname"): schema_name = stmt_node.relation.schemaname # For statements with 'relations' list (like TRUNCATE) elif hasattr(stmt_node, "relations") and stmt_node.relations: for relation in stmt_node.relations: if hasattr(relation, "relname"): object_type = relation.relname if hasattr(relation, "schemaname"): schema_name = relation.schemaname break # Simple approach: Set object_type based on statement type if not already set if object_type is None and stmt_type in self.STATEMENT_TYPE_TO_OBJECT_TYPE: object_type = self.STATEMENT_TYPE_TO_OBJECT_TYPE[stmt_type] # Default schema to public if not set if schema_name is None: schema_name = "public" # Get classification for this statement type classification = self.safety_config.classify_statement(stmt_type, stmt_node) logger.debug( f"Statement category classified as: {classification.get('category', 'UNKNOWN')} - risk level: {classification.get('risk_level', 'UNKNOWN')}" ) logger.debug(f"DEBUGGING QUERY EXTRACTION LOCATION: {stmt.stmt_location} - {stmt.stmt_len}") # Create validation result query_result = ValidatedStatement( category=classification["category"], command=self._map_to_command(stmt_type), risk_level=classification["risk_level"], needs_migration=classification["needs_migration"], object_type=object_type, schema_name=schema_name, query=original_query[stmt.stmt_location : stmt.stmt_location + stmt.stmt_len] if hasattr(stmt, "stmt_location") and hasattr(stmt, "stmt_len") else None, ) # logger.debug(f"Isolated query: {query_result.query}") logger.debug( "Query validation result:", { "statement_category": query_result.category, "risk_level": query_result.risk_level, "needs_migration": query_result.needs_migration, "object_type": query_result.object_type, "schema_name": query_result.schema_name, "query": query_result.query, }, ) # Add result to the batch result.statements.append(query_result) # Update highest risk level if query_result.risk_level > result.highest_risk_level: result.highest_risk_level = query_result.risk_level logger.debug(f"Updated batch validation result to: {query_result.risk_level}") if len(result.statements) == 0: logger.debug("No valid statements found in the query") raise ValidationError("No queries were parsed - please check correctness of your query") logger.debug( f"Validated a total of {len(result.statements)} with the highest risk level of: {result.highest_risk_level}" ) return result except AttributeError as e: # Handle attempting to access missing attributes in the parse tree raise ValidationError(f"Error accessing parse tree structure: {str(e)}") from e except KeyError as e: # Handle missing keys in classification dictionary raise ValidationError(f"Missing classification key: {str(e)}") from e