Merge branch 'feat/TG-1117-mcp-filter-runs-by-schedule' into 'enterprise'

feat(mcp): filter test-run and profiling-run lists by schedule (TG-1117)

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

Author: ci bot

testgen/common/models/job_execution.py

@@ -5,6 +5,7 @@
 
 from sqlalchemy import Column, String, Text, case, delete, func, select, text, update
 from sqlalchemy.dialects import postgresql
+from sqlalchemy.sql.elements import ColumnElement
 
 from testgen.common.enums import JobSource, JobStatus
 from testgen.common.models import Base, database_session, get_current_session
@@ -97,6 +98,7 @@ def claim_actionable(cls, limit: int = 5) -> list[Self]:
     @classmethod
     def select_active_by_kwargs(
         cls,
+        *clauses: ColumnElement[bool],
         project_code: str,
         job_key: str,
         kwargs_match: dict[str, str | list[str]],
@@ -105,13 +107,15 @@ def select_active_by_kwargs(
         """Find JE rows whose ``kwargs`` JSONB matches the given (key, value) pairs.
 
         Values may be a single string or a list of strings (which becomes an ``IN`` filter).
+        Additional caller-supplied WHERE expressions may be passed as ``*clauses``.
         Defaults to active (non-terminal) statuses.
         """
         statuses = statuses or cls._ACTIVE_STATUSES
         query = select(cls).where(
             cls.project_code == project_code,
             cls.job_key == job_key,
             cls.status.in_(statuses),
+            *clauses,
         )
         for k, v in kwargs_match.items():
             if isinstance(v, list):

testgen/common/models/profiling_run.py

@@ -208,13 +208,15 @@ def select_summary(
         project_code: str | None = None,
         table_group_id: str | UUID | None = None,
         job_execution_id: str | UUID | None = None,
+        schedule_id: str | None = None,
         statuses: list[JobStatus] | None = None,
         page: int = 1,
         page_size: int = 20,
     ) -> tuple[list[ProfilingRunSummary], int]:
         if (
             (table_group_id and not is_uuid4(table_group_id))
             or (job_execution_id and not is_uuid4(job_execution_id))
+            or (schedule_id and not is_uuid4(schedule_id))
         ):
             return [], 0
 
@@ -271,6 +273,7 @@ def select_summary(
             {" AND je.project_code = :project_code" if project_code else ""}
             {" AND tg.id = :table_group_id" if table_group_id else ""}
             {" AND je.id = :job_execution_id" if job_execution_id else ""}
+            {" AND je.job_schedule_id = :schedule_id" if schedule_id else ""}
             {" AND je.status IN :statuses" if statuses else ""}
         ORDER BY je.created_at DESC
         LIMIT :limit OFFSET :offset;
@@ -279,6 +282,7 @@ def select_summary(
             "project_code": project_code,
             "table_group_id": str(table_group_id) if table_group_id else None,
             "job_execution_id": str(job_execution_id) if job_execution_id else None,
+            "schedule_id": schedule_id,
             "statuses": tuple(statuses) if statuses else (),
             "limit": page_size,
             "offset": (page - 1) * page_size,

testgen/common/models/test_run.py

@@ -219,6 +219,7 @@ def select_summary(
         test_suite_id: str | None = None,
         test_run_ids: list[str | UUID] | None = None,
         job_execution_id: str | UUID | None = None,
+        schedule_id: str | None = None,
         statuses: list[JobStatus] | None = None,
         page: int = 1,
         page_size: int = 20,
@@ -228,6 +229,7 @@ def select_summary(
             or (test_suite_id and not is_uuid4(test_suite_id))
             or (test_run_ids and not all(is_uuid4(run_id) for run_id in test_run_ids))
             or (job_execution_id and not is_uuid4(job_execution_id))
+            or (schedule_id and not is_uuid4(schedule_id))
         ):
             return [], 0
 
@@ -290,6 +292,7 @@ def select_summary(
             {" AND ts.id = :test_suite_id" if test_suite_id else ""}
             {" AND tr.id IN :test_run_ids" if test_run_ids else ""}
             {" AND je.id = :job_execution_id" if job_execution_id else ""}
+            {" AND je.job_schedule_id = :schedule_id" if schedule_id else ""}
             {" AND je.status IN :statuses" if statuses else ""}
         ORDER BY je.created_at DESC
         LIMIT :limit OFFSET :offset;
@@ -300,6 +303,7 @@ def select_summary(
             "test_suite_id": test_suite_id,
             "test_run_ids": tuple(test_run_ids or []),
             "job_execution_id": str(job_execution_id) if job_execution_id else None,
+            "schedule_id": schedule_id,
             "statuses": tuple(statuses) if statuses else (),
             "limit": page_size,
             "offset": (page - 1) * page_size,

testgen/mcp/tools/profiling.py

@@ -414,6 +414,7 @@ def _render_column_profile_row(c: ColumnProfileSummary) -> list:
 @mcp_permission("catalog")
 def list_profiling_runs(
     table_group_id: str,
+    schedule_id: str | None = None,
     status: str | None = None,
     limit: int = 10,
     page: int = 1,
@@ -423,6 +424,8 @@ def list_profiling_runs(
 
     Args:
         table_group_id: UUID of the table group, e.g. from `get_data_inventory`.
+        schedule_id: Optional UUID of a schedule, e.g. from `list_schedules`. Returns only runs
+            triggered by that schedule.
         status: Optional run status filter. One of: Pending, Running, Completed, Canceled, Error.
         limit: Page size (default 10, max 100).
         page: Page number starting at 1 (default 1).
@@ -431,11 +434,14 @@ def list_profiling_runs(
     validate_page(page)
 
     statuses = parse_run_status_filter(status) if status else None
+    if schedule_id:
+        parse_uuid(schedule_id, "schedule_id")
     tg = resolve_table_group(table_group_id)
 
     summaries, total = ProfilingRun.select_summary(
         project_code=tg.project_code,
         table_group_id=tg.id,
+        schedule_id=schedule_id,
         statuses=statuses,
         page=page,
         page_size=limit,
@@ -445,15 +451,22 @@ def list_profiling_runs(
     # joined-run queries. Surface them as a separate "Pending" section on page 1.
     pending_jes: list[JobExecution] = []
     if page == 1:
+        clauses = [JobExecution.job_schedule_id == schedule_id] if schedule_id else []
         pending_jes = JobExecution.select_active_by_kwargs(
+            *clauses,
             project_code=tg.project_code,
             job_key=RUN_PROFILE_JOB_KEY,
             kwargs_match={"table_group_id": str(tg.id)},
             statuses=statuses,
         )
 
     doc = MdDoc()
-    scope = f" — status `{status}`" if status else ""
+    scope_parts = []
+    if schedule_id:
+        scope_parts.append(f"schedule `{schedule_id}`")
+    if status:
+        scope_parts.append(f"status `{status}`")
+    scope = f" — {', '.join(scope_parts)}" if scope_parts else ""
     doc.heading(1, f"Profiling runs for `{tg.table_groups_name}`{scope}")
 
     next_run = next_scheduled_run(

testgen/mcp/tools/test_runs.py

@@ -30,6 +30,7 @@ def list_test_runs(
     project_code: str | None = None,
     test_suite: str | None = None,
     table_group_id: str | None = None,
+    schedule_id: str | None = None,
     status: str | None = None,
     limit: int = 10,
     page: int = 1,
@@ -43,6 +44,8 @@ def list_test_runs(
         test_suite: Optional test suite name to filter by (case-sensitive).
         table_group_id: Optional UUID of a table group, e.g. from `get_data_inventory`. Returns
             runs for any suite in the group.
+        schedule_id: Optional UUID of a schedule, e.g. from `list_schedules`. Returns only runs
+            triggered by that schedule.
         status: Optional run status filter. One of: Pending, Running, Completed, Canceled, Error.
         limit: Page size (default 10, max 100).
         page: Page number starting at 1 (default 1).
@@ -51,6 +54,8 @@ def list_test_runs(
     validate_page(page)
 
     statuses = parse_run_status_filter(status) if status else None
+    if schedule_id:
+        parse_uuid(schedule_id, "schedule_id")
 
     if not project_code and not table_group_id:
         raise MCPUserError("Provide either `project_code` or `table_group_id`.")
@@ -86,6 +91,7 @@ def list_test_runs(
         project_code=project_code,
         table_group_id=str(table_group.id) if table_group else None,
         test_suite_id=test_suite_id,
+        schedule_id=schedule_id,
         statuses=statuses,
         page=page,
         page_size=limit,
@@ -99,10 +105,11 @@ def list_test_runs(
             project_code=project_code,
             test_suite_id=test_suite_id,
             table_group_id=str(table_group.id) if table_group else None,
+            schedule_id=schedule_id,
             statuses=statuses,
         )
 
-    scope_descriptor = _scope_descriptor(project_code, test_suite, table_group_id, status)
+    scope_descriptor = _scope_descriptor(project_code, test_suite, table_group_id, schedule_id, status)
     doc = MdDoc()
     doc.heading(1, f"Test runs{scope_descriptor}")
 
@@ -199,6 +206,7 @@ def _scope_descriptor(
     project_code: str | None,
     test_suite: str | None,
     table_group_id: str | None,
+    schedule_id: str | None,
     status: str | None,
 ) -> str:
     parts: list[str] = []
@@ -208,6 +216,8 @@ def _scope_descriptor(
         parts.append(f"suite `{test_suite}`")
     if table_group_id:
         parts.append(f"table group `{table_group_id}`")
+    if schedule_id:
+        parts.append(f"schedule `{schedule_id}`")
     if status:
         parts.append(f"status `{status}`")
     return f" — {', '.join(parts)}" if parts else ""
@@ -246,6 +256,7 @@ def _select_pending_test_jes(
     project_code: str,
     test_suite_id: str | None,
     table_group_id: str | None,
+    schedule_id: str | None,
     statuses,
 ) -> list[JobExecution]:
     """Find queued/in-flight test-run JEs for a given suite or table group scope. For a
@@ -267,7 +278,9 @@ def _select_pending_test_jes(
             return []
     else:
         return []
+    clauses = [JobExecution.job_schedule_id == schedule_id] if schedule_id else []
     return JobExecution.select_active_by_kwargs(
+        *clauses,
         project_code=project_code,
         job_key=RUN_TESTS_JOB_KEY,
         kwargs_match={"test_suite_id": suite_ids},

tests/unit/mcp/test_tools_profiling.py

@@ -583,6 +583,53 @@ def test_list_profiling_runs_invalid_status(mock_tg_cls, mock_run_cls, mock_next
         list_profiling_runs(table_group_id=str(uuid4()), status="Bogus")
 
 
+@patch("testgen.mcp.tools.profiling.JobExecution")
+@patch("testgen.mcp.tools.profiling.next_scheduled_run", return_value=None)
+@patch("testgen.mcp.tools.profiling.ProfilingRun")
+@patch("testgen.mcp.tools.common.TableGroup")
+def test_list_profiling_runs_schedule_filter(mock_tg_cls, mock_run_cls, mock_next, mock_je, db_session_mock):
+    mock_je.select_active_by_kwargs.return_value = []
+    mock_tg_cls.get.return_value = _mock_table_group()
+    mock_run_cls.select_summary.return_value = ([], 0)
+    schedule_id = str(uuid4())
+
+    from testgen.mcp.tools.profiling import list_profiling_runs
+    list_profiling_runs(table_group_id=str(uuid4()), schedule_id=schedule_id, status="Completed")
+
+    call_kwargs = mock_run_cls.select_summary.call_args.kwargs
+    assert call_kwargs["schedule_id"] == schedule_id
+    assert call_kwargs["statuses"] == [JobStatus.COMPLETED]
+    # The schedule clause is forwarded to the pending-JE query (positional *clauses arg).
+    assert mock_je.select_active_by_kwargs.call_args.args
+
+
+@patch("testgen.mcp.tools.profiling.JobExecution")
+@patch("testgen.mcp.tools.profiling.next_scheduled_run", return_value=None)
+@patch("testgen.mcp.tools.profiling.ProfilingRun")
+@patch("testgen.mcp.tools.common.TableGroup")
+def test_list_profiling_runs_unknown_schedule_returns_empty_envelope(mock_tg_cls, mock_run_cls, mock_next, mock_je, db_session_mock):
+    mock_je.select_active_by_kwargs.return_value = []
+    mock_tg_cls.get.return_value = _mock_table_group()
+    mock_run_cls.select_summary.return_value = ([], 0)
+
+    from testgen.mcp.tools.profiling import list_profiling_runs
+    result = list_profiling_runs(table_group_id=str(uuid4()), schedule_id=str(uuid4()))
+
+    assert "No profiling runs" in result
+
+
+@patch("testgen.mcp.tools.profiling.next_scheduled_run", return_value=None)
+@patch("testgen.mcp.tools.profiling.ProfilingRun")
+@patch("testgen.mcp.tools.common.TableGroup")
+def test_list_profiling_runs_malformed_schedule_raises(mock_tg_cls, mock_run_cls, mock_next, db_session_mock):
+    mock_tg_cls.get.return_value = _mock_table_group()
+
+    from testgen.mcp.tools.profiling import list_profiling_runs
+    with pytest.raises(MCPUserError):
+        list_profiling_runs(table_group_id=str(uuid4()), schedule_id="not-a-uuid")
+    mock_run_cls.select_summary.assert_not_called()
+
+
 # ----------------------------------------------------------------------
 # get_profiling_run
 # ----------------------------------------------------------------------

tests/unit/mcp/test_tools_test_runs.py

@@ -43,6 +43,7 @@ def test_list_test_runs_default(mock_suite, mock_run, mock_next, db_session_mock
         project_code="demo",
         table_group_id=None,
         test_suite_id=None,
+        schedule_id=None,
         statuses=None,
         page=1,
         page_size=10,
@@ -67,6 +68,66 @@ def test_list_test_runs_with_status_filter(mock_suite, mock_run, mock_next, db_s
     assert call_kwargs["statuses"] == [JobStatus.PENDING, JobStatus.CLAIMED]
 
 
+@patch("testgen.mcp.tools.test_runs.next_scheduled_run", return_value=None)
+@patch("testgen.mcp.tools.test_runs.TestRun")
+@patch("testgen.mcp.tools.test_runs.TestSuite")
+def test_list_test_runs_with_schedule_filter(mock_suite, mock_run, mock_next, db_session_mock):
+    mock_run.select_summary.return_value = ([], 0)
+    schedule_id = str(uuid4())
+
+    from testgen.mcp.tools.test_runs import list_test_runs
+
+    list_test_runs(project_code="demo", schedule_id=schedule_id, status="Completed")
+
+    call_kwargs = mock_run.select_summary.call_args.kwargs
+    assert call_kwargs["schedule_id"] == schedule_id
+    assert call_kwargs["statuses"] == [JobStatus.COMPLETED]
+
+
+@patch("testgen.mcp.tools.test_runs.next_scheduled_run", return_value=None)
+@patch("testgen.mcp.tools.test_runs.TestRun")
+@patch("testgen.mcp.tools.test_runs.TestSuite")
+def test_list_test_runs_unknown_schedule_returns_empty_envelope(mock_suite, mock_run, mock_next, db_session_mock):
+    # Unknown/inaccessible/wrong-kind schedule yields no rows — standard empty envelope, not an error.
+    mock_run.select_summary.return_value = ([], 0)
+
+    from testgen.mcp.tools.test_runs import list_test_runs
+
+    result = list_test_runs(project_code="demo", schedule_id=str(uuid4()))
+
+    assert "No test runs" in result
+
+
+@patch("testgen.mcp.tools.test_runs.next_scheduled_run", return_value=None)
+@patch("testgen.mcp.tools.test_runs.TestRun")
+@patch("testgen.mcp.tools.test_runs.TestSuite")
+def test_list_test_runs_malformed_schedule_raises(mock_suite, mock_run, mock_next, db_session_mock):
+    from testgen.mcp.tools.test_runs import list_test_runs
+
+    with pytest.raises(MCPUserError):
+        list_test_runs(project_code="demo", schedule_id="not-a-uuid")
+    mock_run.select_summary.assert_not_called()
+
+
+@patch("testgen.mcp.tools.test_runs.JobExecution")
+@patch("testgen.mcp.tools.test_runs.next_scheduled_run", return_value=None)
+@patch("testgen.mcp.tools.test_runs.TestRun")
+@patch("testgen.mcp.tools.test_runs.TestSuite")
+def test_list_test_runs_schedule_filters_pending(mock_suite, mock_run, mock_next, mock_je, db_session_mock):
+    mock_je.select_active_by_kwargs.return_value = []
+    suite_id = uuid4()
+    mock_suite.select_minimal_where.return_value = [MagicMock(id=suite_id)]
+    mock_run.select_summary.return_value = ([], 0)
+    schedule_id = str(uuid4())
+
+    from testgen.mcp.tools.test_runs import list_test_runs
+
+    list_test_runs(project_code="demo", test_suite="Quality", schedule_id=schedule_id)
+
+    # A schedule clause is forwarded to the pending-JE query (positional *clauses arg).
+    assert mock_je.select_active_by_kwargs.call_args.args
+
+
 @patch("testgen.mcp.tools.test_runs.JobExecution")
 @patch("testgen.mcp.tools.test_runs.next_scheduled_run", return_value=None)
 @patch("testgen.mcp.tools.test_runs.TestRun")