sql-safe-mcp 1.2.0__py3-none-any.whl

sql_safe_mcp/security/parser.py

@@ -0,0 +1,183 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+import sqlglot
+from sqlglot import exp
+
+from sql_safe_mcp.config import RuntimeConfig
+from sql_safe_mcp.errors import DomainError, ErrorCode
+from sql_safe_mcp.security.dialect import SQLSERVER, SqlDialect
+from sql_safe_mcp.security.reasons import Reason
+
+
+@dataclass(frozen=True, slots=True)
+class ParserLimits:
+    max_sql_chars: int
+    max_ast_nodes: int
+    max_joins: int
+    max_in_list_items: int
+
+    @classmethod
+    def from_runtime(cls, runtime: RuntimeConfig) -> ParserLimits:
+        return cls(
+            runtime.max_sql_chars,
+            runtime.max_ast_nodes,
+            runtime.max_joins,
+            runtime.max_in_list_items,
+        )
+
+
+REJECT_HINT = "Rewrite the operation as a single, simpler SELECT query."
+
+
+def reject(reason: Reason, detail: str | None = None) -> DomainError:
+    """Build the QUERY_REJECTED error; only a fixed Reason (plus an optional name) is accepted."""
+    if not isinstance(reason, Reason):
+        raise TypeError("reject() takes a Reason")
+    text = str(reason) if detail is None else f"{reason}: {detail}"
+    return DomainError(ErrorCode.QUERY_REJECTED, f"Query rejected: {text}.", REJECT_HINT)
+
+
+def check_limits(query: exp.Expression, limits: ParserLimits) -> None:
+    """Enforce AST size limits; also used again after transformations."""
+    for count, _ in enumerate(query.walk(), start=1):
+        if count > limits.max_ast_nodes:
+            raise reject(Reason.LIMIT_NODES)
+    if sum(1 for _ in query.find_all(exp.Join)) > limits.max_joins:
+        raise reject(Reason.LIMIT_JOINS)
+    for node in query.find_all(exp.In):
+        if len(node.expressions) > limits.max_in_list_items:
+            raise reject(Reason.LIMIT_IN_LIST)
+
+
+_COMPARISONS = (exp.EQ, exp.NEQ, exp.GT, exp.GTE, exp.LT, exp.LTE)
+
+# node type -> argument names that may be non-empty; everything else is rejected.
+_ALLOWED_ARGS: dict[type[exp.Expr], frozenset[str]] = {
+    exp.Select: frozenset({"expressions", "from_", "joins", "where", "order", "limit"}),
+    exp.From: frozenset({"this"}),
+    exp.Table: frozenset({"this", "db", "alias"}),
+    exp.TableAlias: frozenset({"this"}),
+    exp.Identifier: frozenset({"this", "quoted"}),
+    exp.Column: frozenset({"this", "table"}),
+    exp.Star: frozenset(),
+    exp.Count: frozenset({"this", "big_int"}),  # big_int: a flag the mysql parser sets
+    exp.Literal: frozenset({"this", "is_string"}),
+    exp.Null: frozenset(),
+    exp.Neg: frozenset({"this"}),
+    exp.Alias: frozenset({"this", "alias"}),
+    exp.Join: frozenset({"this", "on", "side", "kind"}),
+    exp.And: frozenset({"this", "expression"}),
+    exp.Or: frozenset({"this", "expression"}),
+    exp.Paren: frozenset({"this"}),
+    exp.In: frozenset({"this", "expressions"}),
+    exp.Is: frozenset({"this", "expression"}),
+    exp.Where: frozenset({"this"}),
+    exp.Order: frozenset({"expressions"}),
+    exp.Ordered: frozenset({"this", "desc", "nulls_first"}),
+    exp.Limit: frozenset({"expression"}),
+    **{comparison: frozenset({"this", "expression"}) for comparison in _COMPARISONS},
+}
+_PLACEHOLDER_ARGS = frozenset({"this"})
+_PROJECTIONS = (exp.Column, exp.Star, exp.Literal, exp.Null, exp.Neg, exp.Alias, exp.Count)
+
+
+def _is_numeric_literal(node: object) -> bool:
+    return isinstance(node, exp.Literal) and not node.is_string
+
+
+def _is_empty(value: object) -> bool:
+    return value is None or value is False or value == []
+
+
+def _check_node(node: exp.Expr, allow_placeholders: bool) -> None:
+    kind = type(node)
+    if kind is exp.Placeholder and allow_placeholders:
+        allowed_args = _PLACEHOLDER_ARGS
+    elif kind in _ALLOWED_ARGS:
+        allowed_args = _ALLOWED_ARGS[kind]
+    else:
+        raise reject(Reason.NODE_UNSUPPORTED, kind.__name__)
+    for name, value in node.args.items():
+        if name not in allowed_args and not _is_empty(value):
+            raise reject(Reason.OPTION_UNSUPPORTED, f"{kind.__name__}.{name}")
+
+    if isinstance(node, exp.Table):
+        name = node.this
+        if not isinstance(name, exp.Identifier) or str(name.this).startswith(("#", "@")):
+            raise reject(Reason.TABLE_UNSUPPORTED)
+    elif isinstance(node, exp.Column):
+        if not isinstance(node.this, (exp.Identifier, exp.Star)):
+            raise reject(Reason.COLUMN_REFERENCE_UNSUPPORTED)
+    elif isinstance(node, exp.Count):
+        if not isinstance(node.this, exp.Star):
+            raise reject(Reason.COUNT_STAR_ONLY)
+    elif isinstance(node, exp.Star):
+        if not isinstance(node.parent, (exp.Select, exp.Column, exp.Count)):
+            raise reject(Reason.STAR_PROJECTION_ONLY)
+    elif isinstance(node, exp.Join):
+        if not isinstance(node.this, exp.Table) or node.args.get("on") is None:
+            raise reject(Reason.JOIN_NEEDS_ON)
+        if node.side not in {"", "LEFT"} or node.kind not in {"", "INNER"}:
+            raise reject(Reason.JOIN_TYPE_UNSUPPORTED)
+    elif isinstance(node, exp.Is):
+        if not isinstance(node.expression, exp.Null):
+            raise reject(Reason.IS_NULL_ONLY)
+    elif isinstance(node, exp.Neg):
+        if not _is_numeric_literal(node.this):
+            raise reject(Reason.NEGATION_NUMERIC_ONLY)
+    elif isinstance(node, exp.Limit):
+        value = node.expression
+        if not _is_numeric_literal(value) or not str(value.this).isdigit():
+            raise reject(Reason.TOP_INTEGER_ONLY)
+    elif isinstance(node, exp.Ordered):
+        if not isinstance(node.this, exp.Column) or isinstance(node.this.this, exp.Star):
+            raise reject(Reason.ORDER_DIRECT_COLUMNS_ONLY)
+    elif isinstance(node, exp.In):
+        if not isinstance(node.this, exp.Column) or not all(
+            isinstance(item, (exp.Literal, exp.Neg, exp.Placeholder)) for item in node.expressions
+        ):
+            raise reject(Reason.IN_LITERALS_ONLY)
+    elif isinstance(node, _COMPARISONS):
+        for side in (node.this, node.expression):
+            if isinstance(side, exp.Star):
+                raise reject(Reason.STAR_COMPARISON)
+
+
+def validate_allowlist(query: exp.Expression, *, allow_placeholders: bool = False) -> None:
+    """Fail closed on any node type, argument, or shape outside the supported subset."""
+    if not isinstance(query, exp.Select):
+        raise reject(Reason.SELECT_ONLY)
+    for node in query.walk():
+        _check_node(node, allow_placeholders)
+    for projection in query.expressions:
+        if not isinstance(projection, _PROJECTIONS):
+            raise reject(Reason.PROJECTION_UNSUPPORTED)
+
+
+def strip_comments(query: exp.Expression) -> None:
+    """Drop caller comments so no caller-controlled text is ever emitted into generated SQL."""
+    for node in query.walk():
+        node.pop_comments()
+
+
+def parse_select(sql: str, limits: ParserLimits, dialect: SqlDialect = SQLSERVER) -> exp.Select:
+    """Parse exactly one root SELECT; any failure becomes QUERY_REJECTED."""
+    if len(sql) > limits.max_sql_chars:
+        raise reject(Reason.LIMIT_CHARS)
+    try:
+        statements = [s for s in sqlglot.parse(sql, dialect=dialect.name) if s is not None]
+        if len(statements) != 1:
+            raise reject(Reason.ONE_STATEMENT)
+        query = statements[0]
+        if not isinstance(query, exp.Select):
+            raise reject(Reason.SELECT_ONLY)
+        check_limits(query, limits)
+        validate_allowlist(query)
+        strip_comments(query)
+    except DomainError:
+        raise
+    except Exception as exc:
+        raise reject(Reason.PARSE_FAILED) from exc
+    return query

