Merge branch 'feat/TG-1124-mcp-enterprise-gating' into 'enterprise'

feat(mcp): gate connection write tools behind multi_connections plugin (TG-1124)

See merge request dkinternal/testgen/dataops-testgen!565

Author: ci bot

testgen/mcp/server.py

@@ -146,10 +146,10 @@ def build_mcp_server(
         table_health,
     )
     from testgen.mcp.tools.connections import (
-            get_connection,
+        get_connection,
         list_connections,
         test_connection,
-        )
+    )
     from testgen.mcp.tools.data_catalog import update_catalog_metadata
     from testgen.mcp.tools.discovery import (
         get_data_inventory,
@@ -378,6 +378,17 @@ def safe_prompt(fn):
     safe_prompt(profiling_overview)
     safe_prompt(hygiene_triage)
 
+    # Register plugin-provided tools through the same safe_tool wrapper as the core tools above.
+    from testgen.utils.plugins import discover
+
+    for plugin in discover():
+        try:
+            spec = plugin.load()
+            for tool in spec.get_mcp_tools():
+                safe_tool(tool)
+        except Exception:
+            LOG.warning("Plugin %s failed to load; skipping its MCP tools", plugin.package)
+
     return mcp
 
 

testgen/mcp/tools/common.py

@@ -42,6 +42,7 @@
 from testgen.common.models.test_suite import TestSuite
 from testgen.mcp.exceptions import MCPResourceNotAccessible, MCPUserError
 from testgen.mcp.permissions import get_project_permissions
+from testgen.mcp.tools.markdown import MdDoc
 
 # User-facing label for ``Disposition.INACTIVE`` is "Muted" — accept that label on input.
 _DISPOSITION_USER_TO_DB: dict[str, Disposition] = {
@@ -1374,3 +1375,54 @@ def validate_connection_fields(connection: Connection) -> list[str]:
         errors.append(f"`max_query_chars` must be between {_MAX_QUERY_CHARS_MIN} and {_MAX_QUERY_CHARS_MAX}.")
 
     return errors
+
+
+def effective_mode(connection: Connection, connection_mode: str | None) -> str | None:
+    """Mode label to apply: the explicit override, else the connection's current mode."""
+    if connection_mode is not None:
+        return connection_mode
+    inferred = infer_mode(connection)
+    return str(inferred) if inferred is not None else None
+
+
+def raise_validation_error(errors: list[str], header: str) -> None:
+    bullets = "\n".join(f"- {err}" for err in errors)
+    raise MCPUserError(f"{header}\n\n{bullets}")
+
+
+def render_connection_body(doc: MdDoc, connection: Connection) -> None:
+    """Render every non-secret connection field below the heading.
+
+    Encrypted columns are filtered out via ``ConnField.secret``.
+    """
+    doc.field("ID", connection.connection_id, code=True)
+    doc.field("Project", connection.project_code, code=True)
+    doc.field("Type", format_flavor_label(connection.sql_flavor_code))
+
+    # Each populated, non-secret field under its flavor-specific label
+    # (e.g. "Catalog" for Databricks, "Login URL" for Salesforce).
+    for fld in connection_display_fields(connection):
+        if fld.secret:
+            continue
+        value = getattr(connection, fld.column, None)
+        if value in (None, ""):
+            continue
+        doc.field(fld.label, value, code=fld.column != "project_port")
+
+    doc.field("Authentication", authentication_label(connection))
+    if connection.max_threads is not None:
+        doc.field("Max Threads", connection.max_threads)
+    if connection.max_query_chars is not None:
+        doc.field("Max Expression Length", connection.max_query_chars)
+
+
+def authentication_label(connection: Connection) -> str:
+    """The connection's auth method: the active connection mode for multi-mode
+    flavors, else the implicit method (service account key, else password).
+    """
+    mode = infer_mode(connection)
+    if mode is not None:
+        return str(mode)
+    if connection.service_account_key:
+        return "Service Account Key"
+    return "Password"

testgen/mcp/tools/connections.py

@@ -1,43 +1,42 @@
-"""MCP tools for database connection — create, update, and test database connections.
+"""MCP tools for database connections — list, get, and test database connections.
 
-Each tool gates on the ``administer`` permission. The per-flavor connection shape
-(which auth modes exist and which ``connection_params`` keys each needs) lives in
-``testgen.common.database.connection_service`` and is exposed to the model through
-the ``testgen://connection-parameters/{flavor}`` resource. Validation and
-auth-path normalization are delegated to that same module so the rules stay in
-one place.
+The per-flavor connection shape (which auth modes exist and which ``connection_params`` keys
+each needs) lives in ``testgen.common.database.connection_service`` and is exposed to the
+model through the ``testgen://connection-parameters/{flavor}`` resource. Validation and
+auth-path normalization are delegated to that same module so the rules stay in one place.
 """
 
 from __future__ import annotations
 
-from typing import Any
-
 from testgen.common.database.connection_service import (
     apply_connection_defaults,
     normalize_auth_fields,
     test_connection_status,
 )
-from testgen.common.models import get_current_session, with_database_session
+from testgen.common.models import with_database_session
 from testgen.common.models.connection import Connection
 from testgen.mcp.exceptions import MCPResourceNotAccessible, MCPUserError
 from testgen.mcp.permissions import get_project_permissions, mcp_permission
 from testgen.mcp.tools.common import (
     SQL_FLAVOR_CODE_TO_LABEL,
+    DocGroup,
     apply_connection_params,
-    connection_display_fields,
-    connection_field_labels,
+    effective_mode,
     format_flavor_label,
     format_page_footer,
     format_page_info,
-    infer_mode,
     parse_sql_flavor,
+    raise_validation_error,
+    render_connection_body,
     resolve_connection,
     validate_connection_fields,
     validate_limit,
     validate_page,
 )
 from testgen.mcp.tools.markdown import MdDoc
 
+_DOC_GROUP = DocGroup.MANAGE
+
 
 @with_database_session
 @mcp_permission("view")
@@ -104,7 +103,7 @@ def get_connection(connection_id: int) -> str:
     connection = resolve_connection(connection_id)
     doc = MdDoc()
     doc.heading(1, f"Connection `{connection.connection_name}`")
-    _render_connection_body(doc, connection)
+    render_connection_body(doc, connection)
     return doc.render()
 
 
@@ -156,7 +155,7 @@ def test_connection(
         connection = Connection(sql_flavor=family, sql_flavor_code=code)
 
     if connection_params is not None or connection_mode is not None:
-        mode = connection_mode if inline else _effective_mode(connection, connection_mode)
+        mode = connection_mode if inline else effective_mode(connection, connection_mode)
         apply_connection_params(connection, connection.sql_flavor_code, mode, connection_params or {})
 
     normalize_auth_fields(connection)
@@ -165,7 +164,7 @@ def test_connection(
 
     errors = validate_connection_fields(connection)
     if errors:
-        _raise_validation_error(errors, "Cannot test connection. Required fields missing or invalid.")
+        raise_validation_error(errors, "Cannot test connection. Required fields missing or invalid.")
 
     status = test_connection_status(connection)
 
@@ -186,132 +185,3 @@ def test_connection(
     if status.details:
         doc.code_block(status.details)
     return doc.render()
-
-
-# ---------------------------------------------------------------------------
-# Helpers
-# ---------------------------------------------------------------------------
-
-
-def _effective_mode(connection: Connection, connection_mode: str | None) -> str | None:
-    """Mode label to apply: the explicit override, else the connection's current mode."""
-    if connection_mode is not None:
-        return connection_mode
-    inferred = infer_mode(connection)
-    return str(inferred) if inferred is not None else None
-
-
-def _raise_validation_error(errors: list[str], header: str) -> None:
-    bullets = "\n".join(f"- {err}" for err in errors)
-    raise MCPUserError(f"{header}\n\n{bullets}")
-
-
-def _render_created_connection(connection: Connection) -> str:
-    doc = MdDoc()
-    doc.heading(1, f"Connection `{connection.connection_name}` created")
-    _render_connection_body(doc, connection)
-    return doc.render()
-
-
-def _render_connection_body(doc: MdDoc, connection: Connection) -> None:
-    """Render every non-secret connection field below the heading.
-
-    Shared by ``create_connection`` (the response after a successful create) and
-    ``get_connection`` (read tool) so the surfaced field set stays consistent.
-    Encrypted columns are filtered out via ``ConnField.secret``.
-    """
-    doc.field("ID", connection.connection_id, code=True)
-    doc.field("Project", connection.project_code, code=True)
-    doc.field("Type", format_flavor_label(connection.sql_flavor_code))
-
-    # Each populated, non-secret field under its flavor-specific label
-    # (e.g. "Catalog" for Databricks, "Login URL" for Salesforce).
-    for fld in connection_display_fields(connection):
-        if fld.secret:
-            continue
-        value = getattr(connection, fld.column, None)
-        if value in (None, ""):
-            continue
-        doc.field(fld.label, value, code=fld.column != "project_port")
-
-    doc.field("Authentication", _authentication_label(connection))
-    if connection.max_threads is not None:
-        doc.field("Max Threads", connection.max_threads)
-    if connection.max_query_chars is not None:
-        doc.field("Max Expression Length", connection.max_query_chars)
-
-
-def _authentication_label(connection: Connection) -> str:
-    """The connection's auth method: the active connection mode for multi-mode
-    flavors, else the implicit method (service account key, else password).
-    """
-    mode = infer_mode(connection)
-    if mode is not None:
-        return str(mode)
-    if connection.service_account_key:
-        return "Service Account Key"
-    return "Password"
-
-
-def _snapshot(connection: Connection) -> dict[str, Any]:
-    return {attr: getattr(connection, attr, None) for attr in _DIFF_ATTRS}
-
-
-def _render_field_value(attr: str, value: Any) -> str | None:
-    if attr == "sql_flavor_code" and value is not None:
-        return SQL_FLAVOR_CODE_TO_LABEL.get(value, value).value
-    if isinstance(value, bool):
-        return "Yes" if value else "No"
-    if value is None or value == "":
-        return None
-    return str(value)
-
-
-_DIFF_ATTRS: tuple[str, ...] = (
-    "connection_name",
-    "sql_flavor_code",
-    "project_host",
-    "project_port",
-    "project_db",
-    "project_user",
-    "project_pw_encrypted",
-    "url",
-    "connect_by_url",
-    "connect_by_key",
-    "private_key",
-    "private_key_passphrase",
-    "connect_with_identity",
-    "warehouse",
-    "http_path",
-    "service_account_key",
-    "max_threads",
-    "max_query_chars",
-)
-
-_DIFF_LABELS: dict[str, str] = {
-    "connection_name": "Name",
-    "sql_flavor_code": "Type",
-    "project_host": "Host",
-    "project_port": "Port",
-    "project_db": "Database",
-    "project_user": "Username",
-    "project_pw_encrypted": "Password",
-    "url": "URL",
-    "connect_by_url": "Connect by URL",
-    "connect_by_key": "Connect by Key-Pair",
-    "private_key": "Private Key",
-    "private_key_passphrase": "Private Key Passphrase",
-    "connect_with_identity": "Connect with Managed Identity",
-    "warehouse": "Warehouse",
-    "http_path": "HTTP Path",
-    "service_account_key": "Service Account Key",
-    "max_threads": "Max Threads",
-    "max_query_chars": "Max Expression Length",
-}
-
-_ATTR_IS_SECRET: dict[str, bool] = {
-    "project_pw_encrypted": True,
-    "private_key": True,
-    "private_key_passphrase": True,
-    "service_account_key": True,
-}

testgen/mcp/tools/table_groups.py

@@ -21,6 +21,7 @@
 from testgen.mcp.exceptions import MCPPermissionDenied, MCPResourceNotAccessible, MCPUserError
 from testgen.mcp.permissions import get_project_permissions, mcp_permission
 from testgen.mcp.tools.common import (
+    DocGroup,
     format_flavor_label,
     format_page_footer,
     format_page_info,
@@ -32,6 +33,8 @@
 from testgen.mcp.tools.markdown import MdDoc
 from testgen.utils import friendly_score
 
+_DOC_GROUP = DocGroup.MANAGE
+
 _DUPLICATE_NAME_MESSAGE = "A Table Group with the same name already exists."
 _PII_FLAG_DENIED_MESSAGE = (
     "Changing PII detection requires permission to view PII. "

testgen/testing/__init__.py

@@ -0,0 +1,7 @@
+"""Testing support importable by the core test suite and any plugin test suite.
+
+Fixtures live in ``testgen.testing.fixtures``; import the ones you need into a
+``conftest.py`` to register them (pytest registers fixtures that are merely imported
+into a conftest). An importable module is the only way to share fixtures across separate
+test trees — sibling test directories do not share conftest fixtures.
+"""