|
1 | | -"""MCP tools for database connection — create, update, and test database connections. |
2 | | -
|
3 | | -Each tool gates on the ``administer`` permission. The per-flavor connection shape |
4 | | -(which auth modes exist and which ``connection_params`` keys each needs) lives in |
5 | | -``testgen.common.database.connection_service`` and is exposed to the model through |
6 | | -the ``testgen://connection-parameters/{flavor}`` resource. Validation and |
7 | | -auth-path normalization are delegated to that same module so the rules stay in |
8 | | -one place. |
| 1 | +"""MCP tools for database connections — list, get, and test database connections. |
| 2 | +
|
| 3 | +The per-flavor connection shape (which auth modes exist and which ``connection_params`` keys |
| 4 | +each needs) lives in ``testgen.common.database.connection_service`` and is exposed to the |
| 5 | +model through the ``testgen://connection-parameters/{flavor}`` resource. Validation and |
| 6 | +auth-path normalization are delegated to that same module so the rules stay in one place. |
9 | 7 | """ |
10 | 8 |
|
11 | 9 | from __future__ import annotations |
12 | 10 |
|
13 | | -from typing import Any |
14 | | - |
15 | 11 | from testgen.common.database.connection_service import ( |
16 | 12 | apply_connection_defaults, |
17 | 13 | normalize_auth_fields, |
18 | 14 | test_connection_status, |
19 | 15 | ) |
20 | | -from testgen.common.models import get_current_session, with_database_session |
| 16 | +from testgen.common.models import with_database_session |
21 | 17 | from testgen.common.models.connection import Connection |
22 | 18 | from testgen.mcp.exceptions import MCPResourceNotAccessible, MCPUserError |
23 | 19 | from testgen.mcp.permissions import get_project_permissions, mcp_permission |
24 | 20 | from testgen.mcp.tools.common import ( |
25 | 21 | SQL_FLAVOR_CODE_TO_LABEL, |
26 | 22 | apply_connection_params, |
27 | 23 | connection_display_fields, |
28 | | - connection_field_labels, |
29 | 24 | format_flavor_label, |
30 | 25 | format_page_footer, |
31 | 26 | format_page_info, |
@@ -206,18 +201,9 @@ def _raise_validation_error(errors: list[str], header: str) -> None: |
206 | 201 | raise MCPUserError(f"{header}\n\n{bullets}") |
207 | 202 |
|
208 | 203 |
|
209 | | -def _render_created_connection(connection: Connection) -> str: |
210 | | - doc = MdDoc() |
211 | | - doc.heading(1, f"Connection `{connection.connection_name}` created") |
212 | | - _render_connection_body(doc, connection) |
213 | | - return doc.render() |
214 | | - |
215 | | - |
216 | 204 | def _render_connection_body(doc: MdDoc, connection: Connection) -> None: |
217 | 205 | """Render every non-secret connection field below the heading. |
218 | 206 |
|
219 | | - Shared by ``create_connection`` (the response after a successful create) and |
220 | | - ``get_connection`` (read tool) so the surfaced field set stays consistent. |
221 | 207 | Encrypted columns are filtered out via ``ConnField.secret``. |
222 | 208 | """ |
223 | 209 | doc.field("ID", connection.connection_id, code=True) |
@@ -251,67 +237,3 @@ def _authentication_label(connection: Connection) -> str: |
251 | 237 | if connection.service_account_key: |
252 | 238 | return "Service Account Key" |
253 | 239 | return "Password" |
254 | | - |
255 | | - |
256 | | -def _snapshot(connection: Connection) -> dict[str, Any]: |
257 | | - return {attr: getattr(connection, attr, None) for attr in _DIFF_ATTRS} |
258 | | - |
259 | | - |
260 | | -def _render_field_value(attr: str, value: Any) -> str | None: |
261 | | - if attr == "sql_flavor_code" and value is not None: |
262 | | - return SQL_FLAVOR_CODE_TO_LABEL.get(value, value).value |
263 | | - if isinstance(value, bool): |
264 | | - return "Yes" if value else "No" |
265 | | - if value is None or value == "": |
266 | | - return None |
267 | | - return str(value) |
268 | | - |
269 | | - |
270 | | -_DIFF_ATTRS: tuple[str, ...] = ( |
271 | | - "connection_name", |
272 | | - "sql_flavor_code", |
273 | | - "project_host", |
274 | | - "project_port", |
275 | | - "project_db", |
276 | | - "project_user", |
277 | | - "project_pw_encrypted", |
278 | | - "url", |
279 | | - "connect_by_url", |
280 | | - "connect_by_key", |
281 | | - "private_key", |
282 | | - "private_key_passphrase", |
283 | | - "connect_with_identity", |
284 | | - "warehouse", |
285 | | - "http_path", |
286 | | - "service_account_key", |
287 | | - "max_threads", |
288 | | - "max_query_chars", |
289 | | -) |
290 | | - |
291 | | -_DIFF_LABELS: dict[str, str] = { |
292 | | - "connection_name": "Name", |
293 | | - "sql_flavor_code": "Type", |
294 | | - "project_host": "Host", |
295 | | - "project_port": "Port", |
296 | | - "project_db": "Database", |
297 | | - "project_user": "Username", |
298 | | - "project_pw_encrypted": "Password", |
299 | | - "url": "URL", |
300 | | - "connect_by_url": "Connect by URL", |
301 | | - "connect_by_key": "Connect by Key-Pair", |
302 | | - "private_key": "Private Key", |
303 | | - "private_key_passphrase": "Private Key Passphrase", |
304 | | - "connect_with_identity": "Connect with Managed Identity", |
305 | | - "warehouse": "Warehouse", |
306 | | - "http_path": "HTTP Path", |
307 | | - "service_account_key": "Service Account Key", |
308 | | - "max_threads": "Max Threads", |
309 | | - "max_query_chars": "Max Expression Length", |
310 | | -} |
311 | | - |
312 | | -_ATTR_IS_SECRET: dict[str, bool] = { |
313 | | - "project_pw_encrypted": True, |
314 | | - "private_key": True, |
315 | | - "private_key_passphrase": True, |
316 | | - "service_account_key": True, |
317 | | -} |
0 commit comments