""" Snowflake MCP Server Exposes Snowflake as a set of MCP tools so any MCP client (VS Code Copilot, Claude Desktop, etc.) can query Snowflake in natural language. Tools ----- run_query — execute a read-only SQL statement and return results as markdown list_databases — SHOW DATABASES list_schemas — SHOW SCHEMAS IN DATABASE list_tables — SHOW TABLES IN SCHEMA . describe_table — DESCRIBE TABLE .. Auth ---- Authenticates via SSO (externalbrowser). A browser window opens once at server startup; the token is cached for the lifetime of the process. Security -------- - Only SELECT / SHOW / DESCRIBE / EXPLAIN / WITH are permitted in run_query. - Multi-statement SQL (semicolons) is rejected. - SQL comments are stripped before analysis to prevent bypass. - Catalog identifiers are validated against a strict allowlist ([A-Za-z0-9_$]) to prevent identifier injection. - Result rows are capped at MAX_ROWS (500); markdown output capped at MAX_OUTPUT_CHARS. """ from __future__ import annotations import logging import os import re import secrets import threading import time from collections.abc import AsyncIterator from contextlib import asynccontextmanager from dataclasses import dataclass, field from pathlib import Path from typing import Any import snowflake.connector import snowflake.connector.errors from dotenv import dotenv_values from mcp.server.fastmcp import Context, FastMCP logger = logging.getLogger(__name__) # Suppress chatty INFO-level request logs from the FastMCP framework # (VS Code surfaces all stderr as warnings, making them noisy). logging.getLogger("mcp").setLevel(logging.WARNING) # --------------------------------------------------------------------------- # Security constants # --------------------------------------------------------------------------- # Hard cap on rows returned by run_query (even if the query has no LIMIT) MAX_ROWS = 500 # Hard cap on total markdown output characters to avoid flooding the client MAX_OUTPUT_CHARS = 100_000 # Default LIMIT injected into bare SELECTs from run_query DEFAULT_LIMIT = 100 # Strict allowlist for Snowflake identifiers: letters, # digits, underscore, $. # Max 255 chars per component; dots separate db.schema.table. _IDENT_PART = r'[A-Za-z0-9_$]{1,255}' _IDENT_RE = re.compile(rf'^{_IDENT_PART}$') _QUALIFIED_IDENT_RE = re.compile( rf'^{_IDENT_PART}(\.{_IDENT_PART}){{0,2}}$' ) # Strip SQL line comments (--) and block comments (/* */) _STRIP_COMMENTS_RE = re.compile(r'--[^\n]*|/\*.*?\*/', re.DOTALL) # Allowlist: only these statement types are permitted in run_query _ALLOWED_STMT_RE = re.compile( r'^\s*(SELECT|SHOW|DESCRIBE|DESC|EXPLAIN|WITH)\b', re.IGNORECASE, ) # Allowlist for preview_mutation: DML + scoped table DDL. # CREATE is permitted only when the second meaningful token # is TABLE (including OR REPLACE … TABLE and TEMP/TEMPORARY/ # TRANSIENT variants). # CREATE SCHEMA/DATABASE/ROLE/STAGE etc. remain blocked. _MUTATION_STMT_RE = re.compile( r'^\s*(INSERT|UPDATE|DELETE|MERGE|CALL' r'|CREATE\s+(OR\s+REPLACE\s+)?((?:TEMPORARY|TRANSIENT|TEMP)\s+)?TABLE' r')\b', re.IGNORECASE, ) # Staged mutations: token → (sql, expires_at_unix, database, schema) # Tokens are single-use and expire after MUTATION_TOKEN_TTL seconds. # database/schema are the optional per-call context overrides. MUTATION_TOKEN_TTL = 120 # 2 minutes _PENDING_MUTATIONS: dict[ str, tuple[str, float, str | None, str | None] ] = {} # Protects all reads/writes on _PENDING_MUTATIONS to prevent # concurrent purge-vs-pop races. _mutations_lock = threading.Lock() def _validate_identifier( value: str, label: str = "identifier" ) -> str: """Raise ValueError if not a safe Snowflake identifier.""" if not _IDENT_RE.match(value): raise ValueError( f"Invalid {label} {value!r}: only letters," " digits, _ and $ are permitted." ) return value def _validate_qualified_identifier( value: str, label: str = "identifier" ) -> str: """Raise ValueError if not a safe qualified identifier.""" if not _QUALIFIED_IDENT_RE.match(value): raise ValueError( f"Invalid {label} {value!r}: use only letters, digits, _ and $ " "with optional dot-separated qualifications." ) return value def _validate_context_ids( database: str | None, schema: str | None, ) -> tuple[str | None, str | None]: """Validate optional database/schema override identifiers. Returns (db, sc) as validated strings (or None). Raises ValueError with a user-facing message on invalid input. """ db: str | None = None sc: str | None = None if database is not None: db = _validate_identifier(database, "database") if schema is not None: sc = _validate_identifier(schema, "schema") return db, sc def _guard_sql(sql: str) -> str: """Validate and normalise a user-supplied SQL string. Returns the cleaned SQL ready for execution, or raises ValueError. Rules enforced: - Strip comments first (prevents --//* */ bypass) - Exactly one statement (no semicolons between statements) - Statement must start with an allowlisted verb: SELECT, SHOW, DESCRIBE, DESC, EXPLAIN, or WITH (CTE → must ultimately be a read-only query) """ # Remove comments cleaned = _STRIP_COMMENTS_RE.sub(' ', sql).strip() # Split on ; — reject more than one non-empty part parts = [p.strip() for p in cleaned.split(';') if p.strip()] if len(parts) > 1: raise ValueError( "Multi-statement SQL is not permitted." " Submit one statement at a time." ) if not parts: raise ValueError("Empty SQL statement.") stmt = parts[0] if not _ALLOWED_STMT_RE.match(stmt): first_token = ( stmt.split()[0].upper() if stmt.split() else '(empty)' ) raise ValueError( f"Statement type '{first_token}' is not permitted." " Only SELECT, SHOW, DESCRIBE, EXPLAIN, and WITH" " (CTE) are allowed." ) return stmt def _guard_mutation_sql(sql: str) -> str: """Validate a user-supplied mutation SQL string. Applies the same comment-stripping and single-statement rules as _guard_sql, but against the write allowlist: DML : INSERT, UPDATE, DELETE, MERGE, CALL DDL : CREATE [OR REPLACE] [TEMP|TEMPORARY|TRANSIENT] TABLE only Still blocked: DROP, TRUNCATE, ALTER, CREATE SCHEMA/DATABASE/ROLE/STAGE/…, GRANT, REVOKE. """ cleaned = _STRIP_COMMENTS_RE.sub(' ', sql).strip() parts = [p.strip() for p in cleaned.split(';') if p.strip()] if len(parts) > 1: raise ValueError( "Multi-statement SQL is not permitted." " Submit one statement at a time." ) if not parts: raise ValueError("Empty SQL statement.") stmt = parts[0] if not _MUTATION_STMT_RE.match(stmt): first_token = ( stmt.split()[0].upper() if stmt.split() else '(empty)' ) raise ValueError( f"Statement type '{first_token}' is not permitted" " for mutation. Allowed: INSERT, UPDATE, DELETE," " MERGE, CALL, CREATE [OR REPLACE]" " [TEMP|TEMPORARY|TRANSIENT] TABLE." " Blocked: DROP, TRUNCATE, ALTER," " CREATE SCHEMA/DATABASE/ROLE/…, GRANT, REVOKE." ) return stmt # --------------------------------------------------------------------------- # .env loading — walk up from this file to find the workspace root .env # --------------------------------------------------------------------------- def _find_dotenv() -> Path | None: here = Path(__file__).resolve().parent for candidate in [here, here.parent, here.parent.parent]: p = candidate / ".env" if p.exists(): return p return None def _get_config( key: str, _dotenv: dict[str, str | None] | None = None, ) -> str | None: """Return a stripped config value. Checks os.environ first, then falls back to *_dotenv* when supplied (avoids repeated filesystem reads when the caller has already loaded the .env dict), or re-reads the .env file itself when *_dotenv* is None (polling behaviour for standalone callers). Returns None when both sources are blank or absent. """ val = (os.getenv(key) or "").strip() if not val: dotenv: dict[str, str | None] if _dotenv is None: env_path = _find_dotenv() dotenv = dotenv_values(env_path) if env_path else {} else: dotenv = _dotenv val = (dotenv.get(key) or "").strip() return val or None # --------------------------------------------------------------------------- # Connection helpers # --------------------------------------------------------------------------- def _build_connect_kwargs() -> dict[str, Any]: # Read .env exactly once per connection attempt so a .env file # added after startup is honoured without N filesystem reads. _env_path = _find_dotenv() _dotenv = dotenv_values(_env_path) if _env_path else {} account = _get_config("SNOWFLAKE_ACCOUNT", _dotenv) user = _get_config("SNOWFLAKE_USER", _dotenv) warehouse = _get_config("SNOWFLAKE_WAREHOUSE", _dotenv) role = _get_config("SNOWFLAKE_ROLE", _dotenv) database = _get_config("SNOWFLAKE_DATABASE", _dotenv) schema = _get_config("SNOWFLAKE_SCHEMA", _dotenv) missing = [ name for name, val in [ ("SNOWFLAKE_ACCOUNT", account), ("SNOWFLAKE_USER", user), ("SNOWFLAKE_WAREHOUSE", warehouse), ("SNOWFLAKE_ROLE", role), ("SNOWFLAKE_DATABASE", database), ("SNOWFLAKE_SCHEMA", schema), ] if not val ] if missing: raise RuntimeError(f"{', '.join(missing)} must be set.") # Only SSO is supported. password and private_key are # intentionally never read or passed to the connector. _timeout_raw = ( _get_config("SNOWFLAKE_BROWSER_TIMEOUT", _dotenv) or "120" ) return dict( account=account, user=user, warehouse=warehouse, role=role, database=database, schema=schema, authenticator="externalbrowser", external_browser_timeout=int(_timeout_raw), ) # --------------------------------------------------------------------------- # Lifespan — single persistent connection for the server's lifetime # --------------------------------------------------------------------------- @dataclass class AppState: conn: snowflake.connector.SnowflakeConnection | None = None default_database: str | None = None default_schema: str | None = None # Held during USE→execute→restore to prevent interleaving # when context overrides are used on a shared connection. lock: threading.Lock = field(default_factory=threading.Lock) def _get_or_create_conn( state: AppState, ) -> snowflake.connector.SnowflakeConnection: """Return the cached connection, creating it if needed. Called lazily on the first tool invocation so the server starts successfully even when credentials are not yet set. Re-reads config whenever establishing a new connection so a .env file added after startup is picked up on the next (re)connect. Caches the configured default database/schema on AppState so per-call context overrides can be restored afterwards. """ if state.conn is None or state.conn.is_closed(): kwargs = _build_connect_kwargs() state.conn = snowflake.connector.connect(**kwargs) # Validate the configured defaults before caching so they # are safe to interpolate into USE statements during restore. raw_db = kwargs.get("database") raw_sc = kwargs.get("schema") try: state.default_database = ( _validate_identifier(raw_db, "database") if raw_db else None ) state.default_schema = ( _validate_identifier(raw_sc, "schema") if raw_sc else None ) except ValueError as exc: raise RuntimeError(str(exc)) from exc return state.conn @asynccontextmanager async def lifespan(server: FastMCP) -> AsyncIterator[AppState]: state = AppState() try: yield state finally: if state.conn is not None: try: state.conn.close() except Exception: pass # --------------------------------------------------------------------------- # MCP server # --------------------------------------------------------------------------- mcp = FastMCP( "snowflake", instructions=( "Query Snowflake using SQL. " "Use list_databases / list_schemas / list_tables /" " describe_table to explore the catalog, then" " run_query to execute SELECT statements. " "Always fully-qualify table names as" " DATABASE.SCHEMA.TABLE when not in a default context." ), lifespan=lifespan, ) # --------------------------------------------------------------------------- # Formatting helper # --------------------------------------------------------------------------- def _to_markdown(columns: list[str], rows: list[tuple]) -> str: if not rows: return "_No rows returned._" col_widths = [ max(len(c), max((len(str(r[i])) for r in rows), default=0)) for i, c in enumerate(columns) ] sep = ( "| " + " | ".join("-" * w for w in col_widths) + " |" ) header = ( "| " + " | ".join( c.ljust(w) for c, w in zip(columns, col_widths) ) + " |" ) lines = [header, sep] for row in rows: lines.append( "| " + " | ".join( str(v).ljust(w) for v, w in zip(row, col_widths) ) + " |" ) lines.append( f"\n_{len(rows)} row{'s' if len(rows) != 1 else ''}_" ) output = "\n".join(lines) if len(output) > MAX_OUTPUT_CHARS: output = ( output[:MAX_OUTPUT_CHARS] + f"\n\n_[output truncated at" f" {MAX_OUTPUT_CHARS:,} chars]_" ) return output def _execute( ctx: Context, sql: str, database: str | None = None, schema: str | None = None, ) -> tuple[list[str], list[tuple]]: """Execute *sql* on the connection, optionally switching database/schema context first and restoring defaults after. database and schema must already be validated identifiers. The connection lock is always held for the full USE→execute→restore cycle so that non-override calls cannot interleave with override calls that have temporarily changed session context on the shared connection. """ state: AppState = ctx.request_context.lifespan_context conn = _get_or_create_conn(state) # Always hold the lock so non-override calls cannot execute SQL # while an override call has temporarily changed session context. with state.lock: try: with conn.cursor() as cur: if database is not None: cur.execute(f"USE DATABASE {database}") if schema is not None: cur.execute(f"USE SCHEMA {schema}") cur.execute(sql) if cur.description: columns = [d[0] for d in cur.description] rows = cur.fetchmany(MAX_ROWS) else: columns = ["status"] rows = [( f"OK — {cur.rowcount} row" f"{'s' if cur.rowcount != 1 else ''} affected", )] return columns, rows finally: # Restore configured defaults so context doesn't leak # into subsequent tool calls. Always restore both defaults # when either is overridden — switching database can # implicitly change schema in Snowflake. if database is not None or schema is not None: try: with conn.cursor() as cur: if state.default_database is not None: cur.execute( f"USE DATABASE" f" {state.default_database}" ) if state.default_schema is not None: cur.execute( f"USE SCHEMA" f" {state.default_schema}" ) except Exception: logger.exception( "Failed to restore Snowflake session" " defaults after query with" " database=%r schema=%r", database, schema, ) # --------------------------------------------------------------------------- # Tools # --------------------------------------------------------------------------- @mcp.tool() def run_query( sql: str, ctx: Context, limit: int = DEFAULT_LIMIT, database: str | None = None, schema: str | None = None, ) -> str: """Execute a read-only Snowflake SQL query. Returns results as a markdown table. Only SELECT, SHOW, DESCRIBE, EXPLAIN, and WITH (CTE) statements are permitted. Multi-statement SQL is rejected. SQL comments are stripped before analysis. A LIMIT clause is appended to bare SELECTs, capped at 500 rows maximum. Args: sql: A single read-only SQL statement. limit: Max rows to return (1–500, default 100). database: Override the session database for this call only (optional). Only alphanumeric, _ and $ allowed. schema: Override the session schema for this call only (optional). Only alphanumeric, _ and $ allowed. """ # Clamp limit — user cannot request more than MAX_ROWS limit = max(1, min(limit, MAX_ROWS)) try: stmt = _guard_sql(sql) except ValueError as e: return f"⛔ {e}" try: db, sc = _validate_context_ids(database, schema) except ValueError as e: return f"⛔ {e}" # Auto-inject LIMIT on bare SELECTs upper = stmt.upper() if ( _ALLOWED_STMT_RE.match(stmt) and upper.lstrip().startswith("SELECT") and " LIMIT " not in upper ): stmt = f"{stmt} LIMIT {limit}" try: columns, rows = _execute(ctx, stmt, db, sc) output = _to_markdown(columns, rows) if len(rows) == MAX_ROWS: output += ( f"\n\n_Results capped at {MAX_ROWS} rows." " Add a more specific WHERE clause to narrow" " results._" ) return output except RuntimeError as e: return f"⚙️ Configuration error: {e}" except snowflake.connector.errors.ProgrammingError as e: # Return the Snowflake error message but not the full traceback return f"❌ Query error: {e.msg}" @mcp.tool() def list_databases(ctx: Context) -> str: """List all Snowflake databases accessible to the current role.""" try: columns, rows = _execute(ctx, "SHOW DATABASES") name_idx = next( (i for i, c in enumerate(columns) if c.lower() == "name"), 1 ) owner_idx = next( (i for i, c in enumerate(columns) if c.lower() == "owner"), None ) if owner_idx is not None: return _to_markdown( ["name", "owner"], [(r[name_idx], r[owner_idx]) for r in rows], ) return _to_markdown( ["name"], [(r[name_idx],) for r in rows] ) except RuntimeError as e: return f"⚙️ Configuration error: {e}" except snowflake.connector.errors.ProgrammingError as e: return f"❌ {e.msg}" @mcp.tool() def list_schemas(database: str, ctx: Context) -> str: """List all schemas in a Snowflake database. Args: database: The database name (e.g. INTEGRATION). Only alphanumeric, _ and $ allowed. """ try: db = _validate_identifier(database, "database") except ValueError as e: return f"⛔ {e}" try: columns, rows = _execute( ctx, f"SHOW SCHEMAS IN DATABASE {db}" ) name_idx = next( (i for i, c in enumerate(columns) if c.lower() == "name"), 1 ) return _to_markdown( ["schema"], [(r[name_idx],) for r in rows] ) except RuntimeError as e: return f"⚙️ Configuration error: {e}" except snowflake.connector.errors.ProgrammingError as e: return f"❌ {e.msg}" @mcp.tool() def list_tables(database: str, schema: str, ctx: Context) -> str: """List all tables in a Snowflake schema. Args: database: The database name (e.g. INTEGRATION). Only alphanumeric, _ and $ allowed. schema: The schema name (e.g. PUBLIC). Only alphanumeric, _ and $ allowed. """ try: db = _validate_identifier(database, "database") sc = _validate_identifier(schema, "schema") except ValueError as e: return f"⛔ {e}" try: columns, rows = _execute( ctx, f"SHOW TABLES IN SCHEMA {db}.{sc}" ) name_idx = next( (i for i, c in enumerate(columns) if c.lower() == "name"), 1 ) rows_idx = next( (i for i, c in enumerate(columns) if c.lower() == "rows"), None ) if rows_idx is not None: return _to_markdown( ["table", "rows"], [(r[name_idx], r[rows_idx]) for r in rows], ) return _to_markdown( ["table"], [(r[name_idx],) for r in rows] ) except RuntimeError as e: return f"⚙️ Configuration error: {e}" except snowflake.connector.errors.ProgrammingError as e: return f"❌ {e.msg}" @mcp.tool() def describe_table(table_name: str, ctx: Context) -> str: """Describe the columns of a Snowflake table. Args: table_name: Dot-separated qualified name: DATABASE.SCHEMA.TABLE. Only alphanumeric, _ and $ are permitted in each part. """ try: tbl = _validate_qualified_identifier( table_name, "table name" ) except ValueError as e: return f"⛔ {e}" try: columns, rows = _execute(ctx, f"DESCRIBE TABLE {tbl}") return _to_markdown(columns, rows) except RuntimeError as e: return f"⚙️ Configuration error: {e}" except snowflake.connector.errors.ProgrammingError as e: return f"❌ {e.msg}" # --------------------------------------------------------------------------- # Mutation tools (two-phase commit: preview → confirm) # --------------------------------------------------------------------------- @mcp.tool() def preview_mutation( sql: str, ctx: Context, database: str | None = None, schema: str | None = None, ) -> str: """Stage a write SQL statement for human review before execution. Validates the SQL, then returns a preview of the exact statement that will run along with a single-use approval token valid for 2 minutes. Call confirm_mutation(token) to execute. The token is cryptographically random, single-use, time-limited, and bound to the exact SQL shown — it cannot be reused, substituted, or guessed. Permitted: INSERT, UPDATE, DELETE, MERGE, CALL, CREATE [OR REPLACE] [TEMP|TEMPORARY|TRANSIENT] TABLE. Blocked: DROP, TRUNCATE, ALTER, CREATE SCHEMA/DATABASE/ROLE/…, GRANT, REVOKE (and all read statements — use run_query). Args: sql: The write SQL statement to preview. database: Override the session database for this mutation only (optional). Only alphanumeric, _ and $ allowed. schema: Override the session schema for this mutation only (optional). Only alphanumeric, _ and $ allowed. """ try: stmt = _guard_mutation_sql(sql) except ValueError as e: return f"⛔ {e}" try: db, sc = _validate_context_ids(database, schema) except ValueError as e: return f"⛔ {e}" # Purge stale tokens to avoid unbounded memory growth now = time.time() with _mutations_lock: stale = [ t for t, (_, exp, *_rest) in _PENDING_MUTATIONS.items() if exp < now ] for t in stale: del _PENDING_MUTATIONS[t] token = secrets.token_urlsafe(24) _PENDING_MUTATIONS[token] = ( stmt, now + MUTATION_TOKEN_TTL, db, sc ) ttl_min = MUTATION_TOKEN_TTL // 60 context_note = "" if db or sc: parts = [p for p in [db, sc] if p] context_note = ( f"\n\n_Context: {' / '.join(parts)}_" ) return ( f"## Mutation preview\n\n" f"Review the statement below carefully. If it looks correct, call " f"`confirm_mutation` with the token.\n\n" f"```sql\n{stmt}\n```\n\n" f"**Approval token:** `{token}`\n" f"_(single-use · expires in {ttl_min} minutes" f" · bound to the exact SQL" f"{' and execution context' if (db or sc) else ''}" f" above)_{context_note}" ) @mcp.tool() def confirm_mutation(token: str, ctx: Context) -> str: """Execute a mutation that was staged by preview_mutation. The token is consumed on first use (single-use). Expired or unknown tokens are rejected — call preview_mutation again to stage a new one. Args: token: The approval token returned by preview_mutation. """ now = time.time() with _mutations_lock: entry = _PENDING_MUTATIONS.pop(token, None) if entry is None: return ( "⛔ Invalid or already-used token. " "Call preview_mutation again to stage a new mutation." ) stmt, expires_at, db, sc = entry if now > expires_at: return ( "⛔ Token expired. " "Call preview_mutation again to stage a new mutation." ) try: columns, rows = _execute(ctx, stmt, db, sc) output = _to_markdown(columns, rows) return f"✅ Mutation executed.\n\n{output}" except RuntimeError as e: return f"⚙️ Configuration error: {e}" except snowflake.connector.errors.ProgrammingError as e: return f"❌ Mutation error: {e.msg}" # --------------------------------------------------------------------------- # Entry point # --------------------------------------------------------------------------- def main() -> None: mcp.run(transport="stdio") if __name__ == "__main__": main()