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

sql_safe_mcp/errors.py

@@ -0,0 +1,49 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from enum import StrEnum
+from uuid import uuid4
+
+
+class ErrorCode(StrEnum):
+    INVALID_ARGUMENT = "INVALID_ARGUMENT"
+    UNKNOWN_SERVER = "UNKNOWN_SERVER"
+    NOT_FOUND = "NOT_FOUND"
+    AMBIGUOUS_OBJECT = "AMBIGUOUS_OBJECT"
+    ACCESS_LEVEL_DENIED = "ACCESS_LEVEL_DENIED"
+    CONFIG_ERROR = "CONFIG_ERROR"
+    CONNECTION_FAILED = "CONNECTION_FAILED"
+    ACCESS_DENIED = "ACCESS_DENIED"
+    TIMEOUT = "TIMEOUT"
+    METADATA_UNAVAILABLE = "METADATA_UNAVAILABLE"
+    DEFINITION_TOO_LARGE = "DEFINITION_TOO_LARGE"
+    QUERY_REJECTED = "QUERY_REJECTED"
+    INVALID_PII_TOKEN = "INVALID_PII_TOKEN"
+    RESULT_LIMIT_EXCEEDED = "RESULT_LIMIT_EXCEEDED"
+    DATABASE_ERROR = "DATABASE_ERROR"
+
+
+@dataclass(slots=True)
+class DomainError(Exception):
+    code: ErrorCode
+    public_message: str
+    hint: str | None = None
+    retryable: bool = False
+    correlation_id: str | None = None
+
+    def __str__(self) -> str:
+        parts = [f"[{self.code}] {self.public_message}"]
+        if self.hint:
+            parts.append(f"Hint: {self.hint}")
+        if self.correlation_id:
+            parts.append(f"Reference: {self.correlation_id}")
+        return " ".join(parts)
+
+    @classmethod
+    def unexpected(cls) -> DomainError:
+        return cls(
+            ErrorCode.DATABASE_ERROR,
+            "The database operation failed unexpectedly.",
+            retryable=False,
+            correlation_id=uuid4().hex,
+        )

sql_safe_mcp/mcp_server.py

