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.
- sql_safe_mcp/__init__.py +3 -0
- sql_safe_mcp/__main__.py +43 -0
- sql_safe_mcp/config.py +201 -0
- sql_safe_mcp/db/__init__.py +1 -0
- sql_safe_mcp/db/extras.py +35 -0
- sql_safe_mcp/db/mysql.py +63 -0
- sql_safe_mcp/db/reflection.py +132 -0
- sql_safe_mcp/db/registry.py +123 -0
- sql_safe_mcp/db/sqlserver.py +65 -0
- sql_safe_mcp/errors.py +49 -0
- sql_safe_mcp/mcp_server.py +140 -0
- sql_safe_mcp/models.py +124 -0
- sql_safe_mcp/security/__init__.py +1 -0
- sql_safe_mcp/security/dialect.py +29 -0
- sql_safe_mcp/security/executor.py +129 -0
- sql_safe_mcp/security/lineage.py +224 -0
- sql_safe_mcp/security/parser.py +183 -0
- sql_safe_mcp/security/pipeline.py +43 -0
- sql_safe_mcp/security/policy.py +101 -0
- sql_safe_mcp/security/reasons.py +43 -0
- sql_safe_mcp/security/schema.py +139 -0
- sql_safe_mcp/security/tokens.py +207 -0
- sql_safe_mcp/security/validated_query.py +209 -0
- sql_safe_mcp/service.py +325 -0
- sql_safe_mcp-1.2.0.dist-info/METADATA +259 -0
- sql_safe_mcp-1.2.0.dist-info/RECORD +29 -0
- sql_safe_mcp-1.2.0.dist-info/WHEEL +4 -0
- sql_safe_mcp-1.2.0.dist-info/entry_points.txt +2 -0
- sql_safe_mcp-1.2.0.dist-info/licenses/LICENSE +21 -0
sql_safe_mcp/__init__.py
ADDED
sql_safe_mcp/__main__.py
ADDED
|
@@ -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}")
|
sql_safe_mcp/db/mysql.py
ADDED
|
@@ -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
|
+
)
|