sql_safe_mcp/security/pipeline.py

@@ -0,0 +1,43 @@
+from __future__ import annotations
+
+from sql_safe_mcp.config import PiiConfig, RuntimeConfig
+from sql_safe_mcp.security.dialect import SqlDialect
+from sql_safe_mcp.security.lineage import analyze_query
+from sql_safe_mcp.security.parser import ParserLimits, parse_select
+from sql_safe_mcp.security.policy import PiiPolicy
+from sql_safe_mcp.security.schema import TableCatalog, resolve_tables
+from sql_safe_mcp.security.tokens import TokenCodec
+from sql_safe_mcp.security.validated_query import ValidatedQuery, issue_validated_query
+
+
+def validate_sql(
+    sql: str,
+    *,
+    alias: str,
+    database: str,
+    catalog: TableCatalog,
+    pii_config: PiiConfig,
+    codec: TokenCodec,
+    runtime: RuntimeConfig,
+    max_rows: int,
+    dialect: SqlDialect,
+) -> ValidatedQuery:
+    """Run caller SQL through parse, allowlist, reflection, lineage, and PII policy.
+
+    Any rejection raises a DomainError before the catalog or database is used further.
+    """
+    limits = ParserLimits.from_runtime(runtime)
+    query = parse_select(sql, limits, dialect)
+    tables = resolve_tables(query, catalog, dialect)
+    analyzed = analyze_query(query, tables, limits, dialect)
+    decision = PiiPolicy(database, pii_config, codec).evaluate(analyzed)
+    return issue_validated_query(
+        analyzed,
+        decision,
+        alias=alias,
+        database=database,
+        max_rows=max_rows,
+        runtime=runtime,
+        limits=limits,
+        dialect=dialect,
+    )