@@ -0,0 +1,140 @@
+from __future__ import annotations
+
+from collections.abc import AsyncIterator, Awaitable, Callable
+from contextlib import asynccontextmanager
+from dataclasses import dataclass
+
+from mcp.server.mcpserver import Context, MCPServer
+from mcp.server.mcpserver.exceptions import ToolError
+from mcp_types import ToolAnnotations
+
+from sql_safe_mcp.config import AppConfig
+from sql_safe_mcp.db.registry import EngineRegistry
+from sql_safe_mcp.errors import DomainError
+from sql_safe_mcp.models import (
+    DatabaseList,
+    ServerList,
+    SqlResult,
+    StoredProcedureDefinition,
+    StoredProcedureList,
+    TableDefinition,
+    TableList,
+)
+from sql_safe_mcp.service import DatabaseService
+
+READ_ONLY = ToolAnnotations(read_only_hint=True, open_world_hint=False)
+
+
+@dataclass(slots=True)
+class AppContext:
+    service: DatabaseService
+    registry: EngineRegistry
+
+
+async def _domain_call[T](operation: Callable[[], Awaitable[T]]) -> T:
+    try:
+        return await operation()
+    except DomainError as exc:
+        raise ToolError(str(exc)) from exc
+
+
+def create_server(config: AppConfig) -> MCPServer[AppContext]:
+    @asynccontextmanager
+    async def lifespan(_server: MCPServer[AppContext]) -> AsyncIterator[AppContext]:
+        registry = EngineRegistry(config)
+        try:
+            yield AppContext(service=DatabaseService(config, registry), registry=registry)
+        finally:
+            registry.dispose()
+
+    server: MCPServer[AppContext] = MCPServer(
+        "sql-safe-mcp",
+        description="Minimal, read-only, PII-safe SQL database navigation.",
+        version="1.2.0",
+        lifespan=lifespan,
+    )
+
+    @server.tool(annotations=READ_ONLY)
+    async def list_servers(ctx: Context[AppContext]) -> ServerList:
+        """List explicitly configured database server aliases."""
+        return ctx.request_context.lifespan_context.service.list_servers()
+
+    @server.tool(annotations=READ_ONLY)
+    async def list_databases(
+        server: str,
+        ctx: Context[AppContext],
+        name_contains: str | None = None,
+    ) -> DatabaseList:
+        """List databases visible to the configured credentials."""
+        service = ctx.request_context.lifespan_context.service
+        return await _domain_call(lambda: service.list_databases(server, name_contains))
+
+    @server.tool(annotations=READ_ONLY)
+    async def list_tables(
+        server: str,
+        database: str,
+        ctx: Context[AppContext],
+        schema: str | None = None,
+        name_contains: str | None = None,
+    ) -> TableList:
+        """List base tables, optionally filtering by schema and a name substring."""
+        service = ctx.request_context.lifespan_context.service
+        return await _domain_call(
+            lambda: service.list_tables(server, database, schema, name_contains)
+        )
+
+    @server.tool(annotations=READ_ONLY)
+    async def get_table_definition(
+        server: str,
+        database: str,
+        table: str,
+        ctx: Context[AppContext],
+        schema: str | None = None,
+    ) -> TableDefinition:
+        """Get columns, keys, constraints, and indexes for one table."""
+        service = ctx.request_context.lifespan_context.service
+        return await _domain_call(
+            lambda: service.get_table_definition(server, database, table, schema)
+        )
+
+    @server.tool(annotations=READ_ONLY)
+    async def list_stored_procedures(
+        server: str,
+        database: str,
+        ctx: Context[AppContext],
+        schema: str | None = None,
+        name_contains: str | None = None,
+    ) -> StoredProcedureList:
+        """List stored procedures without expanding their definitions."""
+        service = ctx.request_context.lifespan_context.service
+        return await _domain_call(
+            lambda: service.list_stored_procedures(server, database, schema, name_contains)
+        )
+
+    @server.tool(annotations=READ_ONLY)
+    async def get_stored_procedure(
+        server: str,
+        database: str,
+        name: str,
+        ctx: Context[AppContext],
+        schema: str | None = None,
+    ) -> StoredProcedureDefinition:
+        """Get the original definition of one stored procedure when visible."""
+        service = ctx.request_context.lifespan_context.service
+        return await _domain_call(
+            lambda: service.get_stored_procedure(server, database, name, schema)
+        )
+
+    @server.tool(annotations=READ_ONLY)
+    async def execute_sql(
+        server: str,
+        database: str,
+        sql: str,
+        ctx: Context[AppContext],
+        max_rows: int | None = None,
+    ) -> SqlResult:
+        """Run one restricted SELECT on a pii_safe server; protected columns return tokens."""
+        service = ctx.request_context.lifespan_context.service
+        return await _domain_call(lambda: service.execute_sql(server, database, sql, max_rows))
+
+    return server

sql_safe_mcp/models.py

