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

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
@@ -0,0 +1,3 @@
1
+ """Minimal, PII-safe MCP server for relational databases."""
2
+
3
+ __version__ = "1.2.0"
@@ -0,0 +1,43 @@
1
+ from __future__ import annotations
2
+
3
+ import argparse
4
+ import os
5
+ import sys
6
+ from pathlib import Path
7
+
8
+ from sql_safe_mcp.config import load_config
9
+ from sql_safe_mcp.errors import DomainError
10
+ from sql_safe_mcp.mcp_server import create_server
11
+
12
+
13
+ def _parser() -> argparse.ArgumentParser:
14
+ parser = argparse.ArgumentParser(description="Run the sql-safe-mcp stdio server.")
15
+ parser.add_argument(
16
+ "--config",
17
+ type=Path,
18
+ default=os.environ.get("SQL_SAFE_MCP_CONFIG", "sql-safe-mcp.yaml"),
19
+ help="YAML configuration path (default: SQL_SAFE_MCP_CONFIG or sql-safe-mcp.yaml)",
20
+ )
21
+ parser.add_argument(
22
+ "--check-config",
23
+ action="store_true",
24
+ help="Validate configuration and secrets without connecting to databases.",
25
+ )
26
+ return parser
27
+
28
+
29
+ def main() -> None:
30
+ args = _parser().parse_args()
31
+ try:
32
+ config = load_config(args.config)
33
+ except DomainError as exc:
34
+ print(str(exc), file=sys.stderr)
35
+ raise SystemExit(2) from None
36
+ if args.check_config:
37
+ print(f"Configuration valid: {len(config.servers)} server alias(es).")
38
+ return
39
+ create_server(config).run("stdio")
40
+
41
+
42
+ if __name__ == "__main__":
43
+ main()
sql_safe_mcp/config.py ADDED
@@ -0,0 +1,201 @@
1
+ from __future__ import annotations
2
+
3
+ import base64
4
+ import binascii
5
+ import os
6
+ import re
7
+ from collections.abc import Mapping
8
+ from pathlib import Path
9
+ from typing import Literal
10
+ from urllib.parse import quote
11
+
12
+ import yaml
13
+ from pydantic import BaseModel, ConfigDict, Field, SecretStr, ValidationError, model_validator
14
+ from sqlalchemy.engine import make_url
15
+
16
+ from sql_safe_mcp.errors import DomainError, ErrorCode
17
+
18
+ _ENV_PATTERN = re.compile(r"\$\{([A-Z_][A-Z0-9_]*)\}")
19
+ _FULL_ENV_PATTERN = re.compile(r"^\$\{([A-Z_][A-Z0-9_]*)\}$")
20
+
21
+
22
+ class RuntimeConfig(BaseModel):
23
+ model_config = ConfigDict(extra="forbid")
24
+
25
+ max_concurrent_db_operations: int = Field(default=8, ge=1, le=128)
26
+ engine_cache_size: int = Field(default=32, ge=1, le=1024)
27
+ pool_size: int = Field(default=2, ge=1, le=32)
28
+ max_overflow: int = Field(default=2, ge=0, le=64)
29
+ pool_timeout_seconds: int = Field(default=10, ge=1, le=300)
30
+ statement_timeout_seconds: int = Field(default=30, ge=1, le=3600)
31
+ max_definition_chars: int = Field(default=262_144, ge=1024, le=10_000_000)
32
+ default_max_rows: int = Field(default=200, ge=1, le=100_000)
33
+ hard_max_rows: int = Field(default=1000, ge=1, le=100_000)
34
+ max_sql_chars: int = Field(default=65_536, ge=256, le=1_000_000)
35
+ max_ast_nodes: int = Field(default=2000, ge=10, le=100_000)
36
+ max_joins: int = Field(default=8, ge=0, le=100)
37
+ max_in_list_items: int = Field(default=500, ge=1, le=100_000)
38
+
39
+ @model_validator(mode="after")
40
+ def validate_row_limits(self) -> RuntimeConfig:
41
+ if self.default_max_rows > self.hard_max_rows:
42
+ raise ValueError("default_max_rows cannot exceed hard_max_rows")
43
+ return self
44
+
45
+
46
+ class PiiRule(BaseModel):
47
+ model_config = ConfigDict(extra="forbid")
48
+
49
+ database: str = Field(min_length=1)
50
+ schema_: str | None = Field(default=None, alias="schema")
51
+ table: str = Field(min_length=1)
52
+ columns: list[str] = Field(min_length=1)
53
+
54
+ @model_validator(mode="after")
55
+ def validate_wildcards(self) -> PiiRule:
56
+ if "*" in self.table or (self.schema_ and "*" in self.schema_):
57
+ raise ValueError("wildcards are allowed only in pii rule database")
58
+ if self.database != "*" and "*" in self.database:
59
+ raise ValueError("database must be an exact name or '*'")
60
+ folded = [column.casefold() for column in self.columns]
61
+ if len(folded) != len(set(folded)):
62
+ raise ValueError("pii rule columns must be unique")
63
+ return self
64
+
65
+
66
+ class PiiConfig(BaseModel):
67
+ model_config = ConfigDict(extra="forbid")
68
+
69
+ rules: list[PiiRule] = Field(min_length=1)
70
+
71
+
72
+ _DRIVERNAMES = {
73
+ "sqlserver": "mssql+pyodbc",
74
+ "mysql": "mysql+pymysql",
75
+ "mariadb": "mysql+pymysql",
76
+ }
77
+
78
+
79
+ class ServerConfig(BaseModel):
80
+ model_config = ConfigDict(extra="forbid")
81
+
82
+ engine: Literal["sqlserver", "mysql", "mariadb"]
83
+ access_level: Literal["metadata", "pii_safe"] = "metadata"
84
+ connection_url: SecretStr
85
+ pii_key_env: str | None = Field(default=None, pattern=r"^[A-Z_][A-Z0-9_]*$")
86
+ pii: PiiConfig | None = None
87
+ pii_key: SecretStr | None = Field(default=None, exclude=True)
88
+
89
+ @model_validator(mode="after")
90
+ def validate_security_shape(self) -> ServerConfig:
91
+ if self.access_level == "pii_safe":
92
+ if not self.pii_key_env or self.pii is None:
93
+ raise ValueError("pii_safe servers require pii_key_env and pii rules")
94
+ if self.engine != "sqlserver" and any(rule.schema_ for rule in self.pii.rules):
95
+ raise ValueError(f"pii rules for engine {self.engine!r} cannot set schema")
96
+ elif self.pii_key_env is not None or self.pii is not None:
97
+ raise ValueError("metadata servers cannot configure pii_key_env or pii rules")
98
+
99
+ driver = make_url(self.connection_url.get_secret_value()).drivername
100
+ expected = _DRIVERNAMES[self.engine]
101
+ if driver != expected:
102
+ raise ValueError(f"engine {self.engine!r} requires SQLAlchemy dialect {expected!r}")
103
+ return self
104
+
105
+ def key_bytes(self) -> bytes | None:
106
+ if self.pii_key is None:
107
+ return None
108
+ return base64.b64decode(self.pii_key.get_secret_value(), validate=True)
109
+
110
+
111
+ class AppConfig(BaseModel):
112
+ model_config = ConfigDict(extra="forbid")
113
+
114
+ version: Literal[1]
115
+ runtime: RuntimeConfig = Field(default_factory=RuntimeConfig)
116
+ servers: dict[str, ServerConfig] = Field(min_length=1)
117
+
118
+ @model_validator(mode="after")
119
+ def validate_aliases_and_keys(self) -> AppConfig:
120
+ seen_keys: dict[bytes, str] = {}
121
+ for alias, server in self.servers.items():
122
+ if not re.fullmatch(r"[A-Za-z0-9][A-Za-z0-9._-]*", alias):
123
+ raise ValueError(f"invalid server alias {alias!r}")
124
+ key = server.key_bytes()
125
+ if key is None:
126
+ continue
127
+ if len(key) != 32:
128
+ raise ValueError(f"PII key for server {alias!r} must decode to 32 bytes")
129
+ if previous := seen_keys.get(key):
130
+ raise ValueError(
131
+ f"PII keys must be unique per server alias; {alias!r} duplicates {previous!r}"
132
+ )
133
+ seen_keys[key] = alias
134
+ return self
135
+
136
+
137
+ def _expand_connection_url(template: str, environ: Mapping[str, str]) -> str:
138
+ full_match = _FULL_ENV_PATTERN.fullmatch(template)
139
+ if full_match:
140
+ name = full_match.group(1)
141
+ if name not in environ:
142
+ raise ValueError(f"missing environment variable {name}")
143
+ return environ[name]
144
+
145
+ def replace(match: re.Match[str]) -> str:
146
+ name = match.group(1)
147
+ if name not in environ:
148
+ raise ValueError(f"missing environment variable {name}")
149
+ return quote(environ[name], safe="")
150
+
151
+ expanded = _ENV_PATTERN.sub(replace, template)
152
+ if "${" in expanded:
153
+ raise ValueError("invalid environment placeholder in connection_url")
154
+ return expanded
155
+
156
+
157
+ def _validation_summary(error: ValidationError) -> str:
158
+ issues: list[str] = []
159
+ for issue in error.errors(include_url=False, include_context=False, include_input=False):
160
+ location = ".".join(str(part) for part in issue["loc"])
161
+ issues.append(f"{location}: {issue['msg']}" if location else str(issue["msg"]))
162
+ return "; ".join(issues)
163
+
164
+
165
+ def load_config(path: str | Path, environ: Mapping[str, str] | None = None) -> AppConfig:
166
+ env = os.environ if environ is None else environ
167
+ try:
168
+ raw = yaml.safe_load(Path(path).read_text(encoding="utf-8"))
169
+ if not isinstance(raw, dict):
170
+ raise ValueError("configuration root must be a mapping")
171
+ servers = raw.get("servers")
172
+ if not isinstance(servers, dict):
173
+ raise ValueError("servers must be a mapping")
174
+ for alias, value in servers.items():
175
+ if not isinstance(value, dict):
176
+ raise ValueError(f"server {alias!r} must be a mapping")
177
+ url = value.get("connection_url")
178
+ if not isinstance(url, str):
179
+ raise ValueError(f"server {alias!r} requires connection_url")
180
+ value["connection_url"] = _expand_connection_url(url, env)
181
+ key_env = value.get("pii_key_env")
182
+ if key_env is not None:
183
+ if not isinstance(key_env, str) or key_env not in env:
184
+ raise ValueError(f"missing PII key environment variable {key_env!r}")
185
+ try:
186
+ decoded = base64.b64decode(env[key_env], validate=True)
187
+ except (binascii.Error, ValueError) as exc:
188
+ raise ValueError(f"PII key for server {alias!r} is not valid base64") from exc
189
+ if len(decoded) != 32:
190
+ raise ValueError(f"PII key for server {alias!r} must decode to 32 bytes")
191
+ value["pii_key"] = env[key_env]
192
+ return AppConfig.model_validate(raw)
193
+ except ValidationError as exc:
194
+ raise DomainError(
195
+ ErrorCode.CONFIG_ERROR,
196
+ f"Invalid configuration: {_validation_summary(exc)}",
197
+ ) from exc
198
+ except (OSError, yaml.YAMLError) as exc:
199
+ raise DomainError(ErrorCode.CONFIG_ERROR, "Invalid configuration file.") from exc
200
+ except ValueError as exc:
201
+ raise DomainError(ErrorCode.CONFIG_ERROR, f"Invalid configuration: {exc}") from exc
@@ -0,0 +1 @@
1
+ """Database integration layer."""
@@ -0,0 +1,35 @@
1
+ from __future__ import annotations
2
+
3
+ from typing import Protocol
4
+
5
+ from sqlalchemy import Connection
6
+
7
+ from sql_safe_mcp.models import StoredProcedureDefinition, StoredProcedureSummary
8
+
9
+
10
+ class DatabaseExtras(Protocol):
11
+ def list_databases(self, connection: Connection) -> list[str]: ...
12
+
13
+ def list_stored_procedures(
14
+ self, connection: Connection, database: str
15
+ ) -> list[StoredProcedureSummary]: ...
16
+
17
+ def get_stored_procedure(
18
+ self,
19
+ connection: Connection,
20
+ database: str,
21
+ schema: str | None,
22
+ name: str,
23
+ ) -> StoredProcedureDefinition | None: ...
24
+
25
+
26
+ def extras_for(engine: str) -> DatabaseExtras:
27
+ if engine == "sqlserver":
28
+ from sql_safe_mcp.db.sqlserver import SqlServerExtras
29
+
30
+ return SqlServerExtras()
31
+ if engine in ("mysql", "mariadb"):
32
+ from sql_safe_mcp.db.mysql import MySqlExtras
33
+
34
+ return MySqlExtras()
35
+ raise ValueError(f"unsupported engine {engine!r}")
@@ -0,0 +1,63 @@
1
+ from __future__ import annotations
2
+
3
+ from sqlalchemy import Connection, text
4
+
5
+ from sql_safe_mcp.db.extras import DatabaseExtras
6
+ from sql_safe_mcp.models import StoredProcedureDefinition, StoredProcedureSummary
7
+
8
+ SYSTEM_DATABASES = ("information_schema", "mysql", "performance_schema", "sys")
9
+
10
+
11
+ class MySqlExtras(DatabaseExtras):
12
+ """MySQL and MariaDB: a database is the catalog, there is no separate schema."""
13
+
14
+ def list_databases(self, connection: Connection) -> list[str]:
15
+ rows = connection.execute(
16
+ text(
17
+ "SELECT SCHEMA_NAME AS name FROM information_schema.SCHEMATA "
18
+ "WHERE SCHEMA_NAME NOT IN (:s0, :s1, :s2, :s3) ORDER BY SCHEMA_NAME"
19
+ ),
20
+ {f"s{index}": name for index, name in enumerate(SYSTEM_DATABASES)},
21
+ )
22
+ return [str(row.name) for row in rows]
23
+
24
+ def list_stored_procedures(
25
+ self, connection: Connection, database: str
26
+ ) -> list[StoredProcedureSummary]:
27
+ rows = connection.execute(
28
+ text(
29
+ "SELECT ROUTINE_NAME AS procedure_name FROM information_schema.ROUTINES "
30
+ "WHERE ROUTINE_SCHEMA = :database AND ROUTINE_TYPE = 'PROCEDURE' "
31
+ "ORDER BY ROUTINE_NAME"
32
+ ),
33
+ {"database": database},
34
+ )
35
+ return [StoredProcedureSummary(schema_=None, name=str(row.procedure_name)) for row in rows]
36
+
37
+ def get_stored_procedure(
38
+ self,
39
+ connection: Connection,
40
+ database: str,
41
+ schema: str | None,
42
+ name: str,
43
+ ) -> StoredProcedureDefinition | None:
44
+ if schema is not None:
45
+ return None
46
+ row = connection.execute(
47
+ text(
48
+ "SELECT ROUTINE_NAME AS procedure_name, ROUTINE_DEFINITION AS definition "
49
+ "FROM information_schema.ROUTINES "
50
+ "WHERE ROUTINE_SCHEMA = :database AND ROUTINE_NAME = :name "
51
+ "AND ROUTINE_TYPE = 'PROCEDURE'"
52
+ ),
53
+ {"database": database, "name": name},
54
+ ).first()
55
+ if row is None:
56
+ return None
57
+ definition = None if row.definition is None else str(row.definition)
58
+ return StoredProcedureDefinition(
59
+ schema_=None,
60
+ name=str(row.procedure_name),
61
+ definition=definition,
62
+ definition_available=definition is not None,
63
+ )
@@ -0,0 +1,132 @@
1
+ from __future__ import annotations
2
+
3
+ from collections.abc import Iterable, Mapping
4
+ from typing import Any
5
+
6
+ from sqlalchemy import Connection, inspect
7
+
8
+ from sql_safe_mcp.models import (
9
+ ColumnDefinition,
10
+ ForeignKeyDefinition,
11
+ IndexDefinition,
12
+ PrimaryKeyDefinition,
13
+ TableDefinition,
14
+ TableSummary,
15
+ UniqueConstraintDefinition,
16
+ )
17
+
18
+ _SQLSERVER_SYSTEM_SCHEMAS = {
19
+ "db_accessadmin",
20
+ "db_backupoperator",
21
+ "db_datareader",
22
+ "db_datawriter",
23
+ "db_ddladmin",
24
+ "db_denydatareader",
25
+ "db_denydatawriter",
26
+ "db_owner",
27
+ "db_securityadmin",
28
+ "guest",
29
+ "information_schema",
30
+ "sys",
31
+ }
32
+
33
+
34
+ def list_tables(connection: Connection) -> list[TableSummary]:
35
+ inspector = inspect(connection)
36
+ if getattr(connection.dialect, "name", None) == "mysql":
37
+ # MySQL and MariaDB have no schema level: the connection is bound to one database.
38
+ return [
39
+ TableSummary(schema_=None, name=name) for name in inspector.get_table_names(schema=None)
40
+ ]
41
+ tables: list[TableSummary] = []
42
+ for schema in inspector.get_schema_names():
43
+ if schema.casefold() in _SQLSERVER_SYSTEM_SCHEMAS:
44
+ continue
45
+ tables.extend(
46
+ TableSummary(schema_=schema, name=name)
47
+ for name in inspector.get_table_names(schema=schema)
48
+ )
49
+ return tables
50
+
51
+
52
+ def _dict_or_none(value: Any) -> dict[str, Any] | None:
53
+ return dict(value) if isinstance(value, dict) else None
54
+
55
+
56
+ def _normalize_columns(
57
+ connection: Connection, columns: Iterable[Mapping[str, Any]]
58
+ ) -> list[ColumnDefinition]:
59
+ result: list[ColumnDefinition] = []
60
+ for column in columns:
61
+ type_ = column["type"]
62
+ native_type = str(type_.compile(dialect=connection.dialect))
63
+ result.append(
64
+ ColumnDefinition(
65
+ name=str(column["name"]),
66
+ native_type=native_type,
67
+ nullable=bool(column.get("nullable", True)),
68
+ length=getattr(type_, "length", None),
69
+ precision=getattr(type_, "precision", None),
70
+ scale=getattr(type_, "scale", None),
71
+ autoincrement=column.get("autoincrement"),
72
+ identity=_dict_or_none(column.get("identity")),
73
+ computed=_dict_or_none(column.get("computed")),
74
+ default=None if column.get("default") is None else str(column["default"]),
75
+ )
76
+ )
77
+ return result
78
+
79
+
80
+ def get_table_definition(
81
+ connection: Connection,
82
+ schema: str | None,
83
+ table: str,
84
+ ) -> TableDefinition:
85
+ inspector = inspect(connection)
86
+ primary_key = inspector.get_pk_constraint(table, schema=schema)
87
+ foreign_keys = inspector.get_foreign_keys(table, schema=schema)
88
+ indexes = inspector.get_indexes(table, schema=schema)
89
+ try:
90
+ unique_constraints = inspector.get_unique_constraints(table, schema=schema)
91
+ except NotImplementedError:
92
+ unique_constraints = [
93
+ {"name": item.get("name"), "column_names": item.get("column_names")}
94
+ for item in indexes
95
+ if item.get("unique")
96
+ ]
97
+ return TableDefinition(
98
+ schema_=schema,
99
+ name=table,
100
+ columns=_normalize_columns(connection, inspector.get_columns(table, schema=schema)),
101
+ primary_key=PrimaryKeyDefinition(
102
+ name=primary_key.get("name"),
103
+ columns=[str(name) for name in primary_key.get("constrained_columns") or []],
104
+ ),
105
+ foreign_keys=[
106
+ ForeignKeyDefinition(
107
+ name=item.get("name"),
108
+ columns=[str(name) for name in item.get("constrained_columns") or []],
109
+ referred_schema=item.get("referred_schema"),
110
+ referred_table=str(item["referred_table"]),
111
+ referred_columns=[str(name) for name in item.get("referred_columns") or []],
112
+ options=dict(item.get("options") or {}),
113
+ )
114
+ for item in foreign_keys
115
+ ],
116
+ unique_constraints=[
117
+ UniqueConstraintDefinition(
118
+ name=None if item.get("name") is None else str(item["name"]),
119
+ columns=[str(name) for name in item.get("column_names") or []],
120
+ )
121
+ for item in unique_constraints
122
+ ],
123
+ indexes=[
124
+ IndexDefinition(
125
+ name=item.get("name"),
126
+ columns=list(item.get("column_names") or []),
127
+ unique=bool(item.get("unique", False)),
128
+ expressions=[str(value) for value in item.get("expressions") or []],
129
+ )
130
+ for item in indexes
131
+ ],
132
+ )
@@ -0,0 +1,123 @@
1
+ from __future__ import annotations
2
+
3
+ from collections import OrderedDict
4
+ from collections.abc import Callable
5
+ from threading import RLock
6
+ from typing import Any, TypeVar
7
+
8
+ import anyio
9
+ from sqlalchemy import Engine, create_engine, event
10
+
11
+ from sql_safe_mcp.config import AppConfig
12
+ from sql_safe_mcp.errors import DomainError, ErrorCode
13
+
14
+ T = TypeVar("T")
15
+
16
+
17
+ class EngineRegistry:
18
+ """Thread-safe lazy LRU of SQLAlchemy engines."""
19
+
20
+ def __init__(self, config: AppConfig) -> None:
21
+ self._config = config
22
+ self._engines: OrderedDict[tuple[str, str | None], Engine] = OrderedDict()
23
+ self._lock = RLock()
24
+ self._limiter = anyio.CapacityLimiter(config.runtime.max_concurrent_db_operations)
25
+
26
+ def _create_engine(self, alias: str, database: str | None) -> Engine:
27
+ server = self._config.servers[alias]
28
+ url = server.connection_url.get_secret_value()
29
+ if database is not None:
30
+ from sqlalchemy.engine import make_url
31
+
32
+ url = make_url(url).set(database=database)
33
+ timeout = self._config.runtime.statement_timeout_seconds
34
+ options: dict[str, Any] = {}
35
+ if server.engine in ("mysql", "mariadb"):
36
+ options["connect_args"] = {
37
+ "connect_timeout": self._config.runtime.pool_timeout_seconds,
38
+ "read_timeout": timeout,
39
+ "write_timeout": timeout,
40
+ }
41
+ engine = create_engine(
42
+ url,
43
+ pool_size=self._config.runtime.pool_size,
44
+ max_overflow=self._config.runtime.max_overflow,
45
+ pool_timeout=self._config.runtime.pool_timeout_seconds,
46
+ pool_pre_ping=True,
47
+ pool_use_lifo=True,
48
+ **options,
49
+ )
50
+
51
+ @event.listens_for(engine, "connect")
52
+ def set_connection_timeout(dbapi_connection: Any, _record: Any) -> None:
53
+ if hasattr(dbapi_connection, "timeout"):
54
+ dbapi_connection.timeout = timeout
55
+
56
+ if server.engine in ("mysql", "mariadb"):
57
+
58
+ @event.listens_for(engine, "connect")
59
+ def keep_backslash_escapes(dbapi_connection: Any, _record: Any) -> None:
60
+ # NO_BACKSLASH_ESCAPES would make the generated SQL's string escaping unsafe.
61
+ cursor = dbapi_connection.cursor()
62
+ try:
63
+ cursor.execute("SELECT @@SESSION.sql_mode")
64
+ modes = [
65
+ mode
66
+ for mode in str(cursor.fetchone()[0]).split(",")
67
+ if mode and mode != "NO_BACKSLASH_ESCAPES"
68
+ ]
69
+ cursor.execute("SET SESSION sql_mode = %s", (",".join(modes),))
70
+ finally:
71
+ cursor.close()
72
+
73
+ @event.listens_for(engine, "before_cursor_execute")
74
+ def set_statement_timeout(
75
+ _connection: Any,
76
+ cursor: Any,
77
+ _statement: str,
78
+ _parameters: Any,
79
+ _context: Any,
80
+ _executemany: bool,
81
+ ) -> None:
82
+ if hasattr(cursor, "timeout"):
83
+ cursor.timeout = timeout
84
+
85
+ return engine
86
+
87
+ def get(self, alias: str, database: str | None = None) -> Engine:
88
+ if alias not in self._config.servers:
89
+ raise DomainError(ErrorCode.UNKNOWN_SERVER, f"Unknown configured server {alias!r}.")
90
+ key = (alias, database)
91
+ evicted: Engine | None = None
92
+ with self._lock:
93
+ if engine := self._engines.get(key):
94
+ self._engines.move_to_end(key)
95
+ return engine
96
+ engine = self._create_engine(alias, database)
97
+ self._engines[key] = engine
98
+ if len(self._engines) > self._config.runtime.engine_cache_size:
99
+ _, evicted = self._engines.popitem(last=False)
100
+ if evicted is not None:
101
+ evicted.dispose()
102
+ return engine
103
+
104
+ async def run(
105
+ self,
106
+ alias: str,
107
+ database: str | None,
108
+ operation: Callable[[Engine], T],
109
+ ) -> T:
110
+ engine = self.get(alias, database)
111
+ return await anyio.to_thread.run_sync(operation, engine, limiter=self._limiter)
112
+
113
+ def dispose(self) -> None:
114
+ with self._lock:
115
+ engines = list(self._engines.values())
116
+ self._engines.clear()
117
+ for engine in engines:
118
+ engine.dispose()
119
+
120
+ @property
121
+ def cached_engine_count(self) -> int:
122
+ with self._lock:
123
+ return len(self._engines)
@@ -0,0 +1,65 @@
1
+ from __future__ import annotations
2
+
3
+ from sqlalchemy import Connection, text
4
+
5
+ from sql_safe_mcp.db.extras import DatabaseExtras
6
+ from sql_safe_mcp.models import StoredProcedureDefinition, StoredProcedureSummary
7
+
8
+
9
+ class SqlServerExtras(DatabaseExtras):
10
+ def list_databases(self, connection: Connection) -> list[str]:
11
+ rows = connection.execute(
12
+ text(
13
+ "SELECT name FROM sys.databases "
14
+ "WHERE state = 0 AND HAS_DBACCESS(name) = 1 ORDER BY name"
15
+ )
16
+ )
17
+ return [str(row.name) for row in rows]
18
+
19
+ def list_stored_procedures(
20
+ self, connection: Connection, database: str
21
+ ) -> list[StoredProcedureSummary]:
22
+ del database
23
+ rows = connection.execute(
24
+ text(
25
+ "SELECT s.name AS schema_name, p.name AS procedure_name "
26
+ "FROM sys.procedures AS p "
27
+ "JOIN sys.schemas AS s ON s.schema_id = p.schema_id "
28
+ "WHERE p.is_ms_shipped = 0 "
29
+ "ORDER BY s.name, p.name"
30
+ )
31
+ )
32
+ return [
33
+ StoredProcedureSummary(schema_=str(row.schema_name), name=str(row.procedure_name))
34
+ for row in rows
35
+ ]
36
+
37
+ def get_stored_procedure(
38
+ self,
39
+ connection: Connection,
40
+ database: str,
41
+ schema: str | None,
42
+ name: str,
43
+ ) -> StoredProcedureDefinition | None:
44
+ del database
45
+ if schema is None:
46
+ return None
47
+ row = connection.execute(
48
+ text(
49
+ "SELECT s.name AS schema_name, p.name AS procedure_name, "
50
+ "OBJECT_DEFINITION(p.object_id) AS definition "
51
+ "FROM sys.procedures AS p "
52
+ "JOIN sys.schemas AS s ON s.schema_id = p.schema_id "
53
+ "WHERE s.name = :schema AND p.name = :name AND p.is_ms_shipped = 0"
54
+ ),
55
+ {"schema": schema, "name": name},
56
+ ).first()
57
+ if row is None:
58
+ return None
59
+ definition = None if row.definition is None else str(row.definition)
60
+ return StoredProcedureDefinition(
61
+ schema_=str(row.schema_name),
62
+ name=str(row.procedure_name),
63
+ definition=definition,
64
+ definition_available=definition is not None,
65
+ )