sql_safe_mcp/security/policy.py

@@ -0,0 +1,101 @@
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+
+from sqlglot import exp
+
+from sql_safe_mcp.config import PiiConfig
+from sql_safe_mcp.security.lineage import AnalyzedQuery, ColumnRef, SourceColumn
+from sql_safe_mcp.security.parser import reject
+from sql_safe_mcp.security.reasons import Reason
+from sql_safe_mcp.security.tokens import PREFIX, TokenCodec
+
+
+@dataclass(frozen=True, slots=True)
+class TokenSite:
+    """A token literal in an allowed predicate position and its decrypted bind value."""
+
+    literal: exp.Literal
+    value: object = field(repr=False)
+
+
+@dataclass(frozen=True, slots=True)
+class PolicyDecision:
+    protected: tuple[bool, ...]
+    token_sites: tuple[TokenSite, ...]
+
+
+class PiiPolicy:
+    """Decides protected-column usage from source lineage and configured rules only."""
+
+    def __init__(self, database: str, config: PiiConfig, codec: TokenCodec) -> None:
+        self._database = database.casefold()
+        self._config = config
+        self._codec = codec
+
+    def is_protected(self, source: SourceColumn) -> bool:
+        schema = source.schema.casefold()
+        table = source.table.casefold()
+        column = source.column.casefold()
+        for rule in self._config.rules:
+            if rule.database != "*" and rule.database.casefold() != self._database:
+                continue
+            if rule.schema_ is not None and rule.schema_.casefold() != schema:
+                continue
+            if rule.table.casefold() == table and column in {c.casefold() for c in rule.columns}:
+                return True
+        return False
+
+    def evaluate(self, analyzed: AnalyzedQuery) -> PolicyDecision:
+        sites: list[TokenSite] = []
+        for ref in analyzed.references:
+            if self.is_protected(ref.source):
+                sites.extend(self._check_protected(ref))
+        consumed = {id(site.literal) for site in sites}
+        for literal in analyzed.query.find_all(exp.Literal):
+            if (
+                literal.is_string
+                and str(literal.this).startswith(PREFIX)
+                and id(literal) not in consumed
+            ):
+                raise reject(Reason.TOKEN_OUTSIDE_PROTECTED)
+        protected = tuple(
+            output.source is not None and self.is_protected(output.source)
+            for output in analyzed.outputs
+        )
+        return PolicyDecision(protected, tuple(sites))
+
+    def _check_protected(self, ref: ColumnRef) -> list[TokenSite]:
+        column = ref.node
+        parent = column.parent
+        if isinstance(parent, exp.Select):
+            return []
+        if isinstance(parent, exp.Alias) and isinstance(parent.parent, exp.Select):
+            return []
+        if not _in_where(column):
+            raise reject(Reason.PROTECTED_POSITION)
+        if isinstance(parent, exp.EQ) and parent.this is column:
+            return [self._token_site(parent.expression)]
+        if isinstance(parent, exp.In) and parent.this is column and parent.expressions:
+            return [self._token_site(item) for item in parent.expressions]
+        raise reject(Reason.PROTECTED_POSITION)
+
+    def _token_site(self, node: exp.Expression) -> TokenSite:
+        if (
+            not isinstance(node, exp.Literal)
+            or not node.is_string
+            or not str(node.this).startswith(PREFIX)
+        ):
+            raise reject(Reason.PROTECTED_NEEDS_TOKEN)
+        return TokenSite(node, self._codec.decrypt(str(node.this)))
+
+
+def _in_where(node: exp.Expression) -> bool:
+    current = node.parent
+    while current is not None:
+        if isinstance(current, exp.Where):
+            return True
+        if isinstance(current, exp.Join | exp.Order | exp.Select):
+            return False
+        current = current.parent
+    return False