@@ -0,0 +1,124 @@
+from __future__ import annotations
+
+from typing import Any, Literal
+
+from pydantic import BaseModel, ConfigDict, Field
+
+
+class OutputModel(BaseModel):
+    model_config = ConfigDict(extra="forbid", populate_by_name=True)
+
+
+class ServerSummary(OutputModel):
+    name: str
+    engine: Literal["sqlserver", "mysql", "mariadb"]
+    access_level: Literal["metadata", "pii_safe"]
+
+
+class ServerList(OutputModel):
+    servers: list[ServerSummary]
+
+
+class DatabaseSummary(OutputModel):
+    name: str
+
+
+class DatabaseList(OutputModel):
+    databases: list[DatabaseSummary]
+
+
+class TableSummary(OutputModel):
+    schema_: str | None = Field(alias="schema")
+    name: str
+
+
+class TableList(OutputModel):
+    tables: list[TableSummary]
+
+
+class ColumnDefinition(OutputModel):
+    name: str
+    native_type: str
+    nullable: bool
+    length: int | None = None
+    precision: int | None = None
+    scale: int | None = None
+    autoincrement: bool | str | None = None
+    identity: dict[str, Any] | None = None
+    computed: dict[str, Any] | None = None
+    default: str | None = None
+
+
+class PrimaryKeyDefinition(OutputModel):
+    name: str | None = None
+    columns: list[str]
+
+
+class ForeignKeyDefinition(OutputModel):
+    name: str | None = None
+    columns: list[str]
+    referred_schema: str | None = None
+    referred_table: str
+    referred_columns: list[str]
+    options: dict[str, Any] = Field(default_factory=dict)
+
+
+class UniqueConstraintDefinition(OutputModel):
+    name: str | None = None
+    columns: list[str]
+
+
+class IndexDefinition(OutputModel):
+    name: str | None = None
+    columns: list[str | None]
+    unique: bool = False
+    expressions: list[str] = Field(default_factory=list)
+
+
+class TableDefinition(OutputModel):
+    schema_: str | None = Field(alias="schema")
+    name: str
+    columns: list[ColumnDefinition]
+    primary_key: PrimaryKeyDefinition
+    foreign_keys: list[ForeignKeyDefinition]
+    unique_constraints: list[UniqueConstraintDefinition]
+    indexes: list[IndexDefinition]
+
+
+class StoredProcedureSummary(OutputModel):
+    schema_: str | None = Field(alias="schema")
+    name: str
+
+
+class StoredProcedureList(OutputModel):
+    stored_procedures: list[StoredProcedureSummary]
+
+
+class StoredProcedureDefinition(OutputModel):
+    schema_: str | None = Field(alias="schema")
+    name: str
+    definition: str | None
+    definition_available: bool
+
+
+ResultEncoding = Literal["json", "token", "decimal", "date", "time", "datetime", "uuid", "base64"]
+
+
+class ColumnSource(OutputModel):
+    schema_: str | None = Field(alias="schema")
+    table: str
+    column: str
+
+
+class SqlColumn(OutputModel):
+    name: str
+    source: ColumnSource | None
+    protected: bool
+    encoding: ResultEncoding
+
+
+class SqlResult(OutputModel):
+    columns: list[SqlColumn]
+    rows: list[list[Any]]
+    row_count: int
+    truncated: bool

sql_safe_mcp/security/__init__.py

@@ -0,0 +1 @@
+"""SQL validation and PII protection pipeline for execute_sql."""

sql_safe_mcp/security/dialect.py

@@ -0,0 +1,29 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+
+@dataclass(frozen=True, slots=True)
+class SqlDialect:
+    """Everything the shared validation pipeline needs to know about one SQL dialect.
+
+    Policy, tokens, and ValidatedQuery are dialect-independent; only these facts differ. The
+    dialect always comes from the trusted server configuration, never from the caller.
+    """
+
+    name: str  # SQLGlot dialect used to parse and to generate SQL
+    has_schema: bool  # False when a database is the catalog and objects have no schema level
+    bind_marker: str  # DBAPI positional marker written into generated SQL
+    escape_percent: bool  # the driver applies %-formatting, so a literal % must be doubled
+
+
+SQLSERVER = SqlDialect(name="tsql", has_schema=True, bind_marker="?", escape_percent=False)
+MYSQL = SqlDialect(name="mysql", has_schema=False, bind_marker="%s", escape_percent=True)
+
+
+def dialect_for(engine: str) -> SqlDialect:
+    if engine == "sqlserver":
+        return SQLSERVER
+    if engine in ("mysql", "mariadb"):
+        return MYSQL
+    raise ValueError(f"unsupported engine {engine!r}")

sql_safe_mcp/security/executor.py

