feat(mcp): accept hygiene issue id in source data tools

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

Author: luis-dk

testgen/mcp/tools/common.py

@@ -17,7 +17,7 @@
     ProfileMetric,
     SuggestedDataType,
 )
-from testgen.common.models.hygiene_issue import HygieneIssueType
+from testgen.common.models.hygiene_issue import HygieneIssue, HygieneIssueType
 from testgen.common.models.notification_settings import (
     MonitorNotificationTrigger,
     NotificationEvent,
@@ -524,6 +524,16 @@ def resolve_table_group(table_group_id: str) -> TableGroup:
     return tg
 
 
+def resolve_hygiene_issue(issue_id: str) -> HygieneIssue:
+    """Resolve a hygiene issue ID, collapsing missing-or-inaccessible into one error path."""
+    issue_uuid = parse_uuid(issue_id, "issue_id")
+    perms = get_project_permissions()
+    issue = HygieneIssue.get(issue_uuid, HygieneIssue.project_code.in_(perms.allowed_codes))
+    if issue is None:
+        raise MCPResourceNotAccessible("Hygiene issue", issue_id)
+    return issue
+
+
 def resolve_test_suite(test_suite_id: str) -> TestSuite:
     """Resolve a regular (non-monitor) test suite ID, collapsing missing-or-inaccessible into one error path."""
     suite_uuid = parse_uuid(test_suite_id, "test_suite_id")

testgen/mcp/tools/hygiene_issues.py

@@ -25,6 +25,7 @@
     parse_quality_dimension,
     parse_since_arg,
     parse_uuid,
+    resolve_hygiene_issue,
     resolve_issue_type,
     resolve_table_group,
     validate_limit,
@@ -211,17 +212,9 @@ def update_hygiene_issue(*, issue_id: str, disposition: str) -> str:
         issue_id: UUID of the hygiene issue.
         disposition: New disposition. Valid values: 'Confirmed', 'Dismissed', 'Muted'.
     """
-    issue_uuid = parse_uuid(issue_id, "issue_id")
     db_disposition = parse_disposition(disposition)
-    perms = get_project_permissions()
-
-    updated = HygieneIssue.update_disposition(
-        issue_uuid,
-        db_disposition,
-        HygieneIssue.project_code.in_(perms.allowed_codes),
-    )
-    if not updated:
-        raise MCPResourceNotAccessible("Hygiene issue", issue_id)
+    issue = resolve_hygiene_issue(issue_id)
+    issue.disposition = db_disposition
 
     doc = MdDoc()
     doc.text(f"Updated hygiene issue {MdDoc.code(issue_id)} disposition to **{disposition}**.")

testgen/mcp/tools/source_data.py

@@ -1,21 +1,34 @@
 from datetime import datetime
 
 from testgen.common.models import with_database_session
+from testgen.common.models.profiling_run import ProfilingRun
 from testgen.common.models.test_definition import TestDefinition
 from testgen.common.source_data_service import (
     SourceDataResult,
+    build_hygiene_query,
     build_test_result_query,
+    fetch_hygiene_source_data,
     fetch_test_result_source_data,
 )
 from testgen.mcp.exceptions import MCPResourceNotAccessible, MCPUserError
 from testgen.mcp.permissions import get_project_permissions, mcp_permission
-from testgen.mcp.tools.common import DocGroup, parse_uuid, validate_limit
+from testgen.mcp.tools.common import DocGroup, parse_uuid, resolve_hygiene_issue, validate_limit
 from testgen.mcp.tools.markdown import MdDoc
 
 _DOC_GROUP = DocGroup.INVESTIGATE
 
 
-def _resolve_context(test_definition_id: str, reference_date: str | None) -> dict:
+def _validate_source_args(test_definition_id: str | None, issue_id: str | None, reference_date: str | None) -> None:
+    """Enforce 'exactly one entity' and the reference_date-only-with-test-definition rule."""
+    if bool(test_definition_id) == bool(issue_id):
+        raise MCPUserError("Provide exactly one of test_definition_id or issue_id.")
+    if issue_id and reference_date:
+        raise MCPUserError(
+            "reference_date applies only to test_definition_id; omit it when looking up a hygiene issue."
+        )
+
+
+def _resolve_test_definition_context(test_definition_id: str, reference_date: str | None) -> dict:
     """Look up the test definition context and validate permissions."""
     td_uuid = parse_uuid(test_definition_id, "test_definition_id")
     perms = get_project_permissions()
@@ -40,40 +53,86 @@ def _resolve_context(test_definition_id: str, reference_date: str | None) -> dic
     return context
 
 
+def _resolve_hygiene_context(issue_id: str) -> dict:
+    """Resolve a hygiene issue (permission-scoped) into the lookup context the service expects.
+
+    The source profiling run is intrinsic to the issue, so ``profiling_starttime`` comes from the
+    issue's ``ProfilingRun`` — there is no caller-supplied reference date.
+    """
+    issue = resolve_hygiene_issue(issue_id)
+    run = ProfilingRun.get(issue.profile_run_id)
+    return {
+        "table_groups_id": issue.table_groups_id,
+        "anomaly_id": issue.type_id,
+        "detail": issue.detail,
+        "schema_name": issue.schema_name,
+        "table_name": issue.table_name,
+        "column_name": issue.column_name,
+        "profiling_starttime": run.profiling_starttime if run else None,
+        "project_code": issue.project_code,
+    }
+
+
+def _render_header_fields(doc: MdDoc, context: dict) -> None:
+    """Render the entity-neutral location fields shared by both tools."""
+    if context.get("test_type"):
+        doc.field("Test type", context.get("test_type"), code=True)
+    doc.field("Table", f"{context.get('schema_name')}.{context.get('table_name')}", code=True)
+    column = context.get("column_names") or context.get("column_name")
+    if column:
+        doc.field("Column", column, code=True)
+
+
 @with_database_session
 @mcp_permission("view")
 def get_source_data_query(
-    test_definition_id: str,
+    test_definition_id: str | None = None,
+    issue_id: str | None = None,
     reference_date: str | None = None,
     limit: int = 100,
 ) -> str:
-    """Get the SQL query that would be used to look up source data for a test definition, without executing it.
+    """Get the SQL query that would be used to look up source data, without executing it.
 
-    Builds a lookup query using current test definition parameters (thresholds, conditions).
+    Builds a lookup query using the current criteria of a test definition or a hygiene issue.
     The query targets the connected database.
     Some test types (e.g. Freshness Trend, Schema Drift) do not have source data lookups.
 
+    Provide exactly one of ``test_definition_id`` or ``issue_id``.
+
     Args:
         test_definition_id: UUID of a test definition, e.g. from ``list_test_results``.
-        reference_date: ISO 8601 date used as the test reference point (default: now).
+        issue_id: UUID of a hygiene issue, e.g. from ``list_hygiene_issues``. Mutually exclusive
+            with ``test_definition_id``.
+        reference_date: ISO 8601 date used as the test reference point (default: now). Applies only
+            to ``test_definition_id``.
         limit: Maximum rows the query would return (default 100, max 500).
     """
+    _validate_source_args(test_definition_id, issue_id, reference_date)
     validate_limit(limit, 500)
-    context = _resolve_context(test_definition_id, reference_date)
 
-    query = build_test_result_query(context, limit)
+    if test_definition_id:
+        context = _resolve_test_definition_context(test_definition_id, reference_date)
+        entity_label, entity_id = "Test Definition", test_definition_id
+        query = build_test_result_query(context, limit)
+    else:
+        context = _resolve_hygiene_context(issue_id)
+        entity_label, entity_id = "Hygiene Issue", issue_id
+        query = build_hygiene_query(context, limit)
+
     if not query:
+        if test_definition_id:
+            return (
+                f"Source data lookup is not available for test type `{context.get('test_type', 'unknown')}`.\n\n"
+                "This test type does not have a defined lookup query."
+            )
         return (
-            f"Source data lookup is not available for test type `{context.get('test_type', 'unknown')}`.\n\n"
-            "This test type does not have a defined lookup query."
+            "Source data lookup is not available for this hygiene issue.\n\n"
+            "This hygiene issue type does not have a defined lookup query."
         )
 
     doc = MdDoc()
-    doc.heading(1, f"Source Data Query for Test Definition `{test_definition_id}`")
-    doc.field("Test type", context.get("test_type"), code=True)
-    doc.field("Table", f"{context.get('schema_name')}.{context.get('table_name')}", code=True)
-    if context.get("column_names"):
-        doc.field("Column", context["column_names"], code=True)
+    doc.heading(1, f"Source Data Query for {entity_label} `{entity_id}`")
+    _render_header_fields(doc, context)
     doc.field("Limit", limit)
     doc.code_block(query, language="sql")
 
@@ -83,34 +142,46 @@ def get_source_data_query(
 @with_database_session
 @mcp_permission("view")
 def get_source_data(
-    test_definition_id: str,
+    test_definition_id: str | None = None,
+    issue_id: str | None = None,
     reference_date: str | None = None,
     limit: int = 100,
 ) -> str:
-    """Look up rows from the connected database that match or violate a test definition's criteria.
+    """Look up rows from the connected database that match or violate a test or hygiene issue's criteria.
 
     Executes the source data query against the connected database and returns matching rows.
-    Shows CURRENT data — rows may have changed since the test last ran.
+    Shows CURRENT data — rows may have changed since the test or profiling run.
     Some test types (e.g. Freshness Trend, Schema Drift) do not have source data lookups.
 
+    Provide exactly one of ``test_definition_id`` or ``issue_id``.
+
     Args:
         test_definition_id: UUID of a test definition, e.g. from ``list_test_results``.
-        reference_date: ISO 8601 date used as the test reference point (default: now).
+        issue_id: UUID of a hygiene issue, e.g. from ``list_hygiene_issues``. Mutually exclusive
+            with ``test_definition_id``.
+        reference_date: ISO 8601 date used as the test reference point (default: now). Applies only
+            to ``test_definition_id``.
         limit: Maximum rows to return (default 100, max 500).
     """
+    _validate_source_args(test_definition_id, issue_id, reference_date)
     validate_limit(limit, 500)
-    context = _resolve_context(test_definition_id, reference_date)
+
+    if test_definition_id:
+        context = _resolve_test_definition_context(test_definition_id, reference_date)
+        entity_label, entity_id = "Test Definition", test_definition_id
+        fetch = fetch_test_result_source_data
+    else:
+        context = _resolve_hygiene_context(issue_id)
+        entity_label, entity_id = "Hygiene Issue", issue_id
+        fetch = fetch_hygiene_source_data
 
     mask_pii = not get_project_permissions().has_permission("view_pii", context.get("project_code"))
 
-    result: SourceDataResult = fetch_test_result_source_data(context, limit, mask_pii)
+    result: SourceDataResult = fetch(context, limit, mask_pii)
 
     doc = MdDoc()
-    doc.heading(1, f"Source Data for Test Definition `{test_definition_id}`")
-    doc.field("Test type", context.get("test_type"), code=True)
-    doc.field("Table", f"{context.get('schema_name')}.{context.get('table_name')}", code=True)
-    if context.get("column_names"):
-        doc.field("Column", context["column_names"], code=True)
+    doc.heading(1, f"Source Data for {entity_label} `{entity_id}`")
+    _render_header_fields(doc, context)
 
     if result.status == "OK":
         row_count = len(result.df) if result.df is not None else 0

tests/unit/mcp/test_tools_common.py

@@ -29,6 +29,7 @@
     parse_score_group_by,
     parse_score_type,
     parse_uuid,
+    resolve_hygiene_issue,
     resolve_issue_type,
     resolve_profiling_run,
     resolve_test_note,
@@ -379,6 +380,35 @@ def test_resolve_test_note_invalid_uuid():
         resolve_test_note("not-a-uuid")
 
 
+# --- resolve_hygiene_issue ---
+
+
+@patch("testgen.mcp.tools.common.get_project_permissions")
+@patch("testgen.mcp.tools.common.HygieneIssue")
+def test_resolve_hygiene_issue_happy_path(mock_hi_cls, mock_get_perms, db_session_mock):
+    issue = MagicMock()
+    mock_hi_cls.get.return_value = issue
+    mock_get_perms.return_value = _mock_perms()
+
+    assert resolve_hygiene_issue(str(uuid4())) is issue
+
+
+@patch("testgen.mcp.tools.common.get_project_permissions")
+@patch("testgen.mcp.tools.common.HygieneIssue")
+def test_resolve_hygiene_issue_missing_or_inaccessible(mock_hi_cls, mock_get_perms, db_session_mock):
+    """Missing issue and forbidden-project issue both collapse to one error (project scoped in the query)."""
+    mock_hi_cls.get.return_value = None
+    mock_get_perms.return_value = _mock_perms()
+
+    with pytest.raises(MCPResourceNotAccessible, match=r"Hygiene issue .* not found or not accessible"):
+        resolve_hygiene_issue(str(uuid4()))
+
+
+def test_resolve_hygiene_issue_invalid_uuid():
+    with pytest.raises(MCPUserError, match="Invalid issue_id"):
+        resolve_hygiene_issue("not-a-uuid")
+
+
 # --- parse_pii_category ---
 
 

tests/unit/mcp/test_tools_hygiene_issues.py

@@ -727,22 +727,35 @@ def test_update_hygiene_issue_invalid_disposition(db_session_mock, disposition_p
         update_hygiene_issue(issue_id=str(uuid4()), disposition="Bogus")
 
 
-@patch.object(HygieneIssue, "update_disposition")
-def test_update_hygiene_issue_muted_maps_to_inactive(mock_update, db_session_mock, disposition_perms):
+@patch("testgen.mcp.tools.hygiene_issues.resolve_hygiene_issue")
+def test_update_hygiene_issue_muted_maps_to_inactive(mock_resolve, db_session_mock, disposition_perms):
     from testgen.mcp.tools.hygiene_issues import update_hygiene_issue
 
-    mock_update.return_value = True
+    issue = MagicMock()
+    mock_resolve.return_value = issue
     update_hygiene_issue(issue_id=str(uuid4()), disposition="Muted")
 
-    args = mock_update.call_args.args
-    assert args[1] == Disposition.INACTIVE
+    assert issue.disposition == Disposition.INACTIVE
 
 
-@patch.object(HygieneIssue, "update_disposition")
-def test_update_hygiene_issue_returns_success_markdown(mock_update, db_session_mock, disposition_perms):
+@patch("testgen.mcp.tools.hygiene_issues.resolve_hygiene_issue")
+def test_update_hygiene_issue_sets_disposition_on_resolved_issue(
+    mock_resolve, db_session_mock, disposition_perms,
+):
+    from testgen.mcp.tools.hygiene_issues import update_hygiene_issue
+
+    issue = MagicMock()
+    mock_resolve.return_value = issue
+    update_hygiene_issue(issue_id=str(uuid4()), disposition="Dismissed")
+
+    assert issue.disposition == Disposition.DISMISSED
+
+
+@patch("testgen.mcp.tools.hygiene_issues.resolve_hygiene_issue")
+def test_update_hygiene_issue_returns_success_markdown(mock_resolve, db_session_mock, disposition_perms):
     from testgen.mcp.tools.hygiene_issues import update_hygiene_issue
 
-    mock_update.return_value = True
+    mock_resolve.return_value = MagicMock()
     issue_id = str(uuid4())
     result = update_hygiene_issue(issue_id=issue_id, disposition="Dismissed")
 
@@ -751,30 +764,30 @@ def test_update_hygiene_issue_returns_success_markdown(mock_update, db_session_m
     assert "Dismissed" in result
 
 
-@patch.object(HygieneIssue, "update_disposition")
-def test_update_hygiene_issue_not_updated_collapses_to_not_accessible(
-    mock_update, db_session_mock, disposition_perms,
+@patch("testgen.mcp.tools.hygiene_issues.resolve_hygiene_issue")
+def test_update_hygiene_issue_not_accessible_propagates(
+    mock_resolve, db_session_mock, disposition_perms,
 ):
     from testgen.mcp.tools.hygiene_issues import update_hygiene_issue
 
-    mock_update.return_value = False
+    mock_resolve.side_effect = MCPResourceNotAccessible("Hygiene issue", "x")
     with pytest.raises(MCPResourceNotAccessible):
         update_hygiene_issue(issue_id=str(uuid4()), disposition="Confirmed")
 
 
-@patch.object(HygieneIssue, "update_disposition")
-def test_update_hygiene_issue_passes_project_scope_clause(
-    mock_update, db_session_mock, disposition_perms,
+@patch("testgen.mcp.tools.hygiene_issues.resolve_hygiene_issue")
+def test_update_hygiene_issue_delegates_scope_to_resolver(
+    mock_resolve, db_session_mock, disposition_perms,
 ):
+    """Project scoping lives in resolve_hygiene_issue (covered in test_tools_common);
+    the tool must route the issue id through it."""
     from testgen.mcp.tools.hygiene_issues import update_hygiene_issue
 
-    mock_update.return_value = True
-    update_hygiene_issue(issue_id=str(uuid4()), disposition="Confirmed")
+    mock_resolve.return_value = MagicMock()
+    issue_id = str(uuid4())
+    update_hygiene_issue(issue_id=issue_id, disposition="Confirmed")
 
-    # Trailing args after (issue_uuid, db_disposition) are the *clauses
-    clauses = mock_update.call_args.args[2:]
-    sql = "\n".join(str(c.compile(dialect=postgresql.dialect())) for c in clauses)
-    assert "profile_anomaly_results.project_code IN" in sql
+    mock_resolve.assert_called_once_with(issue_id)
 
 
 def test_update_hygiene_issue_uses_disposition_permission():