sql_safe_mcp/security/reasons.py

@@ -0,0 +1,43 @@
+"""Fixed, reviewable reasons for QUERY_REJECTED. Free text never reaches reject()."""
+
+from enum import StrEnum
+
+
+class Reason(StrEnum):
+    TABLE_RESOLUTION_MISMATCH = "table resolution does not match the query"
+    DUPLICATE_BINDING = "table names and aliases must be unique"
+    QUALIFIER_UNKNOWN = "a column qualifier does not match a table in the query"
+    COLUMN_NOT_FOUND = "a referenced column was not found or is ambiguous"
+    COLUMN_AMBIGUOUS = "a referenced column was not found or is ambiguous; qualify it"
+    ORDER_OUTPUT_AMBIGUOUS = "ORDER BY refers to an ambiguous or non-column output"
+    QUALIFY_FAILED = "the query references unknown or ambiguous columns"
+    LIMIT_NODES = "the query exceeds the configured complexity limit"
+    LIMIT_JOINS = "the query has too many joins"
+    LIMIT_IN_LIST = "an IN list exceeds the configured item limit"
+    TABLE_UNSUPPORTED = "only base tables of the current database are supported"
+    COLUMN_REFERENCE_UNSUPPORTED = "unsupported column reference"
+    COUNT_STAR_ONLY = "only COUNT(*) is supported"
+    STAR_PROJECTION_ONLY = "* is supported only as a projection"
+    JOIN_NEEDS_ON = "JOIN requires a table and an explicit ON expression"
+    JOIN_TYPE_UNSUPPORTED = "only INNER JOIN and LEFT JOIN are supported"
+    IS_NULL_ONLY = "IS is supported only with NULL"
+    NEGATION_NUMERIC_ONLY = "negation is supported only for numeric literals"
+    TOP_INTEGER_ONLY = "TOP/LIMIT must be a non-negative integer literal"
+    ORDER_DIRECT_COLUMNS_ONLY = "ORDER BY supports only direct columns"
+    IN_LITERALS_ONLY = "IN supports a column and a list of literals"
+    STAR_COMPARISON = "* cannot be compared"
+    SELECT_ONLY = "only SELECT statements are supported"
+    PROJECTION_UNSUPPORTED = "unsupported projection"
+    LIMIT_CHARS = "the SQL exceeds the configured length limit"
+    ONE_STATEMENT = "exactly one statement is required"
+    PARSE_FAILED = "the SQL could not be parsed"
+    TOKEN_OUTSIDE_PROTECTED = "a PII token can only be compared with a protected column"
+    PROTECTED_POSITION = "a protected column can only be projected or compared with a token"
+    PROTECTED_NEEDS_TOKEN = "a protected column can only be compared with a PII token"
+    TABLE_GONE = "a referenced table no longer exists"
+    TABLE_NOT_FOUND = "a referenced table was not found or is ambiguous; qualify it with a schema"
+    TABLE_SCOPE = "only tables of the current database are supported"
+    BIND_UNSAFE = "bind parameters could not be generated safely"
+    TOKEN_NOT_REPLACED = "a PII token could not be replaced by a bind parameter"
+    NODE_UNSUPPORTED = "unsupported construct"
+    OPTION_UNSUPPORTED = "unsupported option"

sql_safe_mcp/security/schema.py