@@ -0,0 +1,129 @@
+from __future__ import annotations
+
+import base64
+import math
+from collections.abc import Sequence
+from dataclasses import dataclass
+from datetime import date, datetime, time, timedelta
+from decimal import Decimal
+from typing import Any, Protocol
+from uuid import UUID
+
+from sql_safe_mcp.errors import DomainError, ErrorCode
+from sql_safe_mcp.models import ResultEncoding
+from sql_safe_mcp.security.lineage import SourceColumn
+from sql_safe_mcp.security.tokens import TokenCodec
+from sql_safe_mcp.security.validated_query import ValidatedQuery
+
+
+class _Result(Protocol):
+    def keys(self) -> Any: ...
+
+    def fetchmany(self, size: int) -> Sequence[Any]: ...
+
+    def close(self) -> None: ...
+
+
+class SqlConnection(Protocol):
+    """The part of a SQLAlchemy Core connection the executor uses."""
+
+    def exec_driver_sql(self, statement: str, parameters: tuple[Any, ...]) -> _Result: ...
+
+
+@dataclass(frozen=True, slots=True)
+class ResultColumn:
+    name: str
+    source: SourceColumn | None
+    protected: bool
+    encoding: ResultEncoding
+
+
+@dataclass(frozen=True, slots=True)
+class ExecutionResult:
+    columns: tuple[ResultColumn, ...]
+    rows: tuple[tuple[Any, ...], ...]
+    truncated: bool
+
+
+def _unsupported_result() -> DomainError:
+    return DomainError(
+        ErrorCode.DATABASE_ERROR,
+        "The result contains a value type that is not supported.",
+        "Select different columns.",
+    )
+
+
+def _encode(value: Any) -> tuple[Any, ResultEncoding]:
+    """Return the JSON-safe form and encoding name of one non-null, unprotected value."""
+    kind = type(value)
+    if kind in (str, bool, int):
+        return value, "json"
+    if kind is float:
+        if not math.isfinite(value):
+            raise _unsupported_result()
+        return value, "json"
+    if kind is Decimal:
+        if not value.is_finite():
+            raise _unsupported_result()
+        return str(value), "decimal"
+    if kind is datetime:
+        return value.isoformat(), "datetime"
+    if kind is date:
+        return value.isoformat(), "date"
+    if kind is time:
+        return value.isoformat(), "time"
+    if kind is timedelta:  # MySQL TIME; only a time of day fits the `time` encoding
+        if not timedelta(0) <= value < timedelta(days=1):
+            raise _unsupported_result()
+        return (datetime.min + value).time().isoformat(), "time"
+    if kind is UUID:
+        return str(value), "uuid"
+    if kind in (bytes, bytearray, memoryview):
+        return base64.b64encode(bytes(value)).decode("ascii"), "base64"
+    raise _unsupported_result()
+
+
+def execute_validated(
+    connection: SqlConnection, query: ValidatedQuery, codec: TokenCodec
+) -> ExecutionResult:
+    """Run a ValidatedQuery: generated SQL and binds only, at most max_rows + 1 rows buffered."""
+    if not isinstance(query, ValidatedQuery):
+        raise TypeError("the executor accepts only ValidatedQuery")
+    if codec.alias != query.alias:
+        raise TypeError("the token codec belongs to a different server alias")
+
+    result = connection.exec_driver_sql(query.sql, query.parameters)
+    try:
+        if len(result.keys()) != len(query.outputs):
+            raise _unsupported_result()
+        fetched = result.fetchmany(query.max_rows + 1)
+        truncated = len(fetched) > query.max_rows
+        encodings: list[ResultEncoding | None] = [None] * len(query.outputs)
+        rows: list[tuple[Any, ...]] = []
+        for raw in fetched[: query.max_rows]:
+            cells: list[Any] = []
+            for index, (plan, value) in enumerate(zip(query.outputs, tuple(raw), strict=True)):
+                if value is None:
+                    cells.append(None)
+                elif plan.protected:
+                    cells.append(codec.encrypt(value))
+                else:
+                    encoded, encoding = _encode(value)
+                    if encodings[index] not in (None, encoding):
+                        raise _unsupported_result()
+                    encodings[index] = encoding
+                    cells.append(encoded)
+            rows.append(tuple(cells))
+    finally:
+        result.close()
+
+    columns = tuple(
+        ResultColumn(
+            plan.label,
+            plan.source,
+            plan.protected,
+            "token" if plan.protected else (encodings[index] or "json"),
+        )
+        for index, plan in enumerate(query.outputs)
+    )
+    return ExecutionResult(columns, tuple(rows), truncated)

sql_safe_mcp/security/lineage.py

@@ -0,0 +1,224 @@
+from __future__ import annotations
+
+from collections.abc import Sequence
+from dataclasses import dataclass
+from typing import Literal, cast
+
+from sqlglot import exp
+from sqlglot.errors import OptimizeError, SqlglotError
+from sqlglot.optimizer.qualify import qualify
+
+from sql_safe_mcp.security.dialect import SqlDialect
+from sql_safe_mcp.security.parser import (
+    ParserLimits,
+    check_limits,
+    reject,
+    validate_allowlist,
+)
+from sql_safe_mcp.security.reasons import Reason
+from sql_safe_mcp.security.schema import TableSchema
+
+
+@dataclass(frozen=True, slots=True)
+class SourceColumn:
+    schema: str
+    table: str
+    column: str
+
+
+@dataclass(frozen=True, slots=True)
+class OutputColumn:
+    label: str
+    kind: Literal["column", "count", "literal"]
+    source: SourceColumn | None
+    binding: str | None = None  # table alias/name the source column was read through
+
+
+@dataclass(frozen=True, slots=True)
+class ColumnRef:
+    node: exp.Column
+    source: SourceColumn
+
+
+@dataclass(frozen=True, slots=True)
+class AnalyzedQuery:
+    query: exp.Select
+    outputs: tuple[OutputColumn, ...]
+    references: tuple[ColumnRef, ...]
+
+
+@dataclass(frozen=True, slots=True)
+class _Binding:
+    name: str
+    table: TableSchema
+
+
+def _identifier(name: str) -> exp.Identifier:
+    return exp.to_identifier(name, quoted=True)
+
+
+def _collect(query: exp.Select, schemas: Sequence[TableSchema]) -> list[_Binding]:
+    nodes: list[exp.Table] = []
+    from_ = query.args.get("from_")
+    if from_ is not None:
+        nodes.append(from_.this)
+    nodes += [join.this for join in query.args.get("joins") or []]
+    if len(nodes) != len(schemas):
+        raise reject(Reason.TABLE_RESOLUTION_MISMATCH)
+    bindings: list[_Binding] = []
+    seen: set[str] = set()
+    for node, table in zip(nodes, schemas, strict=True):
+        name = node.alias or table.name
+        if name.casefold() in seen:
+            raise reject(Reason.DUPLICATE_BINDING)
+        seen.add(name.casefold())
+        bindings.append(_Binding(name, table))
+    return bindings
+
+
+def _canonical_column(binding: _Binding, name: str) -> str | None:
+    matches = [column for column in binding.table.columns if column.casefold() == name.casefold()]
+    return matches[0] if len(matches) == 1 else None
+
+
+def _resolve(column: exp.Column, bindings: Sequence[_Binding]) -> tuple[_Binding, str]:
+    name = column.name
+    qualifier = column.table
+    if qualifier:
+        matching = [b for b in bindings if b.name.casefold() == qualifier.casefold()]
+        if len(matching) != 1:
+            raise reject(Reason.QUALIFIER_UNKNOWN)
+        canonical = _canonical_column(matching[0], name)
+        if canonical is None:
+            raise reject(Reason.COLUMN_NOT_FOUND)
+        return matching[0], canonical
+    candidates = [
+        (binding, canonical)
+        for binding in bindings
+        if (canonical := _canonical_column(binding, name)) is not None
+    ]
+    if len(candidates) != 1:
+        raise reject(Reason.COLUMN_AMBIGUOUS)
+    return candidates[0]
+
+
+def _rewrite(column: exp.Column, binding: _Binding, canonical: str) -> SourceColumn:
+    column.set("this", _identifier(canonical))
+    column.set("table", _identifier(binding.name))
+    return SourceColumn(binding.table.schema, binding.table.name, canonical)
+
+
+def _expand_stars(query: exp.Select, bindings: Sequence[_Binding]) -> None:
+    expanded: list[exp.Expression] = []
+    for projection in query.expressions:
+        if isinstance(projection, exp.Star):
+            chosen = list(bindings)
+        elif isinstance(projection, exp.Column) and isinstance(projection.this, exp.Star):
+            chosen = [b for b in bindings if b.name.casefold() == projection.table.casefold()]
+            if len(chosen) != 1:
+                raise reject(Reason.QUALIFIER_UNKNOWN)
+        else:
+            expanded.append(projection)
+            continue
+        expanded.extend(
+            exp.Column(this=_identifier(name), table=_identifier(binding.name))
+            for binding in chosen
+            for name in binding.table.columns
+        )
+    query.set("expressions", expanded)
+
+
+def _analyze_projections(
+    query: exp.Select, bindings: Sequence[_Binding], refs: list[ColumnRef]
+) -> list[OutputColumn]:
+    outputs: list[OutputColumn] = []
+    for index, projection in enumerate(query.expressions, start=1):
+        value = projection.this if isinstance(projection, exp.Alias) else projection
+        alias = projection.alias if isinstance(projection, exp.Alias) else ""
+        if isinstance(value, exp.Column):
+            binding, canonical = _resolve(value, bindings)
+            source = _rewrite(value, binding, canonical)
+            refs.append(ColumnRef(value, source))
+            outputs.append(OutputColumn(alias or canonical, "column", source, binding.name))
+        elif isinstance(value, exp.Count):
+            outputs.append(OutputColumn(alias or f"column_{index}", "count", None))
+        else:
+            outputs.append(OutputColumn(alias or f"column_{index}", "literal", None))
+    return outputs
+
+
+def _resolve_order_column(
+    column: exp.Column,
+    bindings: Sequence[_Binding],
+    outputs: Sequence[OutputColumn],
+) -> SourceColumn:
+    if not column.table:
+        matches = [
+            output for output in outputs if output.label.casefold() == column.name.casefold()
+        ]
+        if matches:
+            # SQL Server resolves an unqualified ORDER BY name against output labels first.
+            if len(matches) != 1 or matches[0].source is None:
+                raise reject(Reason.ORDER_OUTPUT_AMBIGUOUS)
+            source = matches[0].source
+            binding = next(b for b in bindings if b.name == matches[0].binding)
+            return _rewrite(column, binding, source.column)
+    binding, canonical = _resolve(column, bindings)
+    return _rewrite(column, binding, canonical)
+
+
+def _cross_check(query: exp.Select, schemas: Sequence[TableSchema], dialect: SqlDialect) -> None:
+    catalog: dict[str, dict] = {}
+    for table in schemas:
+        columns = dict.fromkeys(table.columns, "varchar")
+        if dialect.has_schema:
+            catalog.setdefault(table.schema, {})[table.name] = columns
+        else:
+            catalog[table.name] = columns
+    try:
+        qualify(
+            query.copy(),
+            schema=cast(dict[str, object], catalog),
+            dialect=f"{dialect.name}, normalization_strategy = case_sensitive",
+            validate_qualify_columns=True,
+            expand_stars=False,
+        )
+    except (OptimizeError, SqlglotError) as exc:
+        raise reject(Reason.QUALIFY_FAILED) from exc
+
+
+def analyze_query(
+    query: exp.Select,
+    schemas: Sequence[TableSchema],
+    limits: ParserLimits,
+    dialect: SqlDialect,
+) -> AnalyzedQuery:
+    """Expand stars, resolve every column to a reflected source, and rebuild canonical names.
+
+    The input query is not modified. Unknown or ambiguous lineage raises QUERY_REJECTED.
+    """
+    working = query.copy()
+    bindings = _collect(working, schemas)
+    _expand_stars(working, bindings)
+    check_limits(working, limits)
+
+    refs: list[ColumnRef] = []
+    outputs = _analyze_projections(working, bindings, refs)
+
+    for join in working.args.get("joins") or []:
+        for column in join.args["on"].find_all(exp.Column):
+            binding, canonical = _resolve(column, bindings)
+            refs.append(ColumnRef(column, _rewrite(column, binding, canonical)))
+    where = working.args.get("where")
+    if where is not None:
+        for column in where.find_all(exp.Column):
+            binding, canonical = _resolve(column, bindings)
+            refs.append(ColumnRef(column, _rewrite(column, binding, canonical)))
+    order = working.args.get("order")
+    if order is not None:
+        for column in order.find_all(exp.Column):
+            refs.append(ColumnRef(column, _resolve_order_column(column, bindings, outputs)))
+
+    validate_allowlist(working)
+    _cross_check(working, schemas, dialect)
+    return AnalyzedQuery(working, tuple(outputs), tuple(refs))