@@ -0,0 +1,139 @@
+from __future__ import annotations
+
+import time
+from collections import OrderedDict
+from collections.abc import Callable, Hashable, Sequence
+from dataclasses import dataclass
+from threading import Lock
+from typing import Protocol
+
+from sqlalchemy import Connection, inspect
+from sqlalchemy.exc import NoSuchTableError
+from sqlglot import exp
+
+from sql_safe_mcp.db.reflection import list_tables as reflect_tables
+from sql_safe_mcp.security.dialect import SQLSERVER, SqlDialect
+from sql_safe_mcp.security.parser import reject
+from sql_safe_mcp.security.reasons import Reason
+
+
+@dataclass(frozen=True, slots=True)
+class TableSchema:
+    schema: str
+    name: str
+    columns: tuple[str, ...]
+
+
+class TableCatalog(Protocol):
+    """Reflected user base tables of one alias and database."""
+
+    def list_tables(self) -> Sequence[tuple[str, str]]: ...
+
+    def columns(self, schema: str, table: str) -> Sequence[str]: ...
+
+
+class SchemaCache:
+    """Small thread-safe LRU with a time-to-live for immutable reflection results."""
+
+    def __init__(
+        self,
+        max_entries: int,
+        ttl_seconds: float,
+        clock: Callable[[], float] = time.monotonic,
+    ) -> None:
+        self._max_entries = max_entries
+        self._ttl_seconds = ttl_seconds
+        self._clock = clock
+        self._entries: OrderedDict[Hashable, tuple[float, tuple]] = OrderedDict()
+        self._lock = Lock()
+
+    def get(self, key: Hashable) -> tuple | None:
+        with self._lock:
+            entry = self._entries.get(key)
+            if entry is None:
+                return None
+            stored_at, value = entry
+            if self._clock() - stored_at >= self._ttl_seconds:
+                del self._entries[key]
+                return None
+            self._entries.move_to_end(key)
+            return value
+
+    def put(self, key: Hashable, value: tuple) -> None:
+        with self._lock:
+            self._entries[key] = (self._clock(), value)
+            self._entries.move_to_end(key)
+            while len(self._entries) > self._max_entries:
+                self._entries.popitem(last=False)
+
+
+class ReflectedCatalog:
+    """Inspector-backed catalog; cache keys always include server alias and database."""
+
+    def __init__(
+        self, connection: Connection, cache: SchemaCache, alias: str, database: str
+    ) -> None:
+        self._connection = connection
+        self._cache = cache
+        self._scope = (alias, database)
+
+    def list_tables(self) -> list[tuple[str, str]]:
+        key = (*self._scope, "tables")
+        cached = self._cache.get(key)
+        if cached is None:
+            cached = tuple(
+                (table.schema_ or "", table.name) for table in reflect_tables(self._connection)
+            )
+            self._cache.put(key, cached)
+        return list(cached)
+
+    def columns(self, schema: str, table: str) -> tuple[str, ...]:
+        key = (*self._scope, "columns", schema, table)
+        cached = self._cache.get(key)
+        if cached is None:
+            try:
+                reflected = inspect(self._connection).get_columns(table, schema=schema or None)
+            except NoSuchTableError as exc:
+                raise reject(Reason.TABLE_GONE) from exc
+            cached = tuple(str(column["name"]) for column in reflected)
+            self._cache.put(key, cached)
+        return cached
+
+
+def _match(tables: Sequence[tuple[str, str]], schema: str | None, name: str) -> tuple[str, str]:
+    folded_name = name.casefold()
+    folded_schema = None if schema is None else schema.casefold()
+    matches = [
+        (table_schema, table_name)
+        for table_schema, table_name in tables
+        if table_name.casefold() == folded_name
+        and (folded_schema is None or table_schema.casefold() == folded_schema)
+    ]
+    if len(matches) != 1:
+        raise reject(Reason.TABLE_NOT_FOUND)
+    return matches[0]
+
+
+def resolve_tables(
+    query: exp.Select, catalog: TableCatalog, dialect: SqlDialect = SQLSERVER
+) -> tuple[TableSchema, ...]:
+    """Resolve every raw-AST table against reflection and rewrite it to its reflected name.
+
+    Raises QUERY_REJECTED for missing, ambiguous, system, and cross-database objects.
+    """
+    available = list(catalog.list_tables())
+    resolved: list[TableSchema] = []
+    for table in query.find_all(exp.Table):
+        if table.args.get("catalog") or not isinstance(table.this, exp.Identifier):
+            raise reject(Reason.TABLE_SCOPE)
+        if not dialect.has_schema and table.db:
+            raise reject(Reason.TABLE_SCOPE)  # db.table is a cross-database reference
+        schema = table.db or None
+        schema_name, table_name = _match(available, schema, table.name)
+        resolved.append(
+            TableSchema(schema_name, table_name, tuple(catalog.columns(schema_name, table_name)))
+        )
+        table.set("this", exp.to_identifier(table_name, quoted=True))
+        if dialect.has_schema:
+            table.set("db", exp.to_identifier(schema_name, quoted=True))
+    return tuple(resolved)