From f4628fc719b58b4e97c91b1cea8d5d1c5fc03d69 Mon Sep 17 00:00:00 2001 From: Luis Date: Thu, 28 May 2026 16:17:51 -0400 Subject: [PATCH 01/78] feat(mcp): add tools to manage database connections - create_connection - update_connection - test_connection --- testgen/common/database/connection_service.py | 145 ++++ testgen/common/flavors.py | 71 ++ testgen/mcp/server.py | 6 + testgen/mcp/services/inventory_service.py | 4 +- testgen/mcp/tools/common.py | 447 ++++++++++++ testgen/mcp/tools/connections.py | 235 +++++++ testgen/mcp/tools/reference.py | 135 +++- testgen/ui/views/connections.py | 170 +---- tests/unit/common/test_connection_service.py | 280 ++++++++ tests/unit/common/test_flavors.py | 35 + tests/unit/mcp/conftest.py | 1 + tests/unit/mcp/test_connection_schema.py | 649 ++++++++++++++++++ tests/unit/mcp/test_tools_common.py | 31 + tests/unit/mcp/test_tools_connections.py | 233 +++++++ tests/unit/mcp/test_tools_reference.py | 91 +++ 15 files changed, 2391 insertions(+), 142 deletions(-) create mode 100644 testgen/common/database/connection_service.py create mode 100644 testgen/common/flavors.py create mode 100644 testgen/mcp/tools/connections.py create mode 100644 tests/unit/common/test_connection_service.py create mode 100644 tests/unit/common/test_flavors.py create mode 100644 tests/unit/mcp/test_connection_schema.py create mode 100644 tests/unit/mcp/test_tools_connections.py diff --git a/testgen/common/database/connection_service.py b/testgen/common/database/connection_service.py new file mode 100644 index 00000000..c7cec344 --- /dev/null +++ b/testgen/common/database/connection_service.py @@ -0,0 +1,145 @@ +"""Shared connection domain helpers — moved out of the Streamlit connections page +so MCP and any future caller can reuse identical test/normalize logic without +forking SQL templates or per-flavor rules. + +Two pieces live here, both flavor-aware and both operating on a ``Connection`` by +column / flag (no user-facing labels — those are presentation): + +* ``ConnectionStatus`` + ``test_connection_status`` — open an external connection + and report success / failure with driver-text details. +* ``normalize_auth_fields`` — pre-save scrub of mutually-exclusive credential + columns so flipping auth modes doesn't leave orphan values behind. +* ``apply_connection_defaults`` — pre-save fill of flavor-dependent defaults the + caller didn't supply. + +The MCP-facing connection-parameter contract (field labels, modes, requirement +rules, validation) is presentation/input vocabulary and lives in +``mcp/tools/common.py``; flavor display labels live in ``common/flavors.py``. +""" + +from __future__ import annotations + +import logging +import random +from dataclasses import dataclass, field + +from sqlalchemy.exc import DatabaseError, DBAPIError + +from testgen.common.database.database_service import empty_cache, get_flavor_service +from testgen.common.models.connection import Connection +from testgen.ui.services.database_service import fetch_from_target_db + +try: + from pyodbc import Error as PyODBCError +except ImportError: + PyODBCError = None + +LOG = logging.getLogger("testgen") + + +# --------------------------------------------------------------------------- +# ConnectionStatus + test runner +# --------------------------------------------------------------------------- + + +@dataclass(frozen=True, slots=True) +class ConnectionStatus: + message: str + successful: bool + details: str | None = field(default=None) + # Streamlit's reactive system suppresses re-renders when a frozen dataclass + # compares equal to its previous value. A random per-instance field forces + # ``__eq__`` to differ so repeated "Test Connection" clicks producing the + # same error message still re-render the alert. This field has no consumer + # outside the UI re-render path; do NOT delete it as cleanup. + _: float = field(default_factory=random.random) + + +def is_open_ssl_error(error: Exception) -> bool: + return ( + bool(error.args) + and len(error.args) > 1 + and isinstance(error.args[1], list) + and len(error.args[1]) > 0 + and type(error.args[1][0]).__name__ == "OpenSSLError" + ) + + +def test_connection_status(connection: Connection) -> ConnectionStatus: + """Open the connection, run the flavor's smoke query, classify the outcome.""" + # Drop pooled engines so the next checkout reflects any credential changes. + empty_cache() + try: + flavor_service = get_flavor_service(connection.sql_flavor) + results = fetch_from_target_db(connection, flavor_service.test_query) + connection_successful = len(results) == 1 and next(iter(results[0].values())) == 1 + + if not connection_successful: + return ConnectionStatus(message="Error completing a query to the database server.", successful=False) + return ConnectionStatus(message="The connection was successful.", successful=True) + except KeyError: + return ConnectionStatus( + message="Error attempting the connection. ", + details="Complete all the required fields.", + successful=False, + ) + except DatabaseError as error: + LOG.exception("Error testing database connection") + return ConnectionStatus(message="Error attempting the connection.", details=str(error.orig), successful=False) + except DBAPIError as error: + LOG.exception("Error testing database connection") + details = str(error.orig) + if PyODBCError and isinstance(error.orig, PyODBCError) and error.orig.args: + details = error.orig.args[1] + return ConnectionStatus(message="Error attempting the connection.", details=details, successful=False) + except (TypeError, ValueError) as error: + LOG.exception("Error testing database connection") + details = str(error) + if is_open_ssl_error(error): + details = error.args[0] + return ConnectionStatus(message="Error attempting the connection.", details=details, successful=False) + except Exception: + details = "Something went wrong while testing the connection." + if connection.connect_by_key and not connection.private_key: + details = "The private key is missing." + LOG.exception("Error testing database connection") + return ConnectionStatus(message="Error attempting the connection.", details=details, successful=False) + + +# --------------------------------------------------------------------------- +# Auth-path normalization +# --------------------------------------------------------------------------- + + +def normalize_auth_fields(connection: Connection) -> None: + """Clear credential columns the active auth mode doesn't use. + + Mirrors the UI's pre-save scrub at ``ui/views/connections.py`` so flipping + auth modes doesn't leave stale opposite-mode values behind. Databricks OAuth + stores the client secret in ``project_pw_encrypted`` despite + ``connect_by_key=True``, so it sticks with the password path. + """ + code = connection.sql_flavor_code + + uses_private_key = bool(connection.connect_by_key) and code != "databricks" + if uses_private_key: + connection.project_pw_encrypted = None + else: + connection.private_key = None + connection.private_key_passphrase = None + + if connection.connect_with_identity: + connection.project_user = None + connection.project_pw_encrypted = None + + +# --------------------------------------------------------------------------- +# Flavor-dependent defaults +# --------------------------------------------------------------------------- + + +def apply_connection_defaults(connection: Connection) -> None: + """Fill flavor-dependent defaults for fields the caller didn't supply.""" + if connection.max_query_chars is None: + # Salesforce Data 360's Hyper engine has a lower expression-depth limit + connection.max_query_chars = 15000 if connection.sql_flavor_code == "salesforce_data360" else 20000 diff --git a/testgen/common/flavors.py b/testgen/common/flavors.py new file mode 100644 index 00000000..277bc7da --- /dev/null +++ b/testgen/common/flavors.py @@ -0,0 +1,71 @@ +"""Flavor identity — the single source of truth for database-flavor display labels +and the code→family mapping. + +Flavor *type* labels (e.g. "PostgreSQL", "Snowflake") are shared presentation: the +Streamlit connections page feeds them to its JS component via ``FLAVOR_OPTIONS`` and +the MCP tools accept/return them as the ``sql_flavor`` value. Both surfaces derive +from here so the labels are written exactly once. + +This module is intentionally pure data (no I/O, no Streamlit, no SQLAlchemy) so it +can be imported from the UI, the MCP layer, and common services without cycles. + +(The per-flavor *connection-parameter* field labels and requirement rules are a +different, MCP-only concern and live in ``mcp/tools/common.py`` — they're hardcoded +in the JS form, not sourced from Python.) +""" + +from __future__ import annotations + +from enum import StrEnum + + +class SqlFlavorLabel(StrEnum): + """User-facing database-flavor labels (the ``sql_flavor`` value space).""" + + REDSHIFT = "Amazon Redshift" + REDSHIFT_SPECTRUM = "Amazon Redshift Spectrum" + AZURE_MSSQL = "Azure SQL Database" + SYNAPSE_MSSQL = "Azure Synapse Analytics" + DATABRICKS = "Databricks" + BIGQUERY = "Google BigQuery" + MSSQL = "Microsoft SQL Server" + ORACLE = "Oracle" + POSTGRESQL = "PostgreSQL" + SAP_HANA = "SAP HANA" + SALESFORCE_DATA360 = "Salesforce Data 360" + SNOWFLAKE = "Snowflake" + + +# code → label. Label strings are written once (on the enum); this only pairs them +# with flavor codes. +FLAVOR_CODE_TO_LABEL: dict[str, SqlFlavorLabel] = { + "redshift": SqlFlavorLabel.REDSHIFT, + "redshift_spectrum": SqlFlavorLabel.REDSHIFT_SPECTRUM, + "azure_mssql": SqlFlavorLabel.AZURE_MSSQL, + "synapse_mssql": SqlFlavorLabel.SYNAPSE_MSSQL, + "databricks": SqlFlavorLabel.DATABRICKS, + "bigquery": SqlFlavorLabel.BIGQUERY, + "mssql": SqlFlavorLabel.MSSQL, + "oracle": SqlFlavorLabel.ORACLE, + "postgresql": SqlFlavorLabel.POSTGRESQL, + "sap_hana": SqlFlavorLabel.SAP_HANA, + "salesforce_data360": SqlFlavorLabel.SALESFORCE_DATA360, + "snowflake": SqlFlavorLabel.SNOWFLAKE, +} + +# code → family (multiple codes can share one engine family, e.g. the two Azure +# variants both map to "mssql"). +FLAVOR_CODE_TO_FAMILY: dict[str, str] = { + "redshift": "redshift", + "redshift_spectrum": "redshift_spectrum", + "azure_mssql": "mssql", + "synapse_mssql": "mssql", + "databricks": "databricks", + "bigquery": "bigquery", + "mssql": "mssql", + "oracle": "oracle", + "postgresql": "postgresql", + "sap_hana": "sap_hana", + "salesforce_data360": "salesforce_data360", + "snowflake": "snowflake", +} diff --git a/testgen/mcp/server.py b/testgen/mcp/server.py index 5e91b2f1..46683db5 100644 --- a/testgen/mcp/server.py +++ b/testgen/mcp/server.py @@ -137,6 +137,7 @@ def build_mcp_server( profiling_overview, table_health, ) + from testgen.mcp.tools.connections import test_connection from testgen.mcp.tools.discovery import get_data_inventory, list_projects, list_tables, list_test_suites from testgen.mcp.tools.execution import ( cancel_profiling_run, @@ -184,6 +185,8 @@ def build_mcp_server( ) from testgen.mcp.tools.reference import ( column_profile_fields_resource, + connection_parameters_index_resource, + connection_parameters_resource, get_test_type, glossary_resource, hygiene_issue_types_resource, @@ -310,12 +313,15 @@ def safe_prompt(fn): safe_tool(create_notification) safe_tool(update_notification) safe_tool(delete_notification) + safe_tool(test_connection) # Resources safe_resource("testgen://test-types", test_types_resource) safe_resource("testgen://hygiene-issue-types", hygiene_issue_types_resource) safe_resource("testgen://column-profile-fields", column_profile_fields_resource) safe_resource("testgen://glossary", glossary_resource) + safe_resource("testgen://connection-parameters", connection_parameters_index_resource) + safe_resource("testgen://connection-parameters/{flavor}", connection_parameters_resource) # Prompts safe_prompt(health_check) diff --git a/testgen/mcp/services/inventory_service.py b/testgen/mcp/services/inventory_service.py index cee1c7a6..6d504ed3 100644 --- a/testgen/mcp/services/inventory_service.py +++ b/testgen/mcp/services/inventory_service.py @@ -117,9 +117,9 @@ def get_inventory( lines.append("_No table groups._\n") continue - for _conn_id, conn in proj["connections"].items(): + for conn_id, conn in proj["connections"].items(): if can_view: - lines.append(f"### Connection: {conn['name']}\n") + lines.append(f"### Connection: {conn['name']} (id: `{conn_id}`)\n") if not conn["groups"]: if can_view: diff --git a/testgen/mcp/tools/common.py b/testgen/mcp/tools/common.py index e08ca861..990a759d 100644 --- a/testgen/mcp/tools/common.py +++ b/testgen/mcp/tools/common.py @@ -1,3 +1,4 @@ +from dataclasses import dataclass, field from datetime import date, datetime from enum import StrEnum from uuid import UUID @@ -6,7 +7,9 @@ from testgen.common.date_service import parse_since from testgen.common.enums import Disposition, ImpactDimension, IssueLikelihood, JobStatus, PiiRisk, QualityDimension +from testgen.common.flavors import FLAVOR_CODE_TO_FAMILY, FLAVOR_CODE_TO_LABEL, SqlFlavorLabel from testgen.common.models import get_current_session +from testgen.common.models.connection import Connection from testgen.common.models.data_column import ( GENERAL_TYPE_TO_CODE, ColumnOrderBy, @@ -724,3 +727,447 @@ def format_notification_trigger(event: NotificationEvent | str, settings: dict | if event_enum is NotificationEvent.monitor_run: return _MONITOR_TRIGGER_INTERNAL_TO_LABEL[MonitorNotificationTrigger(raw)].value return None + + +def resolve_connection(connection_id: int) -> Connection: + """Resolve a connection ID, collapsing missing-or-inaccessible into one error path.""" + perms = get_project_permissions() + conn = Connection.get( + connection_id, + Connection.project_code.in_(perms.allowed_codes), + ) + if conn is None: + raise MCPResourceNotAccessible("Connection", str(connection_id)) + return conn + + +# Flavor display labels are the single source of truth in ``common/flavors.py`` +# (shared with the UI page). These maps just re-shape them for the MCP layer. +SQL_FLAVOR_CODE_TO_LABEL: dict[str, SqlFlavorLabel] = dict(FLAVOR_CODE_TO_LABEL) +SQL_FLAVOR_LABEL_TO_CODE: dict[SqlFlavorLabel, str] = { + label: code for code, label in FLAVOR_CODE_TO_LABEL.items() +} +SQL_FLAVOR_CODE_TO_FAMILY: dict[str, str] = dict(FLAVOR_CODE_TO_FAMILY) + + +def parse_sql_flavor(value: str) -> tuple[SqlFlavorLabel, str, str]: + """Validate a user-facing ``sql_flavor`` value and return ``(label, code, family)``.""" + try: + label = SqlFlavorLabel(value) + except ValueError as err: + valid = ", ".join(f.value for f in SqlFlavorLabel) + raise MCPUserError(f"Invalid sql_flavor `{value}`. Valid values: {valid}") from err + code = SQL_FLAVOR_LABEL_TO_CODE[label] + return label, code, SQL_FLAVOR_CODE_TO_FAMILY[code] + + +# =========================================================================== +# Connection-parameter contract (MCP input vocabulary) +# +# The per-flavor connection shape: which auth modes exist and which fields each +# needs, keyed by their UI labels (which double as ``connection_params`` keys). +# This is MCP-only input vocabulary + parsing — it mirrors the per-flavor JS +# forms in ``ui/static/js/components/connection_form.js`` (NOT sourced from +# Python), and drives both the connection tools and the +# ``testgen://connection-parameters/{flavor}`` resource. Field labels map to the +# (often leaky) ``Connection`` columns here so the tool's arg surface speaks the +# target-DB vocabulary. +# =========================================================================== + + +class ConnectionMode(StrEnum): + PASSWORD = "Password" # noqa: S105 — auth-mode label, not a credential + KEY_PAIR = "Key-Pair" + MANAGED_IDENTITY = "Managed Identity" + ACCESS_TOKEN = "Access Token" # noqa: S105 — auth-mode label, not a credential + SERVICE_PRINCIPAL = "Service Principal (OAuth)" + JWT_BEARER = "JWT Bearer Flow" + CLIENT_CREDENTIALS = "Client Credentials Flow" + SERVICE_ACCOUNT = "Service Account Key" + + +class Req(StrEnum): + REQUIRED = "required" # always required in this mode + REQUIRED_UNLESS_URL = "required_unless_url" # required only in host mode + OPTIONAL = "optional" + + +@dataclass(frozen=True) +class ConnField: + label: str # exact UI label == connection_params key + column: str # Connection ORM attribute it maps to + requirement: Req + secret: bool = False + + +@dataclass(frozen=True) +class FlavorMode: + mode: ConnectionMode | None # None for single-auth flavors + fields: tuple[ConnField, ...] # excludes the URL field + supports_url: bool # whether the URL alternative is offered + # Columns forced regardless of supplied params (auth flags, Databricks PAT user). + sets: dict[str, object] = field(default_factory=dict) + + +@dataclass(frozen=True) +class FlavorSchema: + code: str + label: str + modes: tuple[FlavorMode, ...] + url_field: ConnField | None # the shared URL field when any mode supports_url + + +# -- reusable field definitions --------------------------------------------- + +_HOST = ConnField("Host", "project_host", Req.REQUIRED_UNLESS_URL) +_PORT = ConnField("Port", "project_port", Req.REQUIRED_UNLESS_URL) +_DATABASE = ConnField("Database", "project_db", Req.REQUIRED_UNLESS_URL) +_SERVICE_NAME = ConnField("Service Name", "project_db", Req.REQUIRED_UNLESS_URL) +_USERNAME = ConnField("Username", "project_user", Req.REQUIRED) +_PASSWORD_OPT = ConnField("Password", "project_pw_encrypted", Req.OPTIONAL, secret=True) +_PASSWORD_REQ = ConnField("Password", "project_pw_encrypted", Req.REQUIRED, secret=True) +_WAREHOUSE = ConnField("Warehouse", "warehouse", Req.OPTIONAL) +_PRIVATE_KEY = ConnField("Private Key", "private_key", Req.REQUIRED, secret=True) +_PRIVATE_KEY_PASSPHRASE = ConnField("Private Key Passphrase", "private_key_passphrase", Req.OPTIONAL, secret=True) +_URL = ConnField("URL", "url", Req.REQUIRED) + +# Databricks +_HTTP_PATH_RU = ConnField("HTTP Path", "http_path", Req.REQUIRED_UNLESS_URL) +_HTTP_PATH_REQ = ConnField("HTTP Path", "http_path", Req.REQUIRED) +_HOST_REQ = ConnField("Host", "project_host", Req.REQUIRED) +_PORT_REQ = ConnField("Port", "project_port", Req.REQUIRED) +_CATALOG = ConnField("Catalog", "project_db", Req.REQUIRED_UNLESS_URL) +_CATALOG_REQ = ConnField("Catalog", "project_db", Req.REQUIRED) +_ACCESS_TOKEN = ConnField("Access Token", "project_pw_encrypted", Req.REQUIRED, secret=True) +_CLIENT_ID = ConnField("Client ID", "project_user", Req.REQUIRED) +_CLIENT_SECRET = ConnField("Client Secret", "project_pw_encrypted", Req.REQUIRED, secret=True) + +# BigQuery / Salesforce +_SERVICE_ACCOUNT_KEY = ConnField("Service Account Key", "service_account_key", Req.REQUIRED, secret=True) +_LOGIN_URL = ConnField("Login URL", "project_host", Req.REQUIRED) +_CONSUMER_KEY = ConnField("Consumer Key", "project_user", Req.REQUIRED) +_SF_USERNAME = ConnField("Username", "project_db", Req.REQUIRED) +_CONSUMER_SECRET = ConnField("Consumer Secret", "project_pw_encrypted", Req.REQUIRED, secret=True) + + +def _host_auth_schema(code: str, *, db_field: ConnField = _DATABASE) -> FlavorSchema: + """Single-mode host/URL flavors (PostgreSQL, Redshift, MSSQL, Oracle, SAP HANA).""" + return FlavorSchema( + code=code, + label=FLAVOR_CODE_TO_LABEL[code], + modes=( + FlavorMode(mode=None, fields=(_HOST, _PORT, db_field, _USERNAME, _PASSWORD_OPT), supports_url=True), + ), + url_field=_URL, + ) + + +def _azure_schema(code: str) -> FlavorSchema: + return FlavorSchema( + code=code, + label=FLAVOR_CODE_TO_LABEL[code], + modes=( + FlavorMode( + mode=ConnectionMode.PASSWORD, + fields=(_HOST, _PORT, _DATABASE, _USERNAME, _PASSWORD_OPT), + supports_url=True, + sets={"connect_with_identity": False}, + ), + FlavorMode( + mode=ConnectionMode.MANAGED_IDENTITY, + fields=(_HOST, _PORT, _DATABASE), + supports_url=True, + sets={"connect_with_identity": True}, + ), + ), + url_field=_URL, + ) + + +FLAVOR_CONNECTION_SCHEMA: dict[str, FlavorSchema] = { + "postgresql": _host_auth_schema("postgresql"), + "redshift": _host_auth_schema("redshift"), + "redshift_spectrum": _host_auth_schema("redshift_spectrum"), + "mssql": _host_auth_schema("mssql"), + "oracle": _host_auth_schema("oracle", db_field=_SERVICE_NAME), + "sap_hana": _host_auth_schema("sap_hana"), + "azure_mssql": _azure_schema("azure_mssql"), + "synapse_mssql": _azure_schema("synapse_mssql"), + "snowflake": FlavorSchema( + code="snowflake", + label=FLAVOR_CODE_TO_LABEL["snowflake"], + modes=( + FlavorMode( + mode=ConnectionMode.KEY_PAIR, + fields=(_HOST, _PORT, _DATABASE, _USERNAME, _WAREHOUSE, _PRIVATE_KEY, _PRIVATE_KEY_PASSPHRASE), + supports_url=True, + sets={"connect_by_key": True}, + ), + FlavorMode( + mode=ConnectionMode.PASSWORD, + fields=(_HOST, _PORT, _DATABASE, _USERNAME, _WAREHOUSE, _PASSWORD_REQ), + supports_url=True, + sets={"connect_by_key": False}, + ), + ), + url_field=_URL, + ), + "databricks": FlavorSchema( + code="databricks", + label=FLAVOR_CODE_TO_LABEL["databricks"], + modes=( + FlavorMode( + mode=ConnectionMode.ACCESS_TOKEN, + fields=(_HOST, _PORT, _HTTP_PATH_RU, _CATALOG, _ACCESS_TOKEN), + supports_url=True, + # PAT auth: the username is always the literal 'token'. + sets={"connect_by_key": False, "project_user": "token"}, + ), + FlavorMode( + mode=ConnectionMode.SERVICE_PRINCIPAL, + fields=(_HOST_REQ, _PORT_REQ, _HTTP_PATH_REQ, _CATALOG_REQ, _CLIENT_ID, _CLIENT_SECRET), + supports_url=False, + sets={"connect_by_key": True}, + ), + ), + url_field=_URL, + ), + "bigquery": FlavorSchema( + code="bigquery", + label=FLAVOR_CODE_TO_LABEL["bigquery"], + modes=(FlavorMode(mode=None, fields=(_SERVICE_ACCOUNT_KEY,), supports_url=False),), + url_field=None, + ), + "salesforce_data360": FlavorSchema( + code="salesforce_data360", + label=FLAVOR_CODE_TO_LABEL["salesforce_data360"], + modes=( + FlavorMode( + mode=ConnectionMode.JWT_BEARER, + fields=(_LOGIN_URL, _CONSUMER_KEY, _SF_USERNAME, _PRIVATE_KEY), + supports_url=False, + sets={"connect_by_key": True}, + ), + FlavorMode( + mode=ConnectionMode.CLIENT_CREDENTIALS, + fields=(_LOGIN_URL, _CONSUMER_KEY, _CONSUMER_SECRET), + supports_url=False, + sets={"connect_by_key": False}, + ), + ), + url_field=None, + ), +} + + +def schema_for(code: str) -> FlavorSchema: + """Return the connection schema for a flavor code. Raises ``KeyError`` if unknown.""" + return FLAVOR_CONNECTION_SCHEMA[code] + + +def resolve_mode(code: str, mode_label: str | None) -> FlavorMode: + """Resolve the ``FlavorMode`` for a flavor + supplied ``connection_mode`` label. + + Single-mode flavors take no ``connection_mode`` (passing one is an error). + Multi-mode flavors require a valid one. + """ + schema = schema_for(code) + if len(schema.modes) == 1 and schema.modes[0].mode is None: + if mode_label is not None: + raise MCPUserError(f"{schema.label} does not take a connection_mode.") + return schema.modes[0] + + valid = [str(m.mode) for m in schema.modes if m.mode is not None] + if mode_label is None: + raise MCPUserError(f"{schema.label} requires a connection_mode. Valid values: {', '.join(valid)}.") + for fmode in schema.modes: + if fmode.mode is not None and str(fmode.mode) == mode_label: + return fmode + raise MCPUserError( + f"Invalid connection_mode `{mode_label}` for {schema.label}. Valid values: {', '.join(valid)}." + ) + + +def infer_mode(connection: Connection) -> ConnectionMode | None: + """Reverse of a mode's ``sets`` flags — derive the active mode from a stored + connection so update/validation can pick the right field set without the + caller re-supplying ``connection_mode``. + """ + code = connection.sql_flavor_code + schema = FLAVOR_CONNECTION_SCHEMA.get(code) + if schema is None or (len(schema.modes) == 1 and schema.modes[0].mode is None): + return None + + if code in {"azure_mssql", "synapse_mssql"}: + return ConnectionMode.MANAGED_IDENTITY if connection.connect_with_identity else ConnectionMode.PASSWORD + if code == "snowflake": + return ConnectionMode.KEY_PAIR if connection.connect_by_key else ConnectionMode.PASSWORD + if code == "databricks": + return ConnectionMode.SERVICE_PRINCIPAL if connection.connect_by_key else ConnectionMode.ACCESS_TOKEN + if code == "salesforce_data360": + return ConnectionMode.JWT_BEARER if connection.connect_by_key else ConnectionMode.CLIENT_CREDENTIALS + return None + + +def _mode_for_connection(connection: Connection) -> FlavorMode: + """Pick the FlavorMode matching a connection's current flags (for validation).""" + schema = schema_for(connection.sql_flavor_code) + if len(schema.modes) == 1: + return schema.modes[0] + active = infer_mode(connection) + for fmode in schema.modes: + if fmode.mode == active: + return fmode + return schema.modes[0] + + +def connection_display_fields(connection: Connection) -> list[ConnField]: + """Active-mode fields (plus the URL field when in URL mode), in schema order. + + For rendering a connection back to the user with the flavor-specific UI label + for each populated column (e.g. ``Catalog`` for Databricks, ``Login URL`` for + Salesforce). Callers skip secrets and empty values. + """ + schema = schema_for(connection.sql_flavor_code) + fields = list(_mode_for_connection(connection).fields) + if schema.url_field is not None and getattr(connection, "connect_by_url", False): + fields.append(schema.url_field) + return fields + + +def connection_field_labels(connection: Connection) -> dict[str, str]: + """Map each ``Connection`` column to its flavor/mode-specific UI label. + + Used to label diff output. Columns outside the active mode fall back to the + caller's generic labels. + """ + schema = schema_for(connection.sql_flavor_code) + labels = {fld.column: fld.label for fld in _mode_for_connection(connection).fields} + if schema.url_field is not None: + labels[schema.url_field.column] = schema.url_field.label + return labels + + +def apply_connection_params( + connection: Connection, + code: str, + mode_label: str | None, + params: dict[str, object], +) -> None: + """Map a label-keyed ``connection_params`` dict onto a ``Connection``. + + * Resolves the mode and applies its forced ``sets`` columns (auth flags, + Databricks PAT ``project_user='token'``). + * Maps each supplied label to its ``Connection`` column (casting ``Port``). + * Toggles ``connect_by_url`` from the presence of ``URL`` vs the host-group + fields, rejecting an ambiguous mix. + + Raises ``MCPUserError`` on unknown keys, an unsupported / conflicting ``URL``, + or an invalid / missing mode. + """ + fmode = resolve_mode(code, mode_label) + fields_by_label = {f.label: f for f in fmode.fields} + url_fields = {f.label for f in fmode.fields if f.requirement is Req.REQUIRED_UNLESS_URL} + + valid_keys = set(fields_by_label) + if fmode.supports_url: + valid_keys.add(_URL.label) + unknown = [key for key in params if key not in valid_keys] + if unknown: + raise MCPUserError( + f"Unknown connection_params for {schema_for(code).label}: {', '.join(sorted(unknown))}. " + f"Allowed: {', '.join(sorted(valid_keys))}." + ) + + has_url = _URL.label in params + provided_url_fields = [key for key in params if key in url_fields] + if has_url and not fmode.supports_url: + raise MCPUserError(f"{schema_for(code).label} does not support URL connections in this mode.") + if has_url and provided_url_fields: + raise MCPUserError( + f"Provide either a `URL` or host fields ({', '.join(sorted(url_fields))}), not both." + ) + + # Forced columns first so explicit params can't be clobbered by sets. + for attr, value in fmode.sets.items(): + setattr(connection, attr, value) + + for label, value in params.items(): + if label == _URL.label: + continue + column = fields_by_label[label].column + setattr(connection, column, str(value) if column == "project_port" and value is not None else value) + + if has_url: + connection.connect_by_url = True + connection.url = str(params[_URL.label]) + elif provided_url_fields: + connection.connect_by_url = False + + +# -- field-requirement validation ------------------------------------------- + +_CONNECTION_NAME_MIN = 3 +_CONNECTION_NAME_MAX = 40 +_MAX_THREADS_MIN = 1 +_MAX_THREADS_MAX = 8 +_MAX_QUERY_CHARS_MIN = 500 +_MAX_QUERY_CHARS_MAX = 50000 + + +def _missing(value: object) -> bool: + return value is None or (isinstance(value, str) and not value.strip()) + + +def _required_fields_for(connection: Connection) -> list[ConnField]: + """Schema fields that must be non-empty for this connection's flavor + active + mode, accounting for the URL alternative. + """ + fmode = _mode_for_connection(connection) + using_url = fmode.supports_url and bool(connection.connect_by_url) + + required: list[ConnField] = [] + for fld in fmode.fields: + if fld.requirement is Req.REQUIRED: + required.append(fld) + elif fld.requirement is Req.REQUIRED_UNLESS_URL and not using_url: + required.append(fld) + + schema = schema_for(connection.sql_flavor_code) + if using_url and schema.url_field is not None: + required.append(schema.url_field) + return required + + +def validate_connection_fields(connection: Connection) -> list[str]: + """Return every validation error (empty list = valid). + + Mirrors the per-flavor JS form validators in + ``ui/static/js/components/connection_form.js``. The connection tools call this + and raise their user-facing error containing the bullets. Field errors use the + UI label (e.g. ``Host``) so the LLM sees the same wording as a UI user. + """ + errors: list[str] = [] + flavor_label = FLAVOR_CODE_TO_LABEL.get(connection.sql_flavor_code, connection.sql_flavor_code) + + name = connection.connection_name + if _missing(name): + errors.append("`connection_name` is required.") + elif not (_CONNECTION_NAME_MIN <= len(name.strip()) <= _CONNECTION_NAME_MAX): + errors.append( + f"`connection_name` must be between {_CONNECTION_NAME_MIN} and {_CONNECTION_NAME_MAX} characters." + ) + + for fld in _required_fields_for(connection): + if _missing(getattr(connection, fld.column, None)): + errors.append(f"`{fld.label}` is required for {flavor_label}.") + + threads = connection.max_threads + if threads is not None and not (_MAX_THREADS_MIN <= threads <= _MAX_THREADS_MAX): + errors.append(f"`max_threads` must be between {_MAX_THREADS_MIN} and {_MAX_THREADS_MAX}.") + + query_chars = connection.max_query_chars + if query_chars is not None and not (_MAX_QUERY_CHARS_MIN <= query_chars <= _MAX_QUERY_CHARS_MAX): + errors.append(f"`max_query_chars` must be between {_MAX_QUERY_CHARS_MIN} and {_MAX_QUERY_CHARS_MAX}.") + + return errors diff --git a/testgen/mcp/tools/connections.py b/testgen/mcp/tools/connections.py new file mode 100644 index 00000000..7558ae55 --- /dev/null +++ b/testgen/mcp/tools/connections.py @@ -0,0 +1,235 @@ +"""MCP tools for database connection — create, update, 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. +""" + +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.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, + apply_connection_params, + connection_display_fields, + connection_field_labels, + infer_mode, + parse_sql_flavor, + resolve_connection, + validate_connection_fields, +) +from testgen.mcp.tools.markdown import MdDoc + + +@with_database_session +@mcp_permission("administer") +def test_connection( + connection_id: int | None = None, + sql_flavor: str | None = None, + connection_params: dict | None = None, + connection_mode: str | None = None, +) -> str: + """Test connectivity against a stored or inline-supplied connection. + + Two call shapes: pass ``connection_id`` to test a stored connection + (``connection_params`` / ``connection_mode`` override the stored values for + the test only, nothing is saved); or pass ``sql_flavor`` plus + ``connection_params`` (and ``connection_mode`` for multi-mode flavors) to + test inline without persisting. See the + ``testgen://connection-parameters/{flavor}`` resource for each flavor's + ``connection_mode`` values and ``connection_params`` keys. + + Args: + connection_id: Stored connection ID, e.g. from ``get_data_inventory``. + Omit for inline tests. + sql_flavor: SQL Database flavor. Required for inline tests; not accepted + when ``connection_id`` is set. + connection_params: Connection field values. For a stored connection, + omitted fields keep their stored value. + connection_mode: Authentication mode for multi-mode flavors. + """ + if connection_id is None and sql_flavor is None: + raise MCPUserError( + "Provide `connection_id` to test an existing connection, " + "or `sql_flavor` plus `connection_params` to test inline." + ) + + connection: Connection | None = None + if connection_id is not None: + if sql_flavor is not None: + raise MCPUserError( + "`sql_flavor` cannot be overridden when `connection_id` is set " + "— database type is immutable on an existing connection." + ) + connection = resolve_connection(connection_id) + + inline = connection is None + if connection is None: + _, code, family = parse_sql_flavor(sql_flavor) # type: ignore[arg-type] + 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) + apply_connection_params(connection, connection.sql_flavor_code, mode, connection_params or {}) + + normalize_auth_fields(connection) + # No-op for stored connections; keeps the inline-test contract identical to create. + apply_connection_defaults(connection) + + errors = validate_connection_fields(connection) + if errors: + _raise_validation_error(errors, "Cannot test connection. Required fields missing or invalid.") + + status = test_connection_status(connection) + + doc = MdDoc() + heading = "Connection test succeeded" if status.successful else "Connection test failed" + doc.heading(1, heading) + + if connection_id is not None: + doc.field("ID", connection.connection_id, code=True) + if connection.connection_name: + doc.field("Name", connection.connection_name, code=True) + label = SQL_FLAVOR_CODE_TO_LABEL.get(connection.sql_flavor_code) + doc.field("Type", label.value if label else connection.sql_flavor_code) + if connection.project_host: + doc.field("Host", connection.project_host, code=True) + + doc.text(status.message) + 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") + doc.field("ID", connection.connection_id, code=True) + doc.field("Project", connection.project_code, code=True) + label = SQL_FLAVOR_CODE_TO_LABEL.get(connection.sql_flavor_code) + doc.field("Type", label.value if label else connection.sql_flavor_code) + + # Render 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) + + return doc.render() + + +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, +} diff --git a/testgen/mcp/tools/reference.py b/testgen/mcp/tools/reference.py index 210e4b89..a41704f6 100644 --- a/testgen/mcp/tools/reference.py +++ b/testgen/mcp/tools/reference.py @@ -1,7 +1,16 @@ from testgen.common.models import with_database_session from testgen.common.models.hygiene_issue import HygieneIssueType from testgen.common.models.test_definition import TestType -from testgen.mcp.tools.common import DocGroup +from testgen.mcp.exceptions import MCPUserError +from testgen.mcp.tools.common import ( + FLAVOR_CONNECTION_SCHEMA, + ConnField, + DocGroup, + FlavorMode, + FlavorSchema, + Req, + schema_for, +) from testgen.mcp.tools.markdown import MdDoc _DOC_GROUP = DocGroup.DISCOVER @@ -326,3 +335,127 @@ def glossary_resource() -> str: - **referential** — Tests relationships between tables (e.g., foreign key match). - **custom** — User-defined SQL tests. """ + + +# Doc-only hints for the connection-parameters resource: value type / format +# guidance that helps the model but isn't part of the validation/mapping schema. +_FIELD_NOTES: dict[str, str] = { + "Port": "Integer.", + "Service Account Key": "Service-account key JSON, passed as a parsed object.", + "Private Key": "Private key as PEM text.", + "Private Key Passphrase": "Passphrase for the encrypted private key. Omit if the key is unencrypted.", + "Login URL": "My Domain URL of the Salesforce org.", + "Consumer Key": "Consumer key from the Salesforce external client app.", + "Consumer Secret": "Consumer secret from the Salesforce external client app.", + "Access Token": "Databricks personal access token.", + "Client ID": "Service-principal client ID.", + "Client Secret": "Service-principal OAuth secret.", + "HTTP Path": "Databricks SQL warehouse HTTP path.", + "Catalog": "Databricks catalog (the database/namespace).", + "Service Name": "Oracle service name.", + "Warehouse": "Snowflake warehouse name.", + "URL": "Full connection URL. Provide instead of the host fields.", +} + +# Doc-only: the conventional default port per flavor, so the model can fill the +# Port when the user doesn't specify one. +_DEFAULT_PORTS: dict[str, str] = { + "redshift": "5439", + "redshift_spectrum": "5439", + "azure_mssql": "1433", + "synapse_mssql": "1433", + "mssql": "1433", + "postgresql": "5432", + "snowflake": "443", + "databricks": "443", + "oracle": "1521", + "sap_hana": "39015", +} + + +def _requirement_label(field: ConnField) -> str: + if field.requirement is Req.REQUIRED: + return "Required" + if field.requirement is Req.REQUIRED_UNLESS_URL: + return "Required (host mode)" + return "Optional" + + +def _field_note(field: ConnField, schema: FlavorSchema) -> str | None: + parts = [] + if field.secret: + parts.append("Secret — encrypted at rest, never echoed back.") + if note := _FIELD_NOTES.get(field.label): + parts.append(note) + if field.column == "project_port" and (port := _DEFAULT_PORTS.get(schema.code)): + parts.append(f"Default for {schema.label} is {port} — use it unless the user specifies otherwise.") + return " ".join(parts) or None + + +def _append_mode(doc: MdDoc, mode: FlavorMode, schema: FlavorSchema, *, url_offered: bool) -> None: + if mode.mode is not None: + doc.heading(2, f"Mode: {mode.mode}") + doc.text(f'Pass `connection_mode="{mode.mode}"`.') + doc.table( + headers=["Field", "Required", "Notes"], + rows=[[field.label, _requirement_label(field), _field_note(field, schema)] for field in mode.fields], + code=[0], + ) + if mode.supports_url and url_offered: + doc.text( + "Alternatively, provide `URL` instead of the host fields " + "(the fields marked _Required (host mode)_) to connect by URL." + ) + + +def connection_parameters_resource(flavor: str) -> str: + """Per-flavor connection parameter shapes: the auth modes and the exact + ``connection_params`` keys (with required/optional + secret notes) used by + ``create_connection`` / ``update_connection`` / ``test_connection``. + + Args: + flavor: Flavor code, e.g. ``snowflake``, ``azure_mssql``, ``salesforce_data360``. + """ + if flavor not in FLAVOR_CONNECTION_SCHEMA: + valid = ", ".join(sorted(FLAVOR_CONNECTION_SCHEMA)) + raise MCPUserError(f"Unknown flavor `{flavor}`. Valid flavor codes: {valid}.") + + schema = schema_for(flavor) + doc = MdDoc() + doc.heading(1, f"{schema.label} Connection Parameters") + doc.text( + f'Create with `sql_flavor="{schema.label}"` and a `connection_params` dict keyed by the ' + "field labels below. Secrets are encrypted at rest and never returned." + ) + + multi_mode = len([m for m in schema.modes if m.mode is not None]) > 1 + if multi_mode: + modes = ", ".join(f"`{m.mode}`" for m in schema.modes if m.mode is not None) + doc.text(f"Set `connection_mode` to one of: {modes}.") + + for mode in schema.modes: + _append_mode(doc, mode, schema, url_offered=schema.url_field is not None) + + return doc.render() + + +def connection_parameters_index_resource() -> str: + """Supported database flavors for the connection tools: the accepted + ``sql_flavor`` values and, for each, the resource that documents its + connection modes and fields. + """ + doc = MdDoc() + doc.heading(1, "Connection Flavors") + doc.text( + "Accepted `sql_flavor` values. Read the per-flavor resource for a flavor's " + "connection modes and the fields each needs." + ) + doc.table( + headers=["sql_flavor", "Parameters resource"], + rows=[ + [schema.label, f"testgen://connection-parameters/{code}"] + for code, schema in FLAVOR_CONNECTION_SCHEMA.items() + ], + code=[1], + ) + return doc.render() diff --git a/testgen/ui/views/connections.py b/testgen/ui/views/connections.py index 619e10ed..53fa5f7d 100644 --- a/testgen/ui/views/connections.py +++ b/testgen/ui/views/connections.py @@ -1,25 +1,17 @@ import base64 import logging -import random import typing -from dataclasses import asdict, dataclass, field +from dataclasses import asdict, dataclass import streamlit as st -from testgen.commands.test_generation import run_monitor_generation -from testgen.ui.queries import table_group_queries - -try: - from pyodbc import Error as PyODBCError -except ImportError: - PyODBCError = None -from sqlalchemy.exc import DatabaseError, DBAPIError - -import testgen.ui.services.database_service as db from testgen import settings -from testgen.common.database.database_service import empty_cache, get_flavor_service +from testgen.commands.test_generation import run_monitor_generation +from testgen.common.database.connection_service import ConnectionStatus, test_connection_status +from testgen.common.database.database_service import get_flavor_service from testgen.common.database.flavor.flavor_service import resolve_connection_params from testgen.common.enums import JobSource +from testgen.common.flavors import FLAVOR_CODE_TO_FAMILY, FLAVOR_CODE_TO_LABEL from testgen.common.models import get_current_session, with_database_session from testgen.common.models.connection import Connection, ConnectionMinimal from testgen.common.models.job_execution import JobExecution @@ -30,6 +22,7 @@ from testgen.ui.components import widgets as testgen from testgen.ui.navigation.menu import MenuItem from testgen.ui.navigation.page import Page +from testgen.ui.queries import table_group_queries from testgen.ui.services.query_cache import ( get_connection, select_connections_where, @@ -256,43 +249,8 @@ def _format_connection(self, connection: Connection, should_test: bool = False) formatted_connection["status"] = asdict(self.test_connection(connection)) return formatted_connection - def test_connection(self, connection: Connection) -> "ConnectionStatus": - empty_cache() - try: - flavor_service = get_flavor_service(connection.sql_flavor) - results = db.fetch_from_target_db(connection, flavor_service.test_query) - connection_successful = len(results) == 1 and next(iter(results[0].values())) == 1 - - if not connection_successful: - return ConnectionStatus(message="Error completing a query to the database server.", successful=False) - return ConnectionStatus(message="The connection was successful.", successful=True) - except KeyError: - return ConnectionStatus( - message="Error attempting the connection. ", - details="Complete all the required fields.", - successful=False, - ) - except DatabaseError as error: - LOG.exception("Error testing database connection") - return ConnectionStatus(message="Error attempting the connection.", details=str(error.orig), successful=False) - except DBAPIError as error: - LOG.exception("Error testing database connection") - details = str(error.orig) - if PyODBCError and isinstance(error.orig, PyODBCError) and error.orig.args: - details = error.orig.args[1] - return ConnectionStatus(message="Error attempting the connection.", details=details, successful=False) - except (TypeError, ValueError) as error: - LOG.exception("Error testing database connection") - details = str(error) - if is_open_ssl_error(error): - details = error.args[0] - return ConnectionStatus(message="Error attempting the connection.", details=details, successful=False) - except Exception as error: - details = "Something went wrong while testing the connection." - if connection.connect_by_key and not connection.private_key: - details = "The private key is missing." - LOG.exception("Error testing database connection") - return ConnectionStatus(message="Error attempting the connection.", details=details, successful=False) + def test_connection(self, connection: Connection) -> ConnectionStatus: + return test_connection_status(connection) @with_database_session def setup_data_configuration(self, project_code: str, connection_id: str) -> None: @@ -539,24 +497,6 @@ def on_close_clicked(_params: dict) -> None: return wizard_data, wizard_handlers -@dataclass(frozen=True, slots=True) -class ConnectionStatus: - message: str - successful: bool - details: str | None = field(default=None) - _: float = field(default_factory=random.random) - - -def is_open_ssl_error(error: Exception): - return ( - error.args - and len(error.args) > 1 - and isinstance(error.args[1], list) - and len(error.args[1]) > 0 - and type(error.args[1][0]).__name__ == "OpenSSLError" - ) - - def format_connection(connection: Connection | ConnectionMinimal) -> dict: formatted_connection = connection.to_dict(json_safe=True) @@ -582,79 +522,31 @@ class ConnectionFlavor: flavor: str +# Labels and families come from the shared source in ``common/flavors.py`` (the +# same labels MCP uses); only the icon + display order are UI-specific. +_FLAVOR_ICONS: dict[str, str] = { + "redshift": "flavors/redshift.svg", + "redshift_spectrum": "flavors/redshift.svg", + "azure_mssql": "flavors/azure_sql.svg", + "synapse_mssql": "flavors/azure_synapse_table.svg", + "databricks": "flavors/databricks.svg", + "bigquery": "flavors/bigquery.svg", + "mssql": "flavors/mssql.svg", + "oracle": "flavors/oracle.svg", + "postgresql": "flavors/postgresql.svg", + "sap_hana": "flavors/sap_hana.svg", + "salesforce_data360": "flavors/salesforce_data360.svg", + "snowflake": "flavors/snowflake.svg", +} + FLAVOR_OPTIONS = [ ConnectionFlavor( - label="Amazon Redshift", - value="redshift", - flavor="redshift", - icon=get_asset_data_url("flavors/redshift.svg"), - ), - ConnectionFlavor( - label="Amazon Redshift Spectrum", - value="redshift_spectrum", - flavor="redshift_spectrum", - icon=get_asset_data_url("flavors/redshift.svg"), - ), - ConnectionFlavor( - label="Azure SQL Database", - value="azure_mssql", - flavor="mssql", - icon=get_asset_data_url("flavors/azure_sql.svg"), - ), - ConnectionFlavor( - label="Azure Synapse Analytics", - value="synapse_mssql", - flavor="mssql", - icon=get_asset_data_url("flavors/azure_synapse_table.svg"), - ), - ConnectionFlavor( - label="Databricks", - value="databricks", - flavor="databricks", - icon=get_asset_data_url("flavors/databricks.svg"), - ), - ConnectionFlavor( - label="Google BigQuery", - value="bigquery", - flavor="bigquery", - icon=get_asset_data_url("flavors/bigquery.svg"), - ), - ConnectionFlavor( - label="Microsoft SQL Server", - value="mssql", - flavor="mssql", - icon=get_asset_data_url("flavors/mssql.svg"), - ), - ConnectionFlavor( - label="Oracle", - value="oracle", - flavor="oracle", - icon=get_asset_data_url("flavors/oracle.svg"), - ), - ConnectionFlavor( - label="PostgreSQL", - value="postgresql", - flavor="postgresql", - icon=get_asset_data_url("flavors/postgresql.svg"), - ), - ConnectionFlavor( - label="SAP HANA", - value="sap_hana", - flavor="sap_hana", - icon=get_asset_data_url("flavors/sap_hana.svg"), - ), - ConnectionFlavor( - label="Salesforce Data 360", - value="salesforce_data360", - flavor="salesforce_data360", - icon=get_asset_data_url("flavors/salesforce_data360.svg"), - ), - ConnectionFlavor( - label="Snowflake", - value="snowflake", - flavor="snowflake", - icon=get_asset_data_url("flavors/snowflake.svg"), - ), + value=code, + label=str(FLAVOR_CODE_TO_LABEL[code]), + flavor=FLAVOR_CODE_TO_FAMILY[code], + icon=get_asset_data_url(icon), + ) + for code, icon in _FLAVOR_ICONS.items() ] # SAP HANA is hidden in the Docker image because pyhdbcli is glibc-only and fails to load on Alpine/musl. diff --git a/tests/unit/common/test_connection_service.py b/tests/unit/common/test_connection_service.py new file mode 100644 index 00000000..824fce6a --- /dev/null +++ b/tests/unit/common/test_connection_service.py @@ -0,0 +1,280 @@ +"""Tests for the connection_service common-layer module. + +Domain only: ``ConnectionStatus`` (with its load-bearing random field), the +``test_connection_status`` runner, and auth-path normalization. The label-bearing +connection-parameter schema / validation lives in ``mcp/tools/common.py`` and is +tested in ``tests/unit/mcp/test_connection_schema.py``. +""" + +from unittest.mock import patch + +import pytest +from sqlalchemy.exc import DatabaseError, DBAPIError + +from testgen.common.database.connection_service import ( + ConnectionStatus, + apply_connection_defaults, + normalize_auth_fields, +) + +# Aliased on import so pytest doesn't try to collect the ``test_*`` function as a test. +from testgen.common.database.connection_service import test_connection_status as run_connection_test +from testgen.common.models.connection import Connection + +pytestmark = pytest.mark.unit + +MODULE = "testgen.common.database.connection_service" + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +def _conn(**overrides) -> Connection: + """Build a Connection without touching the DB. Defaults to a complete PG config.""" + defaults = { + "sql_flavor": "postgresql", + "sql_flavor_code": "postgresql", + "connection_name": "My DB", + "project_host": "localhost", + "project_port": "5432", + "project_db": "testgen_local", + "project_user": "testgen", + "project_pw_encrypted": "pw", + "connect_by_url": False, + "connect_by_key": False, + "connect_with_identity": False, + "max_threads": 4, + "max_query_chars": 20000, + } + defaults.update(overrides) + return Connection(**defaults) + + +# --------------------------------------------------------------------------- +# ConnectionStatus dataclass +# --------------------------------------------------------------------------- + + +def test_connection_status_random_field_breaks_equality(): + """Two instances with identical (message, successful, details) MUST compare unequal. + + The random ``_`` field is required so Streamlit's reactive system re-renders + on repeated failed-test clicks producing the same error. Removing the field + would silently swallow the second click. + """ + a = ConnectionStatus(message="Error attempting the connection.", successful=False, details="boom") + b = ConnectionStatus(message="Error attempting the connection.", successful=False, details="boom") + assert a != b + + +# --------------------------------------------------------------------------- +# test_connection_status — happy path and exception branches +# --------------------------------------------------------------------------- + + +@patch(f"{MODULE}.empty_cache") +@patch(f"{MODULE}.fetch_from_target_db", return_value=[{"col": 1}]) +@patch(f"{MODULE}.get_flavor_service") +def test_connection_status_success(mock_flavor, mock_fetch, mock_empty_cache): + mock_flavor.return_value.test_query = "SELECT 1" + + status = run_connection_test(_conn()) + + assert status.successful is True + assert status.message == "The connection was successful." + assert status.details is None + mock_empty_cache.assert_called_once() # service owns the cache reset + + +@patch(f"{MODULE}.empty_cache") +@patch(f"{MODULE}.fetch_from_target_db", return_value=[{"col": 0}]) +@patch(f"{MODULE}.get_flavor_service") +def test_connection_status_query_returns_wrong_result(mock_flavor, mock_fetch, mock_empty_cache): + """``SELECT 1`` returns something other than 1 → 'Error completing a query'.""" + mock_flavor.return_value.test_query = "SELECT 1" + + status = run_connection_test(_conn()) + + assert status.successful is False + assert status.message == "Error completing a query to the database server." + + +@patch(f"{MODULE}.empty_cache") +@patch(f"{MODULE}.fetch_from_target_db", side_effect=KeyError("host")) +@patch(f"{MODULE}.get_flavor_service") +def test_connection_status_key_error(mock_flavor, mock_fetch, mock_empty_cache): + """Missing required field → 'Complete all the required fields.'""" + mock_flavor.return_value.test_query = "SELECT 1" + + status = run_connection_test(_conn()) + + assert status.successful is False + assert status.message == "Error attempting the connection. " + assert status.details == "Complete all the required fields." + + +@patch(f"{MODULE}.empty_cache") +@patch(f"{MODULE}.get_flavor_service") +@patch(f"{MODULE}.fetch_from_target_db") +def test_connection_status_database_error(mock_fetch, mock_flavor, mock_empty_cache): + mock_flavor.return_value.test_query = "SELECT 1" + orig = Exception("FATAL: password authentication failed") + mock_fetch.side_effect = DatabaseError("stmt", {}, orig) + + status = run_connection_test(_conn()) + + assert status.successful is False + assert status.message == "Error attempting the connection." + assert "password authentication failed" in str(status.details) + + +@patch(f"{MODULE}.empty_cache") +@patch(f"{MODULE}.get_flavor_service") +@patch(f"{MODULE}.fetch_from_target_db") +def test_connection_status_dbapi_error(mock_fetch, mock_flavor, mock_empty_cache): + mock_flavor.return_value.test_query = "SELECT 1" + orig = Exception("driver-level failure") + mock_fetch.side_effect = DBAPIError("stmt", {}, orig) + + status = run_connection_test(_conn()) + + assert status.successful is False + assert status.message == "Error attempting the connection." + assert "driver-level failure" in str(status.details) + + +@patch(f"{MODULE}.empty_cache") +@patch(f"{MODULE}.get_flavor_service") +@patch(f"{MODULE}.fetch_from_target_db") +def test_connection_status_open_ssl_error(mock_fetch, mock_flavor, mock_empty_cache): + """A TypeError whose args[1][0] is an OpenSSLError uses args[0] as details.""" + mock_flavor.return_value.test_query = "SELECT 1" + + class OpenSSLError: # name matches what is_open_ssl_error sniffs for + pass + + err = TypeError("bad key", [OpenSSLError()]) + mock_fetch.side_effect = err + + status = run_connection_test(_conn()) + + assert status.successful is False + assert status.message == "Error attempting the connection." + assert status.details == "bad key" + + +@patch(f"{MODULE}.empty_cache") +@patch(f"{MODULE}.get_flavor_service") +@patch(f"{MODULE}.fetch_from_target_db") +def test_connection_status_missing_private_key(mock_fetch, mock_flavor, mock_empty_cache): + """connect_by_key=True with no private_key → 'The private key is missing.'""" + mock_flavor.return_value.test_query = "SELECT 1" + mock_fetch.side_effect = RuntimeError("something") + + conn = _conn(sql_flavor="snowflake", sql_flavor_code="snowflake", connect_by_key=True, private_key=None) + status = run_connection_test(conn) + + assert status.successful is False + assert status.message == "Error attempting the connection." + assert status.details == "The private key is missing." + + +@patch(f"{MODULE}.empty_cache") +@patch(f"{MODULE}.get_flavor_service") +@patch(f"{MODULE}.fetch_from_target_db") +def test_connection_status_generic_exception(mock_fetch, mock_flavor, mock_empty_cache): + mock_flavor.return_value.test_query = "SELECT 1" + mock_fetch.side_effect = RuntimeError("unexpected") + + status = run_connection_test(_conn()) + + assert status.successful is False + assert status.message == "Error attempting the connection." + assert status.details == "Something went wrong while testing the connection." + + +# --------------------------------------------------------------------------- +# normalize_auth_fields +# --------------------------------------------------------------------------- + + +def test_normalize_clears_password_when_connect_by_key_non_databricks(): + """Snowflake connect_by_key=True → project_pw_encrypted cleared.""" + conn = _conn( + sql_flavor_code="snowflake", + sql_flavor="snowflake", + connect_by_key=True, + project_pw_encrypted="old_pw", + private_key="key", + ) + normalize_auth_fields(conn) + assert conn.project_pw_encrypted in (None, "") + assert conn.private_key == "key" # untouched + + +def test_normalize_keeps_password_for_databricks_connect_by_key(): + """Databricks OAuth uses connect_by_key but stores the secret in project_pw_encrypted.""" + conn = _conn( + sql_flavor_code="databricks", + sql_flavor="databricks", + connect_by_key=True, + project_pw_encrypted="client_secret_xyz", + ) + normalize_auth_fields(conn) + assert conn.project_pw_encrypted == "client_secret_xyz" + assert not conn.private_key + assert not conn.private_key_passphrase + + +def test_normalize_clears_private_key_fields_when_password_auth(): + """connect_by_key=False → private_key + passphrase cleared.""" + conn = _conn( + sql_flavor_code="snowflake", + sql_flavor="snowflake", + connect_by_key=False, + private_key="old_key", + private_key_passphrase="old_phrase", # noqa: S106 + ) + normalize_auth_fields(conn) + assert not conn.private_key + assert not conn.private_key_passphrase + + +def test_normalize_clears_user_password_when_identity(): + """connect_with_identity=True → project_user + project_pw_encrypted cleared.""" + conn = _conn( + sql_flavor_code="azure_mssql", + sql_flavor="mssql", + connect_with_identity=True, + project_user="leftover_user", + project_pw_encrypted="leftover_pw", + ) + normalize_auth_fields(conn) + assert not conn.project_user + assert not conn.project_pw_encrypted + + +# --------------------------------------------------------------------------- +# apply_connection_defaults +# --------------------------------------------------------------------------- + + +def test_defaults_fill_max_query_chars(): + conn = _conn(max_query_chars=None) + apply_connection_defaults(conn) + assert conn.max_query_chars == 20000 + + +def test_defaults_salesforce_lower_max_query_chars(): + """Salesforce Data 360's Hyper engine gets the lower 15000 default.""" + conn = _conn(sql_flavor="salesforce_data360", sql_flavor_code="salesforce_data360", max_query_chars=None) + apply_connection_defaults(conn) + assert conn.max_query_chars == 15000 + + +def test_defaults_keep_explicit_max_query_chars(): + conn = _conn(max_query_chars=30000) + apply_connection_defaults(conn) + assert conn.max_query_chars == 30000 diff --git a/tests/unit/common/test_flavors.py b/tests/unit/common/test_flavors.py new file mode 100644 index 00000000..4f334817 --- /dev/null +++ b/tests/unit/common/test_flavors.py @@ -0,0 +1,35 @@ +"""Tests for the shared flavor-identity source of truth.""" + +import pytest + +from testgen.common.flavors import ( + FLAVOR_CODE_TO_FAMILY, + FLAVOR_CODE_TO_LABEL, + SqlFlavorLabel, +) + +pytestmark = pytest.mark.unit + + +def test_label_and_family_maps_cover_the_same_codes(): + assert set(FLAVOR_CODE_TO_LABEL) == set(FLAVOR_CODE_TO_FAMILY) + + +def test_every_label_is_a_known_enum_member(): + assert set(FLAVOR_CODE_TO_LABEL.values()) == set(SqlFlavorLabel) + + +def test_labels_are_unique_per_code(): + labels = list(FLAVOR_CODE_TO_LABEL.values()) + assert len(labels) == len(set(labels)) + + +def test_azure_variants_share_the_mssql_family(): + assert FLAVOR_CODE_TO_FAMILY["azure_mssql"] == "mssql" + assert FLAVOR_CODE_TO_FAMILY["synapse_mssql"] == "mssql" + assert FLAVOR_CODE_TO_FAMILY["mssql"] == "mssql" + + +def test_label_renders_as_plain_string(): + assert f"{FLAVOR_CODE_TO_LABEL['postgresql']}" == "PostgreSQL" + assert str(FLAVOR_CODE_TO_LABEL["azure_mssql"]) == "Azure SQL Database" diff --git a/tests/unit/mcp/conftest.py b/tests/unit/mcp/conftest.py index cc932043..53e23a5d 100644 --- a/tests/unit/mcp/conftest.py +++ b/tests/unit/mcp/conftest.py @@ -10,6 +10,7 @@ "view": ["role_a", "role_b"], "catalog": ["role_a", "role_b", "role_c"], "edit": ["role_a"], + "administer": ["role_a"], } diff --git a/tests/unit/mcp/test_connection_schema.py b/tests/unit/mcp/test_connection_schema.py new file mode 100644 index 00000000..75919cfc --- /dev/null +++ b/tests/unit/mcp/test_connection_schema.py @@ -0,0 +1,649 @@ +"""Tests for the connection-parameter contract in ``mcp/tools/common.py``: +the per-flavor schema (``schema_for`` / ``resolve_mode``), label-keyed param +mapping (``apply_connection_params``), schema-driven validation +(``validate_connection_fields``), and mode inference (``infer_mode``). + +The schema mirrors the per-flavor JS forms in +``ui/static/js/components/connection_form.js`` — these tests encode that contract. +""" + +import pytest + +from testgen.common.models.connection import Connection +from testgen.mcp.exceptions import MCPUserError +from testgen.mcp.tools.common import ( + ConnectionMode, + Req, + apply_connection_params, + infer_mode, + resolve_mode, + schema_for, + validate_connection_fields, +) + +pytestmark = pytest.mark.unit + + +def _conn(**overrides) -> Connection: + """Build a Connection without touching the DB. Defaults to a complete PG config.""" + defaults = { + "sql_flavor": "postgresql", + "sql_flavor_code": "postgresql", + "connection_name": "My DB", + "project_host": "localhost", + "project_port": "5432", + "project_db": "testgen_local", + "project_user": "testgen", + "project_pw_encrypted": "pw", + "connect_by_url": False, + "connect_by_key": False, + "connect_with_identity": False, + "max_threads": 4, + "max_query_chars": 20000, + } + defaults.update(overrides) + return Connection(**defaults) + + +# --------------------------------------------------------------------------- +# schema_for +# --------------------------------------------------------------------------- + + +def test_schema_for_unknown_code_raises(): + with pytest.raises(KeyError): + schema_for("not_a_flavor") + + +def test_schema_for_postgresql_single_mode_host_fields(): + schema = schema_for("postgresql") + assert schema.label == "PostgreSQL" + assert len(schema.modes) == 1 + mode = schema.modes[0] + assert mode.mode is None + labels = {f.label: f for f in mode.fields} + assert set(labels) == {"Host", "Port", "Database", "Username", "Password"} + assert labels["Host"].requirement is Req.REQUIRED_UNLESS_URL + assert labels["Username"].requirement is Req.REQUIRED + assert labels["Password"].requirement is Req.OPTIONAL + assert labels["Password"].secret is True + assert mode.supports_url is True + assert schema.url_field is not None and schema.url_field.label == "URL" + assert labels["Host"].column == "project_host" + assert labels["Username"].column == "project_user" + assert labels["Password"].column == "project_pw_encrypted" + + +def test_schema_for_oracle_uses_service_name_label(): + labels = {f.label: f for f in schema_for("oracle").modes[0].fields} + assert "Service Name" in labels + assert "Database" not in labels + assert labels["Service Name"].column == "project_db" + + +def test_schema_for_sap_hana_uses_database_label(): + labels = {f.label for f in schema_for("sap_hana").modes[0].fields} + assert "Database" in labels + + +def test_schema_for_snowflake_two_modes_with_warehouse(): + schema = schema_for("snowflake") + modes = {m.mode: m for m in schema.modes} + assert set(modes) == {ConnectionMode.KEY_PAIR, ConnectionMode.PASSWORD} + key_labels = {f.label: f for f in modes[ConnectionMode.KEY_PAIR].fields} + pw_labels = {f.label: f for f in modes[ConnectionMode.PASSWORD].fields} + assert "Warehouse" in key_labels and key_labels["Warehouse"].requirement is Req.OPTIONAL + assert key_labels["Username"].requirement is Req.REQUIRED + assert key_labels["Private Key"].requirement is Req.REQUIRED + assert key_labels["Private Key"].secret is True + assert key_labels["Private Key Passphrase"].requirement is Req.OPTIONAL + assert "Password" not in key_labels + assert pw_labels["Password"].requirement is Req.REQUIRED + assert "Private Key" not in pw_labels + assert all(m.supports_url for m in schema.modes) + + +def test_schema_for_databricks_modes_and_url_support(): + schema = schema_for("databricks") + modes = {m.mode: m for m in schema.modes} + assert set(modes) == {ConnectionMode.ACCESS_TOKEN, ConnectionMode.SERVICE_PRINCIPAL} + pat = modes[ConnectionMode.ACCESS_TOKEN] + oauth = modes[ConnectionMode.SERVICE_PRINCIPAL] + pat_labels = {f.label: f for f in pat.fields} + oauth_labels = {f.label: f for f in oauth.fields} + assert pat_labels["Catalog"].column == "project_db" + assert pat_labels["Catalog"].requirement is Req.REQUIRED_UNLESS_URL + assert pat_labels["HTTP Path"].requirement is Req.REQUIRED_UNLESS_URL + assert pat_labels["Access Token"].column == "project_pw_encrypted" + assert pat_labels["Access Token"].requirement is Req.REQUIRED + assert "Username" not in pat_labels # auto-set to 'token' + assert pat.supports_url is True + assert oauth_labels["Client ID"].column == "project_user" + assert oauth_labels["Client Secret"].column == "project_pw_encrypted" + assert oauth_labels["Host"].requirement is Req.REQUIRED + assert oauth.supports_url is False + + +def test_schema_for_bigquery_single_field_no_url(): + schema = schema_for("bigquery") + assert len(schema.modes) == 1 + labels = {f.label: f for f in schema.modes[0].fields} + assert set(labels) == {"Service Account Key"} + assert labels["Service Account Key"].column == "service_account_key" + assert labels["Service Account Key"].secret is True + assert schema.modes[0].supports_url is False + assert schema.url_field is None + + +def test_schema_for_salesforce_two_modes_field_mapping(): + schema = schema_for("salesforce_data360") + modes = {m.mode: m for m in schema.modes} + assert set(modes) == {ConnectionMode.JWT_BEARER, ConnectionMode.CLIENT_CREDENTIALS} + jwt = {f.label: f for f in modes[ConnectionMode.JWT_BEARER].fields} + cc = {f.label: f for f in modes[ConnectionMode.CLIENT_CREDENTIALS].fields} + assert jwt["Login URL"].column == "project_host" + assert jwt["Consumer Key"].column == "project_user" + assert jwt["Username"].column == "project_db" + assert jwt["Private Key"].column == "private_key" + assert "Consumer Secret" not in jwt + assert cc["Consumer Secret"].column == "project_pw_encrypted" + assert "Username" not in cc and "Private Key" not in cc + assert all(not m.supports_url for m in schema.modes) + + +# --------------------------------------------------------------------------- +# resolve_mode +# --------------------------------------------------------------------------- + + +def test_resolve_mode_single_mode_no_label(): + assert resolve_mode("postgresql", None).mode is None + + +def test_resolve_mode_single_mode_rejects_label(): + with pytest.raises(MCPUserError): + resolve_mode("postgresql", "Password") + + +def test_resolve_mode_multi_mode_requires_label(): + with pytest.raises(MCPUserError) as exc: + resolve_mode("snowflake", None) + assert "Key-Pair" in str(exc.value) and "Password" in str(exc.value) + + +def test_resolve_mode_multi_mode_invalid_label(): + with pytest.raises(MCPUserError) as exc: + resolve_mode("snowflake", "Bogus") + assert "Key-Pair" in str(exc.value) + + +def test_resolve_mode_multi_mode_valid_label(): + assert resolve_mode("snowflake", "Key-Pair").mode is ConnectionMode.KEY_PAIR + + +# --------------------------------------------------------------------------- +# apply_connection_params +# --------------------------------------------------------------------------- + + +def test_apply_params_postgresql_maps_labels_to_columns(): + conn = Connection(sql_flavor="postgresql", sql_flavor_code="postgresql") + apply_connection_params( + conn, + "postgresql", + None, + {"Host": "h", "Port": 5432, "Database": "d", "Username": "u", "Password": "p"}, + ) + assert conn.project_host == "h" + assert conn.project_port == "5432" # cast to str + assert conn.project_db == "d" + assert conn.project_user == "u" + assert conn.project_pw_encrypted == "p" + assert conn.connect_by_url is False + + +def test_apply_params_url_sets_connect_by_url(): + conn = Connection(sql_flavor="postgresql", sql_flavor_code="postgresql") + apply_connection_params(conn, "postgresql", None, {"URL": "host:5432/db", "Username": "u"}) + assert conn.connect_by_url is True + assert conn.url == "host:5432/db" + assert conn.project_user == "u" + + +def test_apply_params_url_conflicts_with_host_group(): + conn = Connection(sql_flavor="postgresql", sql_flavor_code="postgresql") + with pytest.raises(MCPUserError): + apply_connection_params(conn, "postgresql", None, {"URL": "x", "Host": "h"}) + + +def test_apply_params_url_on_unsupported_flavor_rejected(): + conn = Connection(sql_flavor="bigquery", sql_flavor_code="bigquery") + with pytest.raises(MCPUserError): + apply_connection_params(conn, "bigquery", None, {"URL": "x"}) + + +def test_apply_params_unknown_key_rejected(): + conn = Connection(sql_flavor="postgresql", sql_flavor_code="postgresql") + with pytest.raises(MCPUserError) as exc: + apply_connection_params(conn, "postgresql", None, {"Hostname": "h"}) + assert "Hostname" in str(exc.value) + + +def test_apply_params_snowflake_key_pair_sets_flag(): + conn = Connection(sql_flavor="snowflake", sql_flavor_code="snowflake") + apply_connection_params( + conn, + "snowflake", + "Key-Pair", + {"Host": "h", "Port": 443, "Database": "d", "Username": "u", "Private Key": "KEY"}, + ) + assert conn.connect_by_key is True + assert conn.private_key == "KEY" + + +def test_apply_params_databricks_pat_auto_token_username(): + conn = Connection(sql_flavor="databricks", sql_flavor_code="databricks") + apply_connection_params( + conn, + "databricks", + "Access Token", + {"Host": "h", "Port": 443, "Catalog": "main", "HTTP Path": "/sql/1.0/abc", "Access Token": "tok"}, + ) + assert conn.project_user == "token" + assert conn.connect_by_key is False + assert conn.project_db == "main" + assert conn.http_path == "/sql/1.0/abc" + assert conn.project_pw_encrypted == "tok" + + +def test_apply_params_databricks_oauth_client_id_secret(): + conn = Connection(sql_flavor="databricks", sql_flavor_code="databricks") + apply_connection_params( + conn, + "databricks", + "Service Principal (OAuth)", + {"Host": "h", "Port": 443, "Catalog": "main", "HTTP Path": "/p", "Client ID": "cid", "Client Secret": "csec"}, + ) + assert conn.connect_by_key is True + assert conn.project_user == "cid" + assert conn.project_pw_encrypted == "csec" + + +def test_apply_params_databricks_oauth_rejects_url(): + conn = Connection(sql_flavor="databricks", sql_flavor_code="databricks") + with pytest.raises(MCPUserError): + apply_connection_params(conn, "databricks", "Service Principal (OAuth)", {"URL": "x"}) + + +def test_apply_params_azure_managed_identity_sets_flag(): + conn = Connection(sql_flavor="mssql", sql_flavor_code="azure_mssql") + apply_connection_params(conn, "azure_mssql", "Managed Identity", {"Host": "h", "Port": 1433, "Database": "d"}) + assert conn.connect_with_identity is True + + +def test_apply_params_salesforce_jwt_field_mapping(): + conn = Connection(sql_flavor="salesforce_data360", sql_flavor_code="salesforce_data360") + apply_connection_params( + conn, + "salesforce_data360", + "JWT Bearer Flow", + {"Login URL": "https://my.salesforce.com", "Consumer Key": "ck", "Username": "u@x.com", "Private Key": "PK"}, + ) + assert conn.project_host == "https://my.salesforce.com" + assert conn.project_user == "ck" + assert conn.project_db == "u@x.com" + assert conn.private_key == "PK" + assert conn.connect_by_key is True + + +# --------------------------------------------------------------------------- +# validate_connection_fields — happy paths +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize("flavor_code", ["postgresql", "redshift", "redshift_spectrum", "mssql", "oracle", "sap_hana"]) +def test_validate_passes_basic_host_auth_flavors(flavor_code): + conn = _conn(sql_flavor_code=flavor_code, sql_flavor=flavor_code) + assert validate_connection_fields(conn) == [] + + +def test_validate_passes_azure_mssql_with_user_password(): + assert validate_connection_fields(_conn(sql_flavor_code="azure_mssql", sql_flavor="mssql")) == [] + + +def test_validate_passes_azure_mssql_with_identity(): + conn = _conn( + sql_flavor_code="azure_mssql", + sql_flavor="mssql", + connect_with_identity=True, + project_user=None, + project_pw_encrypted=None, + ) + assert validate_connection_fields(conn) == [] + + +def test_validate_passes_snowflake_password_auth(): + conn = _conn(sql_flavor_code="snowflake", sql_flavor="snowflake", project_port="443", connect_by_key=False) + assert validate_connection_fields(conn) == [] + + +def test_validate_passes_snowflake_key_pair_auth(): + conn = _conn( + sql_flavor_code="snowflake", + sql_flavor="snowflake", + project_port="443", + project_pw_encrypted=None, + connect_by_key=True, + private_key="-----BEGIN PRIVATE KEY-----\n...", + ) + assert validate_connection_fields(conn) == [] + + +def test_validate_passes_databricks_pat(): + conn = _conn( + sql_flavor_code="databricks", + sql_flavor="databricks", + project_port="443", + project_db="main", + project_user="token", + http_path="/sql/1.0/warehouses/abc", + ) + assert validate_connection_fields(conn) == [] + + +def test_validate_passes_bigquery(): + conn = _conn( + sql_flavor_code="bigquery", + sql_flavor="bigquery", + project_host=None, + project_port=None, + project_db=None, + project_user=None, + project_pw_encrypted=None, + service_account_key={"type": "service_account", "project_id": "demo"}, + ) + assert validate_connection_fields(conn) == [] + + +def test_validate_passes_salesforce_jwt(): + conn = _conn( + sql_flavor_code="salesforce_data360", + sql_flavor="salesforce_data360", + project_host="https://my.salesforce.com", + project_port=None, + project_db="user@x.com", + project_user="consumer_key", + project_pw_encrypted=None, + connect_by_key=True, + private_key="-----BEGIN PRIVATE KEY-----\n...", + ) + assert validate_connection_fields(conn) == [] + + +def test_validate_passes_salesforce_client_credentials(): + conn = _conn( + sql_flavor_code="salesforce_data360", + sql_flavor="salesforce_data360", + project_host="https://my.salesforce.com", + project_port=None, + project_db=None, + project_user="consumer_key", + project_pw_encrypted="consumer_secret", + connect_by_key=False, + private_key=None, + ) + assert validate_connection_fields(conn) == [] + + +def test_validate_passes_connect_by_url_keeps_username(): + """URL mode: host/port/db not required, but Username STILL required (matches UI).""" + conn = _conn(connect_by_url=True, url="localhost:5432/mydb", project_host=None, project_port=None, project_db=None) + assert validate_connection_fields(conn) == [] + + +# --------------------------------------------------------------------------- +# validate_connection_fields — divergence fix + per-flavor errors +# --------------------------------------------------------------------------- + + +def test_validate_url_mode_missing_username_fails(): + """The PR-flagged divergence: URL mode must still require Username.""" + conn = _conn( + connect_by_url=True, + url="localhost:5432/mydb", + project_host=None, + project_port=None, + project_db=None, + project_user=None, + ) + assert "`Username` is required for PostgreSQL." in validate_connection_fields(conn) + + +@pytest.mark.parametrize( + "label,model_attr", + [("Host", "project_host"), ("Port", "project_port"), ("Database", "project_db"), ("Username", "project_user")], +) +def test_validate_postgresql_missing_field(label, model_attr): + conn = _conn(**{model_attr: None}) + assert f"`{label}` is required for PostgreSQL." in validate_connection_fields(conn) + + +def test_validate_oracle_missing_service_name(): + conn = _conn(sql_flavor_code="oracle", sql_flavor="oracle", project_db=None) + assert "`Service Name` is required for Oracle." in validate_connection_fields(conn) + + +def test_validate_postgresql_missing_url_when_connect_by_url(): + conn = _conn(connect_by_url=True, url=None, project_host=None, project_port=None) + assert "`URL` is required for PostgreSQL." in validate_connection_fields(conn) + + +def test_validate_snowflake_missing_password_when_not_connect_by_key(): + conn = _conn( + sql_flavor_code="snowflake", + sql_flavor="snowflake", + project_port="443", + connect_by_key=False, + project_pw_encrypted=None, + ) + assert "`Password` is required for Snowflake." in validate_connection_fields(conn) + + +def test_validate_snowflake_missing_private_key_when_connect_by_key(): + conn = _conn( + sql_flavor_code="snowflake", + sql_flavor="snowflake", + project_port="443", + project_pw_encrypted=None, + connect_by_key=True, + private_key=None, + ) + assert "`Private Key` is required for Snowflake." in validate_connection_fields(conn) + + +def test_validate_databricks_pat_missing_catalog(): + conn = _conn( + sql_flavor_code="databricks", + sql_flavor="databricks", + project_port="443", + project_user="token", + project_db=None, + http_path="/sql/1.0/warehouses/abc", + ) + assert "`Catalog` is required for Databricks." in validate_connection_fields(conn) + + +def test_validate_databricks_pat_missing_http_path(): + conn = _conn( + sql_flavor_code="databricks", + sql_flavor="databricks", + project_port="443", + project_user="token", + project_db="main", + http_path=None, + ) + assert "`HTTP Path` is required for Databricks." in validate_connection_fields(conn) + + +def test_validate_databricks_pat_missing_access_token(): + conn = _conn( + sql_flavor_code="databricks", + sql_flavor="databricks", + project_port="443", + project_user="token", + project_db="main", + http_path="/p", + project_pw_encrypted=None, + ) + assert "`Access Token` is required for Databricks." in validate_connection_fields(conn) + + +def test_validate_databricks_oauth_missing_client_id(): + conn = _conn( + sql_flavor_code="databricks", + sql_flavor="databricks", + project_port="443", + project_db="main", + http_path="/p", + connect_by_key=True, + project_user=None, + project_pw_encrypted="secret", + ) + assert "`Client ID` is required for Databricks." in validate_connection_fields(conn) + + +def test_validate_bigquery_missing_service_account_key(): + conn = _conn( + sql_flavor_code="bigquery", + sql_flavor="bigquery", + project_host=None, + project_port=None, + project_db=None, + project_user=None, + project_pw_encrypted=None, + service_account_key=None, + ) + assert "`Service Account Key` is required for Google BigQuery." in validate_connection_fields(conn) + + +def test_validate_salesforce_jwt_missing_private_key(): + conn = _conn( + sql_flavor_code="salesforce_data360", + sql_flavor="salesforce_data360", + project_host="https://my.salesforce.com", + project_port=None, + project_db="user@x.com", + project_user="ck", + project_pw_encrypted=None, + connect_by_key=True, + private_key=None, + ) + assert "`Private Key` is required for Salesforce Data 360." in validate_connection_fields(conn) + + +def test_validate_salesforce_jwt_missing_login_url(): + conn = _conn( + sql_flavor_code="salesforce_data360", + sql_flavor="salesforce_data360", + project_host=None, + project_port=None, + project_db="user@x.com", + project_user="ck", + project_pw_encrypted=None, + connect_by_key=True, + private_key="key", + ) + assert "`Login URL` is required for Salesforce Data 360." in validate_connection_fields(conn) + + +def test_validate_salesforce_client_credentials_missing_secret(): + conn = _conn( + sql_flavor_code="salesforce_data360", + sql_flavor="salesforce_data360", + project_host="https://my.salesforce.com", + project_port=None, + project_db=None, + project_user="ck", + project_pw_encrypted=None, + connect_by_key=False, + private_key=None, + ) + assert "`Consumer Secret` is required for Salesforce Data 360." in validate_connection_fields(conn) + + +def test_validate_azure_mssql_password_auth_missing_user(): + conn = _conn(sql_flavor_code="azure_mssql", sql_flavor="mssql", connect_with_identity=False, project_user=None) + assert "`Username` is required for Azure SQL Database." in validate_connection_fields(conn) + + +# --------------------------------------------------------------------------- +# validate_connection_fields — name / threads / query-chars +# --------------------------------------------------------------------------- + + +def test_validate_missing_connection_name(): + assert "`connection_name` is required." in validate_connection_fields(_conn(connection_name=None)) + + +def test_validate_connection_name_too_short(): + assert "`connection_name` must be between 3 and 40 characters." in validate_connection_fields(_conn(connection_name="ab")) + + +def test_validate_connection_name_too_long(): + assert "`connection_name` must be between 3 and 40 characters." in validate_connection_fields(_conn(connection_name="a" * 41)) + + +@pytest.mark.parametrize("bad", [0, -1, 9, 100]) +def test_validate_max_threads_out_of_range(bad): + assert "`max_threads` must be between 1 and 8." in validate_connection_fields(_conn(max_threads=bad)) + + +@pytest.mark.parametrize("bad", [499, 0, 50001, 100000]) +def test_validate_max_query_chars_out_of_range(bad): + assert "`max_query_chars` must be between 500 and 50000." in validate_connection_fields(_conn(max_query_chars=bad)) + + +def test_validate_aggregates_all_errors(): + conn = _conn(connection_name="", project_host=None, project_port=None, max_threads=99) + joined = "\n".join(validate_connection_fields(conn)) + assert "`connection_name` is required." in joined + assert "`Host` is required for PostgreSQL." in joined + assert "`Port` is required for PostgreSQL." in joined + assert "`max_threads` must be between 1 and 8." in joined + + +# --------------------------------------------------------------------------- +# infer_mode +# --------------------------------------------------------------------------- + + +def test_infer_mode_single_mode_flavor_is_none(): + assert infer_mode(_conn(sql_flavor_code="postgresql")) is None + + +def test_infer_mode_snowflake_key_pair(): + assert infer_mode(_conn(sql_flavor_code="snowflake", sql_flavor="snowflake", connect_by_key=True)) is ConnectionMode.KEY_PAIR + + +def test_infer_mode_snowflake_password(): + assert infer_mode(_conn(sql_flavor_code="snowflake", sql_flavor="snowflake", connect_by_key=False)) is ConnectionMode.PASSWORD + + +def test_infer_mode_azure_identity(): + conn = _conn(sql_flavor_code="azure_mssql", sql_flavor="mssql", connect_with_identity=True) + assert infer_mode(conn) is ConnectionMode.MANAGED_IDENTITY + + +def test_infer_mode_databricks_oauth(): + conn = _conn(sql_flavor_code="databricks", sql_flavor="databricks", connect_by_key=True) + assert infer_mode(conn) is ConnectionMode.SERVICE_PRINCIPAL + + +def test_infer_mode_databricks_pat(): + conn = _conn(sql_flavor_code="databricks", sql_flavor="databricks", connect_by_key=False) + assert infer_mode(conn) is ConnectionMode.ACCESS_TOKEN + + +def test_infer_mode_salesforce_jwt(): + conn = _conn(sql_flavor_code="salesforce_data360", sql_flavor="salesforce_data360", connect_by_key=True) + assert infer_mode(conn) is ConnectionMode.JWT_BEARER diff --git a/tests/unit/mcp/test_tools_common.py b/tests/unit/mcp/test_tools_common.py index 3b5e7f32..2abf85c0 100644 --- a/tests/unit/mcp/test_tools_common.py +++ b/tests/unit/mcp/test_tools_common.py @@ -711,3 +711,34 @@ def test_score_chain_leaf_field_values(): def test_score_chain_leaf_to_column_mapping(): assert SCORE_CHAIN_LEAF_TO_COLUMN[ScoreChainLeafField.TABLE] == "table_name" assert SCORE_CHAIN_LEAF_TO_COLUMN[ScoreChainLeafField.COLUMN] == "column_name" + + +# --- SqlFlavorLabel --- + + +def test_sql_flavor_label_set_matches_common_layer(): + """Codes covered by the MCP enum and the common-layer maps must stay in sync.""" + from testgen.common.flavors import FLAVOR_CODE_TO_FAMILY, FLAVOR_CODE_TO_LABEL + from testgen.mcp.tools.common import SQL_FLAVOR_CODE_TO_LABEL, SQL_FLAVOR_LABEL_TO_CODE + + assert set(SQL_FLAVOR_CODE_TO_LABEL) == set(FLAVOR_CODE_TO_LABEL) + assert set(SQL_FLAVOR_LABEL_TO_CODE.values()) == set(FLAVOR_CODE_TO_FAMILY.keys()) + + +def test_parse_sql_flavor_returns_label_code_family(): + from testgen.mcp.tools.common import SqlFlavorLabel, parse_sql_flavor + + label, code, family = parse_sql_flavor("Azure SQL Database") + assert label == SqlFlavorLabel.AZURE_MSSQL + assert code == "azure_mssql" + assert family == "mssql" + + +def test_parse_sql_flavor_invalid_lists_display_values(): + from testgen.mcp.tools.common import SqlFlavorLabel, parse_sql_flavor + + with pytest.raises(MCPUserError, match="Invalid sql_flavor `bogus`") as exc: + parse_sql_flavor("bogus") + msg = str(exc.value) + for member in SqlFlavorLabel: + assert member.value in msg diff --git a/tests/unit/mcp/test_tools_connections.py b/tests/unit/mcp/test_tools_connections.py new file mode 100644 index 00000000..3533023e --- /dev/null +++ b/tests/unit/mcp/test_tools_connections.py @@ -0,0 +1,233 @@ +"""Tests for the MCP connection CRUD tools — create / update / test. + +The tools take a flavor-shaped ``connection_params`` dict (keyed by UI label) +plus an explicit ``connection_mode``; mapping + validation are delegated to +``connection_service``. +""" + +from unittest.mock import MagicMock, patch + +import pytest + +from testgen.common.database.connection_service import ConnectionStatus +from testgen.mcp.exceptions import MCPPermissionDenied, MCPResourceNotAccessible, MCPUserError +from testgen.mcp.permissions import ProjectPermissions + +pytestmark = pytest.mark.unit + +MODULE = "testgen.mcp.tools.connections" + + +# --------------------------------------------------------------------------- +# Fixtures / helpers +# --------------------------------------------------------------------------- + + +def _patch_perms(allowed=("demo",), memberships=None, permission="administer"): + """Inject a ProjectPermissions for the given access set.""" + memberships = memberships or dict.fromkeys(allowed, "role_a") + return patch( + "testgen.mcp.permissions._compute_project_permissions", + return_value=ProjectPermissions( + memberships=memberships, permission=permission, username="test_user", + ), + ) + + +def _mock_connection(**overrides) -> MagicMock: + """Build a MagicMock matching the Connection model surface used by the tools.""" + conn = MagicMock() + conn.connection_id = overrides.get("connection_id", 42) + conn.project_code = overrides.get("project_code", "demo") + conn.connection_name = overrides.get("connection_name", "Local PG") + conn.sql_flavor = overrides.get("sql_flavor", "postgresql") + conn.sql_flavor_code = overrides.get("sql_flavor_code", "postgresql") + conn.project_host = overrides.get("project_host", "localhost") + conn.project_port = overrides.get("project_port", "5432") + conn.project_db = overrides.get("project_db", "testgen_local") + conn.project_user = overrides.get("project_user", "testgen") + conn.project_pw_encrypted = overrides.get("project_pw_encrypted", "stored_pw") + conn.connect_by_url = overrides.get("connect_by_url", False) + conn.url = overrides.get("url", None) + conn.connect_by_key = overrides.get("connect_by_key", False) + conn.private_key = overrides.get("private_key", None) + conn.private_key_passphrase = overrides.get("private_key_passphrase", None) + conn.connect_with_identity = overrides.get("connect_with_identity", False) + conn.http_path = overrides.get("http_path", None) + conn.warehouse = overrides.get("warehouse", None) + conn.service_account_key = overrides.get("service_account_key", None) + conn.max_threads = overrides.get("max_threads", 4) + conn.max_query_chars = overrides.get("max_query_chars", None) + return conn + + +# --------------------------------------------------------------------------- +# test_connection +# --------------------------------------------------------------------------- + + +def test_test_connection_neither_id_nor_flavor(db_session_mock): + from testgen.mcp.tools.connections import test_connection + + with _patch_perms(), pytest.raises(MCPUserError) as exc: + test_connection() + msg = str(exc.value) + assert "`connection_id`" in msg + assert "`sql_flavor`" in msg + + +@patch(f"{MODULE}.resolve_connection") +def test_test_connection_rejects_sql_flavor_with_id(mock_resolve, db_session_mock): + """sql_flavor override is meaningless on a stored connection — must reject loud.""" + mock_resolve.return_value = _mock_connection() + + from testgen.mcp.tools.connections import test_connection + + with _patch_perms(), pytest.raises(MCPUserError) as exc: + test_connection(connection_id=7, sql_flavor="Snowflake") + assert "sql_flavor" in str(exc.value) + + +@patch(f"{MODULE}.test_connection_status", return_value=ConnectionStatus(message="The connection was successful.", successful=True)) +@patch(f"{MODULE}.resolve_connection") +def test_test_connection_with_id_only(mock_resolve, mock_runner, db_session_mock): + conn = _mock_connection(connection_id=7, connection_name="Local PG", project_host="localhost") + mock_resolve.return_value = conn + + from testgen.mcp.tools.connections import test_connection + + with _patch_perms(): + out = test_connection(connection_id=7) + + assert "Connection test succeeded" in out + assert "**ID:** `7`" in out + assert "**Name:** `Local PG`" in out + assert "**Type:** PostgreSQL" in out + mock_runner.assert_called_once() + + +@patch(f"{MODULE}.test_connection_status", return_value=ConnectionStatus(message="The connection was successful.", successful=True)) +@patch(f"{MODULE}.resolve_connection") +def test_test_connection_with_id_and_overrides(mock_resolve, mock_runner, db_session_mock): + """Overrides win: assigned to the loaded connection before status runs.""" + conn = _mock_connection(connection_id=7, project_host="old.host") + mock_resolve.return_value = conn + + from testgen.mcp.tools.connections import test_connection + + with _patch_perms(): + test_connection(connection_id=7, connection_params={"Host": "new.host"}) + + assert conn.project_host == "new.host" + + +@patch(f"{MODULE}.test_connection_status", return_value=ConnectionStatus(message="The connection was successful.", successful=True)) +@patch(f"{MODULE}.Connection") +def test_test_connection_inline_only(mock_conn_cls, mock_runner, db_session_mock): + """No connection_id supplied → builds an inline Connection from args.""" + inline_conn = _mock_connection() + mock_conn_cls.return_value = inline_conn + + from testgen.mcp.tools.connections import test_connection + + with _patch_perms(): + out = test_connection( + sql_flavor="PostgreSQL", + connection_params={ + "Host": "localhost", + "Port": 5432, + "Database": "d", + "Username": "u", + "Password": "p", + }, + ) + + assert "Connection test succeeded" in out + # No ID/Name lines on inline tests (no entity). + assert "**ID:**" not in out + assert "**Name:**" not in out + mock_runner.assert_called_once() + + +@patch(f"{MODULE}.test_connection_status", return_value=ConnectionStatus(message="OK", successful=True)) +@patch(f"{MODULE}.Connection") +def test_test_connection_inline_applies_defaults(mock_conn_cls, mock_runner, db_session_mock): + """Inline tests get the same flavor defaults as create.""" + inline = _mock_connection(max_query_chars=None) + mock_conn_cls.return_value = inline + + from testgen.mcp.tools.connections import test_connection + + with _patch_perms(): + test_connection( + sql_flavor="PostgreSQL", + connection_params={"Host": "h", "Port": 5432, "Database": "d", "Username": "u", "Password": "p"}, + ) + + assert inline.max_query_chars == 20000 + + +@patch(f"{MODULE}.test_connection_status") +@patch(f"{MODULE}.resolve_connection") +def test_test_connection_failure_renders_details(mock_resolve, mock_runner, db_session_mock): + """Failure with detail string renders the detail verbatim in a code block.""" + mock_resolve.return_value = _mock_connection() + driver_text = "FATAL: password authentication failed for user 'dq'" + mock_runner.return_value = ConnectionStatus( + message="Error attempting the connection.", + successful=False, + details=driver_text, + ) + + from testgen.mcp.tools.connections import test_connection + + with _patch_perms(): + out = test_connection(connection_id=7) + + assert "Connection test failed" in out + assert driver_text in out # verbatim, no scrubbing + + +def test_test_connection_inline_validation_failure(db_session_mock): + """Inline test with missing required field rejects before opening any DB connection.""" + from testgen.mcp.tools.connections import test_connection + + # PostgreSQL inline with no params -> uses a real Connection (no patch) with empty fields. + with _patch_perms(), pytest.raises(MCPUserError) as exc: + test_connection(sql_flavor="PostgreSQL", connection_params={}) + msg = str(exc.value) + assert "Cannot test connection" in msg + assert "`Host` is required for PostgreSQL." in msg + + +def test_test_connection_inline_multi_mode_requires_mode(db_session_mock): + """Inline test of a multi-mode flavor without connection_mode is rejected (not defaulted).""" + from testgen.mcp.tools.connections import test_connection + + with _patch_perms(), pytest.raises(MCPUserError) as exc: + test_connection( + sql_flavor="Salesforce Data 360", + connection_params={"Login URL": "https://my.salesforce.com", "Consumer Key": "ck"}, + ) + msg = str(exc.value) + assert "requires a connection_mode" in msg + assert "JWT Bearer Flow" in msg and "Client Credentials Flow" in msg + + +@patch(f"{MODULE}.test_connection_status") +@patch(f"{MODULE}.Connection") +def test_test_connection_inline_does_not_save(mock_conn_cls, mock_runner, db_session_mock): + """Inline test never persists — the constructed connection is not saved.""" + inline = _mock_connection() + mock_conn_cls.return_value = inline + mock_runner.return_value = ConnectionStatus(message="OK", successful=True) + + from testgen.mcp.tools.connections import test_connection + + with _patch_perms(): + test_connection( + sql_flavor="PostgreSQL", + connection_params={"Host": "h", "Port": 5432, "Database": "d", "Username": "u", "Password": "p"}, + ) + + inline.save.assert_not_called() diff --git a/tests/unit/mcp/test_tools_reference.py b/tests/unit/mcp/test_tools_reference.py index 96ead51f..97bb03bd 100644 --- a/tests/unit/mcp/test_tools_reference.py +++ b/tests/unit/mcp/test_tools_reference.py @@ -1,5 +1,7 @@ from unittest.mock import MagicMock, patch +import pytest + @patch("testgen.mcp.tools.reference.TestType") def test_get_test_type_found(mock_tt_cls, db_session_mock): @@ -263,3 +265,92 @@ def test_server_instructions_reference_column_profile_fields_resource(): # Sanity check the existing references are still present. assert "testgen://test-types" in SERVER_INSTRUCTIONS assert "testgen://hygiene-issue-types" in SERVER_INSTRUCTIONS + + +# --------------------------------------------------------------------------- +# connection_parameters_resource +# --------------------------------------------------------------------------- + + +def test_connection_parameters_index_lists_flavors(): + from testgen.mcp.tools.reference import connection_parameters_index_resource + + out = connection_parameters_index_resource() + # Accepted sql_flavor labels. + assert "PostgreSQL" in out and "Azure SQL Database" in out and "Salesforce Data 360" in out + # Each links to its per-flavor resource (keyed by code). + assert "testgen://connection-parameters/postgresql" in out + assert "testgen://connection-parameters/salesforce_data360" in out + + +def test_connection_parameters_resource_unknown_flavor(): + from testgen.mcp.exceptions import MCPUserError + from testgen.mcp.tools.reference import connection_parameters_resource + + with pytest.raises(MCPUserError) as exc: + connection_parameters_resource("not_a_flavor") + msg = str(exc.value) + assert "snowflake" in msg and "salesforce_data360" in msg + + +def test_connection_parameters_resource_postgresql_single_mode(): + from testgen.mcp.tools.reference import connection_parameters_resource + + out = connection_parameters_resource("postgresql") + assert "PostgreSQL Connection Parameters" in out + assert "Host" in out and "Username" in out + assert "Required (host mode)" in out # Host/Port/Database + # URL alternative is advertised. + assert "connect by URL" in out + # Single-mode flavor: no connection_mode instruction. + assert "Set `connection_mode`" not in out + + +def test_connection_parameters_resource_snowflake_both_modes(): + from testgen.mcp.tools.reference import connection_parameters_resource + + out = connection_parameters_resource("snowflake") + assert "Mode: Key-Pair" in out + assert "Mode: Password" in out + assert "Private Key" in out + assert "Warehouse" in out + assert "Set `connection_mode`" in out + assert 'connection_mode="Key-Pair"' in out + + +def test_connection_parameters_resource_databricks_pat_fields(): + from testgen.mcp.tools.reference import connection_parameters_resource + + out = connection_parameters_resource("databricks") + assert "Mode: Access Token" in out + assert "Mode: Service Principal (OAuth)" in out + assert "Catalog" in out + assert "HTTP Path" in out + assert "Client ID" in out + + +def test_connection_parameters_resource_port_default_note(): + """The Port row documents the flavor's conventional default port (doc-only — + the field stays required; the LLM supplies the default when the user doesn't).""" + from testgen.mcp.tools.reference import connection_parameters_resource + + out = connection_parameters_resource("postgresql") + assert "Default for PostgreSQL is 5432" in out + assert "Required (host mode)" in out # Port requirement unchanged + + +def test_connection_parameters_resource_port_default_per_flavor(): + from testgen.mcp.tools.reference import connection_parameters_resource + + out = connection_parameters_resource("sap_hana") + assert "Default for SAP HANA is 39015" in out + + +def test_connection_parameters_resource_marks_secrets(): + from testgen.mcp.tools.reference import connection_parameters_resource + + out = connection_parameters_resource("bigquery") + assert "Service Account Key" in out + assert "Secret" in out + # No URL alternative for BigQuery. + assert "connect by URL" not in out From 35157f28d5ab49a5a522b86ee296449ee74ecc9f Mon Sep 17 00:00:00 2001 From: Luis Date: Mon, 1 Jun 2026 14:11:50 -0400 Subject: [PATCH 02/78] feat(mcp): accept hygiene issue id in source data tools Co-Authored-By: Claude Opus 4.8 (1M context) --- testgen/mcp/tools/common.py | 12 +- testgen/mcp/tools/hygiene_issues.py | 13 +- testgen/mcp/tools/source_data.py | 123 ++++++++--- tests/unit/mcp/test_tools_common.py | 30 +++ tests/unit/mcp/test_tools_hygiene_issues.py | 55 +++-- tests/unit/mcp/test_tools_source_data.py | 214 +++++++++++++++++++- 6 files changed, 388 insertions(+), 59 deletions(-) diff --git a/testgen/mcp/tools/common.py b/testgen/mcp/tools/common.py index 990a759d..c1b563d6 100644 --- a/testgen/mcp/tools/common.py +++ b/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") diff --git a/testgen/mcp/tools/hygiene_issues.py b/testgen/mcp/tools/hygiene_issues.py index 19042588..01084fbc 100644 --- a/testgen/mcp/tools/hygiene_issues.py +++ b/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}**.") diff --git a/testgen/mcp/tools/source_data.py b/testgen/mcp/tools/source_data.py index 1b3fb0bb..bcab0191 100644 --- a/testgen/mcp/tools/source_data.py +++ b/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 diff --git a/tests/unit/mcp/test_tools_common.py b/tests/unit/mcp/test_tools_common.py index 2abf85c0..79c83f2e 100644 --- a/tests/unit/mcp/test_tools_common.py +++ b/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 --- diff --git a/tests/unit/mcp/test_tools_hygiene_issues.py b/tests/unit/mcp/test_tools_hygiene_issues.py index 741c9b4b..5dcbeb24 100644 --- a/tests/unit/mcp/test_tools_hygiene_issues.py +++ b/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(): diff --git a/tests/unit/mcp/test_tools_source_data.py b/tests/unit/mcp/test_tools_source_data.py index 0a888b46..1f1f26b1 100644 --- a/tests/unit/mcp/test_tools_source_data.py +++ b/tests/unit/mcp/test_tools_source_data.py @@ -1,4 +1,4 @@ -from unittest.mock import patch +from unittest.mock import MagicMock, patch from uuid import uuid4 import pandas as pd @@ -21,6 +21,21 @@ def _make_context(**overrides): return base +def _make_hygiene_issue(**overrides): + issue = MagicMock() + issue.table_groups_id = uuid4() + issue.type_id = "1001" + issue.detail = "Empty String: 12" + issue.schema_name = "public" + issue.table_name = "users" + issue.column_name = "address" + issue.profile_run_id = uuid4() + issue.project_code = "demo" + for key, value in overrides.items(): + setattr(issue, key, value) + return issue + + # --- get_source_data_query --- @@ -281,3 +296,200 @@ def test_get_source_data_passes_project_codes(mock_compute, mock_td, mock_fetch, call_kwargs = mock_td.get_source_data_context.call_args.kwargs assert call_kwargs["project_codes"] == ["proj_a"] + + +# --- argument mutual-exclusion + reference_date rule --- + + +@pytest.mark.parametrize("tool_name", ["get_source_data", "get_source_data_query"]) +def test_source_data_both_ids_rejected(tool_name, db_session_mock): + import testgen.mcp.tools.source_data as mod + + tool = getattr(mod, tool_name) + with pytest.raises(MCPUserError, match="Provide exactly one of test_definition_id or issue_id"): + tool(test_definition_id=str(uuid4()), issue_id=str(uuid4())) + + +@pytest.mark.parametrize("tool_name", ["get_source_data", "get_source_data_query"]) +def test_source_data_neither_id_rejected(tool_name, db_session_mock): + import testgen.mcp.tools.source_data as mod + + tool = getattr(mod, tool_name) + with pytest.raises(MCPUserError, match="Provide exactly one of test_definition_id or issue_id"): + tool() + + +@pytest.mark.parametrize("tool_name", ["get_source_data", "get_source_data_query"]) +def test_source_data_reference_date_with_issue_id_rejected(tool_name, db_session_mock): + import testgen.mcp.tools.source_data as mod + + tool = getattr(mod, tool_name) + with pytest.raises(MCPUserError, match="reference_date"): + tool(issue_id=str(uuid4()), reference_date="2026-01-01") + + +# --- get_source_data_query: hygiene issue --- + + +@patch("testgen.mcp.tools.source_data.build_hygiene_query") +@patch("testgen.mcp.tools.source_data.ProfilingRun") +@patch("testgen.mcp.tools.source_data.resolve_hygiene_issue") +def test_get_source_data_query_hygiene_basic(mock_resolve, mock_pr, mock_build, db_session_mock): + issue_id = str(uuid4()) + mock_resolve.return_value = _make_hygiene_issue() + mock_pr.get.return_value = MagicMock(profiling_starttime="2026-05-12 00:00:00") + mock_build.return_value = "SELECT * FROM users WHERE address = ''" + + from testgen.mcp.tools.source_data import get_source_data_query + + result = get_source_data_query(issue_id=issue_id) + + assert f"# Source Data Query for Hygiene Issue `{issue_id}`" in result + assert "public.users" in result + assert "`address`" in result + assert "SELECT * FROM users" in result + assert "Test Definition" not in result + + +@patch("testgen.mcp.tools.source_data.build_hygiene_query") +@patch("testgen.mcp.tools.source_data.ProfilingRun") +@patch("testgen.mcp.tools.source_data.resolve_hygiene_issue") +def test_get_source_data_query_hygiene_no_query_available(mock_resolve, mock_pr, mock_build, db_session_mock): + mock_resolve.return_value = _make_hygiene_issue() + mock_pr.get.return_value = MagicMock(profiling_starttime="2026-05-12 00:00:00") + mock_build.return_value = None + + from testgen.mcp.tools.source_data import get_source_data_query + + result = get_source_data_query(issue_id=str(uuid4())) + + assert "not available" in result + + +@patch("testgen.mcp.tools.source_data.resolve_hygiene_issue") +def test_get_source_data_query_hygiene_not_accessible(mock_resolve, db_session_mock): + mock_resolve.side_effect = MCPResourceNotAccessible("Hygiene issue", "x") + + from testgen.mcp.tools.source_data import get_source_data_query + + with pytest.raises(MCPResourceNotAccessible, match="Hygiene issue"): + get_source_data_query(issue_id=str(uuid4())) + + +def test_get_source_data_query_hygiene_invalid_uuid(db_session_mock): + from testgen.mcp.tools.source_data import get_source_data_query + + with pytest.raises(MCPUserError, match="not a valid UUID"): + get_source_data_query(issue_id="bad-uuid") + + +# --- get_source_data: hygiene issue --- + + +@patch("testgen.mcp.tools.source_data.fetch_hygiene_source_data") +@patch("testgen.mcp.tools.source_data.ProfilingRun") +@patch("testgen.mcp.tools.source_data.resolve_hygiene_issue") +def test_get_source_data_hygiene_ok(mock_resolve, mock_pr, mock_fetch, db_session_mock): + issue_id = str(uuid4()) + mock_resolve.return_value = _make_hygiene_issue() + mock_pr.get.return_value = MagicMock(profiling_starttime="2026-05-12 00:00:00") + df = pd.DataFrame({"address": ["", ""], "count": [12, 3]}) + mock_fetch.return_value = SourceDataResult(status="OK", message=None, query="SELECT ...", df=df) + + from testgen.mcp.tools.source_data import get_source_data + + result = get_source_data(issue_id=issue_id) + + assert f"# Source Data for Hygiene Issue `{issue_id}`" in result + assert "**Rows returned:** 2" in result + assert "public.users" in result + assert "`address`" in result + assert "SELECT ..." in result + assert "Test Definition" not in result + + +@patch("testgen.mcp.tools.source_data.fetch_hygiene_source_data") +@patch("testgen.mcp.tools.source_data.ProfilingRun") +@patch("testgen.mcp.tools.source_data.resolve_hygiene_issue") +def test_get_source_data_hygiene_na(mock_resolve, mock_pr, mock_fetch, db_session_mock): + mock_resolve.return_value = _make_hygiene_issue() + mock_pr.get.return_value = MagicMock(profiling_starttime="2026-05-12 00:00:00") + mock_fetch.return_value = SourceDataResult( + status="NA", message="Source data lookup is not available for this hygiene issue.", query=None, df=None, + ) + + from testgen.mcp.tools.source_data import get_source_data + + result = get_source_data(issue_id=str(uuid4())) + + assert "not available for this hygiene issue" in result + + +@patch("testgen.mcp.tools.source_data.fetch_hygiene_source_data") +@patch("testgen.mcp.tools.source_data.ProfilingRun") +@patch("testgen.mcp.tools.source_data.resolve_hygiene_issue") +def test_get_source_data_hygiene_nd(mock_resolve, mock_pr, mock_fetch, db_session_mock): + mock_resolve.return_value = _make_hygiene_issue() + mock_pr.get.return_value = MagicMock(profiling_starttime="2026-05-12 00:00:00") + mock_fetch.return_value = SourceDataResult( + status="ND", message="No matching data.", query="SELECT * FROM users WHERE 1=0", df=None, + ) + + from testgen.mcp.tools.source_data import get_source_data + + result = get_source_data(issue_id=str(uuid4())) + + assert "No matching data." in result + assert "SELECT * FROM users WHERE 1=0" in result + + +@patch("testgen.mcp.tools.source_data.fetch_hygiene_source_data") +@patch("testgen.mcp.tools.source_data.ProfilingRun") +@patch("testgen.mcp.tools.source_data.resolve_hygiene_issue") +def test_get_source_data_hygiene_err(mock_resolve, mock_pr, mock_fetch, db_session_mock): + mock_resolve.return_value = _make_hygiene_issue() + mock_pr.get.return_value = MagicMock(profiling_starttime="2026-05-12 00:00:00") + mock_fetch.return_value = SourceDataResult( + status="ERR", message="Connection refused", query="SELECT 1", df=None, + ) + + from testgen.mcp.tools.source_data import get_source_data + + result = get_source_data(issue_id=str(uuid4())) + + assert "**Error:** Connection refused" in result + assert "SELECT 1" in result + + +@patch("testgen.mcp.tools.source_data.fetch_hygiene_source_data") +@patch("testgen.mcp.tools.source_data.ProfilingRun") +@patch("testgen.mcp.tools.source_data.resolve_hygiene_issue") +def test_get_source_data_hygiene_mask_pii_passed(mock_resolve, mock_pr, mock_fetch, db_session_mock): + mock_resolve.return_value = _make_hygiene_issue() + mock_pr.get.return_value = MagicMock(profiling_starttime="2026-05-12 00:00:00") + df = pd.DataFrame({"address": [""]}) + mock_fetch.return_value = SourceDataResult(status="OK", message=None, query=None, df=df) + + from testgen.mcp.tools.source_data import get_source_data + + get_source_data(issue_id=str(uuid4())) + + # mask_pii is the third positional arg + assert isinstance(mock_fetch.call_args[0][2], bool) + + +@patch("testgen.mcp.tools.source_data.resolve_hygiene_issue") +def test_get_source_data_hygiene_not_accessible(mock_resolve, db_session_mock): + mock_resolve.side_effect = MCPResourceNotAccessible("Hygiene issue", "x") + + from testgen.mcp.tools.source_data import get_source_data + + with pytest.raises(MCPResourceNotAccessible, match="Hygiene issue"): + get_source_data(issue_id=str(uuid4())) + + +def test_get_source_data_hygiene_invalid_uuid(db_session_mock): + from testgen.mcp.tools.source_data import get_source_data + + with pytest.raises(MCPUserError, match="not a valid UUID"): + get_source_data(issue_id="bad-uuid") From 0e0c35e4e307bf2e5abe2f22bd0ccfb997d92cbb Mon Sep 17 00:00:00 2001 From: Luis Date: Thu, 28 May 2026 17:52:36 -0400 Subject: [PATCH 03/78] feat(mcp): add tools to manage table group - create_table_group - update_table_group - add_tables_to_group - preview_table_group Co-Authored-By: Claude Opus 4.7 (1M context) --- .../common/database/table_group_service.py | 260 ++++++ testgen/mcp/server.py | 8 + testgen/mcp/tools/common.py | 24 +- testgen/mcp/tools/table_groups.py | 618 ++++++++++++++ testgen/ui/queries/table_group_queries.py | 174 +--- tests/unit/common/test_table_group_service.py | 330 ++++++++ tests/unit/mcp/test_tools_table_groups.py | 757 ++++++++++++++++++ 7 files changed, 2025 insertions(+), 146 deletions(-) create mode 100644 testgen/common/database/table_group_service.py create mode 100644 testgen/mcp/tools/table_groups.py create mode 100644 tests/unit/common/test_table_group_service.py create mode 100644 tests/unit/mcp/test_tools_table_groups.py diff --git a/testgen/common/database/table_group_service.py b/testgen/common/database/table_group_service.py new file mode 100644 index 00000000..f4c09882 --- /dev/null +++ b/testgen/common/database/table_group_service.py @@ -0,0 +1,260 @@ +"""Shared table-group helpers — common-layer logic shared by MCP, the UI, and +any future caller. Two pieces live here, both flavor-aware: + +* ``validate_table_group_fields`` — required-field + format checks. Callers + surface the returned bullets verbatim. +* ``preview_table_group`` — schema-introspection preview. Returns + ``(preview, data_chars, sql_generator)``; the ``save_data_chars`` callback + is built by callers via ``make_save_data_chars``. + +None of these belong on the ``TableGroup`` model: preview opens an external +DB connection (I/O against an arbitrary target system), and the field-required +rules are form-shaped per-flavor logic, not model invariants. +""" + +from __future__ import annotations + +from collections.abc import Callable +from datetime import UTC, datetime +from typing import TypedDict +from uuid import UUID + +from testgen.commands.queries.refresh_data_chars_query import RefreshDataCharsSQL +from testgen.commands.run_refresh_data_chars import write_data_chars +from testgen.common.database.column_chars import ColumnChars +from testgen.common.database.flavor.flavor_service import resolve_connection_params +from testgen.common.models.connection import Connection +from testgen.common.models.table_group import TableGroup +from testgen.ui.services.database_service import fetch_from_target_db + +# --------------------------------------------------------------------------- +# Validation +# --------------------------------------------------------------------------- + +_NAME_MIN = 3 +_NAME_MAX = 40 +_SAMPLE_PCT_MIN = 1 +_SAMPLE_PCT_MAX = 100 + + +def _missing(value: object) -> bool: + return value is None or (isinstance(value, str) and not value.strip()) + + +def _coerce_int(value: object) -> int | None: + """Return ``int(value)`` for ints / numeric strings, ``None`` otherwise. + + The model stores ``profiling_delay_days`` and ``profile_sample_percent`` as + ``String`` columns; callers may pass either ints or numeric strings. Return + ``None`` to signal "not parseable" so the validator can emit the right error. + """ + if isinstance(value, bool): + return None + if isinstance(value, int): + return value + if isinstance(value, str): + try: + return int(value.strip()) + except (ValueError, AttributeError): + return None + return None + + +def validate_table_group_fields(table_group: TableGroup) -> list[str]: + """Return every validation error (empty list = valid). + + Mirrors the per-field form validators on the UI ``Add Table Group`` wizard. + The MCP tools call this and raise an ``MCPUserError`` containing the bullets + — they never duplicate a rule themselves. + """ + errors: list[str] = [] + + name = table_group.table_groups_name + if _missing(name): + errors.append("`table_group_name` is required.") + elif not (_NAME_MIN <= len(name.strip()) <= _NAME_MAX): + errors.append(f"`table_group_name` must be between {_NAME_MIN} and {_NAME_MAX} characters.") + + if _missing(table_group.table_group_schema): + errors.append("`schema` is required.") + + delay = _coerce_int(table_group.profiling_delay_days) + if delay is None or delay < 0: + errors.append("`profiling_delay_days` must be a non-negative integer.") + + pct = _coerce_int(table_group.profile_sample_percent) + if pct is None or not (_SAMPLE_PCT_MIN <= pct <= _SAMPLE_PCT_MAX): + errors.append(f"`profile_sample_percent` must be between {_SAMPLE_PCT_MIN} and {_SAMPLE_PCT_MAX}.") + + min_count = table_group.profile_sample_min_count + if not isinstance(min_count, int) or isinstance(min_count, bool) or min_count < 0: + errors.append("`profile_sample_min_count` must be a non-negative integer.") + + return errors + + +# --------------------------------------------------------------------------- +# Preview +# --------------------------------------------------------------------------- + + +class StatsPreview(TypedDict, total=False): + id: UUID | None + table_groups_name: str + table_group_schema: str + table_ct: int | None + column_ct: int | None + approx_record_ct: int | None + approx_data_point_ct: int | None + + +class TablePreview(TypedDict): + column_ct: int + approx_record_ct: int | None + approx_data_point_ct: int | None + can_access: bool | None + + +class TableGroupPreview(TypedDict): + stats: StatsPreview + tables: dict[str, TablePreview] + success: bool + message: str | None + + +_NO_CONNECTION_MESSAGE = "No connection selected. Please select a connection to preview the Table Group." +_NO_TABLES_MESSAGE = ( + "No tables found matching the criteria. Please check the Table Group configuration" + " or the database permissions." +) +_INACCESSIBLE_MESSAGE = "Some tables were not accessible. Please the check the database permissions." + + +def preview_table_group( + table_group: TableGroup, + *, + connection: Connection | None = None, + verify_access: bool = True, +) -> tuple[TableGroupPreview, list[ColumnChars] | None, RefreshDataCharsSQL | None]: + """Probe the connected target DB for tables matching ``table_group``'s filters. + + Returns ``(preview, data_chars, sql_generator)`` — three picklable values. + On the error path (no connection, DDF failure, or empty schema) the second + and third tuple elements are ``None``. + + Use ``make_save_data_chars(data_chars, sql_generator)`` in callers that + need to record the introspected metadata in ``data_chars``. + """ + preview: TableGroupPreview = { + "stats": { + "id": table_group.id, + "table_groups_name": table_group.table_groups_name, + "table_group_schema": table_group.table_group_schema, + }, + "tables": {}, + "success": True, + "message": None, + } + + if not (connection or table_group.connection_id): + preview["success"] = False + preview["message"] = _NO_CONNECTION_MESSAGE + return preview, None, None + + data_chars: list[ColumnChars] | None = None + sql_generator: RefreshDataCharsSQL | None = None + try: + if connection is None: + connection = Connection.get(table_group.connection_id) + preview, data_chars, sql_generator = _build_preview(table_group, connection) + + if verify_access and preview["success"]: + for table_name in list(preview["tables"]): + try: + results = fetch_from_target_db(connection, *sql_generator.verify_access(table_name)) + except Exception: + preview["tables"][table_name]["can_access"] = False + else: + preview["tables"][table_name]["can_access"] = bool(results) and len(results) > 0 + if not all(t["can_access"] for t in preview["tables"].values()): + preview["message"] = _INACCESSIBLE_MESSAGE + except Exception as error: + preview["success"] = False + preview["message"] = error.args[0] if error.args else str(error) + data_chars = None + sql_generator = None + + return preview, data_chars, sql_generator + + +def make_save_data_chars( + data_chars: list[ColumnChars], + sql_generator: RefreshDataCharsSQL, +) -> Callable[[UUID], None]: + """Build the ``save_data_chars(table_group_id)`` callback for the caller. + + Kept out of ``preview_table_group`` so the service's return value stays + picklable; local closures don't pickle. + """ + def save(table_group_id: UUID) -> None: + # Unsaved table groups won't have an ID yet; sync it before writing. + sql_generator.table_group.id = table_group_id + write_data_chars(data_chars, sql_generator, datetime.now(UTC)) + return save + + +def _build_preview( + table_group: TableGroup, + connection: Connection, +) -> tuple[TableGroupPreview, list[ColumnChars], RefreshDataCharsSQL]: + sql_generator = RefreshDataCharsSQL(connection, table_group) + if sql_generator.flavor_service.metadata_via_api: + params = resolve_connection_params(connection.__dict__) + api_columns = sql_generator.flavor_service.get_schema_columns(params, table_group.table_group_schema) or [] + data_chars = sql_generator.filter_schema_columns(api_columns) + else: + rows = fetch_from_target_db(connection, *sql_generator.get_schema_ddf()) + data_chars = [ColumnChars(**column) for column in rows] + + preview: TableGroupPreview = { + "stats": { + "id": table_group.id, + "table_groups_name": table_group.table_groups_name, + "table_group_schema": table_group.table_group_schema, + "table_ct": 0, + "column_ct": 0, + "approx_record_ct": None, + "approx_data_point_ct": None, + }, + "tables": {}, + "success": True, + "message": None, + } + stats = preview["stats"] + tables = preview["tables"] + + for column in data_chars: + if not tables.get(column.table_name): + tables[column.table_name] = { + "column_ct": 0, + "approx_record_ct": column.approx_record_ct, + "approx_data_point_ct": None, + "can_access": None, + } + stats["table_ct"] += 1 + if column.approx_record_ct is not None: + stats["approx_record_ct"] = (stats["approx_record_ct"] or 0) + column.approx_record_ct + + stats["column_ct"] += 1 + tables[column.table_name]["column_ct"] += 1 + if column.approx_record_ct is not None: + stats["approx_data_point_ct"] = (stats["approx_data_point_ct"] or 0) + column.approx_record_ct + tables[column.table_name]["approx_data_point_ct"] = ( + tables[column.table_name]["approx_data_point_ct"] or 0 + ) + column.approx_record_ct + + if len(data_chars) <= 0: + preview["success"] = False + preview["message"] = _NO_TABLES_MESSAGE + + return preview, data_chars, sql_generator diff --git a/testgen/mcp/server.py b/testgen/mcp/server.py index 46683db5..17f555d3 100644 --- a/testgen/mcp/server.py +++ b/testgen/mcp/server.py @@ -201,6 +201,11 @@ def build_mcp_server( update_schedule, ) from testgen.mcp.tools.source_data import get_source_data, get_source_data_query + from testgen.mcp.tools.table_groups import ( + create_table_group, + preview_table_group, + update_table_group, + ) from testgen.mcp.tools.test_definitions import ( bulk_update_tests, create_test, @@ -314,6 +319,9 @@ def safe_prompt(fn): safe_tool(update_notification) safe_tool(delete_notification) safe_tool(test_connection) + safe_tool(create_table_group) + safe_tool(update_table_group) + safe_tool(preview_table_group) # Resources safe_resource("testgen://test-types", test_types_resource) diff --git a/testgen/mcp/tools/common.py b/testgen/mcp/tools/common.py index c1b563d6..c6f649ec 100644 --- a/testgen/mcp/tools/common.py +++ b/testgen/mcp/tools/common.py @@ -514,6 +514,18 @@ def format_page_footer(total: int, page: int, limit: int) -> str: # Extract a new resolve_ here when a second caller needs the same parse-uuid + # perm-scoped lookup + collapsed-error pattern. +def resolve_connection(connection_id: int) -> Connection: + """Resolve a connection ID, collapsing missing-or-inaccessible into one error path.""" + perms = get_project_permissions() + conn = Connection.get( + connection_id, + Connection.project_code.in_(perms.allowed_codes), + ) + if conn is None: + raise MCPResourceNotAccessible("Connection", str(connection_id)) + return conn + + def resolve_table_group(table_group_id: str) -> TableGroup: """Resolve a TG ID, collapsing missing-or-inaccessible into one error path.""" tg_uuid = parse_uuid(table_group_id, "table_group_id") @@ -739,18 +751,6 @@ def format_notification_trigger(event: NotificationEvent | str, settings: dict | return None -def resolve_connection(connection_id: int) -> Connection: - """Resolve a connection ID, collapsing missing-or-inaccessible into one error path.""" - perms = get_project_permissions() - conn = Connection.get( - connection_id, - Connection.project_code.in_(perms.allowed_codes), - ) - if conn is None: - raise MCPResourceNotAccessible("Connection", str(connection_id)) - return conn - - # Flavor display labels are the single source of truth in ``common/flavors.py`` # (shared with the UI page). These maps just re-shape them for the MCP layer. SQL_FLAVOR_CODE_TO_LABEL: dict[str, SqlFlavorLabel] = dict(FLAVOR_CODE_TO_LABEL) diff --git a/testgen/mcp/tools/table_groups.py b/testgen/mcp/tools/table_groups.py new file mode 100644 index 00000000..1485055c --- /dev/null +++ b/testgen/mcp/tools/table_groups.py @@ -0,0 +1,618 @@ +"""MCP tools for table groups — create, update, append tables, preview. + +Each tool gates on the ``edit`` permission. Validation and target-DB +introspection are delegated to ``testgen.common.database.table_group_service`` +so the rules and SQL paths stay in one place. +""" + +from __future__ import annotations + +from typing import Any + +from sqlalchemy.exc import IntegrityError + +from testgen.common.database.table_group_service import ( + preview_table_group as preview_table_group_service, +) +from testgen.common.database.table_group_service import validate_table_group_fields +from testgen.common.models import with_database_session +from testgen.common.models.connection import Connection +from testgen.common.models.table_group import TableGroup +from testgen.mcp.exceptions import MCPUserError +from testgen.mcp.permissions import mcp_permission +from testgen.mcp.tools.common import resolve_connection, resolve_table_group +from testgen.mcp.tools.markdown import MdDoc + +_DUPLICATE_NAME_MESSAGE = "A Table Group with the same name already exists." +_SCHEMA_LOCKED_MESSAGE = ( + "Schema cannot be changed once the table group has been used. " + "Delete and recreate the table group to use a different schema." +) + + +@with_database_session +@mcp_permission("edit") +def create_table_group( + connection_id: int, + table_group_name: str, + schema: str, + *, + description: str | None = None, + table_set: list[str] | None = None, + include_mask: str | None = None, + exclude_mask: str | None = None, + profile_id_column_mask: str | None = None, + profile_sk_column_mask: str | None = None, + profile_use_sampling: bool | None = None, + profile_sample_percent: int | None = None, + profile_sample_min_count: int | None = None, + profiling_delay_days: int | None = None, + profile_flag_cdes: bool | None = None, + profile_flag_pii: bool | None = None, + profile_exclude_xde: bool | None = None, + include_in_dashboard: bool | None = None, + add_scorecard: bool = True, + data_source: str | None = None, + source_system: str | None = None, + source_process: str | None = None, + data_location: str | None = None, + business_domain: str | None = None, + stakeholder_group: str | None = None, + transform_level: str | None = None, + data_product: str | None = None, +) -> str: + """Create a table group on an existing connection. + + The table group inherits its project from the connection. ``include_mask`` + and ``exclude_mask`` are SQL ``LIKE`` patterns (e.g. ``fact_%,dim_%``); + ``table_set`` is an explicit list. All filters compose with ``AND``. + + Args: + connection_id: Bigint connection ID, e.g. from ``get_data_inventory``. + table_group_name: 3-40 character display name. Must be unique within the project. + schema: Schema name on the target database, e.g. ``public``. For Salesforce Data 360 + connections, use the data space name. + description: Optional free-text description. + table_set: Explicit list of table names. Combined with masks if also set. + include_mask: Comma-separated SQL LIKE patterns to include (e.g. ``fact_%,dim_%``). + exclude_mask: Comma-separated SQL LIKE patterns to exclude. + profile_id_column_mask: SQL LIKE pattern marking ID columns. Default ``%id``. + profile_sk_column_mask: SQL LIKE pattern marking surrogate-key columns. Default ``%_sk``. + profile_use_sampling: Whether to sample large tables during profiling. + profile_sample_percent: Sample size as a percent (1-100). + profile_sample_min_count: Minimum row count when sampling. + profiling_delay_days: Number of days to wait before new profiling will be available + to generate tests. + profile_flag_cdes: Whether profiling flags Critical Data Elements. + profile_flag_pii: Whether profiling flags Personally Identifiable Information. + profile_exclude_xde: Whether profiling excludes columns flagged as excluded data elements. + include_in_dashboard: Whether the table group appears on the project dashboard. + add_scorecard: Whether to add a scorecard for the table group to the Quality Dashboard. + data_source: Catalog tag — original source of the dataset. + source_system: Catalog tag — enterprise system source for the dataset. + source_process: Catalog tag — process, program, or data flow that produced the dataset. + data_location: Catalog tag — physical or virtual location of the dataset + (e.g. ``Headquarters``, ``Cloud``). + business_domain: Catalog tag — business division responsible for the dataset + (e.g. ``Finance``, ``Sales``, ``Manufacturing``). + stakeholder_group: Catalog tag — data owners or stakeholders responsible for the dataset. + transform_level: Catalog tag — data warehouse processing stage (e.g. ``Raw``, + ``Conformed``, ``Processed``, ``Reporting``) or Medallion level + (``bronze``, ``silver``, ``gold``). + data_product: Catalog tag — data domain that comprises the dataset. + """ + connection = resolve_connection(connection_id) + + table_group = TableGroup( + project_code=connection.project_code, + connection_id=connection.connection_id, + table_groups_name=table_group_name, + table_group_schema=schema, + **_model_defaults(), + ) + _apply_args_to_table_group( + table_group, + description=description, + table_set=table_set, + include_mask=include_mask, + exclude_mask=exclude_mask, + profile_id_column_mask=profile_id_column_mask, + profile_sk_column_mask=profile_sk_column_mask, + profile_use_sampling=profile_use_sampling, + profile_sample_percent=profile_sample_percent, + profile_sample_min_count=profile_sample_min_count, + profiling_delay_days=profiling_delay_days, + profile_flag_cdes=profile_flag_cdes, + profile_flag_pii=profile_flag_pii, + profile_exclude_xde=profile_exclude_xde, + include_in_dashboard=include_in_dashboard, + data_source=data_source, + source_system=source_system, + source_process=source_process, + data_location=data_location, + business_domain=business_domain, + stakeholder_group=stakeholder_group, + transform_level=transform_level, + data_product=data_product, + ) + + errors = validate_table_group_fields(table_group) + if errors: + _raise_validation_error(errors, "Table group creation rejected. No changes saved.") + + try: + table_group.save(add_scorecard_definition=add_scorecard) + except IntegrityError as err: + _maybe_raise_duplicate_name(err) + raise + + return _render_created_table_group(table_group, connection) + + +@with_database_session +@mcp_permission("edit") +def update_table_group( + table_group_id: str, + *, + table_group_name: str | None = None, + schema: str | None = None, + description: str | None = None, + table_set: list[str] | None = None, + include_mask: str | None = None, + exclude_mask: str | None = None, + profile_id_column_mask: str | None = None, + profile_sk_column_mask: str | None = None, + profile_use_sampling: bool | None = None, + profile_sample_percent: int | None = None, + profile_sample_min_count: int | None = None, + profiling_delay_days: int | None = None, + profile_flag_cdes: bool | None = None, + profile_flag_pii: bool | None = None, + profile_exclude_xde: bool | None = None, + include_in_dashboard: bool | None = None, + data_source: str | None = None, + source_system: str | None = None, + source_process: str | None = None, + data_location: str | None = None, + business_domain: str | None = None, + stakeholder_group: str | None = None, + transform_level: str | None = None, + data_product: str | None = None, +) -> str: + """Update fields on an existing table group. Atomic — no partial save. + + Connection and project are immutable — delete and recreate the table group + to re-parent it. ``schema`` is also immutable once the table group has been + used (profiled or has test suites); supply a different ``schema`` only on + unused table groups. + + Args: + table_group_id: UUID of the table group to update. + table_group_name: New display name (3-40 chars). Must be unique within the project. + schema: New target-DB schema. Rejected if the table group has been used. + description: Free-text description. + table_set: Replacement explicit table list (full replacement of the current list). + include_mask: Comma-separated SQL LIKE patterns to include. + exclude_mask: Comma-separated SQL LIKE patterns to exclude. + profile_id_column_mask: SQL LIKE pattern marking ID columns. + profile_sk_column_mask: SQL LIKE pattern marking surrogate-key columns. + profile_use_sampling: Whether to sample large tables during profiling. + profile_sample_percent: Sample size as a percent (1-100). + profile_sample_min_count: Minimum row count when sampling. + profiling_delay_days: Number of days to wait before new profiling will be available + to generate tests. + profile_flag_cdes: Whether profiling flags CDEs. + profile_flag_pii: Whether profiling flags PII. + profile_exclude_xde: Whether profiling excludes XDE columns. + include_in_dashboard: Whether the table group appears on the project dashboard. + data_source: Catalog tag — original source of the dataset. + source_system: Catalog tag — enterprise system source for the dataset. + source_process: Catalog tag — process, program, or data flow that produced the dataset. + data_location: Catalog tag — physical or virtual location of the dataset + (e.g. ``Headquarters``, ``Cloud``). + business_domain: Catalog tag — business division responsible for the dataset + (e.g. ``Finance``, ``Sales``, ``Manufacturing``). + stakeholder_group: Catalog tag — data owners or stakeholders responsible for the dataset. + transform_level: Catalog tag — data warehouse processing stage (e.g. ``Raw``, + ``Conformed``, ``Processed``, ``Reporting``) or Medallion level + (``bronze``, ``silver``, ``gold``). + data_product: Catalog tag — data domain that comprises the dataset. + """ + supplied = { + "table_group_name": table_group_name, + "schema": schema, + "description": description, + "table_set": table_set, + "include_mask": include_mask, + "exclude_mask": exclude_mask, + "profile_id_column_mask": profile_id_column_mask, + "profile_sk_column_mask": profile_sk_column_mask, + "profile_use_sampling": profile_use_sampling, + "profile_sample_percent": profile_sample_percent, + "profile_sample_min_count": profile_sample_min_count, + "profiling_delay_days": profiling_delay_days, + "profile_flag_cdes": profile_flag_cdes, + "profile_flag_pii": profile_flag_pii, + "profile_exclude_xde": profile_exclude_xde, + "include_in_dashboard": include_in_dashboard, + "data_source": data_source, + "source_system": source_system, + "source_process": source_process, + "data_location": data_location, + "business_domain": business_domain, + "stakeholder_group": stakeholder_group, + "transform_level": transform_level, + "data_product": data_product, + } + if all(value is None for value in supplied.values()): + raise MCPUserError("No fields supplied to update.") + + table_group = resolve_table_group(table_group_id) + + if ( + schema is not None + and schema != table_group.table_group_schema + and TableGroup.is_in_use([table_group.id]) + ): + raise MCPUserError(_SCHEMA_LOCKED_MESSAGE) + + before = _snapshot(table_group) + _apply_args_to_table_group(table_group, **supplied) + + errors = validate_table_group_fields(table_group) + if errors: + _raise_validation_error(errors, "Update rejected. No changes saved.") + + after = _snapshot(table_group) + changed = {attr for attr in before if before[attr] != after[attr]} + + doc = MdDoc() + doc.heading(1, f"Table Group `{table_group.table_groups_name}` updated") + doc.field("ID", str(table_group.id), code=True) + + if not changed: + doc.text("No fields changed — supplied values matched the current state.") + return doc.render() + + try: + table_group.save() + except IntegrityError as err: + _maybe_raise_duplicate_name(err) + raise + + rows: list[list[object]] = [] + for attr in _DIFF_ATTRS: + if attr not in changed: + continue + label = _DIFF_LABELS.get(attr, attr) + rows.append([label, _render_field_value(before[attr]), _render_field_value(after[attr])]) + doc.table(["Field", "Before", "After"], rows, code=[0]) + return doc.render() + + +@with_database_session +@mcp_permission("edit") +def preview_table_group( + table_group_id: str, + verify_access: bool = False, +) -> str: + """Probe the target database for tables matching a table group's filters. + + Returns counts plus a per-table breakdown. Does not save anything to the + application database. + + Args: + table_group_id: UUID of the table group. + verify_access: When True, probe read access on every matched table. + """ + table_group = resolve_table_group(table_group_id) + connection = Connection.get_by_table_group(table_group.id) + if connection is None: + raise MCPUserError("Cannot preview — the table group's connection is unavailable.") + + preview, _data_chars, _sql_generator = preview_table_group_service( + table_group, connection=connection, verify_access=verify_access, + ) + stats = preview["stats"] + name = stats.get("table_groups_name") or table_group.table_groups_name + + doc = MdDoc() + + if not preview["success"]: + if preview.get("message", "").startswith("No tables found matching the criteria"): + doc.heading(1, f"Preview for table group `{name}` returned no tables") + doc.text(preview["message"]) + else: + doc.heading(1, f"Preview failed for table group `{name}`") + doc.text(preview.get("message") or "Preview failed for an unknown reason.") + return doc.render() + + doc.heading(1, f"Preview for table group `{name}`") + doc.field("Table Group ID", str(table_group.id), code=True) + doc.field("Schema", stats.get("table_group_schema"), code=True) + doc.field("Tables matched", stats.get("table_ct") or 0) + doc.field("Total columns", stats.get("column_ct") or 0) + if stats.get("approx_record_ct") is not None: + doc.field("Approx rows", stats.get("approx_record_ct")) + if stats.get("approx_data_point_ct") is not None: + doc.field("Approx data points", stats.get("approx_data_point_ct")) + + headers = ["Table", "Columns", "Approx Rows", "Approx Data Points"] + if verify_access: + headers.append("Read Access") + + rows: list[list[object]] = [] + for table_name, info in preview["tables"].items(): + row: list[object] = [ + table_name, + info.get("column_ct"), + info.get("approx_record_ct"), + info.get("approx_data_point_ct"), + ] + if verify_access: + access = info.get("can_access") + row.append("Yes" if access is True else "No" if access is False else "Unknown") + rows.append(row) + doc.table(headers, rows, code=[0]) + + if verify_access and preview.get("message"): + doc.text(preview["message"]) + + return doc.render() + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +_DIFF_ATTRS: tuple[str, ...] = ( + "table_groups_name", + "table_group_schema", + "profiling_table_set", + "profiling_include_mask", + "profiling_exclude_mask", + "profile_id_column_mask", + "profile_sk_column_mask", + "profile_use_sampling", + "profile_sample_percent", + "profile_sample_min_count", + "profiling_delay_days", + "profile_flag_cdes", + "profile_flag_pii", + "profile_exclude_xde", + "include_in_dashboard", + "description", + "data_source", + "source_system", + "source_process", + "data_location", + "business_domain", + "stakeholder_group", + "transform_level", + "data_product", +) + +_DIFF_LABELS: dict[str, str] = { + "table_groups_name": "Name", + "table_group_schema": "Schema", + "profiling_table_set": "Table set", + "profiling_include_mask": "Include mask", + "profiling_exclude_mask": "Exclude mask", + "profile_id_column_mask": "ID column mask", + "profile_sk_column_mask": "SK column mask", + "profile_use_sampling": "Sampling", + "profile_sample_percent": "Sample %", + "profile_sample_min_count": "Sample min rows", + "profiling_delay_days": "Min profiling age (days)", + "profile_flag_cdes": "Flag CDEs", + "profile_flag_pii": "Flag PII", + "profile_exclude_xde": "Exclude XDE", + "include_in_dashboard": "Include in dashboard", + "description": "Description", + "data_source": "Data source", + "source_system": "Source system", + "source_process": "Source process", + "data_location": "Data location", + "business_domain": "Business domain", + "stakeholder_group": "Stakeholder group", + "transform_level": "Transform level", + "data_product": "Data product", +} + +_CATALOG_ATTRS: tuple[str, ...] = ( + "data_source", + "source_system", + "source_process", + "data_location", + "business_domain", + "stakeholder_group", + "transform_level", + "data_product", +) + + +# Mirror SQLAlchemy ``Column(default=...)`` values into the constructor kwargs. +# Column defaults only fire at flush time, but ``validate_table_group_fields`` +# runs *before* flush. Without seeding these, every create call that omits +# e.g. ``profile_sample_percent`` would fail validation. Computing once at +# import time keeps the values in sync with the model and survives tests that +# patch the ``TableGroup`` class itself. +# +# ``YNString`` columns store raw "Y"/"N" strings as their default but expose +# the attribute as ``bool``; normalize so the in-memory render path doesn't +# treat "N" as truthy before the row is reloaded. +def _normalize_default(column, raw: Any) -> Any: + from testgen.common.models.custom_types import YNString + + if isinstance(column.type, YNString) and isinstance(raw, str): + return raw == "Y" + return raw + + +_MODEL_DEFAULTS: dict[str, Any] = { + column.name: _normalize_default(column, column.default.arg) + for column in TableGroup.__table__.columns + if ( + column.default is not None + and not column.primary_key + and getattr(column.default, "is_scalar", False) + ) +} + + +def _model_defaults() -> dict[str, Any]: + return dict(_MODEL_DEFAULTS) + + +def _apply_args_to_table_group( + table_group: TableGroup, + *, + table_group_name: str | None = None, + schema: str | None = None, + description: str | None = None, + table_set: list[str] | None = None, + include_mask: str | None = None, + exclude_mask: str | None = None, + profile_id_column_mask: str | None = None, + profile_sk_column_mask: str | None = None, + profile_use_sampling: bool | None = None, + profile_sample_percent: int | None = None, + profile_sample_min_count: int | None = None, + profiling_delay_days: int | None = None, + profile_flag_cdes: bool | None = None, + profile_flag_pii: bool | None = None, + profile_exclude_xde: bool | None = None, + include_in_dashboard: bool | None = None, + data_source: str | None = None, + source_system: str | None = None, + source_process: str | None = None, + data_location: str | None = None, + business_domain: str | None = None, + stakeholder_group: str | None = None, + transform_level: str | None = None, + data_product: str | None = None, +) -> None: + """Apply every non-None arg to its model field. + + Casts ``table_set: list[str]`` to comma-joined string, and casts + integer args for columns the model stores as strings. + """ + if table_group_name is not None: + table_group.table_groups_name = table_group_name + if schema is not None: + table_group.table_group_schema = schema + if description is not None: + table_group.description = description + if table_set is not None: + table_group.profiling_table_set = ",".join(table_set) + if include_mask is not None: + table_group.profiling_include_mask = include_mask + if exclude_mask is not None: + table_group.profiling_exclude_mask = exclude_mask + if profile_id_column_mask is not None: + table_group.profile_id_column_mask = profile_id_column_mask + if profile_sk_column_mask is not None: + table_group.profile_sk_column_mask = profile_sk_column_mask + if profile_use_sampling is not None: + table_group.profile_use_sampling = profile_use_sampling + if profile_sample_percent is not None: + table_group.profile_sample_percent = str(profile_sample_percent) + if profile_sample_min_count is not None: + table_group.profile_sample_min_count = profile_sample_min_count + if profiling_delay_days is not None: + table_group.profiling_delay_days = str(profiling_delay_days) + if profile_flag_cdes is not None: + table_group.profile_flag_cdes = profile_flag_cdes + if profile_flag_pii is not None: + table_group.profile_flag_pii = profile_flag_pii + if profile_exclude_xde is not None: + table_group.profile_exclude_xde = profile_exclude_xde + if include_in_dashboard is not None: + table_group.include_in_dashboard = include_in_dashboard + if data_source is not None: + table_group.data_source = data_source + if source_system is not None: + table_group.source_system = source_system + if source_process is not None: + table_group.source_process = source_process + if data_location is not None: + table_group.data_location = data_location + if business_domain is not None: + table_group.business_domain = business_domain + if stakeholder_group is not None: + table_group.stakeholder_group = stakeholder_group + if transform_level is not None: + table_group.transform_level = transform_level + if data_product is not None: + table_group.data_product = data_product + + +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 _maybe_raise_duplicate_name(err: IntegrityError) -> None: + if "table_groups_name_unique" in str(err.orig): + raise MCPUserError(_DUPLICATE_NAME_MESSAGE) from err + + +def _snapshot(table_group: TableGroup) -> dict[str, Any]: + return {attr: getattr(table_group, attr, None) for attr in _DIFF_ATTRS} + + +def _render_field_value(value: Any) -> str | None: + if isinstance(value, bool): + return "Yes" if value else "No" + if value is None or value == "": + return None + return str(value) + + +def _render_created_table_group(table_group: TableGroup, connection: Connection) -> str: + doc = MdDoc() + doc.heading(1, f"Table Group `{table_group.table_groups_name}` created") + doc.field("ID", str(table_group.id), code=True) + doc.field("Project", table_group.project_code, code=True) + doc.field( + "Connection", + f"{connection.connection_name} (`{connection.connection_id}`)", + ) + doc.field("Schema", table_group.table_group_schema, code=True) + if table_group.description: + doc.field("Description", table_group.description) + + if table_group.profiling_table_set: + doc.field("Table set", table_group.profiling_table_set, code=True) + if table_group.profiling_include_mask: + doc.field("Include mask", table_group.profiling_include_mask, code=True) + if table_group.profiling_exclude_mask: + doc.field("Exclude mask", table_group.profiling_exclude_mask, code=True) + + doc.field("Profile sampling", "Yes" if table_group.profile_use_sampling else "No") + if table_group.profile_use_sampling: + doc.field("Sample %", table_group.profile_sample_percent) + doc.field("Sample min rows", table_group.profile_sample_min_count) + + profile_flags = [] + if table_group.profile_flag_cdes: + profile_flags.append("Flag CDEs") + if table_group.profile_flag_pii: + profile_flags.append("Flag PII") + if table_group.profile_exclude_xde: + profile_flags.append("Exclude XDE") + if profile_flags: + doc.field("Profiling flags", ", ".join(profile_flags)) + + doc.field("Min profiling age (days)", table_group.profiling_delay_days) + doc.field("Include in dashboard", "Yes" if table_group.include_in_dashboard else "No") + + if any(getattr(table_group, attr, None) for attr in _CATALOG_ATTRS): + doc.heading(2, "Catalog") + for attr in _CATALOG_ATTRS: + value = getattr(table_group, attr, None) + if value: + doc.field(_DIFF_LABELS[attr], value) + + return doc.render() diff --git a/testgen/ui/queries/table_group_queries.py b/testgen/ui/queries/table_group_queries.py index 0db27e12..1d4cdf6f 100644 --- a/testgen/ui/queries/table_group_queries.py +++ b/testgen/ui/queries/table_group_queries.py @@ -1,98 +1,60 @@ +"""UI cache adapter around the table-group preview service. + +The implementation lives in ``testgen.common.database.table_group_service`` so +MCP tools and the CLI can reuse identical logic without a Streamlit runtime. +This module wraps the service in ``@st.cache_data`` for the Streamlit pages. +""" + from collections.abc import Callable -from datetime import UTC, datetime -from typing import TypedDict from uuid import UUID import streamlit as st -from testgen.commands.queries.refresh_data_chars_query import RefreshDataCharsSQL -from testgen.commands.run_refresh_data_chars import write_data_chars -from testgen.common.database.column_chars import ColumnChars -from testgen.common.database.flavor.flavor_service import resolve_connection_params +from testgen.common.database.table_group_service import ( + TableGroupPreview, + make_save_data_chars, + preview_table_group, +) from testgen.common.models.connection import Connection from testgen.common.models.table_group import TableGroup -from testgen.ui.services.database_service import fetch_from_target_db from testgen.ui.services.query_cache import get_connection -class StatsPreview(TypedDict): - id: UUID - table_groups_name: str - table_group_schema: str - table_ct: int | None - column_ct: int | None - approx_record_ct: int | None - approx_data_point_ct: int | None - -class TablePreview(TypedDict): - column_ct: int - approx_record_ct: int | None - approx_data_point_ct: int | None - can_access: bool | None - - -class TableGroupPreview(TypedDict): - stats: StatsPreview - tables: dict[str, TablePreview] - success: bool - message: str | None - - def get_table_group_preview( table_group: TableGroup, connection: Connection | None = None, verify_table_access: bool = False, -) -> tuple[TableGroupPreview, Callable[[UUID], None]]: - table_group_preview: TableGroupPreview = { - "stats": { - "id": table_group.id, - "table_groups_name": table_group.table_groups_name, - "table_group_schema": table_group.table_group_schema, - }, - "tables": {}, - "success": True, - "message": None, - } - save_data_chars = None - - if connection or table_group.connection_id: - try: - connection = connection or get_connection(table_group.connection_id) - table_group_preview, data_chars, sql_generator = _get_preview(table_group, connection) - - def save_data_chars(table_group_id: UUID) -> None: - # Unsaved table groups will not have an ID, so we have to update it after saving - sql_generator.table_group.id = table_group_id - write_data_chars(data_chars, sql_generator, datetime.now(UTC)) - - if verify_table_access: - tables_preview = table_group_preview["tables"] - for table_name in tables_preview.keys(): - try: - results = fetch_from_target_db(connection, *sql_generator.verify_access(table_name)) - except Exception as error: - tables_preview[table_name]["can_access"] = False - else: - tables_preview[table_name]["can_access"] = results is not None and len(results) > 0 - - if not all(table["can_access"] for table in tables_preview.values()): - table_group_preview["message"] = ( - "Some tables were not accessible. Please the check the database permissions." - ) - except Exception as error: - table_group_preview["success"] = False - table_group_preview["message"] = error.args[0] - else: - table_group_preview["success"] = False - table_group_preview["message"] = ( - "No connection selected. Please select a connection to preview the Table Group." +) -> tuple[TableGroupPreview, Callable[[UUID], None] | None]: + """Streamlit-cached wrapper around ``preview_table_group``. + + The service returns ``(preview, data_chars, sql_generator)`` — all picklable — + so the cache can store the result. The ``save_data_chars`` closure is built + here, outside the cached function, because local closures can't be pickled. + + When the caller does not supply a ``Connection`` and the table group has a + ``connection_id``, the connection is resolved via the Streamlit cache so + repeated previews on the same page don't re-fetch it. + """ + if connection is None and table_group.connection_id: + connection = get_connection(table_group.connection_id) + + if verify_table_access: + preview, data_chars, sql_generator = preview_table_group( + table_group, connection=connection, verify_access=True, ) + else: + preview, data_chars, sql_generator = _cached_preview(table_group, connection) - return table_group_preview, save_data_chars + save = ( + make_save_data_chars(data_chars, sql_generator) + if data_chars is not None and sql_generator is not None + else None + ) + return preview, save def reset_table_group_preview() -> None: - _get_preview.clear() + _cached_preview.clear() @st.cache_data( @@ -107,61 +69,5 @@ def reset_table_group_preview() -> None: Connection: lambda x: x.to_dict(), }, ) -def _get_preview( - table_group: TableGroup, - connection: Connection, -) -> tuple[TableGroupPreview, list[ColumnChars], RefreshDataCharsSQL]: - sql_generator = RefreshDataCharsSQL(connection, table_group) - if sql_generator.flavor_service.metadata_via_api: - params = resolve_connection_params(connection.__dict__) - api_columns = sql_generator.flavor_service.get_schema_columns(params, table_group.table_group_schema) or [] - data_chars = sql_generator.filter_schema_columns(api_columns) - else: - rows = fetch_from_target_db(connection, *sql_generator.get_schema_ddf()) - data_chars = [ColumnChars(**column) for column in rows] - - preview: TableGroupPreview = { - "stats": { - "id": table_group.id, - "table_groups_name": table_group.table_groups_name, - "table_group_schema": table_group.table_group_schema, - "table_ct": 0, - "column_ct": 0, - "approx_record_ct": None, - "approx_data_point_ct": None, - }, - "tables": {}, - "success": True, - "message": None, - } - stats = preview["stats"] - tables = preview["tables"] - - for column in data_chars: - if not tables.get(column.table_name): - tables[column.table_name] = { - "column_ct": 0, - "approx_record_ct": column.approx_record_ct, - "approx_data_point_ct": None, - "can_access": None, - } - stats["table_ct"] += 1 - if column.approx_record_ct is not None: - stats["approx_record_ct"] = (stats["approx_record_ct"] or 0) + column.approx_record_ct - - stats["column_ct"] += 1 - tables[column.table_name]["column_ct"] += 1 - if column.approx_record_ct is not None: - stats["approx_data_point_ct"] = (stats["approx_data_point_ct"] or 0) + column.approx_record_ct - tables[column.table_name]["approx_data_point_ct"] = ( - tables[column.table_name]["approx_data_point_ct"] or 0 - ) + column.approx_record_ct - - if len(data_chars) <= 0: - preview["success"] = False - preview["message"] = ( - "No tables found matching the criteria. Please check the Table Group configuration" - " or the database permissions." - ) - - return preview, data_chars, sql_generator +def _cached_preview(table_group, connection): + return preview_table_group(table_group, connection=connection, verify_access=False) diff --git a/tests/unit/common/test_table_group_service.py b/tests/unit/common/test_table_group_service.py new file mode 100644 index 00000000..b8d051ba --- /dev/null +++ b/tests/unit/common/test_table_group_service.py @@ -0,0 +1,330 @@ +"""Tests for the table_group_service common-layer module. + +Covers ``validate_table_group_fields`` and ``preview_table_group`` (happy path, +verify-access flag, partial-inaccessible footer, no-connection branch, +connection failure, empty result). +""" + +from unittest.mock import patch + +import pytest + +from testgen.common.database.table_group_service import ( + preview_table_group, + validate_table_group_fields, +) +from testgen.common.models.connection import Connection +from testgen.common.models.table_group import TableGroup + +pytestmark = pytest.mark.unit + +MODULE = "testgen.common.database.table_group_service" + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +def _tg(**overrides) -> TableGroup: + """Build a TableGroup without touching the DB. Defaults to a valid baseline.""" + defaults = { + "project_code": "demo", + "connection_id": 7, + "table_groups_name": "Sample TG", + "table_group_schema": "public", + "profile_use_sampling": False, + "profile_sample_percent": "30", + "profile_sample_min_count": 100000, + "profiling_delay_days": "0", + } + defaults.update(overrides) + return TableGroup(**defaults) + + +def _conn(**overrides) -> Connection: + defaults = { + "sql_flavor": "postgresql", + "sql_flavor_code": "postgresql", + "connection_id": 7, + "project_code": "demo", + "connection_name": "Local PG", + "project_host": "localhost", + "project_port": "5432", + "project_db": "demo", + "project_user": "demo", + } + defaults.update(overrides) + return Connection(**defaults) + + +# --------------------------------------------------------------------------- +# validate_table_group_fields — happy paths and per-field errors +# --------------------------------------------------------------------------- + + +def test_validate_passes_baseline(): + assert validate_table_group_fields(_tg()) == [] + + +def test_validate_missing_name(): + errors = validate_table_group_fields(_tg(table_groups_name=None)) + assert "`table_group_name` is required." in errors + + +def test_validate_blank_name(): + errors = validate_table_group_fields(_tg(table_groups_name=" ")) + assert "`table_group_name` is required." in errors + + +def test_validate_name_too_short(): + errors = validate_table_group_fields(_tg(table_groups_name="ab")) + assert "`table_group_name` must be between 3 and 40 characters." in errors + + +def test_validate_name_too_long(): + errors = validate_table_group_fields(_tg(table_groups_name="a" * 41)) + assert "`table_group_name` must be between 3 and 40 characters." in errors + + +def test_validate_name_at_lower_bound_passes(): + assert validate_table_group_fields(_tg(table_groups_name="abc")) == [] + + +def test_validate_name_at_upper_bound_passes(): + assert validate_table_group_fields(_tg(table_groups_name="a" * 40)) == [] + + +def test_validate_missing_schema(): + errors = validate_table_group_fields(_tg(table_group_schema=None)) + assert "`schema` is required." in errors + + +def test_validate_blank_schema(): + errors = validate_table_group_fields(_tg(table_group_schema=" ")) + assert "`schema` is required." in errors + + +@pytest.mark.parametrize("bad_delay", ["-1", "-100"]) +def test_validate_delay_days_negative(bad_delay): + errors = validate_table_group_fields(_tg(profiling_delay_days=bad_delay)) + assert "`profiling_delay_days` must be a non-negative integer." in errors + + +def test_validate_delay_days_non_integer(): + errors = validate_table_group_fields(_tg(profiling_delay_days="abc")) + assert "`profiling_delay_days` must be a non-negative integer." in errors + + +def test_validate_delay_days_int_input_accepted(): + """Validator accepts int input as well — apply-args casts but validators should be tolerant.""" + assert validate_table_group_fields(_tg(profiling_delay_days=0)) == [] + + +@pytest.mark.parametrize("bad_pct", ["0", "101", "200", "-1"]) +def test_validate_sample_percent_out_of_range(bad_pct): + errors = validate_table_group_fields(_tg(profile_sample_percent=bad_pct)) + assert "`profile_sample_percent` must be between 1 and 100." in errors + + +def test_validate_sample_percent_non_integer(): + errors = validate_table_group_fields(_tg(profile_sample_percent="abc")) + assert "`profile_sample_percent` must be between 1 and 100." in errors + + +def test_validate_sample_percent_int_input_accepted(): + assert validate_table_group_fields(_tg(profile_sample_percent=50)) == [] + + +def test_validate_sample_min_count_negative(): + errors = validate_table_group_fields(_tg(profile_sample_min_count=-5)) + assert "`profile_sample_min_count` must be a non-negative integer." in errors + + +def test_validate_aggregates_all_errors(): + """Multiple failures are collected and surfaced together — never one-at-a-time.""" + errors = validate_table_group_fields(_tg( + table_groups_name="", + table_group_schema="", + profiling_delay_days="-1", + profile_sample_percent="200", + profile_sample_min_count=-5, + )) + joined = "\n".join(errors) + assert "`table_group_name` is required." in joined + assert "`schema` is required." in joined + assert "`profiling_delay_days` must be a non-negative integer." in joined + assert "`profile_sample_percent` must be between 1 and 100." in joined + assert "`profile_sample_min_count` must be a non-negative integer." in joined + + +# --------------------------------------------------------------------------- +# preview_table_group — lifted UI logic +# --------------------------------------------------------------------------- + + +def _column_row(table_name: str, column_name: str, approx_record_ct: int | None = None) -> dict: + return { + "schema_name": "public", + "table_name": table_name, + "column_name": column_name, + "ordinal_position": 1, + "general_type": "A", + "column_type": "varchar", + "db_data_type": "VARCHAR", + "is_decimal": False, + "approx_record_ct": approx_record_ct, + "record_ct": None, + } + + +@patch(f"{MODULE}.RefreshDataCharsSQL") +@patch(f"{MODULE}.fetch_from_target_db") +def test_preview_returns_picklable_shape(mock_fetch, mock_sql_cls): + """Contract: return ``(preview, data_chars, sql_generator)`` — three picklable values. + + A local ``save_data_chars`` closure as the second element would fail + pickling (``Can't get local object 'preview_table_group..save_data_chars'``) + in any caller that caches the result. Closures are built by callers via + ``make_save_data_chars``. + """ + from inspect import isfunction + + sql_generator = mock_sql_cls.return_value + sql_generator.flavor_service.metadata_via_api = False + sql_generator.get_schema_ddf.return_value = ("SELECT ...", {}) + mock_fetch.return_value = [_column_row("customer", "id", approx_record_ct=100)] + + result = preview_table_group(_tg(), connection=_conn(), verify_access=False) + + assert isinstance(result, tuple) and len(result) == 3, ( + "preview_table_group must return (preview, data_chars, sql_generator)" + ) + preview, data_chars, sql_gen = result + assert isinstance(preview, dict) + assert not isfunction(data_chars), "data_chars must not be a local closure" + assert not isfunction(sql_gen), "sql_generator must not be a local closure" + + +@patch(f"{MODULE}.RefreshDataCharsSQL") +@patch(f"{MODULE}.fetch_from_target_db") +def test_preview_happy_path_aggregates_stats(mock_fetch, mock_sql_cls): + sql_generator = mock_sql_cls.return_value + sql_generator.flavor_service.metadata_via_api = False + sql_generator.get_schema_ddf.return_value = ("SELECT ...", {}) + mock_fetch.return_value = [ + _column_row("customer", "id", approx_record_ct=100), + _column_row("customer", "email", approx_record_ct=100), + _column_row("rental", "id", approx_record_ct=50), + ] + + preview, data_chars, sql_gen = preview_table_group(_tg(), connection=_conn(), verify_access=False) + + assert preview["success"] is True + assert preview["message"] is None + assert preview["stats"]["table_ct"] == 2 + assert preview["stats"]["column_ct"] == 3 + assert preview["stats"]["approx_record_ct"] == 150 # 100 + 50, counted once per table + assert "customer" in preview["tables"] + assert "rental" in preview["tables"] + assert preview["tables"]["customer"]["column_ct"] == 2 + assert data_chars is not None and len(data_chars) == 3 + assert sql_gen is sql_generator + + +@patch(f"{MODULE}.RefreshDataCharsSQL") +@patch(f"{MODULE}.fetch_from_target_db") +def test_preview_verify_access_marks_all_accessible(mock_fetch, mock_sql_cls): + """When verify_access=True and every table is reachable, no footer message is set.""" + sql_generator = mock_sql_cls.return_value + sql_generator.flavor_service.metadata_via_api = False + sql_generator.get_schema_ddf.return_value = ("SELECT ...", {}) + sql_generator.verify_access.return_value = ("SELECT 1", None) + mock_fetch.side_effect = [ + [_column_row("customer", "id", approx_record_ct=100)], # initial DDF + [{"col": 1}], # verify customer + ] + + preview, _data_chars, _sql_gen = preview_table_group(_tg(), connection=_conn(), verify_access=True) + + assert preview["success"] is True + assert preview["tables"]["customer"]["can_access"] is True + assert preview["message"] is None + + +@patch(f"{MODULE}.RefreshDataCharsSQL") +@patch(f"{MODULE}.fetch_from_target_db") +def test_preview_verify_access_partial_failure_sets_footer(mock_fetch, mock_sql_cls): + """One inaccessible table → can_access=False AND the UI-verbatim footer message.""" + sql_generator = mock_sql_cls.return_value + sql_generator.flavor_service.metadata_via_api = False + sql_generator.get_schema_ddf.return_value = ("SELECT ...", {}) + sql_generator.verify_access.side_effect = lambda name: (f"SELECT 1 FROM {name}", None) + + def fetch_side_effect(_conn_arg, query, *_args, **_kwargs): + if "SELECT ..." in query: + return [ + _column_row("customer", "id", approx_record_ct=100), + _column_row("rental", "id", approx_record_ct=50), + ] + if "SELECT 1 FROM customer" in query: + return [{"col": 1}] + if "SELECT 1 FROM rental" in query: + raise RuntimeError("permission denied") + raise AssertionError(f"unexpected query: {query}") + + mock_fetch.side_effect = fetch_side_effect + + preview, _data_chars, _sql_gen = preview_table_group(_tg(), connection=_conn(), verify_access=True) + + assert preview["success"] is True + assert preview["tables"]["customer"]["can_access"] is True + assert preview["tables"]["rental"]["can_access"] is False + assert preview["message"] == ( + "Some tables were not accessible. Please the check the database permissions." + ) + + +@patch(f"{MODULE}.RefreshDataCharsSQL") +@patch(f"{MODULE}.fetch_from_target_db", side_effect=RuntimeError("connection refused")) +def test_preview_connection_failure_marks_unsuccessful(mock_fetch, mock_sql_cls): + """When the DDF query blows up, preview comes back ``success=False`` carrying the driver text.""" + sql_generator = mock_sql_cls.return_value + sql_generator.flavor_service.metadata_via_api = False + sql_generator.get_schema_ddf.return_value = ("SELECT ...", {}) + + preview, data_chars, sql_gen = preview_table_group(_tg(), connection=_conn(), verify_access=False) + + assert preview["success"] is False + assert "connection refused" in (preview["message"] or "") + assert data_chars is None + assert sql_gen is None + + +@patch(f"{MODULE}.RefreshDataCharsSQL") +@patch(f"{MODULE}.fetch_from_target_db", return_value=[]) +def test_preview_empty_result_uses_ui_verbatim_message(mock_fetch, mock_sql_cls): + """Zero matching tables → UI-verbatim "No tables found matching the criteria." message.""" + sql_generator = mock_sql_cls.return_value + sql_generator.flavor_service.metadata_via_api = False + sql_generator.get_schema_ddf.return_value = ("SELECT ...", {}) + + preview, _data_chars, _sql_gen = preview_table_group(_tg(), connection=_conn(), verify_access=False) + + assert preview["success"] is False + assert preview["message"] == ( + "No tables found matching the criteria. Please check the Table Group configuration" + " or the database permissions." + ) + + +def test_preview_no_connection_returns_failure_no_io(): + """No connection passed AND no ``table_group.connection_id`` → fail fast, no DB calls attempted.""" + tg = _tg(connection_id=None) + preview, data_chars, sql_gen = preview_table_group(tg, connection=None, verify_access=False) + + assert preview["success"] is False + assert preview["message"] is not None + assert "connection" in preview["message"].lower() + assert data_chars is None + assert sql_gen is None diff --git a/tests/unit/mcp/test_tools_table_groups.py b/tests/unit/mcp/test_tools_table_groups.py new file mode 100644 index 00000000..9c78c95a --- /dev/null +++ b/tests/unit/mcp/test_tools_table_groups.py @@ -0,0 +1,757 @@ +"""Tests for the MCP table-group CRUD tools — create / update / preview.""" + +from unittest.mock import MagicMock, patch +from uuid import uuid4 + +import pytest +from sqlalchemy.exc import IntegrityError + +from testgen.mcp.exceptions import MCPPermissionDenied, MCPResourceNotAccessible, MCPUserError +from testgen.mcp.permissions import ProjectPermissions + +pytestmark = pytest.mark.unit + +MODULE = "testgen.mcp.tools.table_groups" + + +# --------------------------------------------------------------------------- +# Fixtures / helpers +# --------------------------------------------------------------------------- + + +def _patch_perms(allowed=("demo",), memberships=None, permission="edit"): + memberships = memberships or dict.fromkeys(allowed, "role_a") + return patch( + "testgen.mcp.permissions._compute_project_permissions", + return_value=ProjectPermissions( + memberships=memberships, permission=permission, username="test_user", + ), + ) + + +def _mock_connection(**overrides) -> MagicMock: + conn = MagicMock() + conn.connection_id = overrides.get("connection_id", 42) + conn.project_code = overrides.get("project_code", "demo") + conn.connection_name = overrides.get("connection_name", "Local PG") + conn.sql_flavor = overrides.get("sql_flavor", "postgresql") + conn.sql_flavor_code = overrides.get("sql_flavor_code", "postgresql") + return conn + + +def _mock_table_group(**overrides) -> MagicMock: + """Build a MagicMock matching the TableGroup model surface used by the tools.""" + tg_id = overrides.get("id", uuid4()) + tg = MagicMock() + tg.id = tg_id + tg.project_code = overrides.get("project_code", "demo") + tg.connection_id = overrides.get("connection_id", 42) + tg.table_groups_name = overrides.get("table_groups_name", "Sample TG") + tg.table_group_schema = overrides.get("table_group_schema", "public") + tg.profiling_table_set = overrides.get("profiling_table_set", None) + tg.profiling_include_mask = overrides.get("profiling_include_mask", None) + tg.profiling_exclude_mask = overrides.get("profiling_exclude_mask", None) + tg.profile_id_column_mask = overrides.get("profile_id_column_mask", "%id") + tg.profile_sk_column_mask = overrides.get("profile_sk_column_mask", "%_sk") + tg.profile_use_sampling = overrides.get("profile_use_sampling", False) + tg.profile_sample_percent = overrides.get("profile_sample_percent", "30") + tg.profile_sample_min_count = overrides.get("profile_sample_min_count", 100000) + tg.profiling_delay_days = overrides.get("profiling_delay_days", "0") + tg.profile_flag_cdes = overrides.get("profile_flag_cdes", True) + tg.profile_flag_pii = overrides.get("profile_flag_pii", True) + tg.profile_exclude_xde = overrides.get("profile_exclude_xde", True) + tg.include_in_dashboard = overrides.get("include_in_dashboard", True) + tg.description = overrides.get("description", None) + tg.data_source = overrides.get("data_source", None) + tg.source_system = overrides.get("source_system", None) + tg.source_process = overrides.get("source_process", None) + tg.data_location = overrides.get("data_location", None) + tg.business_domain = overrides.get("business_domain", None) + tg.stakeholder_group = overrides.get("stakeholder_group", None) + tg.transform_level = overrides.get("transform_level", None) + tg.data_product = overrides.get("data_product", None) + return tg + + +def _integrity_error(message: str) -> IntegrityError: + orig = Exception(message) + return IntegrityError("stmt", {}, orig) + + +# --------------------------------------------------------------------------- +# create_table_group +# --------------------------------------------------------------------------- + + +@patch(f"{MODULE}.TableGroup") +@patch(f"{MODULE}.resolve_connection") +def test_create_table_group_happy_path(mock_resolve, mock_tg_cls, db_session_mock): + mock_resolve.return_value = _mock_connection() + instance = _mock_table_group(table_groups_name="Sample TG") + mock_tg_cls.return_value = instance + + from testgen.mcp.tools.table_groups import create_table_group + + with _patch_perms(): + out = create_table_group( + connection_id=42, + table_group_name="Sample TG", + schema="public", + ) + + instance.save.assert_called_once() + assert "Table Group `Sample TG` created" in out + assert "**Project:** `demo`" in out + assert "**Schema:** `public`" in out + + +def test_model_defaults_normalize_ynstring_columns_to_python_bool(): + """``YNString`` columns store ``"Y"``/``"N"`` but expose ``bool``. + + The defaults dict mirrors ``Column(default=...)`` raw values, so without + normalization a ``"N"`` default would render as truthy (non-empty string) + in the create-tool output until the row is reloaded from the DB. + """ + from testgen.mcp.tools.table_groups import _MODEL_DEFAULTS + + assert _MODEL_DEFAULTS["profile_use_sampling"] is False + assert _MODEL_DEFAULTS["profile_do_pair_rules"] is False + + +@patch(f"{MODULE}.TableGroup") +@patch(f"{MODULE}.resolve_connection") +def test_create_table_group_applies_model_defaults_when_optional_args_omitted( + mock_resolve, mock_tg_cls, db_session_mock, +): + """Optional profiling fields must seed from the model's ``Column(default=...)`` values. + + SQLAlchemy column defaults only fire at flush time, but validation runs + before flush — so the create tool must populate them in-memory or + validation rejects every call that omits them. + """ + mock_resolve.return_value = _mock_connection() + instance = _mock_table_group() + mock_tg_cls.return_value = instance + + from testgen.mcp.tools.table_groups import create_table_group + + with _patch_perms(): + create_table_group( + connection_id=42, + table_group_name="Sample TG", + schema="public", + ) + + kwargs = mock_tg_cls.call_args.kwargs + assert kwargs["profile_sample_percent"] == "30" + assert kwargs["profile_sample_min_count"] == 100000 + assert kwargs["profiling_delay_days"] == "0" + + +@patch(f"{MODULE}.TableGroup") +@patch(f"{MODULE}.resolve_connection") +def test_create_table_group_with_table_set_joins_to_string(mock_resolve, mock_tg_cls, db_session_mock): + """``table_set: list[str]`` is comma-joined into the model's ``profiling_table_set`` column.""" + mock_resolve.return_value = _mock_connection() + instance = _mock_table_group(profiling_table_set="film,actor,customer") + mock_tg_cls.return_value = instance + + from testgen.mcp.tools.table_groups import create_table_group + + with _patch_perms(): + create_table_group( + connection_id=42, + table_group_name="Sample TG", + schema="public", + table_set=["film", "actor", "customer"], + ) + + assert instance.profiling_table_set == "film,actor,customer" + + +@patch(f"{MODULE}.TableGroup") +@patch(f"{MODULE}.resolve_connection") +def test_create_table_group_renders_sample_settings_when_sampling_on( + mock_resolve, mock_tg_cls, db_session_mock, +): + """Sample % and Sample min rows only matter when sampling is on; render them then.""" + mock_resolve.return_value = _mock_connection() + instance = _mock_table_group( + profile_use_sampling=True, + profile_sample_percent="50", + profile_sample_min_count=5000, + ) + mock_tg_cls.return_value = instance + + from testgen.mcp.tools.table_groups import create_table_group + + with _patch_perms(): + out = create_table_group( + connection_id=42, + table_group_name="Sample TG", + schema="public", + profile_use_sampling=True, + profile_sample_percent=50, + profile_sample_min_count=5000, + ) + + assert "Sample %" in out + assert "50" in out + assert "Sample min rows" in out + assert "5000" in out + + +@patch(f"{MODULE}.TableGroup") +@patch(f"{MODULE}.resolve_connection") +def test_create_table_group_skips_sample_settings_when_sampling_off( + mock_resolve, mock_tg_cls, db_session_mock, +): + """When sampling is off, Sample % / min rows are irrelevant — omit them.""" + mock_resolve.return_value = _mock_connection() + instance = _mock_table_group(profile_use_sampling=False) + mock_tg_cls.return_value = instance + + from testgen.mcp.tools.table_groups import create_table_group + + with _patch_perms(): + out = create_table_group( + connection_id=42, + table_group_name="Sample TG", + schema="public", + ) + + assert "Sample %" not in out + assert "Sample min rows" not in out + + +@patch(f"{MODULE}.TableGroup") +@patch(f"{MODULE}.resolve_connection") +def test_create_table_group_with_catalog_tags_rendered(mock_resolve, mock_tg_cls, db_session_mock): + mock_resolve.return_value = _mock_connection() + instance = _mock_table_group( + data_source="Postgres", + business_domain="Sales", + ) + mock_tg_cls.return_value = instance + + from testgen.mcp.tools.table_groups import create_table_group + + with _patch_perms(): + out = create_table_group( + connection_id=42, + table_group_name="Sample TG", + schema="public", + data_source="Postgres", + business_domain="Sales", + ) + + assert "## Catalog" in out + assert "Data source" in out + assert "Postgres" in out + assert "Business domain" in out + assert "Sales" in out + + +@patch(f"{MODULE}.TableGroup") +@patch(f"{MODULE}.resolve_connection") +def test_create_table_group_validation_error_no_save(mock_resolve, mock_tg_cls, db_session_mock): + mock_resolve.return_value = _mock_connection() + instance = _mock_table_group(table_groups_name="ab") + mock_tg_cls.return_value = instance + + from testgen.mcp.tools.table_groups import create_table_group + + with _patch_perms(), pytest.raises(MCPUserError) as exc: + create_table_group( + connection_id=42, + table_group_name="ab", + schema="public", + ) + msg = str(exc.value) + assert "Table group creation rejected" in msg + assert "must be between 3 and 40 characters" in msg + instance.save.assert_not_called() + + +@patch(f"{MODULE}.TableGroup") +@patch(f"{MODULE}.resolve_connection") +def test_create_table_group_duplicate_name_maps_to_user_error(mock_resolve, mock_tg_cls, db_session_mock): + mock_resolve.return_value = _mock_connection() + instance = _mock_table_group() + instance.save.side_effect = _integrity_error( + 'duplicate key value violates unique constraint "table_groups_name_unique"' + ) + mock_tg_cls.return_value = instance + + from testgen.mcp.tools.table_groups import create_table_group + + with _patch_perms(), pytest.raises(MCPUserError) as exc: + create_table_group( + connection_id=42, + table_group_name="Sample TG", + schema="public", + ) + assert str(exc.value) == "A Table Group with the same name already exists." + + +@patch(f"{MODULE}.resolve_connection") +def test_create_table_group_connection_not_accessible(mock_resolve, db_session_mock): + """Connection's project not in allowed_codes → MCPResourceNotAccessible from resolve_connection.""" + mock_resolve.side_effect = MCPResourceNotAccessible("Connection", "42") + + from testgen.mcp.tools.table_groups import create_table_group + + with _patch_perms(allowed=("other",)), pytest.raises(MCPResourceNotAccessible) as exc: + create_table_group( + connection_id=42, + table_group_name="Sample TG", + schema="public", + ) + assert "Connection" in str(exc.value) + + +@patch(f"{MODULE}.TableGroup") +@patch(f"{MODULE}.resolve_connection") +def test_create_table_group_adds_scorecard_by_default(mock_resolve, mock_tg_cls, db_session_mock): + """Mirrors the UI checkbox default — a new table group gets a scorecard unless opted out.""" + mock_resolve.return_value = _mock_connection() + instance = _mock_table_group() + mock_tg_cls.return_value = instance + + from testgen.mcp.tools.table_groups import create_table_group + + with _patch_perms(): + create_table_group( + connection_id=42, + table_group_name="Sample TG", + schema="public", + ) + + instance.save.assert_called_once_with(add_scorecard_definition=True) + + +@patch(f"{MODULE}.TableGroup") +@patch(f"{MODULE}.resolve_connection") +def test_create_table_group_add_scorecard_false_skips_scorecard(mock_resolve, mock_tg_cls, db_session_mock): + mock_resolve.return_value = _mock_connection() + instance = _mock_table_group() + mock_tg_cls.return_value = instance + + from testgen.mcp.tools.table_groups import create_table_group + + with _patch_perms(): + create_table_group( + connection_id=42, + table_group_name="Sample TG", + schema="public", + add_scorecard=False, + ) + + instance.save.assert_called_once_with(add_scorecard_definition=False) + + +def test_create_table_group_requires_edit(db_session_mock): + """Role without 'edit' permission → MCPPermissionDenied.""" + from testgen.mcp.tools.table_groups import create_table_group + + # role_c lacks 'edit' + with _patch_perms(memberships={"demo": "role_c"}), pytest.raises(MCPPermissionDenied): + create_table_group( + connection_id=42, + table_group_name="Sample TG", + schema="public", + ) + + +# --------------------------------------------------------------------------- +# update_table_group +# --------------------------------------------------------------------------- + + +@patch(f"{MODULE}.resolve_table_group") +def test_update_table_group_no_fields_supplied(mock_resolve, db_session_mock): + mock_resolve.return_value = _mock_table_group() + + from testgen.mcp.tools.table_groups import update_table_group + + with _patch_perms(), pytest.raises(MCPUserError) as exc: + update_table_group(table_group_id=str(uuid4())) + assert str(exc.value) == "No fields supplied to update." + + +@patch(f"{MODULE}.TableGroup") +@patch(f"{MODULE}.resolve_table_group") +def test_update_table_group_description_diff_table(mock_resolve, mock_tg_cls, db_session_mock): + tg = _mock_table_group(description=None) + mock_resolve.return_value = tg + mock_tg_cls.is_in_use.return_value = False + + from testgen.mcp.tools.table_groups import update_table_group + + with _patch_perms(): + out = update_table_group(table_group_id=str(tg.id), description="Pulled from postgres tutorial DB") + + tg.save.assert_called_once() + assert "| Field | Before | After |" in out + assert "Description" in out + assert "Pulled from postgres tutorial DB" in out + + +@patch(f"{MODULE}.TableGroup") +@patch(f"{MODULE}.resolve_table_group") +def test_update_table_group_no_op(mock_resolve, mock_tg_cls, db_session_mock): + """Supplying the current value → no-op message, no save call.""" + tg = _mock_table_group(description="existing") + mock_resolve.return_value = tg + mock_tg_cls.is_in_use.return_value = False + + from testgen.mcp.tools.table_groups import update_table_group + + with _patch_perms(): + out = update_table_group(table_group_id=str(tg.id), description="existing") + + tg.save.assert_not_called() + assert "No fields changed" in out + + +@patch(f"{MODULE}.TableGroup") +@patch(f"{MODULE}.resolve_table_group") +def test_update_table_group_schema_locked_when_in_use(mock_resolve, mock_tg_cls, db_session_mock): + """``is_in_use=True`` + new schema → MCPUserError, no save.""" + tg = _mock_table_group(table_group_schema="public") + mock_resolve.return_value = tg + mock_tg_cls.is_in_use.return_value = True + + from testgen.mcp.tools.table_groups import update_table_group + + with _patch_perms(), pytest.raises(MCPUserError) as exc: + update_table_group(table_group_id=str(tg.id), schema="staging") + assert "Schema cannot be changed" in str(exc.value) + assert "Delete and recreate" in str(exc.value) + tg.save.assert_not_called() + + +@patch(f"{MODULE}.TableGroup") +@patch(f"{MODULE}.resolve_table_group") +def test_update_table_group_schema_unlocked_when_not_in_use(mock_resolve, mock_tg_cls, db_session_mock): + tg = _mock_table_group(table_group_schema="public") + mock_resolve.return_value = tg + mock_tg_cls.is_in_use.return_value = False + + from testgen.mcp.tools.table_groups import update_table_group + + with _patch_perms(): + out = update_table_group(table_group_id=str(tg.id), schema="staging") + + tg.save.assert_called_once() + assert "Schema" in out + assert "staging" in out + + +@patch(f"{MODULE}.TableGroup") +@patch(f"{MODULE}.resolve_table_group") +def test_update_table_group_same_schema_on_in_use_group_is_noop(mock_resolve, mock_tg_cls, db_session_mock): + """Re-supplying the current schema on an in-use group is a no-op, not a lock violation.""" + tg = _mock_table_group(table_group_schema="public") + mock_resolve.return_value = tg + mock_tg_cls.is_in_use.return_value = True + + from testgen.mcp.tools.table_groups import update_table_group + + with _patch_perms(): + out = update_table_group(table_group_id=str(tg.id), schema="public") + + tg.save.assert_not_called() + assert "No fields changed" in out + + +@patch(f"{MODULE}.TableGroup") +@patch(f"{MODULE}.resolve_table_group") +def test_update_table_group_validation_error_no_save(mock_resolve, mock_tg_cls, db_session_mock): + tg = _mock_table_group(table_groups_name="Sample TG") + mock_resolve.return_value = tg + mock_tg_cls.is_in_use.return_value = False + + from testgen.mcp.tools.table_groups import update_table_group + + with _patch_perms(), pytest.raises(MCPUserError) as exc: + update_table_group(table_group_id=str(tg.id), table_group_name="ab") + msg = str(exc.value) + assert "Update rejected" in msg + assert "must be between 3 and 40 characters" in msg + tg.save.assert_not_called() + + +@patch(f"{MODULE}.TableGroup") +@patch(f"{MODULE}.resolve_table_group") +def test_update_table_group_duplicate_name_maps_to_user_error(mock_resolve, mock_tg_cls, db_session_mock): + tg = _mock_table_group(table_groups_name="Sample TG") + tg.save.side_effect = _integrity_error( + 'duplicate key value violates unique constraint "table_groups_name_unique"' + ) + mock_resolve.return_value = tg + mock_tg_cls.is_in_use.return_value = False + + from testgen.mcp.tools.table_groups import update_table_group + + with _patch_perms(), pytest.raises(MCPUserError) as exc: + update_table_group(table_group_id=str(tg.id), table_group_name="Existing Name") + assert str(exc.value) == "A Table Group with the same name already exists." + + +@patch(f"{MODULE}.resolve_table_group") +def test_update_table_group_not_accessible(mock_resolve, db_session_mock): + mock_resolve.side_effect = MCPResourceNotAccessible("Table group", "abc") + + from testgen.mcp.tools.table_groups import update_table_group + + with _patch_perms(allowed=("other",)), pytest.raises(MCPResourceNotAccessible): + update_table_group(table_group_id=str(uuid4()), description="any") + + +@patch(f"{MODULE}.TableGroup") +@patch(f"{MODULE}.resolve_table_group") +def test_update_table_group_delay_days_int_cast_to_str(mock_resolve, mock_tg_cls, db_session_mock): + """``profiling_delay_days: int`` from the caller gets cast to ``str`` to match the model column.""" + tg = _mock_table_group(profiling_delay_days="0") + mock_resolve.return_value = tg + mock_tg_cls.is_in_use.return_value = False + + from testgen.mcp.tools.table_groups import update_table_group + + with _patch_perms(): + update_table_group(table_group_id=str(tg.id), profiling_delay_days=3) + + assert tg.profiling_delay_days == "3" + + +# --------------------------------------------------------------------------- +# preview_table_group +# --------------------------------------------------------------------------- + + +@patch(f"{MODULE}.preview_table_group_service") +@patch(f"{MODULE}.Connection") +@patch(f"{MODULE}.resolve_table_group") +def test_preview_success_renders_table(mock_resolve, mock_conn_cls, mock_preview_svc, db_session_mock): + tg = _mock_table_group(table_groups_name="Sample TG", table_group_schema="public") + mock_resolve.return_value = tg + mock_conn_cls.get_by_table_group.return_value = _mock_connection() + mock_preview_svc.return_value = ( + { + "stats": { + "id": tg.id, + "table_groups_name": "Sample TG", + "table_group_schema": "public", + "table_ct": 2, + "column_ct": 5, + "approx_record_ct": 150, + "approx_data_point_ct": 350, + }, + "tables": { + "customer": { + "column_ct": 3, + "approx_record_ct": 100, + "approx_data_point_ct": 300, + "can_access": True, + }, + "rental": { + "column_ct": 2, + "approx_record_ct": 50, + "approx_data_point_ct": 50, + "can_access": True, + }, + }, + "success": True, + "message": None, + }, + None, + None, + ) + + from testgen.mcp.tools.table_groups import preview_table_group + + with _patch_perms(): + out = preview_table_group(table_group_id=str(tg.id)) + + assert "Preview for table group" in out + assert "Sample TG" in out + assert "customer" in out + assert "rental" in out + # verify_access defaults to False — no Read Access column + assert "| Table | Columns | Approx Rows | Approx Data Points |" in out + assert "Read Access" not in out + _, kwargs = mock_preview_svc.call_args + assert kwargs["verify_access"] is False + + +@patch(f"{MODULE}.preview_table_group_service") +@patch(f"{MODULE}.Connection") +@patch(f"{MODULE}.resolve_table_group") +def test_preview_verify_access_adds_column(mock_resolve, mock_conn_cls, mock_preview_svc, db_session_mock): + tg = _mock_table_group(table_groups_name="Sample TG", table_group_schema="public") + mock_resolve.return_value = tg + mock_conn_cls.get_by_table_group.return_value = _mock_connection() + mock_preview_svc.return_value = ( + { + "stats": { + "id": tg.id, + "table_groups_name": "Sample TG", + "table_group_schema": "public", + "table_ct": 1, + "column_ct": 3, + "approx_record_ct": 100, + "approx_data_point_ct": 300, + }, + "tables": { + "customer": { + "column_ct": 3, + "approx_record_ct": 100, + "approx_data_point_ct": 300, + "can_access": True, + }, + }, + "success": True, + "message": None, + }, + None, + None, + ) + + from testgen.mcp.tools.table_groups import preview_table_group + + with _patch_perms(): + out = preview_table_group(table_group_id=str(tg.id), verify_access=True) + + assert "| Table | Columns | Approx Rows | Approx Data Points | Read Access |" in out + _, kwargs = mock_preview_svc.call_args + assert kwargs["verify_access"] is True + + +@patch(f"{MODULE}.preview_table_group_service") +@patch(f"{MODULE}.Connection") +@patch(f"{MODULE}.resolve_table_group") +def test_preview_partial_inaccessible_appends_footer(mock_resolve, mock_conn_cls, mock_preview_svc, db_session_mock): + tg = _mock_table_group(table_groups_name="Sample TG", table_group_schema="public") + mock_resolve.return_value = tg + mock_conn_cls.get_by_table_group.return_value = _mock_connection() + mock_preview_svc.return_value = ( + { + "stats": { + "id": tg.id, + "table_groups_name": "Sample TG", + "table_group_schema": "public", + "table_ct": 2, + "column_ct": 5, + "approx_record_ct": 150, + "approx_data_point_ct": 350, + }, + "tables": { + "customer": { + "column_ct": 3, "approx_record_ct": 100, + "approx_data_point_ct": 300, "can_access": True, + }, + "rental": { + "column_ct": 2, "approx_record_ct": 50, + "approx_data_point_ct": 50, "can_access": False, + }, + }, + "success": True, + "message": "Some tables were not accessible. Please the check the database permissions.", + }, + None, + None, + ) + + from testgen.mcp.tools.table_groups import preview_table_group + + with _patch_perms(): + out = preview_table_group(table_group_id=str(tg.id), verify_access=True) + + assert "Some tables were not accessible" in out + + +@patch(f"{MODULE}.preview_table_group_service") +@patch(f"{MODULE}.Connection") +@patch(f"{MODULE}.resolve_table_group") +def test_preview_no_match(mock_resolve, mock_conn_cls, mock_preview_svc, db_session_mock): + tg = _mock_table_group(table_groups_name="Sample TG", table_group_schema="public") + mock_resolve.return_value = tg + mock_conn_cls.get_by_table_group.return_value = _mock_connection() + mock_preview_svc.return_value = ( + { + "stats": { + "id": tg.id, + "table_groups_name": "Sample TG", + "table_group_schema": "public", + "table_ct": 0, + "column_ct": 0, + "approx_record_ct": None, + "approx_data_point_ct": None, + }, + "tables": {}, + "success": False, + "message": ( + "No tables found matching the criteria. Please check the Table Group configuration" + " or the database permissions." + ), + }, + None, + None, + ) + + from testgen.mcp.tools.table_groups import preview_table_group + + with _patch_perms(): + out = preview_table_group(table_group_id=str(tg.id)) + + assert "returned no tables" in out + assert "No tables found matching the criteria" in out + + +@patch(f"{MODULE}.preview_table_group_service") +@patch(f"{MODULE}.Connection") +@patch(f"{MODULE}.resolve_table_group") +def test_preview_failed_returns_text_not_raises(mock_resolve, mock_conn_cls, mock_preview_svc, db_session_mock): + """A failed preview surfaces as a text response — no exception raised.""" + tg = _mock_table_group(table_groups_name="Sample TG") + mock_resolve.return_value = tg + mock_conn_cls.get_by_table_group.return_value = _mock_connection() + mock_preview_svc.return_value = ( + { + "stats": { + "id": tg.id, "table_groups_name": "Sample TG", "table_group_schema": "public", + "table_ct": 0, "column_ct": 0, "approx_record_ct": None, "approx_data_point_ct": None, + }, + "tables": {}, + "success": False, + "message": "Could not connect to target DB: connection refused", + }, + None, + None, + ) + + from testgen.mcp.tools.table_groups import preview_table_group + + with _patch_perms(): + out = preview_table_group(table_group_id=str(tg.id)) + + assert "Preview failed" in out + assert "Could not connect" in out + + +@patch(f"{MODULE}.resolve_table_group") +def test_preview_not_accessible(mock_resolve, db_session_mock): + mock_resolve.side_effect = MCPResourceNotAccessible("Table group", "abc") + + from testgen.mcp.tools.table_groups import preview_table_group + + with _patch_perms(allowed=("other",)), pytest.raises(MCPResourceNotAccessible): + preview_table_group(table_group_id=str(uuid4())) + + +def test_preview_requires_edit(db_session_mock): + from testgen.mcp.tools.table_groups import preview_table_group + + with _patch_perms(memberships={"demo": "role_c"}), pytest.raises(MCPPermissionDenied): + preview_table_group(table_group_id=str(uuid4())) From 91a3fb7b0b1c8c50af2bba82b39c1ca6ca54ab22 Mon Sep 17 00:00:00 2001 From: Luis Date: Thu, 4 Jun 2026 11:17:42 -0400 Subject: [PATCH 04/78] fix(profiling): tolerate NULL max_query_chars in row-count chunking Co-Authored-By: Claude Opus 4.8 (1M context) --- .../commands/queries/refresh_data_chars_query.py | 5 +++-- testgen/common/models/connection.py | 5 +++++ .../queries/test_refresh_data_chars_query.py | 13 +++++++++++++ 3 files changed, 21 insertions(+), 2 deletions(-) diff --git a/testgen/commands/queries/refresh_data_chars_query.py b/testgen/commands/queries/refresh_data_chars_query.py index 762e7261..c6d74a7c 100644 --- a/testgen/commands/queries/refresh_data_chars_query.py +++ b/testgen/commands/queries/refresh_data_chars_query.py @@ -5,7 +5,7 @@ from testgen.common import read_template_sql_file from testgen.common.database.column_chars import ColumnChars from testgen.common.database.database_service import get_flavor_service, replace_params -from testgen.common.models.connection import Connection +from testgen.common.models.connection import DEFAULT_MAX_QUERY_CHARS, Connection from testgen.common.models.table_group import TableGroup from testgen.utils import chunk_queries, to_sql_timestamp @@ -133,7 +133,8 @@ def get_row_counts(self, table_names: Iterable[str]) -> list[tuple[str, None]]: f"SELECT '{table}' AS table_name, COUNT(*) AS row_count FROM {self.flavor_service.get_table_ref(schema, table)}" for table in table_names ] - chunked_queries = chunk_queries(count_queries, " UNION ALL ", self.connection.max_query_chars) + max_query_chars = self.connection.max_query_chars or DEFAULT_MAX_QUERY_CHARS + chunked_queries = chunk_queries(count_queries, " UNION ALL ", max_query_chars) return [ (query, None) for query in chunked_queries ] def verify_access(self, table_name: str) -> tuple[str, None]: diff --git a/testgen/common/models/connection.py b/testgen/common/models/connection.py index a940edba..2004a67f 100644 --- a/testgen/common/models/connection.py +++ b/testgen/common/models/connection.py @@ -28,6 +28,11 @@ SQLFlavorCode = Literal["redshift", "redshift_spectrum", "snowflake", "mssql", "azure_mssql", "synapse_mssql", "postgresql", "databricks", "bigquery", "oracle", "sap_hana", "salesforce_data360"] +# Fallback when a connection row has a NULL max_query_chars (no DB default; older +# rows / non-UI insert paths may not seed it). Matches the UI's non-Salesforce +# default and the 0160 migration. +DEFAULT_MAX_QUERY_CHARS = 20000 + @dataclass class ConnectionMinimal(EntityMinimal): diff --git a/tests/unit/commands/queries/test_refresh_data_chars_query.py b/tests/unit/commands/queries/test_refresh_data_chars_query.py index 5dc3ff16..427587e8 100644 --- a/tests/unit/commands/queries/test_refresh_data_chars_query.py +++ b/tests/unit/commands/queries/test_refresh_data_chars_query.py @@ -217,3 +217,16 @@ def test_filter_schema_columns_no_filters_returns_all(): filtered = sql_generator.filter_schema_columns(columns) assert {c.table_name for c in filtered} == {"users", "orders"} + + +def test_get_row_counts_handles_null_max_query_chars(): + """A connection with NULL max_query_chars must not crash chunking — it falls + back to DEFAULT_MAX_QUERY_CHARS so the UNION ALL count queries still build.""" + connection = Connection(sql_flavor="postgresql", max_query_chars=None) + table_group = TableGroup(table_group_schema="test_schema") + sql_generator = RefreshDataCharsSQL(connection, table_group) + + result = sql_generator.get_row_counts(["orders", "customers"]) + + assert result + assert all(isinstance(query, str) and query for query, _ in result) From b6e8f29e6976e2f837108271c769a0ac398fd0c3 Mon Sep 17 00:00:00 2001 From: Luis Date: Thu, 4 Jun 2026 11:19:14 -0400 Subject: [PATCH 05/78] fix(test-execution): tolerate NULL max_query_chars in CAT batching Co-Authored-By: Claude Opus 4.8 (1M context) --- testgen/commands/queries/execute_tests_query.py | 4 ++-- .../queries/test_execute_tests_query.py | 17 +++++++++++++++++ 2 files changed, 19 insertions(+), 2 deletions(-) diff --git a/testgen/commands/queries/execute_tests_query.py b/testgen/commands/queries/execute_tests_query.py index 03eab489..751fd4d5 100644 --- a/testgen/commands/queries/execute_tests_query.py +++ b/testgen/commands/queries/execute_tests_query.py @@ -21,7 +21,7 @@ is_excluded_day, resolve_holiday_dates, ) -from testgen.common.models.connection import Connection +from testgen.common.models.connection import DEFAULT_MAX_QUERY_CHARS, Connection from testgen.common.models.scheduler import JobSchedule from testgen.common.models.table_group import TableGroup from testgen.common.models.test_definition import TestRunType, TestScope @@ -534,7 +534,7 @@ def aggregate_cat_tests( null_value=self.null_value, ) - max_query_chars = self.connection.max_query_chars - 400 + max_query_chars = (self.connection.max_query_chars or DEFAULT_MAX_QUERY_CHARS) - 400 groups = group_cat_tests(test_defs, max_query_chars, concat_operator, single) aggregate_queries: list[tuple[str, None]] = [] diff --git a/tests/unit/commands/queries/test_execute_tests_query.py b/tests/unit/commands/queries/test_execute_tests_query.py index 4839f99b..ca1c9555 100644 --- a/tests/unit/commands/queries/test_execute_tests_query.py +++ b/tests/unit/commands/queries/test_execute_tests_query.py @@ -11,6 +11,8 @@ group_cat_tests, parse_cat_results, ) +from testgen.common.database.database_service import get_flavor_service +from testgen.common.models.connection import Connection pytestmark = pytest.mark.unit @@ -478,3 +480,18 @@ def test_resolve_cat_no_freshness_result_uses_band_check(_mock_changed): assert operator == "NOT BETWEEN" +def test_aggregate_cat_tests_handles_null_max_query_chars(): + """A connection with NULL max_query_chars must not crash CAT batching — the + `- 400` headroom subtraction falls back to DEFAULT_MAX_QUERY_CHARS.""" + instance = _make_execution_sql() + instance.connection = Connection(sql_flavor="postgresql", max_query_chars=None) + instance.flavor = "postgresql" + instance.flavor_service = get_flavor_service("postgresql") + + td = _make_td(measure_expression="m_expr", condition_expression="c_expr") + queries, grouped_defs = instance.aggregate_cat_tests([td], single=True) + + assert len(queries) == 1 + assert grouped_defs == [[td]] + + From f6abfa3c1b2ebb75bd6e40df4b92a600856ea6a6 Mon Sep 17 00:00:00 2001 From: Luis Date: Thu, 4 Jun 2026 11:19:41 -0400 Subject: [PATCH 06/78] refactor(connection): reuse DEFAULT_MAX_QUERY_CHARS constant Co-Authored-By: Claude Opus 4.8 (1M context) --- testgen/common/database/connection_service.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/testgen/common/database/connection_service.py b/testgen/common/database/connection_service.py index c7cec344..879f6208 100644 --- a/testgen/common/database/connection_service.py +++ b/testgen/common/database/connection_service.py @@ -26,7 +26,7 @@ from sqlalchemy.exc import DatabaseError, DBAPIError from testgen.common.database.database_service import empty_cache, get_flavor_service -from testgen.common.models.connection import Connection +from testgen.common.models.connection import DEFAULT_MAX_QUERY_CHARS, Connection from testgen.ui.services.database_service import fetch_from_target_db try: @@ -142,4 +142,4 @@ def apply_connection_defaults(connection: Connection) -> None: """Fill flavor-dependent defaults for fields the caller didn't supply.""" if connection.max_query_chars is None: # Salesforce Data 360's Hyper engine has a lower expression-depth limit - connection.max_query_chars = 15000 if connection.sql_flavor_code == "salesforce_data360" else 20000 + connection.max_query_chars = 15000 if connection.sql_flavor_code == "salesforce_data360" else DEFAULT_MAX_QUERY_CHARS From 71fa2855bd5a8c95125f7197e3283758f69be935 Mon Sep 17 00:00:00 2001 From: Ricardo Boni Date: Thu, 4 Jun 2026 13:30:10 -0400 Subject: [PATCH 07/78] refactor(mcp,api): run-kind labels, shared status labels, API schema hygiene - Label run identifiers by kind (Test Run / Profiling Run) in MCP output instead of "Job ID" - Single JOB_STATUS_LABEL in common.enums; run-summary models alias it - Strip enum descriptions from the OpenAPI schema so enum docstrings don't leak into the public docs - Add PublicJobKey so the API exposes only externally-triggerable job kinds; remove dead backfill job source - Generalize an OAuth route docstring; drop an incidental JobKey re-export TG-1116 Co-Authored-By: Claude Opus 4.8 (1M context) --- testgen/api/deps.py | 7 ++-- testgen/api/jobs.py | 8 ++--- testgen/api/oauth/metadata.py | 2 +- testgen/api/runs.py | 2 +- testgen/api/schemas.py | 4 +-- testgen/common/enums.py | 21 +++++++++++- testgen/common/models/job_execution.py | 11 +------ testgen/common/models/profiling_run.py | 12 ++----- testgen/common/models/test_run.py | 12 ++----- testgen/mcp/tools/execution.py | 4 +-- testgen/mcp/tools/profiling.py | 6 ++-- testgen/mcp/tools/schedules.py | 7 ++-- testgen/mcp/tools/test_runs.py | 6 ++-- testgen/server/__init__.py | 7 ++++ tests/unit/api/test_jobs.py | 44 ++++++++++++++++++++++++++ 15 files changed, 99 insertions(+), 54 deletions(-) diff --git a/testgen/api/deps.py b/testgen/api/deps.py index 1807d68d..e1f834ec 100644 --- a/testgen/api/deps.py +++ b/testgen/api/deps.py @@ -7,8 +7,9 @@ from sqlalchemy import select from testgen.common.auth import authorize_token, decode_jwt_token +from testgen.common.enums import PublicJobKey from testgen.common.models import Session, _current_session_wrapper, get_current_session -from testgen.common.models.job_execution import PUBLIC_JOB_KEYS, JobExecution +from testgen.common.models.job_execution import JobExecution from testgen.common.models.project_membership import ProjectMembership from testgen.common.models.table_group import TableGroup from testgen.common.models.test_suite import TestSuite @@ -122,7 +123,7 @@ def dependency(test_suite_id: UUID, user: User = _require_user) -> TestSuite: def resolve_job(permission: str, *extra_filters): """Resolve a JobExecution by ``job_id`` path param and verify project permission. - Only jobs whose ``job_key`` is in ``PUBLIC_JOB_KEYS`` are exposed via the API. + Only externally-triggerable jobs (``PublicJobKey``) are exposed via the API. Internal kinds (score rollups, recalculations, monitor runs) are filtered out by construction. Extra ORM clauses are appended to the WHERE clause to further restrict by job_key when a caller wants a single kind. @@ -130,7 +131,7 @@ def resolve_job(permission: str, *extra_filters): def dependency(job_id: UUID, user: User = _require_user) -> JobExecution: query = select(JobExecution).where( JobExecution.id == job_id, - JobExecution.job_key.in_(PUBLIC_JOB_KEYS), + JobExecution.job_key.in_(list(PublicJobKey)), *extra_filters, ) return _check_access(get_current_session().scalars(query).first(), user, permission) diff --git a/testgen/api/jobs.py b/testgen/api/jobs.py index 3ab291cf..c4f28f2f 100644 --- a/testgen/api/jobs.py +++ b/testgen/api/jobs.py @@ -11,8 +11,8 @@ resolve_test_suite, ) from testgen.api.schemas import ErrorResponse, JobListResponse, JobResponse, JobSubmittedResponse -from testgen.common.enums import JobKey, JobSource, JobStatus -from testgen.common.models.job_execution import PUBLIC_JOB_KEYS, JobExecution +from testgen.common.enums import JobKey, JobSource, JobStatus, PublicJobKey +from testgen.common.models.job_execution import JobExecution from testgen.common.models.table_group import TableGroup from testgen.common.models.test_suite import TestSuite @@ -98,7 +98,7 @@ def cancel_job(job: JobExecution = resolve_job("edit")): # noqa: B008 ) def list_jobs( project_code: str = resolve_project_code("view"), - job_key: JobKey | None = Query(default=None), # noqa: B008 + job_key: PublicJobKey | None = Query(default=None), # noqa: B008 status: JobStatus | None = Query(default=None), # noqa: B008 page: int = Query(default=1, ge=1), limit: int = Query(default=20, ge=1, le=100), @@ -106,7 +106,7 @@ def list_jobs( """List job executions for a project, with optional filters and pagination.""" items, total = JobExecution.list_for_project( project_code, - JobExecution.job_key.in_(PUBLIC_JOB_KEYS), + JobExecution.job_key.in_(list(PublicJobKey)), job_key=job_key, status=status, page=page, diff --git a/testgen/api/oauth/metadata.py b/testgen/api/oauth/metadata.py index 55d6fb28..d54d3dfe 100644 --- a/testgen/api/oauth/metadata.py +++ b/testgen/api/oauth/metadata.py @@ -12,7 +12,7 @@ def authorization_server_metadata(): """Return OAuth 2.1 Authorization Server Metadata per RFC 8414. - MCP clients use this for server discovery. + OAuth clients use this for server discovery. """ base_url = settings.BASE_URL.rstrip("/") diff --git a/testgen/api/runs.py b/testgen/api/runs.py index fcbe0445..be2f62f1 100644 --- a/testgen/api/runs.py +++ b/testgen/api/runs.py @@ -7,13 +7,13 @@ from testgen.api.schemas import ( ErrorResponse, IssueBreakdown, - JobKey, ProfilingRunResponse, ProfilingRunResult, TestBreakdown, TestRunResponse, TestRunResult, ) +from testgen.common.enums import JobKey from testgen.common.models import get_current_session from testgen.common.models.hygiene_issue import HygieneIssue from testgen.common.models.job_execution import JobExecution diff --git a/testgen/api/schemas.py b/testgen/api/schemas.py index 2543f227..fbe56584 100644 --- a/testgen/api/schemas.py +++ b/testgen/api/schemas.py @@ -6,7 +6,7 @@ from pydantic import BaseModel, field_validator -from testgen.common.enums import JobKey, JobSource, JobStatus +from testgen.common.enums import JobSource, JobStatus, PublicJobKey # --- Jobs --- @@ -24,7 +24,7 @@ class JobResponse(BaseModel): """Full job execution record returned by status and cancel endpoints.""" id: UUID - job_key: JobKey + job_key: PublicJobKey status: JobStatus source: JobSource created_at: datetime diff --git a/testgen/common/enums.py b/testgen/common/enums.py index 3679f709..f315023d 100644 --- a/testgen/common/enums.py +++ b/testgen/common/enums.py @@ -40,6 +40,14 @@ class JobKey(StrEnum): run_data_cleanup = "run-data-cleanup" +class PublicJobKey(StrEnum): + """``job_key`` values exposed through the public API — the externally-triggerable + subset of ``JobKey``. Internal maintenance kinds are intentionally absent.""" + run_profile = JobKey.run_profile.value + run_tests = JobKey.run_tests.value + run_test_generation = JobKey.run_test_generation.value + + class JobSource(StrEnum): """``source`` column values for ``job_executions``.""" api = "api" @@ -47,7 +55,6 @@ class JobSource(StrEnum): scheduler = "scheduler" mcp = "mcp" cli = "cli" - backfill = "backfill" system = "system" @@ -63,6 +70,18 @@ class JobStatus(StrEnum): CANCELED = "canceled" +# User-facing display labels for JobStatus values. +JOB_STATUS_LABEL: dict[str, str] = { + JobStatus.COMPLETED: "Completed", + JobStatus.CANCELED: "Canceled", + JobStatus.CANCEL_REQUESTED: "Canceling", + JobStatus.PENDING: "Pending", + JobStatus.CLAIMED: "Starting", + JobStatus.RUNNING: "Running", + JobStatus.ERROR: "Error", +} + + class Disposition(StrEnum): """Stored disposition values for ``profile_anomaly_results.disposition`` and ``test_results.disposition``. The user-facing label for ``INACTIVE`` is "Muted".""" diff --git a/testgen/common/models/job_execution.py b/testgen/common/models/job_execution.py index 8d8160c2..9ae65325 100644 --- a/testgen/common/models/job_execution.py +++ b/testgen/common/models/job_execution.py @@ -6,21 +6,12 @@ from sqlalchemy import Column, String, Text, case, delete, func, select, text, update from sqlalchemy.dialects import postgresql -from testgen.common.enums import JobKey, JobSource, JobStatus +from testgen.common.enums import JobSource, JobStatus from testgen.common.models import Base, database_session, get_current_session LOG = logging.getLogger("testgen") -# Job kinds that are externally triggerable. Internal kinds (run-score-update, -# recalculate-project-scores, ...) are absent and filtered out by construction. -PUBLIC_JOB_KEYS: frozenset[JobKey] = frozenset({ - JobKey.run_profile, - JobKey.run_tests, - JobKey.run_test_generation, -}) - - _VALID_TRANSITIONS: dict[JobStatus, frozenset[JobStatus]] = { JobStatus.PENDING: frozenset({JobStatus.CLAIMED, JobStatus.CANCEL_REQUESTED}), JobStatus.CLAIMED: frozenset({JobStatus.RUNNING, JobStatus.ERROR, JobStatus.CANCEL_REQUESTED}), diff --git a/testgen/common/models/profiling_run.py b/testgen/common/models/profiling_run.py index 225c9942..39bace1c 100644 --- a/testgen/common/models/profiling_run.py +++ b/testgen/common/models/profiling_run.py @@ -10,7 +10,7 @@ from sqlalchemy.orm.attributes import flag_modified from sqlalchemy.sql.expression import case -from testgen.common.enums import Disposition, JobStatus +from testgen.common.enums import JOB_STATUS_LABEL, Disposition, JobStatus from testgen.common.models import database_session, get_current_session from testgen.common.models.connection import Connection from testgen.common.models.entity import Entity, EntityMinimal @@ -71,15 +71,7 @@ class ProfilingRunSummary(EntityMinimal): dq_score_profiling: float | None total_count: int - STATUS_LABEL: ClassVar[dict[str, str]] = { - JobStatus.COMPLETED: "Completed", - JobStatus.CANCELED: "Canceled", - JobStatus.CANCEL_REQUESTED: "Canceling", - JobStatus.PENDING: "Pending", - JobStatus.CLAIMED: "Starting", - JobStatus.RUNNING: "Running", - JobStatus.ERROR: "Error", - } + STATUS_LABEL: ClassVar[dict[str, str]] = JOB_STATUS_LABEL @property def status_label(self) -> str: diff --git a/testgen/common/models/test_run.py b/testgen/common/models/test_run.py index 2bb9bc51..385b429b 100644 --- a/testgen/common/models/test_run.py +++ b/testgen/common/models/test_run.py @@ -8,7 +8,7 @@ from sqlalchemy.orm.attributes import flag_modified from sqlalchemy.sql.expression import case -from testgen.common.enums import JobStatus +from testgen.common.enums import JOB_STATUS_LABEL, JobStatus from testgen.common.models import database_session, get_current_session from testgen.common.models.connection import Connection from testgen.common.models.entity import Entity, EntityMinimal @@ -70,15 +70,7 @@ class TestRunSummary(EntityMinimal): dq_score_testing: float | None total_count: int - STATUS_LABEL: ClassVar[dict[str, str]] = { - JobStatus.COMPLETED: "Completed", - JobStatus.CANCELED: "Canceled", - JobStatus.CANCEL_REQUESTED: "Canceling", - JobStatus.PENDING: "Pending", - JobStatus.CLAIMED: "Starting", - JobStatus.RUNNING: "Running", - JobStatus.ERROR: "Error", - } + STATUS_LABEL: ClassVar[dict[str, str]] = JOB_STATUS_LABEL @property def status_label(self) -> str: diff --git a/testgen/mcp/tools/execution.py b/testgen/mcp/tools/execution.py index 3ed02e0e..2847e77a 100644 --- a/testgen/mcp/tools/execution.py +++ b/testgen/mcp/tools/execution.py @@ -133,7 +133,7 @@ def _render_submission( ) -> str: doc = MdDoc() doc.heading(1, f"{kind} submitted for `{scope_name}`") - doc.field("Job ID", job.id, code=True) + doc.field(kind.title(), job.id, code=True) doc.field(scope_label, scope_name) doc.field("Status", "Pending") doc.text(f"Use `{poll_tool}` {poll_hint}.") @@ -147,7 +147,7 @@ def _render_cancel(job: JobExecution, kind: str, poll_tool: str) -> str: ) doc = MdDoc() doc.heading(1, f"{kind} cancellation requested") - doc.field("Job ID", job.id, code=True) + doc.field(kind.title(), job.id, code=True) doc.field("Status", job.status) doc.text(f"Use `{poll_tool}` to confirm cancellation.") return doc.render() diff --git a/testgen/mcp/tools/profiling.py b/testgen/mcp/tools/profiling.py index 79ed638c..c861e802 100644 --- a/testgen/mcp/tools/profiling.py +++ b/testgen/mcp/tools/profiling.py @@ -509,7 +509,7 @@ def get_profiling_run(job_execution_id: str) -> str: doc = MdDoc() tg_label = summary.table_groups_name or "—" doc.heading(1, f"Profiling run: {tg_label}") - doc.field("Job ID", summary.job_execution_id, code=True) + doc.field("Profiling Run", summary.job_execution_id, code=True) if summary.table_groups_name: doc.field("Table group", summary.table_groups_name) if summary.table_group_schema: @@ -561,7 +561,7 @@ def get_profiling_run(job_execution_id: str) -> str: def _render_pending_profiling_je(doc: MdDoc, je: JobExecution, label: str) -> None: status_label = ProfilingRunSummary.STATUS_LABEL.get(je.status, je.status) doc.heading(3, f"{label} — {status_label}") - doc.field("Job ID", je.id, code=True) + doc.field("Profiling Run", je.id, code=True) if je.job_schedule_id is not None: doc.field("Schedule", je.job_schedule_id, code=True) doc.field("Submitted", je.created_at) @@ -572,7 +572,7 @@ def _render_pending_profiling_je(doc: MdDoc, je: JobExecution, label: str) -> No def _render_profiling_run_section(doc: MdDoc, run: ProfilingRunSummary) -> None: title = run.table_groups_name or run.profiling_run_id or run.job_execution_id doc.heading(2, f"{title} — {run.status_label}") - doc.field("Job ID", run.job_execution_id, code=True) + doc.field("Profiling Run", run.job_execution_id, code=True) if run.job_schedule_id is not None: doc.field("Schedule", run.job_schedule_id, code=True) doc.field("Submitted", run.created_at) diff --git a/testgen/mcp/tools/schedules.py b/testgen/mcp/tools/schedules.py index 9c1ed1d0..3aa0a021 100644 --- a/testgen/mcp/tools/schedules.py +++ b/testgen/mcp/tools/schedules.py @@ -6,12 +6,11 @@ from sqlalchemy import select from testgen.common.cron_service import describe_cron, get_cron_sample -from testgen.common.enums import JobKey +from testgen.common.enums import JOB_STATUS_LABEL, JobKey from testgen.common.models import get_current_session, with_database_session from testgen.common.models.job_execution import JobExecution from testgen.common.models.scheduler import JobSchedule from testgen.common.models.table_group import TableGroup -from testgen.common.models.test_run import TestRunSummary # STATUS_LABEL is shared with ProfilingRunSummary from testgen.common.models.test_suite import TestSuite from testgen.mcp.exceptions import MCPResourceNotAccessible, MCPUserError from testgen.mcp.permissions import get_project_permissions, mcp_permission @@ -421,13 +420,13 @@ def get_schedule(schedule_id: str) -> str: for je in history: rows.append([ je.id, - TestRunSummary.STATUS_LABEL.get(je.status, je.status), + JOB_STATUS_LABEL.get(je.status, je.status), je.started_at, je.completed_at, format_run_duration(je.started_at, je.completed_at), ]) doc.table( - ["Job ID", "Status", "Started", "Completed", "Duration"], + [_kind_display(sched.key), "Status", "Started", "Completed", "Duration"], rows, code=[0], ) diff --git a/testgen/mcp/tools/test_runs.py b/testgen/mcp/tools/test_runs.py index 415571e6..2303a500 100644 --- a/testgen/mcp/tools/test_runs.py +++ b/testgen/mcp/tools/test_runs.py @@ -160,7 +160,7 @@ def get_test_run(job_execution_id: str) -> str: doc = MdDoc() suite_label = summary.test_suite or "—" doc.heading(1, f"Test run: {suite_label}") - doc.field("Job ID", summary.job_execution_id, code=True) + doc.field("Test Run", summary.job_execution_id, code=True) doc.field("Test suite", suite_label) if summary.table_groups_name: doc.field("Table group", summary.table_groups_name) @@ -278,7 +278,7 @@ def _select_pending_test_jes( def _render_pending_je(doc: MdDoc, je: JobExecution, label: str) -> None: status_label = TestRunSummary.STATUS_LABEL.get(je.status, je.status) doc.heading(3, f"{label} — {status_label}") - doc.field("Job ID", je.id, code=True) + doc.field("Test Run", je.id, code=True) if je.job_schedule_id is not None: doc.field("Schedule", je.job_schedule_id, code=True) doc.field("Submitted", je.created_at) @@ -289,7 +289,7 @@ def _render_pending_je(doc: MdDoc, je: JobExecution, label: str) -> None: def _render_test_run_section(doc: MdDoc, run: TestRunSummary) -> None: title = run.test_suite or run.project_code doc.heading(2, f"{title} — {run.status_label}") - doc.field("Job ID", run.job_execution_id, code=True) + doc.field("Test Run", run.job_execution_id, code=True) if run.job_schedule_id is not None: doc.field("Schedule", run.job_schedule_id, code=True) if run.test_suite: diff --git a/testgen/server/__init__.py b/testgen/server/__init__.py index d0867e9d..a1c5afef 100644 --- a/testgen/server/__init__.py +++ b/testgen/server/__init__.py @@ -44,6 +44,11 @@ def _patch_openapi_schema(app: FastAPI) -> None: - Top-level component schema titles (shown in Redoc sidebar) - Path/query parameter schemas + It also strips ``description`` from enum component schemas. A ``StrEnum`` used + as a field type surfaces its class docstring as the schema description; enums + are internal primitives free to carry technical docstrings, so the public + schema must not echo them. Per-field API text comes from ``Field(description=...)``. + FastAPI caches the schema after the first call, so the patching runs once. """ _original = app.openapi @@ -56,6 +61,8 @@ def patched_openapi() -> dict: for branch in prop.get("anyOf", []): branch.pop("title", None) model_schema.pop("title", None) + if "enum" in model_schema: + model_schema.pop("description", None) for methods in schema.get("paths", {}).values(): for details in methods.values(): if isinstance(details, dict): diff --git a/tests/unit/api/test_jobs.py b/tests/unit/api/test_jobs.py index 6af04eb4..0659bb04 100644 --- a/tests/unit/api/test_jobs.py +++ b/tests/unit/api/test_jobs.py @@ -7,6 +7,7 @@ import pytest from fastapi import FastAPI, HTTPException from fastapi.testclient import TestClient +from sqlalchemy.dialects import postgresql from testgen.api.deps import db_session, get_authorized_user from testgen.api.jobs import ( @@ -18,6 +19,8 @@ submit_test_generation, submit_test_run, ) +from testgen.common.enums import JobKey, PublicJobKey +from testgen.common.models.job_execution import JobExecution pytestmark = pytest.mark.unit @@ -236,3 +239,44 @@ def test_list_jobs_accepts_valid_status(mock_je_cls, _mock_perm): assert resp.status_code == 200 # Verify the status string was forwarded to the model layer. assert mock_je_cls.list_for_project.call_args.kwargs["status"] == "completed" + + +# --- internal job-kind exclusion --- +# JobResponse.job_key is typed PublicJobKey, which is only safe because every path +# that builds a JobResponse excludes internal job kinds. These pin that invariant. + + +def test_public_job_key_excludes_internal_kinds(): + public = {member.value for member in PublicJobKey} + assert public == {"run-profile", "run-tests", "run-test-generation"} + internal = {"run-monitors", "run-score-update", "recalculate-project-scores", "run-data-cleanup"} + assert public.isdisjoint(internal) + # Every public value must be a real JobKey. + assert public <= {member.value for member in JobKey} + + +@patch.object(JobExecution, "list_for_project", return_value=([], 0)) +def test_list_jobs_filters_to_public_job_kinds(_mock_list): + list_jobs(project_code="DEFAULT", job_key=None, status=None, page=1, limit=20) + + clause = _mock_list.call_args.args[1] + sql = str(clause.compile(dialect=postgresql.dialect(), compile_kwargs={"literal_binds": True})) + assert "run-profile" in sql + assert "run-tests" in sql + assert "run-test-generation" in sql + for internal in ("run-monitors", "run-score-update", "recalculate-project-scores", "run-data-cleanup"): + assert internal not in sql + + +@patch("testgen.api.deps.has_project_permission", return_value=True) +@patch(f"{MODULE}.JobExecution") +def test_list_jobs_rejects_internal_job_key(mock_je_cls, _mock_perm): + mock_je_cls.list_for_project.return_value = ([], 0) + client = TestClient(_client_with_overrides()) + + resp = client.get("/api/v1/projects/DEFAULT/jobs?job_key=run-data-cleanup") + + assert resp.status_code == 422 + body = resp.json() + assert body["detail"][0]["loc"] == ["query", "job_key"] + assert body["detail"][0]["type"] == "enum" From f30a932edf81077157be3dc821012ba2fa97e492 Mon Sep 17 00:00:00 2001 From: Ricardo Boni Date: Thu, 4 Jun 2026 15:26:13 -0400 Subject: [PATCH 08/78] refactor(runs): read run lifecycle through a job_execution relationship Add a job_execution relationship to TestRun/ProfilingRun and migrate the entity-instance lifecycle readers (notification trigger logic, monitor end-time, CLI run summary) to read status/timestamps through it. Unify the profiling-run email onto the job-execution-sourced context (lowercase JobStatus + status_label) and retire the now-dead format_status helper. No behavior change; prepares for run-table deduplication (TG-1047). TG-1115 Co-Authored-By: Claude Opus 4.8 --- testgen/commands/job_runner.py | 4 +-- testgen/common/models/profiling_run.py | 9 ++++- testgen/common/models/test_run.py | 8 +++++ testgen/common/notifications/monitor_run.py | 2 +- testgen/common/notifications/notifications.py | 8 ----- testgen/common/notifications/profiling_run.py | 34 +++++++++++-------- testgen/common/notifications/test_run.py | 3 +- .../test_monitor_run_notifications.py | 30 ++++++---------- .../test_profiling_run_notifications.py | 20 +++++++---- .../test_test_run_notifications.py | 29 +++++++++------- tests/unit/common/test_enums.py | 18 ++++++++++ 11 files changed, 98 insertions(+), 67 deletions(-) create mode 100644 tests/unit/common/test_enums.py diff --git a/testgen/commands/job_runner.py b/testgen/commands/job_runner.py index f95584e6..d1b7b068 100644 --- a/testgen/commands/job_runner.py +++ b/testgen/commands/job_runner.py @@ -73,12 +73,12 @@ def _print_run_summary(job_id: UUID, job_key: str) -> None: case "run-profile": run = session.scalars(select(ProfilingRun).where(ProfilingRun.job_execution_id == job_id)).first() if run: - status_msg = "Profiling encountered an error. Check log for details." if run.status == "Error" else "Profiling completed." + status_msg = "Profiling encountered an error. Check log for details." if run.job_execution.status == JobStatus.ERROR else "Profiling completed." click.echo(f"\n {status_msg}\n Run ID: {run.id}\n ") case "run-tests" | "run-monitors": run = session.scalars(select(TestRun).where(TestRun.job_execution_id == job_id)).first() if run: - status_msg = "Test execution encountered an error. Check log for details." if run.status == "Error" else "Test execution completed." + status_msg = "Test execution encountered an error. Check log for details." if run.job_execution.status == JobStatus.ERROR else "Test execution completed." click.echo(f"\n {status_msg}\n Run ID: {run.id}\n ") case "run-test-generation": click.echo("Test generation completed.") diff --git a/testgen/common/models/profiling_run.py b/testgen/common/models/profiling_run.py index 39bace1c..0c898172 100644 --- a/testgen/common/models/profiling_run.py +++ b/testgen/common/models/profiling_run.py @@ -6,7 +6,7 @@ from sqlalchemy import BigInteger, Column, Float, Integer, String, desc, func, select, text, update from sqlalchemy.dialects import postgresql -from sqlalchemy.orm import InstrumentedAttribute +from sqlalchemy.orm import InstrumentedAttribute, foreign, relationship from sqlalchemy.orm.attributes import flag_modified from sqlalchemy.sql.expression import case @@ -117,6 +117,13 @@ class ProfilingRun(Entity): process_id: int = Column(Integer) job_execution_id: UUID | None = Column(postgresql.UUID(as_uuid=True), nullable=True) + job_execution = relationship( + JobExecution, + primaryjoin=foreign(job_execution_id) == JobExecution.id, + uselist=False, + viewonly=True, + ) + _default_order_by = (desc(profiling_starttime),) _minimal_columns = ( id, diff --git a/testgen/common/models/test_run.py b/testgen/common/models/test_run.py index 385b429b..1887d886 100644 --- a/testgen/common/models/test_run.py +++ b/testgen/common/models/test_run.py @@ -5,6 +5,7 @@ from sqlalchemy import BigInteger, Column, Float, ForeignKey, Integer, String, Text, desc, func, select, text, update from sqlalchemy.dialects import postgresql +from sqlalchemy.orm import foreign, relationship from sqlalchemy.orm.attributes import flag_modified from sqlalchemy.sql.expression import case @@ -121,6 +122,13 @@ class TestRun(Entity): process_id: int = Column(Integer) job_execution_id: UUID | None = Column(postgresql.UUID(as_uuid=True), nullable=True) + job_execution = relationship( + JobExecution, + primaryjoin=foreign(job_execution_id) == JobExecution.id, + uselist=False, + viewonly=True, + ) + _default_order_by = (desc(test_starttime),) _minimal_columns = ( id, diff --git a/testgen/common/notifications/monitor_run.py b/testgen/common/notifications/monitor_run.py index 4967216a..9c9a55ec 100644 --- a/testgen/common/notifications/monitor_run.py +++ b/testgen/common/notifications/monitor_run.py @@ -236,7 +236,7 @@ def send_monitor_notifications(test_run: TestRun, result_list_ct=20): notification.recipients, { "summary": { - "test_endtime": test_run.test_endtime, + "test_endtime": test_run.job_execution.completed_at, "table_groups_name": table_group.table_groups_name, "project_name": project.project_name, "table_name": table_name, diff --git a/testgen/common/notifications/notifications.py b/testgen/common/notifications/notifications.py index 68e20732..23a9f896 100644 --- a/testgen/common/notifications/notifications.py +++ b/testgen/common/notifications/notifications.py @@ -1,8 +1,6 @@ import math from datetime import datetime -from testgen.common.models.profiling_run import ProfilingRunStatus -from testgen.common.models.test_definition import TestRunStatus from testgen.common.notifications.base import BaseEmailTemplate from testgen.utils import friendly_score @@ -29,12 +27,6 @@ def format_duration_helper(self, start_time: datetime, end_time: datetime) -> st formatted = " ".join([ f"{unit[0]}{unit[1]}" for unit in units if unit[0] ]) return formatted.strip() or "< 1s" - def format_status_helper(self, status: TestRunStatus | ProfilingRunStatus) -> str: - return { - "Complete": "Completed", - "Cancelled": "Canceled", - }.get(status, status) - def format_score_helper(self, score: float) -> str: return friendly_score(score) diff --git a/testgen/common/notifications/profiling_run.py b/testgen/common/notifications/profiling_run.py index 2a4c254c..4d7b2197 100644 --- a/testgen/common/notifications/profiling_run.py +++ b/testgen/common/notifications/profiling_run.py @@ -4,6 +4,7 @@ from sqlalchemy import select from testgen import settings +from testgen.common.enums import JOB_STATUS_LABEL, JobStatus from testgen.common.models import get_current_session, with_database_session from testgen.common.models.hygiene_issue import HygieneIssue from testgen.common.models.notification_settings import ( @@ -23,7 +24,7 @@ class ProfilingRunEmailTemplate(BaseNotificationTemplate): def get_subject_template(self) -> str: return ( - "[TestGen] Profiling Run {{format_status profiling_run.status}}: {{table_groups_name}}" + "[TestGen] Profiling Run {{profiling_run.status_label}}: {{table_groups_name}}" "{{#if issue_count}}" ' | {{format_number issue_count}} hygiene {{pluralize issue_count "issue" "issues"}}' "{{/if}}" @@ -32,9 +33,9 @@ def get_subject_template(self) -> str: def get_title_template(self): return """ TestGen Profiling Run - {{format_status profiling_run.status}} + {{#if (eq profiling_run.status 'error')}} text-red {{/if}} + {{#if (eq profiling_run.status 'canceled')}} text-purple {{/if}} + ">{{profiling_run.status_label}} """ def get_main_content_template(self): @@ -85,7 +86,7 @@ def get_main_content_template(self): - {{#if (eq profiling_run.status 'Complete')}} + {{#if (eq profiling_run.status 'completed')}} {{#if (eq notification_trigger 'on_changes')}} Profiling run detected new hygiene issues. {{/if}} @@ -95,15 +96,15 @@ def get_main_content_template(self): {{/if}} {{/if}} {{/if}} - {{#if (eq profiling_run.status 'Error')}} + {{#if (eq profiling_run.status 'error')}} Profiling encountered an error. {{/if}} - {{#if (eq profiling_run.status 'Cancelled')}} + {{#if (eq profiling_run.status 'canceled')}} Profiling run was canceled. {{/if}} - {{#if (eq profiling_run.status 'Complete')}} + {{#if (eq profiling_run.status 'completed')}} @@ -131,12 +132,12 @@ def get_main_content_template(self): {{/if}} - {{#if (eq profiling_run.status 'Error')}} + {{#if (eq profiling_run.status 'error')}} {{/if}} - {{#if (eq profiling_run.status 'Complete')}} + {{#if (eq profiling_run.status 'completed')}} - + {{/if}} {{#if (eq test_run.status 'completed')}} @@ -330,7 +330,7 @@ def send_test_run_notifications(test_run: TestRun, result_list_ct=20, result_sta "/test-runs:results?project_code=", str(tr_summary.project_code), "&run_id=", - str(test_run.job_execution_id), + str(test_run.id), "&source=email" ) ) diff --git a/testgen/mcp/tools/common.py b/testgen/mcp/tools/common.py index d05511ce..4469abb5 100644 --- a/testgen/mcp/tools/common.py +++ b/testgen/mcp/tools/common.py @@ -602,7 +602,7 @@ def resolve_profiling_run(job_execution_id: str) -> ProfilingRun: so callers don't leak existence of runs they shouldn't see. """ run_uuid = parse_uuid(job_execution_id, "job_execution_id") - run = ProfilingRun.get_by_id_or_job(run_uuid) + run = ProfilingRun.get(run_uuid) perms = get_project_permissions() if run is None or not perms.has_access(run.project_code): raise MCPResourceNotAccessible("Profiling run", job_execution_id) diff --git a/testgen/mcp/tools/discovery.py b/testgen/mcp/tools/discovery.py index 978cc858..3ff393e7 100644 --- a/testgen/mcp/tools/discovery.py +++ b/testgen/mcp/tools/discovery.py @@ -2,7 +2,6 @@ from testgen.common.models import with_database_session from testgen.common.models.data_table import DataTable from testgen.common.models.project import Project -from testgen.common.models.test_run import TestRun from testgen.common.models.test_suite import TestSuite from testgen.mcp.exceptions import MCPResourceNotAccessible from testgen.mcp.permissions import get_project_permissions, mcp_permission @@ -71,10 +70,6 @@ def list_test_suites(project_code: str) -> str: if not summaries: return f"No test suites found for project `{project_code}`." - # Batch-lookup job_execution_ids for latest runs - run_ids = [s.latest_run_id for s in summaries if s.latest_run_id] - job_exec_map = TestRun.get_job_execution_ids(run_ids) if run_ids else {} - doc = MdDoc() doc.heading(1, f"Test Suites for `{project_code}`") for s in summaries: @@ -86,8 +81,7 @@ def list_test_suites(project_code: str) -> str: doc.field("Test definitions", s.test_ct or 0) if s.latest_run_id: - run_id = job_exec_map.get(s.latest_run_id) or s.latest_run_id - doc.field("Latest run", f"`{run_id}` ({s.latest_run_start})") + doc.field("Latest run", f"`{s.latest_run_id}` ({s.latest_run_start})") results_summary = ( f"{s.last_run_test_ct or 0} tests: " f"{s.last_run_passed_ct or 0} passed, " diff --git a/testgen/mcp/tools/hygiene_issues.py b/testgen/mcp/tools/hygiene_issues.py index 01084fbc..ac7846cf 100644 --- a/testgen/mcp/tools/hygiene_issues.py +++ b/testgen/mcp/tools/hygiene_issues.py @@ -87,12 +87,12 @@ def _resolve_profile_run_je_id( return je_uuid job_uuid = parse_uuid(job_execution_id, "job_execution_id") - run = ProfilingRun.get_by_id_or_job(job_uuid) + run = ProfilingRun.get(job_uuid) perms = get_project_permissions() tg = TableGroup.get(run.table_groups_id) if run else None if run is None or tg is None or not perms.has_access(tg.project_code): raise MCPResourceNotAccessible("Profiling run", job_execution_id) - return run.job_execution_id + return run.id @with_database_session diff --git a/testgen/mcp/tools/profile_history.py b/testgen/mcp/tools/profile_history.py index 17839798..62537cc4 100644 --- a/testgen/mcp/tools/profile_history.py +++ b/testgen/mcp/tools/profile_history.py @@ -203,7 +203,7 @@ def _delta_cell(metric: ProfileMetric, baseline: object | None, target: object | def _require_completed(run: ProfilingRun, label: str) -> None: """Raise if the run's job execution isn't completed.""" - je = get_current_session().get(JobExecution, run.job_execution_id) + je = get_current_session().get(JobExecution, run.id) if je.status != JobStatus.COMPLETED: status_label = ProfilingRunSummary.STATUS_LABEL.get(je.status, je.status) raise MCPUserError( @@ -475,8 +475,8 @@ def _render_run_comparison( ["", "Target", "Baseline"], [ ["Profiling Run", - MdDoc.code(str(target_run.job_execution_id)), - MdDoc.code(str(baseline_run.job_execution_id))], + MdDoc.code(str(target_run.id)), + MdDoc.code(str(baseline_run.id))], ["Started", target_run.profiling_starttime, baseline_run.profiling_starttime], ], ) @@ -792,7 +792,7 @@ def _render_schema_history( baseline_snap=snapshots.get(baseline.id, {}), ) doc.heading(2, f"Run started {_format_run_label(target)}") - doc.field("Profiling Run", target.job_execution_id, code=True) + doc.field("Profiling Run", target.id, code=True) if section_lines: doc.bullets(section_lines) else: diff --git a/testgen/mcp/tools/profiling.py b/testgen/mcp/tools/profiling.py index 874cd5c3..7cce2b34 100644 --- a/testgen/mcp/tools/profiling.py +++ b/testgen/mcp/tools/profiling.py @@ -4,6 +4,7 @@ from sqlalchemy import func, or_ from testgen.common.data_catalog_service import build_create_table_script +from testgen.common.enums import JOB_STATUS_LABEL, JobStatus from testgen.common.models import with_database_session from testgen.common.models.data_column import ( SUGGESTED_DATA_TYPE_TO_PREFIX, @@ -744,7 +745,7 @@ def get_column_profile_detail( "Run profiling for the table group first." ) - if detail.profile_run_status in ("Running", "Error", "Cancelled"): + if detail.profile_run_status in (JobStatus.RUNNING, JobStatus.ERROR, JobStatus.CANCELED): _raise_run_not_ready(detail) payload = dataclasses.asdict(detail) @@ -766,11 +767,12 @@ def _raise_run_not_ready(detail: ColumnProfileDetail) -> None: ended = detail.profile_run_ended_at started_label = started.strftime("%Y-%m-%d %H:%M UTC") if started else "—" ended_label = ended.strftime("%Y-%m-%d %H:%M UTC") if ended else "—" + status_label = JOB_STATUS_LABEL.get(status, status) lines = [ - f"Profiling run `{je}` is in `{status}` state — no profile detail available.", + f"Profiling run `{je}` is in `{status_label}` state — no profile detail available.", f"Started: {started_label}. Ended: {ended_label}.", ] - if status == "Error" and detail.profile_run_log_message: + if status == JobStatus.ERROR and detail.profile_run_log_message: lines.append(f"Error: {detail.profile_run_log_message}") raise MCPUserError("\n".join(lines)) @@ -953,7 +955,7 @@ def get_column_frequent_values( doc = MdDoc() doc.heading(1, f"Frequent values: {table_name}.{column_name}") doc.field("Table group", tg.id, code=True) - doc.field("Profiling Run", profiling_run.job_execution_id, code=True) + doc.field("Profiling Run", profiling_run.id, code=True) doc.field("Row Count", profile.record_ct) doc.field("Distinct values", profile.distinct_value_ct) if pii_flag: @@ -1008,7 +1010,7 @@ def get_column_patterns( doc = MdDoc() doc.heading(1, f"Character patterns: {table_name}.{column_name}") doc.field("Table group", tg.id, code=True) - doc.field("Profiling Run", profiling_run.job_execution_id, code=True) + doc.field("Profiling Run", profiling_run.id, code=True) doc.field("Row Count", profile.record_ct) doc.field("Distinct values", profile.distinct_value_ct) diff --git a/testgen/mcp/tools/test_results.py b/testgen/mcp/tools/test_results.py index d4d215cf..f38c574b 100644 --- a/testgen/mcp/tools/test_results.py +++ b/testgen/mcp/tools/test_results.py @@ -90,14 +90,14 @@ def list_test_results( raise MCPResourceNotAccessible("Test suite", test_suite_id) if suite.last_complete_test_run_id is None: raise MCPUserError(f"No completed test runs found for test suite `{test_suite_id}`.") - test_run = TestRun.get_by_id_or_job(suite.last_complete_test_run_id) + test_run = TestRun.get(suite.last_complete_test_run_id) if test_run is None: raise MCPUserError(f"No completed test runs found for test suite `{test_suite_id}`.") resolved_via_suite = True - run_id_label = str(test_run.job_execution_id) + run_id_label = str(test_run.id) else: job_uuid = parse_uuid(job_execution_id, "job_execution_id") - test_run = TestRun.get_by_id_or_job(job_uuid) + test_run = TestRun.get(job_uuid) suite = TestSuite.get_regular(test_run.test_suite_id) if test_run else None if test_run is None or suite is None or not perms.has_access(suite.project_code): raise MCPResourceNotAccessible("Test run", job_execution_id) @@ -208,7 +208,7 @@ def get_failure_summary( if job_execution_id: job_uuid = parse_uuid(job_execution_id, "job_execution_id") - test_run = TestRun.get_by_id_or_job(job_uuid) + test_run = TestRun.get(job_uuid) suite = TestSuite.get_regular(test_run.test_suite_id) if test_run else None if test_run is None or suite is None or not perms.has_access(suite.project_code): raise MCPResourceNotAccessible("Test run", job_execution_id) @@ -548,7 +548,7 @@ def compare_test_runs( perms = get_project_permissions() def _resolve_accessible(je_id_str: str, je_uuid: UUID) -> TestRun: - run = TestRun.get_by_id_or_job(je_uuid) + run = TestRun.get(je_uuid) if run is None: raise MCPResourceNotAccessible("Test run", je_id_str) suite = TestSuite.get_regular(run.test_suite_id) @@ -557,7 +557,7 @@ def _resolve_accessible(je_id_str: str, je_uuid: UUID) -> TestRun: return run def _require_completed(run: TestRun, label: str) -> None: - je = get_current_session().get(JobExecution, run.job_execution_id) + je = get_current_session().get(JobExecution, run.id) if je.status != JobStatus.COMPLETED: status_label = TestRunSummary.STATUS_LABEL.get(je.status, je.status) raise MCPUserError( @@ -595,8 +595,8 @@ def _require_completed(run: TestRun, label: str) -> None: ["", "Target", "Baseline"], [ ["Test Run", - MdDoc.code(str(target_run.job_execution_id)), - MdDoc.code(str(baseline_run.job_execution_id))], + MdDoc.code(str(target_run.id)), + MdDoc.code(str(baseline_run.id))], ["Started", target_run.test_starttime, baseline_run.test_starttime], ], ) @@ -705,12 +705,12 @@ def bulk_update_test_results( db_disposition = parse_test_result_disposition(disposition) if job_execution_id: - run = TestRun.get_by_id_or_job(parse_uuid(job_execution_id, "job_execution_id")) + run = TestRun.get(parse_uuid(job_execution_id, "job_execution_id")) if run is None or run.test_suite_id != suite.id: raise MCPResourceNotAccessible("Test run", job_execution_id) else: run = ( - TestRun.get_by_id_or_job(suite.last_complete_test_run_id) + TestRun.get(suite.last_complete_test_run_id) if suite.last_complete_test_run_id else None ) @@ -751,7 +751,7 @@ def bulk_update_test_results( doc.heading(1, f"Updated {update.matched} test results in suite `{suite.test_suite}`") doc.field("Disposition", disposition) - doc.field("Run", run.job_execution_id, code=True) + doc.field("Run", run.id, code=True) doc.field("Filter", filter_str) if update.passed_skipped: doc.text( diff --git a/testgen/template/dbsetup/030_initialize_new_schema_structure.sql b/testgen/template/dbsetup/030_initialize_new_schema_structure.sql index 2e381725..a71e6718 100644 --- a/testgen/template/dbsetup/030_initialize_new_schema_structure.sql +++ b/testgen/template/dbsetup/030_initialize_new_schema_structure.sql @@ -144,10 +144,7 @@ CREATE TABLE profiling_runs ( connection_id BIGINT NOT NULL, table_groups_id UUID NOT NULL, profiling_starttime TIMESTAMP, - profiling_endtime TIMESTAMP, - status VARCHAR(100) DEFAULT 'Running', progress JSONB, - log_message VARCHAR, table_ct BIGINT, column_ct BIGINT, record_ct BIGINT, @@ -157,9 +154,7 @@ CREATE TABLE profiling_runs ( anomaly_column_ct BIGINT, dq_affected_data_points BIGINT, dq_total_data_points BIGINT, - dq_score_profiling FLOAT, - process_id INTEGER, - job_execution_id UUID + dq_score_profiling FLOAT ); CREATE TABLE test_suites ( @@ -612,10 +607,7 @@ CREATE TABLE test_runs ( PRIMARY KEY, test_suite_id UUID NOT NULL, test_starttime TIMESTAMP, - test_endtime TIMESTAMP, - status VARCHAR(100) DEFAULT 'Running', progress JSONB, - log_message TEXT, test_ct INTEGER, passed_ct INTEGER, failed_ct INTEGER, @@ -629,8 +621,6 @@ CREATE TABLE test_runs ( dq_affected_data_points BIGINT, dq_total_data_points BIGINT, dq_score_test_run FLOAT, - process_id INTEGER, - job_execution_id UUID, CONSTRAINT test_runs_test_suites_fk FOREIGN KEY (test_suite_id) REFERENCES test_suites ); @@ -1141,3 +1131,26 @@ CREATE TABLE notification_settings ( REFERENCES score_definitions (id) ON DELETE CASCADE ); + +-- Run identity = job execution id, with a cascading delete chain: +-- deleting a job_executions row removes its run and the run's results. +-- Declared here because job_executions is created after the run tables. +ALTER TABLE profiling_runs + ADD CONSTRAINT profiling_runs_job_executions_id_fk + FOREIGN KEY (id) REFERENCES job_executions (id) ON DELETE CASCADE; + +ALTER TABLE test_runs + ADD CONSTRAINT test_runs_job_executions_id_fk + FOREIGN KEY (id) REFERENCES job_executions (id) ON DELETE CASCADE; + +ALTER TABLE test_results + ADD CONSTRAINT test_results_test_runs_id_fk + FOREIGN KEY (test_run_id) REFERENCES test_runs (id) ON DELETE CASCADE; + +ALTER TABLE profile_results + ADD CONSTRAINT profile_results_profiling_runs_id_fk + FOREIGN KEY (profile_run_id) REFERENCES profiling_runs (id) ON DELETE CASCADE; + +ALTER TABLE profile_anomaly_results + ADD CONSTRAINT profile_anomaly_results_profiling_runs_id_fk + FOREIGN KEY (profile_run_id) REFERENCES profiling_runs (id) ON DELETE CASCADE; diff --git a/testgen/template/dbupgrade/0194_incremental_upgrade.sql b/testgen/template/dbupgrade/0194_incremental_upgrade.sql new file mode 100644 index 00000000..0d89f265 --- /dev/null +++ b/testgen/template/dbupgrade/0194_incremental_upgrade.sql @@ -0,0 +1,150 @@ +SET SEARCH_PATH TO {SCHEMA_NAME}; + +-- Deduplicate the run tables against job_executions. +-- +-- The run primary key becomes the job execution id: job_execution_id is +-- promoted to the primary key (which is also a foreign key to +-- job_executions.id with cascade on delete), and the duplicated status, +-- end-time, log-message, and process-id columns are dropped. External +-- references are rewritten from the old run id to the job execution id. +-- Because job_execution_id already holds that value on every row, the +-- run identity is reshaped with pure DDL rather than rewriting primary +-- key values in place. + +-- Standard views join the run tables on their primary key, so they must be +-- dropped before the id column can be reshaped. _refresh_static_metadata +-- recreates them from 060_create_standard_views.sql after the upgrade. +DROP VIEW IF EXISTS v_latest_profile_results CASCADE; +DROP VIEW IF EXISTS v_inactive_anomalies CASCADE; +DROP VIEW IF EXISTS v_queued_observability_results CASCADE; +DROP VIEW IF EXISTS v_dq_profile_scoring_latest_by_column CASCADE; +DROP VIEW IF EXISTS v_dq_profile_scoring_latest_by_dimension CASCADE; +DROP VIEW IF EXISTS v_dq_test_scoring_latest_by_column CASCADE; +DROP VIEW IF EXISTS v_dq_test_scoring_latest_by_dimension CASCADE; +DROP VIEW IF EXISTS v_dq_profile_scoring_latest_by_impact_dimension CASCADE; +DROP VIEW IF EXISTS v_dq_test_scoring_latest_by_impact_dimension CASCADE; +DROP VIEW IF EXISTS v_dq_profile_scoring_history_by_column CASCADE; +DROP VIEW IF EXISTS v_dq_test_scoring_history_by_column CASCADE; + +-- Remove result rows orphaned by an already-deleted run. They are +-- unreachable in the application and would block the value rewrite and the +-- new cascade constraints below. +DO $$ +DECLARE + orphan_ct BIGINT; +BEGIN + DELETE FROM test_results c + WHERE NOT EXISTS (SELECT 1 FROM test_runs r WHERE r.id = c.test_run_id); + GET DIAGNOSTICS orphan_ct = ROW_COUNT; + IF orphan_ct > 0 THEN RAISE NOTICE 'deleted % orphan test_results row(s)', orphan_ct; END IF; + + DELETE FROM profile_results c + WHERE NOT EXISTS (SELECT 1 FROM profiling_runs r WHERE r.id = c.profile_run_id); + GET DIAGNOSTICS orphan_ct = ROW_COUNT; + IF orphan_ct > 0 THEN RAISE NOTICE 'deleted % orphan profile_results row(s)', orphan_ct; END IF; + + DELETE FROM profile_anomaly_results c + WHERE NOT EXISTS (SELECT 1 FROM profiling_runs r WHERE r.id = c.profile_run_id); + GET DIAGNOSTICS orphan_ct = ROW_COUNT; + IF orphan_ct > 0 THEN RAISE NOTICE 'deleted % orphan profile_anomaly_results row(s)', orphan_ct; END IF; +END $$; + +-- Rewrite result-row references from the old run id to the job execution id. +UPDATE test_results c + SET test_run_id = r.job_execution_id + FROM test_runs r + WHERE r.id = c.test_run_id; + +UPDATE profile_results c + SET profile_run_id = r.job_execution_id + FROM profiling_runs r + WHERE r.id = c.profile_run_id; + +UPDATE profile_anomaly_results c + SET profile_run_id = r.job_execution_id + FROM profiling_runs r + WHERE r.id = c.profile_run_id; + +-- Repoint the cache/reference pointers to the job execution id. These are +-- not foreign keys; rows pointing at an already-deleted run are left as-is +-- (a dangling cache that resolves to no run, exactly as before). +UPDATE table_groups tg + SET last_complete_profile_run_id = r.job_execution_id + FROM profiling_runs r + WHERE r.id = tg.last_complete_profile_run_id; + +UPDATE test_suites ts + SET last_complete_test_run_id = r.job_execution_id + FROM test_runs r + WHERE r.id = ts.last_complete_test_run_id; + +UPDATE data_table_chars dtc + SET last_complete_profile_run_id = r.job_execution_id + FROM profiling_runs r + WHERE r.id = dtc.last_complete_profile_run_id; + +UPDATE data_column_chars dcc + SET last_complete_profile_run_id = r.job_execution_id + FROM profiling_runs r + WHERE r.id = dcc.last_complete_profile_run_id; + +UPDATE score_history_latest_runs sr + SET last_test_run_id = r.job_execution_id + FROM test_runs r + WHERE r.id = sr.last_test_run_id; + +UPDATE score_history_latest_runs sr + SET last_profiling_run_id = r.job_execution_id + FROM profiling_runs r + WHERE r.id = sr.last_profiling_run_id; + +UPDATE test_definitions td + SET profile_run_id = r.job_execution_id + FROM profiling_runs r + WHERE r.id = td.profile_run_id; + +-- Promote job_execution_id to the primary key on profiling_runs. The column +-- already holds the job execution id, so this is pure DDL. +ALTER TABLE profiling_runs DROP CONSTRAINT pk_prun_id; +ALTER TABLE profiling_runs DROP COLUMN id; +ALTER TABLE profiling_runs RENAME COLUMN job_execution_id TO id; +ALTER TABLE profiling_runs ADD CONSTRAINT pk_prun_id PRIMARY KEY (id); +ALTER TABLE profiling_runs + ADD CONSTRAINT profiling_runs_job_executions_id_fk + FOREIGN KEY (id) REFERENCES job_executions (id) ON DELETE CASCADE; + +-- Promote job_execution_id to the primary key on test_runs. +ALTER TABLE test_runs DROP CONSTRAINT test_runs_id_pk; +ALTER TABLE test_runs DROP COLUMN id; +ALTER TABLE test_runs RENAME COLUMN job_execution_id TO id; +ALTER TABLE test_runs ADD CONSTRAINT test_runs_id_pk PRIMARY KEY (id); +ALTER TABLE test_runs + ADD CONSTRAINT test_runs_job_executions_id_fk + FOREIGN KEY (id) REFERENCES job_executions (id) ON DELETE CASCADE; + +-- Make the result-row references real cascade foreign keys so deleting a job +-- execution removes its run and the run's results in one DB-level cascade. +ALTER TABLE test_results + ADD CONSTRAINT test_results_test_runs_id_fk + FOREIGN KEY (test_run_id) REFERENCES test_runs (id) ON DELETE CASCADE; + +ALTER TABLE profile_results + ADD CONSTRAINT profile_results_profiling_runs_id_fk + FOREIGN KEY (profile_run_id) REFERENCES profiling_runs (id) ON DELETE CASCADE; + +ALTER TABLE profile_anomaly_results + ADD CONSTRAINT profile_anomaly_results_profiling_runs_id_fk + FOREIGN KEY (profile_run_id) REFERENCES profiling_runs (id) ON DELETE CASCADE; + +-- Drop the columns now owned by job_executions. +ALTER TABLE test_runs + DROP COLUMN status, + DROP COLUMN test_endtime, + DROP COLUMN log_message, + DROP COLUMN process_id; + +ALTER TABLE profiling_runs + DROP COLUMN status, + DROP COLUMN profiling_endtime, + DROP COLUMN log_message, + DROP COLUMN process_id; diff --git a/testgen/template/execution/get_errored_autogen_monitors.sql b/testgen/template/execution/get_errored_autogen_monitors.sql index dd0fb0a5..7d54a266 100644 --- a/testgen/template/execution/get_errored_autogen_monitors.sql +++ b/testgen/template/execution/get_errored_autogen_monitors.sql @@ -1,9 +1,10 @@ WITH prev_run AS ( - SELECT id + SELECT test_runs.id FROM test_runs + INNER JOIN job_executions ON job_executions.id = test_runs.id WHERE test_suite_id = :TEST_SUITE_ID ::UUID - AND id <> :TEST_RUN_ID ::UUID - AND status = 'Complete' + AND test_runs.id <> :TEST_RUN_ID ::UUID + AND job_executions.status = 'completed' ORDER BY test_starttime DESC LIMIT 1 ) diff --git a/testgen/template/get_entities/get_profile_list.sql b/testgen/template/get_entities/get_profile_list.sql index 088bc745..8bbb83e0 100644 --- a/testgen/template/get_entities/get_profile_list.sql +++ b/testgen/template/get_entities/get_profile_list.sql @@ -5,13 +5,15 @@ Optional: table-name*/ SELECT p.id as profile_run_id, p.project_code as project_key, schema_name, p.table_groups_id, - profiling_starttime as start_time, status, + profiling_starttime as start_time, je.status, COUNT(DISTINCT table_name) as tables, COUNT(DISTINCT table_name || '.' || column_name) as columns FROM profiling_runs p INNER JOIN profile_results r ON (p.id = r.profile_run_id) +INNER JOIN job_executions je + ON (je.id = p.id) WHERE p.table_groups_id = :TABLE_GROUP_ID GROUP BY p.id, p.project_code, p.connection_id, schema_name, p.table_groups_id, - profiling_starttime, status + profiling_starttime, je.status ORDER BY profiling_starttime DESC; diff --git a/testgen/template/get_entities/get_test_run_list.sql b/testgen/template/get_entities/get_test_run_list.sql index 50f9ecc7..ea65632e 100644 --- a/testgen/template/get_entities/get_test_run_list.sql +++ b/testgen/template/get_entities/get_test_run_list.sql @@ -5,16 +5,16 @@ Optional: table-name, column-name, from-date, thru-date*/ SELECT ts.test_suite as test_suite_key, tr.test_starttime as test_time, - tr.status, + je.status, tr.id::VARCHAR as test_run_id, COUNT(DISTINCT lower(r.schema_name || '.' || table_name)) as table_ct, COUNT(*) as result_ct, SUM(CASE WHEN r.result_code = 0 THEN 1 END) as fail_ct, - SUM(CASE WHEN r.observability_status = 'Sent' THEN 1 END) as sent_to_obs, - process_id + SUM(CASE WHEN r.observability_status = 'Sent' THEN 1 END) as sent_to_obs FROM test_runs tr INNER JOIN test_results r ON tr.id = r.test_run_id INNER JOIN test_suites ts ON tr.test_suite_id = ts.id +INNER JOIN job_executions je ON je.id = tr.id WHERE ts.project_code = :PROJECT_CODE AND ts.test_suite = :TEST_SUITE AND ts.is_monitor IS NOT TRUE @@ -22,4 +22,4 @@ INNER JOIN test_suites ts ON tr.test_suite_id = ts.id ts.project_code, ts.test_suite, tr.test_starttime, - tr.status; + je.status; diff --git a/testgen/template/rollup_scores/rollup_scores_profile_table_group.sql b/testgen/template/rollup_scores/rollup_scores_profile_table_group.sql index ccd67104..f1d69599 100644 --- a/testgen/template/rollup_scores/rollup_scores_profile_table_group.sql +++ b/testgen/template/rollup_scores/rollup_scores_profile_table_group.sql @@ -2,7 +2,8 @@ WITH last_profile_date AS (SELECT table_groups_id, MAX(profiling_starttime) as last_profile_run_date FROM profiling_runs - WHERE status = 'Complete' + INNER JOIN job_executions ON job_executions.id = profiling_runs.id + WHERE job_executions.status = 'completed' GROUP BY table_groups_id), score_calc AS (SELECT run.table_groups_id, run.id as profile_run_id, diff --git a/testgen/template/rollup_scores/rollup_scores_test_table_group.sql b/testgen/template/rollup_scores/rollup_scores_test_table_group.sql index 6009e5d0..dfe6d7b9 100644 --- a/testgen/template/rollup_scores/rollup_scores_test_table_group.sql +++ b/testgen/template/rollup_scores/rollup_scores_test_table_group.sql @@ -2,7 +2,8 @@ WITH last_test_date AS (SELECT r.test_suite_id, MAX(r.test_starttime) as last_test_run_date FROM test_runs r - WHERE r.status = 'Complete' + INNER JOIN job_executions je ON je.id = r.id + WHERE je.status = 'completed' GROUP BY r.test_suite_id), score_calc AS (SELECT ts.table_groups_id, diff --git a/testgen/template/score_cards/add_latest_runs.sql b/testgen/template/score_cards/add_latest_runs.sql index bf332d8c..4403cb8a 100644 --- a/testgen/template/score_cards/add_latest_runs.sql +++ b/testgen/template/score_cards/add_latest_runs.sql @@ -1,11 +1,12 @@ -- Insert latest profiling runs as of cutoff WITH ranked_profiling - AS (SELECT project_code, table_groups_id, id as profiling_run_id, + AS (SELECT r.project_code, table_groups_id, r.id as profiling_run_id, ROW_NUMBER() OVER (PARTITION BY table_groups_id ORDER BY profiling_starttime DESC) as rank FROM profiling_runs r - WHERE project_code = :project_code + INNER JOIN job_executions je ON je.id = r.id + WHERE r.project_code = :project_code AND profiling_starttime <= :score_history_cutoff_time - AND r.status = 'Complete') + AND je.status = 'completed') INSERT INTO score_history_latest_runs (definition_id, score_history_cutoff_time, table_groups_id, last_profiling_run_id) SELECT :definition_id as definition_id, :score_history_cutoff_time as score_history_cutoff_time, table_groups_id, profiling_run_id @@ -20,9 +21,11 @@ WITH ranked_test_runs FROM test_runs r INNER JOIN test_suites s ON (r.test_suite_id = s.id) + INNER JOIN job_executions je + ON (je.id = r.id) WHERE s.project_code = :project_code AND r.test_starttime <= :score_history_cutoff_time - AND r.status = 'Complete') + AND je.status = 'completed') INSERT INTO score_history_latest_runs (definition_id, score_history_cutoff_time, test_suite_id, last_test_run_id) SELECT :definition_id as definition_id, :score_history_cutoff_time as score_history_cutoff_time, test_suite_id, test_run_id diff --git a/testgen/ui/queries/profiling_queries.py b/testgen/ui/queries/profiling_queries.py index ce7ff242..0c9a64a9 100644 --- a/testgen/ui/queries/profiling_queries.py +++ b/testgen/ui/queries/profiling_queries.py @@ -678,7 +678,7 @@ def get_profiling_anomalies_by_ids(anomaly_ids: list[str]) -> pd.DataFrame: t.anomaly_description, r.detail, t.detail_redactable, t.suggested_action, t.dq_dimension, r.impact_dimension, r.anomaly_id, r.table_groups_id::VARCHAR, r.id::VARCHAR, p.profiling_starttime, r.profile_run_id::VARCHAR, - p.job_execution_id::VARCHAR as job_execution_id, + p.id::VARCHAR as job_execution_id, tg.table_groups_name, tg.project_code, dcc.functional_data_type, dcc.description as column_description, diff --git a/testgen/ui/queries/scoring_queries.py b/testgen/ui/queries/scoring_queries.py index 26ff23a5..019c9bb7 100644 --- a/testgen/ui/queries/scoring_queries.py +++ b/testgen/ui/queries/scoring_queries.py @@ -42,7 +42,7 @@ def get_score_card_issue_reports(selected_issues: list["SelectedIssue"]) -> list groups.table_groups_name, results.disposition, results.profile_run_id::VARCHAR, - runs.job_execution_id::VARCHAR, + runs.id::VARCHAR AS job_execution_id, types.suggested_action, results.table_groups_id::VARCHAR, results.project_code, @@ -108,7 +108,7 @@ def get_score_card_issue_reports(selected_issues: list["SelectedIssue"]) -> list ELSE 'Passed' END as disposition, results.test_run_id::VARCHAR, - test_runs.job_execution_id::VARCHAR, + test_runs.id::VARCHAR AS job_execution_id, types.usage_notes, types.test_type, results.auto_gen, diff --git a/testgen/ui/queries/test_result_queries.py b/testgen/ui/queries/test_result_queries.py index 82ebc23a..4e0b0c58 100644 --- a/testgen/ui/queries/test_result_queries.py +++ b/testgen/ui/queries/test_result_queries.py @@ -199,7 +199,7 @@ def get_test_results_by_ids(test_result_ids: list[str]) -> pd.DataFrame: END as execution_error_ct, p.project_code, r.table_groups_id::VARCHAR, r.id::VARCHAR as test_result_id, r.test_run_id::VARCHAR, - tr.job_execution_id::VARCHAR as job_execution_id, + tr.id::VARCHAR as job_execution_id, c.id::VARCHAR as connection_id, r.test_suite_id::VARCHAR, r.test_definition_id::VARCHAR, r.auto_gen, diff --git a/testgen/ui/views/profiling_runs.py b/testgen/ui/views/profiling_runs.py index 5707838a..9f034615 100644 --- a/testgen/ui/views/profiling_runs.py +++ b/testgen/ui/views/profiling_runs.py @@ -289,9 +289,6 @@ def on_cancel_run(payload: dict) -> None: job_exec = JobExecution.get(job_execution_id) if job_exec and job_exec.request_cancel(): - # Stopgap: also update the run status so the UI reflects cancellation immediately. - if profiling_run_id := payload.get("profiling_run_id"): - ProfilingRun.cancel_run(profiling_run_id) get_profiling_run_summaries.clear() fm.reset_post_updates(str_message=":green[Cancellation requested.]", as_toast=True) else: @@ -307,7 +304,7 @@ def on_delete_runs(job_execution_ids: list[str]) -> None: continue if job_exec.status in (JobStatus.PENDING, JobStatus.CLAIMED, JobStatus.RUNNING, JobStatus.CANCEL_REQUESTED): job_exec.request_cancel() - profiling_run = next(iter(select_profiling_runs_where(ProfilingRun.job_execution_id == je_id)), None) + profiling_run = next(iter(select_profiling_runs_where(ProfilingRun.id == je_id)), None) if profiling_run: ProfilingRun.cascade_delete([str(profiling_run.id)]) get_current_session().delete(job_exec) diff --git a/testgen/ui/views/test_runs.py b/testgen/ui/views/test_runs.py index f6563b94..84bc515b 100644 --- a/testgen/ui/views/test_runs.py +++ b/testgen/ui/views/test_runs.py @@ -301,9 +301,6 @@ def on_cancel_run(payload: dict) -> None: job_exec = JobExecution.get(job_execution_id) if job_exec and job_exec.request_cancel(): - # Stopgap: also update the run status so the UI reflects cancellation immediately. - if test_run_id := payload.get("test_run_id"): - TestRun.cancel_run(test_run_id) get_test_run_summaries.clear() fm.reset_post_updates(str_message=":green[Cancellation requested.]", as_toast=True) else: @@ -319,7 +316,7 @@ def on_delete_runs(job_execution_ids: list[str]) -> None: continue if job_exec.status in (JobStatus.PENDING, JobStatus.CLAIMED, JobStatus.RUNNING, JobStatus.CANCEL_REQUESTED): job_exec.request_cancel() - test_run = next(iter(select_test_runs_where(TestRun.job_execution_id == je_id)), None) + test_run = next(iter(select_test_runs_where(TestRun.id == je_id)), None) if test_run: TestRun.cascade_delete([str(test_run.id)]) get_current_session().delete(job_exec) diff --git a/tests/unit/api/test_runs.py b/tests/unit/api/test_runs.py index 0fe8143c..cdeb4574 100644 --- a/tests/unit/api/test_runs.py +++ b/tests/unit/api/test_runs.py @@ -70,7 +70,7 @@ def _mock_profiling_run(**overrides): @patch(f"{MODULE}.TestRun") def test_get_test_run_completed(mock_tr_cls, mock_result_cls, mock_session): job = _mock_job() - mock_tr_cls.get_by_id_or_job.return_value = _mock_test_run() + mock_tr_cls.get.return_value = _mock_test_run() mock_result_cls.count_by_status.return_value = ResultStatusCounts( passed=90, failed=5, warning=3, error=2, log=0, dismissed=12, ) @@ -92,7 +92,7 @@ def test_get_test_run_completed(mock_tr_cls, mock_result_cls, mock_session): @patch(f"{MODULE}.TestRun") def test_get_test_run_pending_no_run(mock_tr_cls): job = _mock_job(status="pending", started_at=None, completed_at=None) - mock_tr_cls.get_by_id_or_job.return_value = None + mock_tr_cls.get.return_value = None result = get_test_run(job) @@ -110,7 +110,7 @@ def test_get_test_run_pending_no_run(mock_tr_cls): @patch(f"{MODULE}.ProfilingRun") def test_get_profiling_run_completed(mock_pr_cls, mock_issue_cls): job = _mock_job() - mock_pr_cls.get_by_id_or_job.return_value = _mock_profiling_run() + mock_pr_cls.get.return_value = _mock_profiling_run() mock_issue_cls.count_by_likelihood.return_value = IssueLikelihoodCounts( definite=5, likely=3, possible=8, dismissed=2, ) @@ -131,7 +131,7 @@ def test_get_profiling_run_completed(mock_pr_cls, mock_issue_cls): @patch(f"{MODULE}.ProfilingRun") def test_get_profiling_run_pending_no_run(mock_pr_cls): job = _mock_job(status="pending", started_at=None, completed_at=None) - mock_pr_cls.get_by_id_or_job.return_value = None + mock_pr_cls.get.return_value = None result = get_profiling_run(job) diff --git a/tests/unit/commands/test_run_data_cleanup.py b/tests/unit/commands/test_run_data_cleanup.py index 6c6e4304..07b463b2 100644 --- a/tests/unit/commands/test_run_data_cleanup.py +++ b/tests/unit/commands/test_run_data_cleanup.py @@ -51,17 +51,16 @@ def _patch_orchestrator( } started = {name: p.start() for name, p in patches.items()} - started["ProfilingRun"].find_latest_per_table_group.return_value = protected_profiling or set() - # get_job_execution_ids returns dict[run_id, je_id]; orchestrator filters nulls. - started["ProfilingRun"].get_job_execution_ids.return_value = { - uuid4(): je_id for je_id in (protected_profiling_jes or set()) - } + # The run id IS the job execution id, so find_latest_* returns the protected + # JE ids directly. The *_jes params let a test name the same set as JE ids. + started["ProfilingRun"].find_latest_per_table_group.return_value = ( + protected_profiling_jes or protected_profiling or set() + ) started["ProfilingRun"].delete_older_than.return_value = deleted_profiling - started["TestRun"].find_latest_per_test_suite.return_value = protected_tests or set() - started["TestRun"].get_job_execution_ids.return_value = { - uuid4(): je_id for je_id in (protected_test_jes or set()) - } + started["TestRun"].find_latest_per_test_suite.return_value = ( + protected_test_jes or protected_tests or set() + ) started["TestRun"].delete_older_than.return_value = deleted_tests started["JobExecution"].delete_older_than.return_value = deleted_job_executions diff --git a/tests/unit/common/notifications/test_profiling_run_notifications.py b/tests/unit/common/notifications/test_profiling_run_notifications.py index 5e985cd2..16351212 100644 --- a/tests/unit/common/notifications/test_profiling_run_notifications.py +++ b/tests/unit/common/notifications/test_profiling_run_notifications.py @@ -107,7 +107,6 @@ def test_send_profiling_run_notification( ): profiling_run = ProfilingRun( id="pr-id", - job_execution_id="pr-id", table_groups_id="tg-id", project_code="proj", ) diff --git a/tests/unit/common/notifications/test_test_run_notifications.py b/tests/unit/common/notifications/test_test_run_notifications.py index 3b4a035b..d791037e 100644 --- a/tests/unit/common/notifications/test_test_run_notifications.py +++ b/tests/unit/common/notifications/test_test_run_notifications.py @@ -139,7 +139,6 @@ def test_send_test_run_notification( test_run = TestRun( id="tr-id", - job_execution_id="tr-id", test_suite_id="ts-id", failed_ct=failed_ct, warning_ct=warning_ct, diff --git a/tests/unit/mcp/test_model_hygiene_issue.py b/tests/unit/mcp/test_model_hygiene_issue.py index bf90e2b6..74c7de55 100644 --- a/tests/unit/mcp/test_model_hygiene_issue.py +++ b/tests/unit/mcp/test_model_hygiene_issue.py @@ -95,14 +95,14 @@ def _stub_paginate(session_mock, *, total=0, rows=()): # --------------------------------------------------------------------------- -def test_list_for_run_filters_by_je_id_not_run_pk(session_mock): +def test_list_for_run_filters_by_run_id(session_mock): _stub_paginate(session_mock) HygieneIssue.list_for_run(uuid4()) sql = _all_compiled_sql(session_mock) - assert "profiling_runs.job_execution_id =" in sql - # The legacy run PK must NOT be the filter: + # The run id is the job execution id; filter targets the profiling run id directly. + assert "profiling_runs.id =" in sql assert "profile_anomaly_results.profile_run_id =" not in sql diff --git a/tests/unit/mcp/test_model_profiling_run.py b/tests/unit/mcp/test_model_profiling_run.py index 94d88b1b..888c50fa 100644 --- a/tests/unit/mcp/test_model_profiling_run.py +++ b/tests/unit/mcp/test_model_profiling_run.py @@ -53,13 +53,13 @@ def test_get_latest_complete_je_id_orders_desc_limit_1(session_mock): assert "LIMIT" in sql -def test_get_latest_complete_je_id_selects_je_id_not_run_pk(session_mock): +def test_get_latest_complete_je_id_selects_run_id(session_mock): """Pin the docstring contract — does NOT read ``table_groups.last_complete_profile_run_id``, - selects the JE id directly from ``profiling_runs``.""" + selects the run id (which is the JE id) directly from ``profiling_runs``.""" session_mock.scalar.return_value = None ProfilingRun.get_latest_complete_je_id_for_table_group(uuid4()) sql = _compiled_sql(session_mock.scalar.call_args[0][0]) - assert "SELECT profiling_runs.job_execution_id" in sql + assert "SELECT profiling_runs.id" in sql assert "table_groups.last_complete_profile_run_id" not in sql diff --git a/tests/unit/mcp/test_tools_common.py b/tests/unit/mcp/test_tools_common.py index 068d1096..15a7c0b2 100644 --- a/tests/unit/mcp/test_tools_common.py +++ b/tests/unit/mcp/test_tools_common.py @@ -311,7 +311,7 @@ def _mock_perms(allowed_projects=("demo",)): def test_resolve_profiling_run_happy_path(mock_pr_cls, mock_get_perms, db_session_mock): run = MagicMock() run.project_code = "demo" - mock_pr_cls.get_by_id_or_job.return_value = run + mock_pr_cls.get.return_value = run mock_get_perms.return_value = _mock_perms(allowed_projects=("demo",)) result = resolve_profiling_run(str(uuid4())) @@ -322,7 +322,7 @@ def test_resolve_profiling_run_happy_path(mock_pr_cls, mock_get_perms, db_sessio @patch("testgen.mcp.tools.common.get_project_permissions") @patch("testgen.mcp.tools.common.ProfilingRun") def test_resolve_profiling_run_unknown_run_id(mock_pr_cls, mock_get_perms, db_session_mock): - mock_pr_cls.get_by_id_or_job.return_value = None + mock_pr_cls.get.return_value = None mock_get_perms.return_value = _mock_perms() with pytest.raises(MCPResourceNotAccessible, match=r"Profiling run .* not found or not accessible"): @@ -335,7 +335,7 @@ def test_resolve_profiling_run_inaccessible_project(mock_pr_cls, mock_get_perms, """Run exists but caller can't access its project — same unified error as unknown run.""" run = MagicMock() run.project_code = "forbidden" - mock_pr_cls.get_by_id_or_job.return_value = run + mock_pr_cls.get.return_value = run mock_get_perms.return_value = _mock_perms(allowed_projects=("demo",)) with pytest.raises(MCPResourceNotAccessible, match=r"Profiling run .* not found or not accessible"): diff --git a/tests/unit/mcp/test_tools_discovery.py b/tests/unit/mcp/test_tools_discovery.py index 409b20f2..15bfac0b 100644 --- a/tests/unit/mcp/test_tools_discovery.py +++ b/tests/unit/mcp/test_tools_discovery.py @@ -116,11 +116,9 @@ def test_list_projects_filters_for_scoped_user(mock_compute, mock_project, db_se assert "Secret" not in result -@patch("testgen.mcp.tools.discovery.TestRun") @patch("testgen.mcp.tools.discovery.TestSuite") -def test_list_test_suites_returns_stats(mock_suite, mock_test_run, db_session_mock): +def test_list_test_suites_returns_stats(mock_suite, db_session_mock): run_id = uuid4() - job_exec_id = uuid4() summary = MagicMock() summary.id = uuid4() summary.test_suite = "Quality Suite" @@ -137,7 +135,6 @@ def test_list_test_suites_returns_stats(mock_suite, mock_test_run, db_session_mo summary.last_run_error_ct = 0 summary.last_run_dismissed_ct = 0 mock_suite.select_summary.return_value = [summary] - mock_test_run.get_job_execution_ids.return_value = {run_id: job_exec_id} from testgen.mcp.tools.discovery import list_test_suites @@ -146,7 +143,8 @@ def test_list_test_suites_returns_stats(mock_suite, mock_test_run, db_session_mo assert "Quality Suite" in result assert "45 passed" in result assert "3 failed" in result - assert str(job_exec_id) in result + # latest_run_id is already the job execution id — surfaced directly. + assert str(run_id) in result @patch("testgen.mcp.tools.discovery.TestSuite") diff --git a/tests/unit/mcp/test_tools_hygiene_issues.py b/tests/unit/mcp/test_tools_hygiene_issues.py index 5dcbeb24..343aab73 100644 --- a/tests/unit/mcp/test_tools_hygiene_issues.py +++ b/tests/unit/mcp/test_tools_hygiene_issues.py @@ -252,7 +252,7 @@ def test_resolve_je_id_table_group_no_completed_runs(mock_latest, mock_resolve_t @patch("testgen.mcp.tools.hygiene_issues.TableGroup") -@patch.object(ProfilingRun, "get_by_id_or_job") +@patch.object(ProfilingRun, "get") def test_resolve_je_id_je_branch_unknown_run(mock_get, mock_tg_cls, db_session_mock): from testgen.mcp.tools.hygiene_issues import _resolve_profile_run_je_id @@ -267,7 +267,7 @@ def test_resolve_je_id_je_branch_unknown_run(mock_get, mock_tg_cls, db_session_m @patch("testgen.mcp.tools.hygiene_issues.TableGroup") -@patch.object(ProfilingRun, "get_by_id_or_job") +@patch.object(ProfilingRun, "get") def test_resolve_je_id_je_branch_inaccessible_tg(mock_get, mock_tg_cls, db_session_mock): from testgen.mcp.tools.hygiene_issues import _resolve_profile_run_je_id @@ -309,7 +309,7 @@ def test_list_hygiene_issues_invalid_je_uuid(db_session_mock): @patch("testgen.mcp.tools.hygiene_issues.TableGroup") -@patch.object(ProfilingRun, "get_by_id_or_job") +@patch.object(ProfilingRun, "get") @patch.object(HygieneIssue, "list_for_run") def test_list_hygiene_issues_resolves_via_je_id(mock_list, mock_get, mock_tg_cls, db_session_mock): from testgen.mcp.tools.hygiene_issues import list_hygiene_issues @@ -321,7 +321,8 @@ def test_list_hygiene_issues_resolves_via_je_id(mock_list, mock_get, mock_tg_cls list_hygiene_issues(job_execution_id=str(uuid4())) - assert mock_list.call_args.args[0] == run.job_execution_id + # The run id is the job execution id — list_for_run is called with run.id. + assert mock_list.call_args.args[0] == run.id @patch("testgen.mcp.tools.hygiene_issues.resolve_table_group") diff --git a/tests/unit/mcp/test_tools_profiling.py b/tests/unit/mcp/test_tools_profiling.py index 733902f8..989a6108 100644 --- a/tests/unit/mcp/test_tools_profiling.py +++ b/tests/unit/mcp/test_tools_profiling.py @@ -267,7 +267,7 @@ def test_list_column_profiles_with_valid_job_execution_id( pr.project_code = tg.project_code mock_tg_cls.get.return_value = tg - mock_pr_cls.get_by_id_or_job.return_value = pr + mock_pr_cls.get.return_value = pr mock_dcc_cls.list_for_table_group.return_value = ([_column_summary()], 1) from testgen.mcp.tools.profiling import list_column_profiles @@ -289,7 +289,7 @@ def test_list_column_profiles_rejects_je_from_different_tg( pr.project_code = tg.project_code mock_tg_cls.get.return_value = tg - mock_pr_cls.get_by_id_or_job.return_value = pr + mock_pr_cls.get.return_value = pr from testgen.mcp.tools.profiling import list_column_profiles with pytest.raises(MCPResourceNotAccessible, match="Profiling run .* not found or not accessible"): @@ -300,7 +300,7 @@ def test_list_column_profiles_rejects_je_from_different_tg( @patch("testgen.mcp.tools.common.TableGroup") def test_list_column_profiles_rejects_unknown_je(mock_tg_cls, mock_pr_cls, db_session_mock): mock_tg_cls.get.return_value = _mock_table_group() - mock_pr_cls.get_by_id_or_job.return_value = None + mock_pr_cls.get.return_value = None from testgen.mcp.tools.profiling import list_column_profiles with pytest.raises(MCPResourceNotAccessible, match="Profiling run .* not found or not accessible"): @@ -593,7 +593,7 @@ def test_get_profiling_run_returns_detail(mock_run_cls, db_session_mock): summary = _mock_profiling_run() mock_run_cls.select_summary.return_value = ([summary], 1) mock_run = MagicMock(project_code="demo") - mock_run_cls.get_by_id_or_job.return_value = mock_run + mock_run_cls.get.return_value = mock_run mock_run_cls.select_table_breakdown.return_value = [ MagicMock(schema_name="demo", table_name="orders", record_ct=1000, column_ct=5, anomaly_ct=2), ] @@ -624,7 +624,7 @@ def test_get_profiling_run_pending_no_breakdown(mock_run_cls, db_session_mock): anomalies_possible_ct=None, dq_score_profiling=None, ) mock_run_cls.select_summary.return_value = ([summary], 1) - mock_run_cls.get_by_id_or_job.return_value = MagicMock(project_code="demo") + mock_run_cls.get.return_value = MagicMock(project_code="demo") with patch("testgen.mcp.permissions._compute_project_permissions") as mock_compute: mock_compute.return_value = ProjectPermissions( @@ -760,7 +760,7 @@ def _column_detail(**overrides) -> ColumnProfileDetail: # Run identity "profile_run_id": uuid4(), "profile_run_je_id": uuid4(), - "profile_run_status": "Complete", + "profile_run_status": JobStatus.COMPLETED, "profile_run_started_at": datetime(2026, 5, 1, 12, 0, 0), "profile_run_ended_at": datetime(2026, 5, 1, 12, 5, 0), "profile_run_log_message": None, @@ -1026,7 +1026,7 @@ def test_get_column_profile_detail_pinned_run_without_column_rejects( pr.project_code = tg.project_code mock_tg_cls.get.return_value = tg - mock_pr_cls.get_by_id_or_job.return_value = pr + mock_pr_cls.get.return_value = pr mock_dcc_cls.get_column_detail.return_value = _column_detail( profile_run_id=None, profile_run_je_id=None, @@ -1105,7 +1105,7 @@ def test_get_column_profile_detail_pinned_run_passes_id_to_model( pr.project_code = tg.project_code mock_tg_cls.get.return_value = tg - mock_pr_cls.get_by_id_or_job.return_value = pr + mock_pr_cls.get.return_value = pr mock_dcc_cls.get_column_detail.return_value = _column_detail() from testgen.mcp.tools.profiling import get_column_profile_detail @@ -1126,7 +1126,7 @@ def test_get_column_profile_detail_pinned_run_from_different_tg_unified_error( pr.project_code = tg.project_code mock_tg_cls.get.return_value = tg - mock_pr_cls.get_by_id_or_job.return_value = pr + mock_pr_cls.get.return_value = pr from testgen.mcp.tools.profiling import get_column_profile_detail with pytest.raises(MCPResourceNotAccessible, match=r"Profiling run .* not found or not accessible"): @@ -1141,7 +1141,7 @@ def test_get_column_profile_detail_pinned_run_unknown_unified_error( mock_tg_cls, mock_pr_cls, db_session_mock, ): mock_tg_cls.get.return_value = _mock_table_group() - mock_pr_cls.get_by_id_or_job.return_value = None + mock_pr_cls.get.return_value = None from testgen.mcp.tools.profiling import get_column_profile_detail with pytest.raises(MCPResourceNotAccessible, match=r"Profiling run .* not found or not accessible"): @@ -1161,7 +1161,7 @@ def test_get_column_profile_detail_running_run_rejects_with_status( mock_tg_cls.get.return_value = _mock_table_group() je_id = uuid4() mock_dcc_cls.get_column_detail.return_value = _column_detail( - profile_run_status="Running", + profile_run_status=JobStatus.RUNNING, profile_run_je_id=je_id, profile_run_ended_at=None, ) @@ -1183,7 +1183,7 @@ def test_get_column_profile_detail_error_run_includes_log_message( mock_tg_cls.get.return_value = _mock_table_group() je_id = uuid4() mock_dcc_cls.get_column_detail.return_value = _column_detail( - profile_run_status="Error", + profile_run_status=JobStatus.ERROR, profile_run_je_id=je_id, profile_run_log_message="connection timed out", ) @@ -1621,8 +1621,8 @@ def test_get_column_frequent_values_surfaces_job_execution_id_not_profile_run_id from testgen.mcp.tools.profiling import get_column_frequent_values result = get_column_frequent_values(str(uuid4()), "customers", "country") - # The internal profile_run_id PK must not leak; only the job_execution_id is followable. - assert str(run.job_execution_id) in result + # The internal profile_run_id PK must not leak; only the run id (the job execution id) is followable. + assert str(run.id) in result assert str(profile.profile_run_id) not in result diff --git a/tests/unit/mcp/test_tools_test_results.py b/tests/unit/mcp/test_tools_test_results.py index dc4f9084..47e04f6c 100644 --- a/tests/unit/mcp/test_tools_test_results.py +++ b/tests/unit/mcp/test_tools_test_results.py @@ -24,7 +24,7 @@ def _mock_test_run(test_run_id=None): @patch("testgen.mcp.tools.test_results.TestResult") def test_list_test_results_basic(mock_result, mock_tt_cls, mock_test_run_cls, mock_suite_cls, db_session_mock): job_id = str(uuid4()) - mock_test_run_cls.get_by_id_or_job.return_value = _mock_test_run() + mock_test_run_cls.get.return_value = _mock_test_run() mock_suite_cls.get_regular.return_value = _mock_test_suite() r1 = MagicMock() @@ -59,7 +59,7 @@ def test_list_test_results_basic(mock_result, mock_tt_cls, mock_test_run_cls, mo @patch("testgen.mcp.tools.test_results.TestType") @patch("testgen.mcp.tools.test_results.TestResult") def test_list_test_results_emits_test_result_id(mock_result, mock_tt_cls, mock_test_run_cls, mock_suite_cls, db_session_mock): - mock_test_run_cls.get_by_id_or_job.return_value = _mock_test_run() + mock_test_run_cls.get.return_value = _mock_test_run() mock_suite_cls.get_regular.return_value = _mock_test_suite() result_id = uuid4() @@ -93,7 +93,7 @@ def test_list_test_results_emits_test_result_id(mock_result, mock_tt_cls, mock_t @patch("testgen.mcp.tools.test_results.TestType") @patch("testgen.mcp.tools.test_results.TestResult") def test_list_test_results_table_level_title(mock_result, mock_tt_cls, mock_test_run_cls, mock_suite_cls, db_session_mock): - mock_test_run_cls.get_by_id_or_job.return_value = _mock_test_run() + mock_test_run_cls.get.return_value = _mock_test_run() mock_suite_cls.get_regular.return_value = _mock_test_suite() r1 = MagicMock() @@ -124,7 +124,7 @@ def test_list_test_results_table_level_title(mock_result, mock_tt_cls, mock_test @patch("testgen.mcp.tools.test_results.TestRun") @patch("testgen.mcp.tools.test_results.TestResult") def test_list_test_results_empty(mock_result, mock_test_run_cls, mock_suite_cls, db_session_mock): - mock_test_run_cls.get_by_id_or_job.return_value = _mock_test_run() + mock_test_run_cls.get.return_value = _mock_test_run() mock_suite_cls.get_regular.return_value = _mock_test_suite() mock_result.select_results.return_value = [] @@ -143,7 +143,7 @@ def test_list_test_results_empty(mock_result, mock_test_run_cls, mock_suite_cls, def test_list_test_results_with_filters( mock_result, mock_tt_cls, mock_test_run_cls, mock_tt_common, mock_suite_cls, db_session_mock ): - mock_test_run_cls.get_by_id_or_job.return_value = _mock_test_run() + mock_test_run_cls.get.return_value = _mock_test_run() mock_suite_cls.get_regular.return_value = _mock_test_suite() tt = MagicMock() tt.test_type = "Alpha_Trunc" @@ -171,7 +171,7 @@ def test_list_test_results_invalid_uuid(db_session_mock): @patch("testgen.mcp.tools.test_results.TestSuite") @patch("testgen.mcp.tools.test_results.TestRun") def test_list_test_results_invalid_status(mock_test_run_cls, mock_suite_cls, db_session_mock): - mock_test_run_cls.get_by_id_or_job.return_value = _mock_test_run() + mock_test_run_cls.get.return_value = _mock_test_run() mock_suite_cls.get_regular.return_value = _mock_test_suite() from testgen.mcp.tools.test_results import list_test_results @@ -182,7 +182,7 @@ def test_list_test_results_invalid_status(mock_test_run_cls, mock_suite_cls, db_ @patch("testgen.mcp.tools.test_results.TestRun") def test_list_test_results_run_not_found(mock_test_run_cls, db_session_mock): - mock_test_run_cls.get_by_id_or_job.return_value = None + mock_test_run_cls.get.return_value = None from testgen.mcp.tools.test_results import list_test_results @@ -194,7 +194,7 @@ def test_list_test_results_run_not_found(mock_test_run_cls, db_session_mock): @patch("testgen.mcp.tools.test_results.TestRun") def test_list_test_results_run_in_monitor_suite_rejected(mock_test_run_cls, mock_suite_cls, db_session_mock): # Run exists, but the resolved suite is monitor → TestSuite.get_regular returns None. - mock_test_run_cls.get_by_id_or_job.return_value = _mock_test_run() + mock_test_run_cls.get.return_value = _mock_test_run() mock_suite_cls.get_regular.return_value = None from testgen.mcp.tools.test_results import list_test_results @@ -210,7 +210,7 @@ def test_list_test_results_run_in_forbidden_project( mock_compute, mock_test_run_cls, mock_suite_cls, db_session_mock ): mock_compute.return_value = ProjectPermissions(memberships={"proj_a": "role_a"}, permission="view", username="test_user") - mock_test_run_cls.get_by_id_or_job.return_value = _mock_test_run() + mock_test_run_cls.get.return_value = _mock_test_run() mock_suite_cls.get_regular.return_value = _mock_test_suite(project_code="forbidden_project") from testgen.mcp.tools.test_results import list_test_results @@ -231,7 +231,7 @@ def test_list_test_results_passes_project_codes( permission="view", username="test_user", ) - mock_test_run_cls.get_by_id_or_job.return_value = _mock_test_run() + mock_test_run_cls.get.return_value = _mock_test_run() mock_suite_cls.get_regular.return_value = _mock_test_suite(project_code="proj_a") mock_result.select_results.return_value = [] @@ -247,12 +247,12 @@ def test_list_test_results_passes_project_codes( @patch("testgen.mcp.tools.test_results.TestRun") @patch("testgen.mcp.tools.test_results.TestResult") @patch("testgen.mcp.tools.test_results.TestType") -def test_list_test_results_resolves_via_get_by_id_or_job( +def test_list_test_results_resolves_via_get( mock_tt_cls, mock_result, mock_test_run_cls, mock_suite_cls, db_session_mock ): """Verify the resolved test_run.id is passed to select_results.""" resolved_run_id = uuid4() - mock_test_run_cls.get_by_id_or_job.return_value = _mock_test_run(resolved_run_id) + mock_test_run_cls.get.return_value = _mock_test_run(resolved_run_id) mock_suite_cls.get_regular.return_value = _mock_test_suite() mock_result.select_results.return_value = [] @@ -332,10 +332,8 @@ def test_list_test_results_by_suite_id_resolves_latest_run( mock_suite_cls.get_regular.return_value = _mock_test_suite(last_complete_test_run_id=last_run_id) resolved_run_id = uuid4() - resolved_je_id = uuid4() resolved_run = _mock_test_run(resolved_run_id) - resolved_run.job_execution_id = resolved_je_id - mock_test_run_cls.get_by_id_or_job.return_value = resolved_run + mock_test_run_cls.get.return_value = resolved_run r1 = MagicMock() r1.status = TestResultStatus.Failed @@ -358,11 +356,11 @@ def test_list_test_results_by_suite_id_resolves_latest_run( suite_id = str(uuid4()) result = list_test_results(test_suite_id=suite_id) - # Resolution chain: suite.last_complete_test_run_id → TestRun.get_by_id_or_job → test_run.id → select_results - mock_test_run_cls.get_by_id_or_job.assert_called_once_with(last_run_id) + # Resolution chain: suite.last_complete_test_run_id → TestRun.get → test_run.id → select_results + mock_test_run_cls.get.assert_called_once_with(last_run_id) assert mock_result.select_results.call_args.kwargs["test_run_id"] == resolved_run_id - # Output indicates which run the suite was resolved to. - assert str(resolved_je_id) in result + # Output indicates which run the suite was resolved to (run id is the job execution id). + assert str(resolved_run_id) in result assert f"Latest completed run of test suite `{suite_id}`" in result @@ -373,7 +371,7 @@ def test_list_test_results_by_suite_id_resolves_latest_run( def test_get_failure_summary_by_test_type( mock_result, mock_tt_cls, mock_test_run_cls, mock_suite_cls, db_session_mock, ): - mock_test_run_cls.get_by_id_or_job.return_value = _mock_test_run() + mock_test_run_cls.get.return_value = _mock_test_run() mock_suite_cls.get_regular.return_value = _mock_test_suite(project_code="demo") mock_result.select_failures.return_value = [ ("Alpha_Trunc", TestResultStatus.Failed, 5), @@ -405,7 +403,7 @@ def test_get_failure_summary_by_test_type( @patch("testgen.mcp.tools.test_results.TestRun") @patch("testgen.mcp.tools.test_results.TestResult") def test_get_failure_summary_empty(mock_result, mock_test_run_cls, mock_suite_cls, db_session_mock): - mock_test_run_cls.get_by_id_or_job.return_value = _mock_test_run() + mock_test_run_cls.get.return_value = _mock_test_run() mock_suite_cls.get_regular.return_value = _mock_test_suite(project_code="demo") mock_result.select_failures.return_value = [] @@ -420,7 +418,7 @@ def test_get_failure_summary_empty(mock_result, mock_test_run_cls, mock_suite_cl @patch("testgen.mcp.tools.test_results.TestRun") @patch("testgen.mcp.tools.test_results.TestResult") def test_get_failure_summary_by_table(mock_result, mock_test_run_cls, mock_suite_cls, db_session_mock): - mock_test_run_cls.get_by_id_or_job.return_value = _mock_test_run() + mock_test_run_cls.get.return_value = _mock_test_run() mock_suite_cls.get_regular.return_value = _mock_test_suite(project_code="demo") mock_result.select_failures.return_value = [("orders", 10)] @@ -437,7 +435,7 @@ def test_get_failure_summary_by_table(mock_result, mock_test_run_cls, mock_suite @patch("testgen.mcp.tools.test_results.TestRun") @patch("testgen.mcp.tools.test_results.TestResult") def test_get_failure_summary_by_column(mock_result, mock_test_run_cls, mock_suite_cls, db_session_mock): - mock_test_run_cls.get_by_id_or_job.return_value = _mock_test_run() + mock_test_run_cls.get.return_value = _mock_test_run() mock_suite_cls.get_regular.return_value = _mock_test_suite(project_code="demo") mock_result.select_failures.return_value = [("orders", "total_value", 34), ("orders", None, 2)] @@ -460,7 +458,7 @@ def test_get_failure_summary_invalid_uuid(db_session_mock): @patch("testgen.mcp.tools.test_results.TestRun") def test_get_failure_summary_run_not_found(mock_test_run_cls, db_session_mock): - mock_test_run_cls.get_by_id_or_job.return_value = None + mock_test_run_cls.get.return_value = None from testgen.mcp.tools.test_results import get_failure_summary @@ -475,7 +473,7 @@ def test_get_failure_summary_run_in_forbidden_project( mock_compute, mock_test_run_cls, mock_suite_cls, db_session_mock, ): mock_compute.return_value = ProjectPermissions(memberships={"proj_a": "role_a"}, permission="view", username="test_user") - mock_test_run_cls.get_by_id_or_job.return_value = _mock_test_run() + mock_test_run_cls.get.return_value = _mock_test_run() mock_suite_cls.get_regular.return_value = _mock_test_suite(project_code="forbidden_project") from testgen.mcp.tools.test_results import get_failure_summary @@ -488,7 +486,7 @@ def test_get_failure_summary_run_in_forbidden_project( @patch("testgen.mcp.tools.test_results.TestRun") def test_get_failure_summary_run_in_monitor_suite_rejected(mock_test_run_cls, mock_suite_cls, db_session_mock): # Run exists, but the resolved suite is monitor → TestSuite.get_regular returns None. - mock_test_run_cls.get_by_id_or_job.return_value = _mock_test_run() + mock_test_run_cls.get.return_value = _mock_test_run() mock_suite_cls.get_regular.return_value = None from testgen.mcp.tools.test_results import get_failure_summary @@ -509,7 +507,7 @@ def test_get_failure_summary_passes_project_codes( permission="view", username="test_user", ) - mock_test_run_cls.get_by_id_or_job.return_value = _mock_test_run() + mock_test_run_cls.get.return_value = _mock_test_run() mock_suite_cls.get_regular.return_value = _mock_test_suite(project_code="proj_a") mock_result.select_failures.return_value = [] @@ -1034,7 +1032,7 @@ def test_compare_test_runs_happy_path( baseline_run = _mock_run(suite_id) target_run = _mock_run(suite_id) # Tool resolves target first, then baseline. - mock_test_run_cls.get_by_id_or_job.side_effect = [target_run, baseline_run] + mock_test_run_cls.get.side_effect = [target_run, baseline_run] mock_test_suite_cls.get_regular.return_value = _mock_test_suite(suite_id=suite_id, project_code="proj_a") diff = MagicMock() @@ -1134,7 +1132,7 @@ def test_compare_test_runs_single_arg_resolves_previous( target_run = _mock_run(suite_id) baseline_run = _mock_run(suite_id) target_run.get_previous.return_value = baseline_run - mock_test_run_cls.get_by_id_or_job.return_value = target_run + mock_test_run_cls.get.return_value = target_run mock_test_suite_cls.get_regular.return_value = _mock_test_suite(suite_id=suite_id, project_code="proj_a") diff = MagicMock( @@ -1150,8 +1148,8 @@ def test_compare_test_runs_single_arg_resolves_previous( target_run.get_previous.assert_called_once_with() mock_result.diff_with_details.assert_called_once_with(baseline_run.id, target_run.id) - # Rendered Baseline cell shows the resolved JE ID, not an input string. - assert str(baseline_run.job_execution_id) in out + # Rendered Baseline cell shows the resolved run id (the job execution id), not an input string. + assert str(baseline_run.id) in out @patch("testgen.mcp.tools.test_results.TestSuite") @@ -1169,7 +1167,7 @@ def test_compare_test_runs_single_arg_no_previous_raises( suite_id = uuid4() target_run = _mock_run(suite_id) target_run.get_previous.return_value = None - mock_test_run_cls.get_by_id_or_job.return_value = target_run + mock_test_run_cls.get.return_value = target_run mock_test_suite_cls.get_regular.return_value = _mock_test_suite(suite_id=suite_id, project_code="proj_a") from testgen.mcp.tools.test_results import compare_test_runs @@ -1192,7 +1190,7 @@ def test_compare_test_runs_single_arg_inaccessible_target( ) suite_id = uuid4() target_run = _mock_run(suite_id) - mock_test_run_cls.get_by_id_or_job.return_value = target_run + mock_test_run_cls.get.return_value = target_run # Monitor suite or inaccessible project — get_regular returns None either way. mock_test_suite_cls.get_regular.return_value = None @@ -1215,7 +1213,7 @@ def test_compare_test_runs_run_not_found( permission="view", username="test_user", ) - mock_test_run_cls.get_by_id_or_job.return_value = None + mock_test_run_cls.get.return_value = None from testgen.mcp.tools.test_results import compare_test_runs @@ -1237,7 +1235,7 @@ def test_compare_test_runs_rejects_inaccessible_project( ) suite_id = uuid4() run = _mock_run(suite_id) - mock_test_run_cls.get_by_id_or_job.return_value = run + mock_test_run_cls.get.return_value = run mock_test_suite_cls.get_regular.return_value = _mock_test_suite(suite_id=suite_id, project_code="proj_forbidden") from testgen.mcp.tools.test_results import compare_test_runs @@ -1262,7 +1260,7 @@ def test_compare_test_runs_rejects_different_suites( suite_id_baseline = uuid4() target_run = _mock_run(suite_id_target) baseline_run = _mock_run(suite_id_baseline) - mock_test_run_cls.get_by_id_or_job.side_effect = [target_run, baseline_run] + mock_test_run_cls.get.side_effect = [target_run, baseline_run] mock_test_suite_cls.get_regular.side_effect = [ _mock_test_suite(suite_id=suite_id_target, project_code="proj_a"), _mock_test_suite(suite_id=suite_id_baseline, project_code="proj_a"), @@ -1295,7 +1293,7 @@ def test_compare_test_runs_rejects_monitor_suite( ) suite_id = uuid4() run = _mock_run(suite_id) - mock_test_run_cls.get_by_id_or_job.return_value = run + mock_test_run_cls.get.return_value = run mock_test_suite_cls.get_regular.return_value = None from testgen.mcp.tools.test_results import compare_test_runs @@ -1318,7 +1316,7 @@ def test_compare_test_runs_rejects_target_not_completed( ) suite_id = uuid4() target_run = _mock_run(suite_id) - mock_test_run_cls.get_by_id_or_job.return_value = target_run + mock_test_run_cls.get.return_value = target_run mock_test_suite_cls.get_regular.return_value = _mock_test_suite(suite_id=suite_id, project_code="proj_a") from testgen.mcp.tools.test_results import compare_test_runs @@ -1344,7 +1342,7 @@ def test_compare_test_runs_rejects_baseline_not_completed( suite_id = uuid4() target_run = _mock_run(suite_id) baseline_run = _mock_run(suite_id) - mock_test_run_cls.get_by_id_or_job.side_effect = [target_run, baseline_run] + mock_test_run_cls.get.side_effect = [target_run, baseline_run] mock_test_suite_cls.get_regular.return_value = _mock_test_suite(suite_id=suite_id, project_code="proj_a") from testgen.mcp.tools.test_results import compare_test_runs @@ -1456,7 +1454,7 @@ def test_bulk_update_uses_latest_run_when_run_omitted( mock_session.return_value.scalars.return_value.all.return_value = matched_ids mock_set.return_value = DispositionUpdate(matched=2, passed_skipped=0) - with patch("testgen.mcp.tools.test_results.TestRun.get_by_id_or_job", + with patch("testgen.mcp.tools.test_results.TestRun.get", return_value=MagicMock(id=run_id, job_execution_id=run_id, test_suite_id=suite.id)): out = bulk_update_test_results(test_suite_id=str(uuid4()), disposition="Dismissed") @@ -1488,7 +1486,7 @@ def test_bulk_update_reports_passed_exclusions( mock_session.return_value.scalars.return_value.all.return_value = [uuid4(), uuid4(), uuid4()] mock_set.return_value = DispositionUpdate(matched=2, passed_skipped=1) - with patch("testgen.mcp.tools.test_results.TestRun.get_by_id_or_job", + with patch("testgen.mcp.tools.test_results.TestRun.get", return_value=MagicMock(id=run_id, job_execution_id=run_id, test_suite_id=suite.id)): out = bulk_update_test_results(test_suite_id=str(uuid4()), disposition="Muted") @@ -1508,7 +1506,7 @@ def test_bulk_update_no_matches(mock_resolve_suite, mock_session, mock_set, db_s mock_session.return_value.scalars.return_value.all.return_value = [] mock_set.return_value = DispositionUpdate(matched=0, passed_skipped=0) - with patch("testgen.mcp.tools.test_results.TestRun.get_by_id_or_job", + with patch("testgen.mcp.tools.test_results.TestRun.get", return_value=MagicMock(id=run_id, job_execution_id=run_id, test_suite_id=suite.id)): out = bulk_update_test_results(test_suite_id=str(uuid4()), disposition="Confirmed") @@ -1538,7 +1536,7 @@ def test_bulk_update_explicit_run_in_suite( mock_set.return_value = DispositionUpdate(matched=1, passed_skipped=0) with patch( - "testgen.mcp.tools.test_results.TestRun.get_by_id_or_job", + "testgen.mcp.tools.test_results.TestRun.get", return_value=MagicMock(id=uuid4(), job_execution_id=run_id, test_suite_id=suite.id), ): out = bulk_update_test_results( @@ -1559,7 +1557,7 @@ def test_bulk_update_explicit_run_from_other_suite_rejected( mock_resolve_suite.return_value = suite with patch( - "testgen.mcp.tools.test_results.TestRun.get_by_id_or_job", + "testgen.mcp.tools.test_results.TestRun.get", return_value=MagicMock(test_suite_id=uuid4()), # different suite ): with pytest.raises(MCPResourceNotAccessible, match=r"Test run .* not found or not accessible"): From 76739051e688d378287d5d983809a2581dad15f5 Mon Sep 17 00:00:00 2001 From: Aarthy Adityan Date: Tue, 9 Jun 2026 23:40:43 -0400 Subject: [PATCH 22/78] fix(mcp): allow the server's own port-less host in the transport allowlist When BASE_URL carries a non-default port, the netloc and :* wildcard entries match only a Host/Origin that includes a port. Hosted MCP gateways (e.g. Databricks) send a port-less Host/Origin for the server's own host, which was rejected with "invalid Host header" unless the host was re-listed in TG_MCP_EXTRA_ALLOWED_HOSTS. Add the bare base_host and origin to the BASE_URL-derived allowlist so the server accepts its own host regardless of port. Co-Authored-By: Claude Opus 4.8 (1M context) --- testgen/mcp/server.py | 5 +++++ tests/unit/mcp/test_transport_security.py | 14 ++++++++++++++ 2 files changed, 19 insertions(+) diff --git a/testgen/mcp/server.py b/testgen/mcp/server.py index cd673604..70e6f788 100644 --- a/testgen/mcp/server.py +++ b/testgen/mcp/server.py @@ -93,8 +93,12 @@ def _build_transport_security() -> TransportSecuritySettings: netloc = parsed.netloc scheme = parsed.scheme or "http" + # base_host (bare) is allowed alongside netloc: when BASE_URL carries a non-default + # port, netloc and the :* wildcard only match a host that includes a port, but some + # clients (e.g. hosted MCP gateways) send a port-less Host/Origin for their own host. allowed_hosts: set[str] = { netloc, + base_host, f"{base_host}:*", "127.0.0.1:*", "localhost:*", @@ -102,6 +106,7 @@ def _build_transport_security() -> TransportSecuritySettings: } allowed_origins: set[str] = { f"{scheme}://{netloc}", + f"{scheme}://{base_host}", "http://127.0.0.1:*", "https://127.0.0.1:*", "http://localhost:*", "https://localhost:*", "http://[::1]:*", "https://[::1]:*", diff --git a/tests/unit/mcp/test_transport_security.py b/tests/unit/mcp/test_transport_security.py index 94b250f8..a2b6f048 100644 --- a/tests/unit/mcp/test_transport_security.py +++ b/tests/unit/mcp/test_transport_security.py @@ -30,6 +30,20 @@ def test_loopback_and_base_url_always_present(): assert "https://localhost:*" in settings.allowed_origins +def test_ported_base_url_also_allows_bare_host_and_origin(): + """A BASE_URL with a non-default port also allows the bare (port-less) host and origin. + + Clients such as hosted MCP gateways may send a port-less Host/Origin for the server's + own host; the `:*` wildcard requires a port, so the bare forms must be present too. + """ + settings = _build_with("https://tg.example.com:8530") + + assert "tg.example.com:8530" in settings.allowed_hosts + assert "tg.example.com" in settings.allowed_hosts + assert "https://tg.example.com:8530" in settings.allowed_origins + assert "https://tg.example.com" in settings.allowed_origins + + def test_extra_host_without_port_gets_wildcard_and_bare(): """An extras entry without `:` is allowed both with a `:*` port wildcard and bare. From f7ad98c962519a4ab981ee57c0fefae0fdbe005b Mon Sep 17 00:00:00 2001 From: Ricardo Boni Date: Wed, 10 Jun 2026 08:55:36 -0400 Subject: [PATCH 23/78] feat(mcp): filter test-run and profiling-run lists by schedule (TG-1117) Add an optional schedule_id filter to list_test_runs and list_profiling_runs so the followable Schedule ID already surfaced on scheduler-triggered runs closes the loop into "all runs from this schedule". The filter is a soft additive WHERE clause, not a resolve-and-raise: the existing select_summary queries already pin je.job_key and scope by project_code, so an unknown, inaccessible, or wrong-kind schedule yields the standard empty-result envelope rather than an error. JobExecution .select_active_by_kwargs gains a *clauses arg so the pending-JE section stays schedule-scoped too. Co-Authored-By: Claude Opus 4.8 (1M context) --- testgen/common/models/job_execution.py | 4 ++ testgen/common/models/profiling_run.py | 4 ++ testgen/common/models/test_run.py | 4 ++ testgen/mcp/tools/profiling.py | 15 ++++++- testgen/mcp/tools/test_runs.py | 15 ++++++- tests/unit/mcp/test_tools_profiling.py | 47 ++++++++++++++++++++ tests/unit/mcp/test_tools_test_runs.py | 61 ++++++++++++++++++++++++++ 7 files changed, 148 insertions(+), 2 deletions(-) diff --git a/testgen/common/models/job_execution.py b/testgen/common/models/job_execution.py index 9ae65325..bb6d8d83 100644 --- a/testgen/common/models/job_execution.py +++ b/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,6 +107,7 @@ 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 @@ -112,6 +115,7 @@ def select_active_by_kwargs( 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): diff --git a/testgen/common/models/profiling_run.py b/testgen/common/models/profiling_run.py index 0c898172..f672ed0a 100644 --- a/testgen/common/models/profiling_run.py +++ b/testgen/common/models/profiling_run.py @@ -208,6 +208,7 @@ 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, @@ -215,6 +216,7 @@ def select_summary( 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, diff --git a/testgen/common/models/test_run.py b/testgen/common/models/test_run.py index 1887d886..161f14a3 100644 --- a/testgen/common/models/test_run.py +++ b/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, diff --git a/testgen/mcp/tools/profiling.py b/testgen/mcp/tools/profiling.py index 874cd5c3..b5b6f42b 100644 --- a/testgen/mcp/tools/profiling.py +++ b/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,7 +451,9 @@ 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)}, @@ -453,7 +461,12 @@ def list_profiling_runs( ) 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( diff --git a/testgen/mcp/tools/test_runs.py b/testgen/mcp/tools/test_runs.py index 2303a500..8f92d86a 100644 --- a/testgen/mcp/tools/test_runs.py +++ b/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}, diff --git a/tests/unit/mcp/test_tools_profiling.py b/tests/unit/mcp/test_tools_profiling.py index 733902f8..bddee5ac 100644 --- a/tests/unit/mcp/test_tools_profiling.py +++ b/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 # ---------------------------------------------------------------------- diff --git a/tests/unit/mcp/test_tools_test_runs.py b/tests/unit/mcp/test_tools_test_runs.py index 7a71380e..e37f6df9 100644 --- a/tests/unit/mcp/test_tools_test_runs.py +++ b/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") From 2746a700e4c2bf4f45fea5105124973163a4a371 Mon Sep 17 00:00:00 2001 From: Luis Date: Tue, 9 Jun 2026 12:40:17 -0400 Subject: [PATCH 24/78] feat(mcp): add update_catalog_metadata write tool Add a single MCP write tool that applies table/column catalog metadata (description, CDE, XDE, PII, and 8 catalog tags) via a list of per-row update specs, sharing one persistence path with the Data Catalog UI. - Map description + tag columns on DataTableChars and DataColumnChars - Add common/data_catalog_service.py (validate, apply, disable-autoflags) and relocate TAG_FIELDS into it - Rewire on_tags_changed and CSV import to the shared service - Gate on disposition (+ view_pii for PII); rows processed independently with a per-row outcome report Co-Authored-By: Claude Opus 4.8 (1M context) --- testgen/common/data_catalog_service.py | 79 ++++- testgen/common/models/data_column.py | 13 +- testgen/common/models/data_table.py | 17 +- testgen/mcp/server.py | 2 + testgen/mcp/tools/data_catalog.py | 221 +++++++++++++ testgen/ui/queries/profiling_queries.py | 11 +- testgen/ui/views/data_catalog.py | 84 ++--- .../views/dialogs/import_metadata_dialog.py | 68 ++-- .../unit/common/test_data_catalog_service.py | 110 +++++++ tests/unit/mcp/test_model_data_column.py | 25 ++ tests/unit/mcp/test_model_data_table.py | 38 +++ tests/unit/mcp/test_tools_data_catalog.py | 294 ++++++++++++++++++ tests/unit/ui/test_import_metadata.py | 42 +++ 13 files changed, 896 insertions(+), 108 deletions(-) create mode 100644 testgen/mcp/tools/data_catalog.py create mode 100644 tests/unit/mcp/test_tools_data_catalog.py diff --git a/testgen/common/data_catalog_service.py b/testgen/common/data_catalog_service.py index 936e7544..d31a0c55 100644 --- a/testgen/common/data_catalog_service.py +++ b/testgen/common/data_catalog_service.py @@ -1,11 +1,13 @@ -"""Shared data catalog convenience service. +"""Shared Data Catalog service. -Generates CREATE TABLE scripts from profiled column metadata and fetches -sample rows from source tables. Used by both the Streamlit UI and MCP tools. +Generates CREATE TABLE scripts from profiled column metadata, fetches sample rows +from source tables, and reads/writes table/column catalog metadata. Used by both +the Streamlit UI and the MCP tools so they share one code path. """ import logging +from collections.abc import Mapping from dataclasses import dataclass -from typing import Literal +from typing import Any, Literal from uuid import UUID import pandas as pd @@ -14,12 +16,32 @@ from testgen.common.database.flavor.flavor_service import FlavorService from testgen.common.models.connection import Connection from testgen.common.models.data_column import CreateScriptColumn, DataColumnChars +from testgen.common.models.data_table import DataTable +from testgen.common.models.table_group import TableGroup from testgen.common.pii_masking import get_pii_columns, mask_source_data_pii from testgen.ui.services.database_service import fetch_from_target_db from testgen.utils import to_dataframe LOG = logging.getLogger("testgen") +DESCRIPTION_MAX_LENGTH = 1000 +TAG_MAX_LENGTH = 40 + +TAG_FIELDS = [ + "data_source", + "source_system", + "source_process", + "business_domain", + "stakeholder_group", + "transform_level", + "aggregation_level", + "data_product", +] + +# Metadata fields settable per target type, keyed by their data_*_chars column names. +_TABLE_FIELDS = ("description", *TAG_FIELDS, "critical_data_element") +_COLUMN_FIELDS = ("description", *TAG_FIELDS, "critical_data_element", "excluded_data_element", "pii_flag") + @dataclass class TableSampleResult: @@ -117,3 +139,52 @@ def fetch_table_sample( pii_columns = get_pii_columns(str(table_group_id), schema_name, table_name) pii_redacted = mask_source_data_pii(df, pii_columns) return TableSampleResult("OK", df=df, pii_redacted=pii_redacted) + + +def validate_metadata_fields(fields: Mapping[str, Any]) -> list[str]: + """Validate free-text field lengths, returning one message per violation. + + Only checks length: ``description`` and the tag fields. ``None`` values clear the + field and are always valid; absent keys are skipped. + """ + errors: list[str] = [] + + description = fields.get("description") + if description is not None and len(description) > DESCRIPTION_MAX_LENGTH: + errors.append(f"description must be {DESCRIPTION_MAX_LENGTH} characters or fewer.") + + for tag in TAG_FIELDS: + value = fields.get(tag) + if value is not None and len(value) > TAG_MAX_LENGTH: + errors.append(f"{tag} must be {TAG_MAX_LENGTH} characters or fewer.") + + return errors + + +def apply_table_metadata(table: DataTable, fields: Mapping[str, Any]) -> None: + """Set table metadata attributes for every present key (None clears, value sets).""" + for field in _TABLE_FIELDS: + if field in fields: + setattr(table, field, fields[field]) + + +def apply_column_metadata(column: DataColumnChars, fields: Mapping[str, Any]) -> None: + """Set column metadata attributes for every present key (None clears, value sets).""" + for field in _COLUMN_FIELDS: + if field in fields: + setattr(column, field, fields[field]) + + +def disable_autoflags(table_group: TableGroup, *, wrote_cde: bool, wrote_pii: bool) -> list[str]: + """Turn off auto-detect flags so the next profiling pass preserves manual marks. + + Returns the names of the flags that were turned off (skips flags already off). + """ + disabled: list[str] = [] + if wrote_cde and table_group.profile_flag_cdes: + table_group.profile_flag_cdes = False + disabled.append("profile_flag_cdes") + if wrote_pii and table_group.profile_flag_pii: + table_group.profile_flag_pii = False + disabled.append("profile_flag_pii") + return disabled diff --git a/testgen/common/models/data_column.py b/testgen/common/models/data_column.py index ed21e312..2c7adbdb 100644 --- a/testgen/common/models/data_column.py +++ b/testgen/common/models/data_column.py @@ -260,6 +260,15 @@ class DataColumnChars(Entity): critical_data_element: bool | None = Column(Boolean) excluded_data_element: bool | None = Column(Boolean, nullable=True) pii_flag: str | None = Column(String(50), nullable=True) + description: str | None = Column(String(1000)) + data_source: str | None = Column(String(40)) + source_system: str | None = Column(String(40)) + source_process: str | None = Column(String(40)) + business_domain: str | None = Column(String(40)) + stakeholder_group: str | None = Column(String(40)) + transform_level: str | None = Column(String(40)) + aggregation_level: str | None = Column(String(40)) + data_product: str | None = Column(String(40)) drop_date: datetime | None = Column(postgresql.TIMESTAMP) last_complete_profile_run_id: UUID | None = Column(postgresql.UUID(as_uuid=True)) dq_score_profiling: float | None = Column(Float) @@ -267,9 +276,7 @@ class DataColumnChars(Entity): _default_order_by = (asc(ordinal_position), asc(column_name)) - # Unmapped columns: description, data_source, source_system, source_process, - # business_domain, stakeholder_group, transform_level, aggregation_level, - # data_product, add_date, last_mod_date, test_ct, last_test_date, + # Unmapped columns: add_date, last_mod_date, test_ct, last_test_date, # tests_last_run, tests_7_days_prior, tests_30_days_prior, fails_last_run, # fails_7_days_prior, fails_30_days_prior, warnings_last_run, # warnings_7_days_prior, warnings_30_days_prior, valid_profile_issue_ct, diff --git a/testgen/common/models/data_table.py b/testgen/common/models/data_table.py index ce149b94..5b866376 100644 --- a/testgen/common/models/data_table.py +++ b/testgen/common/models/data_table.py @@ -68,14 +68,25 @@ class DataTable(Entity): record_ct: int | None = Column(BigInteger) approx_record_ct: int | None = Column(BigInteger) critical_data_element: bool | None = Column(Boolean) + description: str | None = Column(String(1000)) + data_source: str | None = Column(String(40)) + source_system: str | None = Column(String(40)) + source_process: str | None = Column(String(40)) + business_domain: str | None = Column(String(40)) + stakeholder_group: str | None = Column(String(40)) + transform_level: str | None = Column(String(40)) + aggregation_level: str | None = Column(String(40)) + data_product: str | None = Column(String(40)) drop_date: datetime | None = Column(postgresql.TIMESTAMP) last_complete_profile_run_id: UUID | None = Column(postgresql.UUID(as_uuid=True)) dq_score_profiling: float | None = Column(Float) dq_score_testing: float | None = Column(Float) - # Unmapped columns: functional_table_type, description, data_source, - # source_system, source_process, business_domain, stakeholder_group, - # transform_level, aggregation_level, data_product, add_date, + # The inherited Entity default orders by the textual label "id", but this model's + # primary key column is "table_id" — order by a real column instead. + _default_order_by = (asc(table_name),) + + # Unmapped columns: functional_table_type, add_date, # last_refresh_date, last_profile_record_ct @classmethod diff --git a/testgen/mcp/server.py b/testgen/mcp/server.py index cd673604..c628036a 100644 --- a/testgen/mcp/server.py +++ b/testgen/mcp/server.py @@ -141,6 +141,7 @@ def build_mcp_server( table_health, ) from testgen.mcp.tools.connections import test_connection + from testgen.mcp.tools.data_catalog import update_catalog_metadata from testgen.mcp.tools.discovery import get_data_inventory, list_projects, list_tables, list_test_suites from testgen.mcp.tools.execution import ( cancel_profiling_run, @@ -332,6 +333,7 @@ def safe_prompt(fn): safe_tool(create_table_group) safe_tool(update_table_group) safe_tool(preview_table_group) + safe_tool(update_catalog_metadata) # Resources safe_resource("testgen://test-types", test_types_resource) diff --git a/testgen/mcp/tools/data_catalog.py b/testgen/mcp/tools/data_catalog.py new file mode 100644 index 00000000..334c65bb --- /dev/null +++ b/testgen/mcp/tools/data_catalog.py @@ -0,0 +1,221 @@ +from collections.abc import Mapping + +from testgen.common.data_catalog_service import ( + TAG_FIELDS, + apply_column_metadata, + apply_table_metadata, + disable_autoflags, + validate_metadata_fields, +) +from testgen.common.models import with_database_session +from testgen.common.models.data_column import DataColumnChars +from testgen.common.models.data_table import DataTable +from testgen.mcp.exceptions import MCPUserError +from testgen.mcp.permissions import get_project_permissions, mcp_permission +from testgen.mcp.tools.common import DocGroup, resolve_table_group +from testgen.mcp.tools.markdown import MdDoc + +_DOC_GROUP = DocGroup.MANAGE + +_BOOLEAN_ARGS = ("cde", "xde", "pii") + +# MCP argument name -> data_*_chars column name. +_ARG_TO_COLUMN = {"cde": "critical_data_element", "xde": "excluded_data_element", "pii": "pii_flag"} +_COLUMN_TO_ARG = {v: k for k, v in _ARG_TO_COLUMN.items()} + + +@with_database_session +@mcp_permission("disposition") +def update_catalog_metadata(updates: list[dict]) -> str: + """Apply metadata updates to tables and columns within table groups. + + Each update targets one table, or one column within a table, and sets one or + more metadata fields. Rows are processed independently: a failure on one row + does not stop the others, and the response reports the outcome of every row. + + Omit a field to leave it unchanged, pass null to clear it, or pass a value to set it. + + Args: + updates: List of per-row update specs. Each spec accepts: + table_group_id: UUID of the table group, e.g. from `get_data_inventory`. + table_name: Name of the target table. + column_name: Name of the target column. Omit for a table-level update. + description: Free-text description (max 1000 characters). + cde: Whether this is a critical data element (true/false). Set at the + table level to apply to all its columns unless a column overrides it. + xde: Whether to exclude the column from future profiling and test + generation (true/false). Columns only. + pii: Whether the column contains personally identifiable information + (true/false). Columns only. + data_source, source_system, source_process, business_domain, + stakeholder_group, transform_level, aggregation_level, data_product: + Catalog tags (max 40 characters each). + """ + if not updates: + raise MCPUserError("Provide at least one update.") + + perms = get_project_permissions() + tg_cache: dict[str, object] = {} + flag_writes: dict = {} # tg.id -> {"tg": tg, "cde": bool, "pii": bool} + inheritance_notices: list[str] = [] + exclusion_notices: list[str] = [] + rows: list[tuple[str, str, str]] = [] + + for spec in updates: + target = _target_label(spec) + try: + outcome, detail = _apply_update(spec, perms, tg_cache, flag_writes, inheritance_notices, exclusion_notices) + rows.append((target, outcome, detail)) + except MCPUserError as err: + rows.append((target, "Failed", str(err))) + + autoflag_notices: list[str] = [] + for entry in flag_writes.values(): + disabled = disable_autoflags(entry["tg"], wrote_cde=entry["cde"], wrote_pii=entry["pii"]) + for flag in disabled: + autoflag_notices.append( + f"Auto-disabled {flag} on table group `{entry['tg'].table_groups_name}` to preserve manual marks." + ) + + return _render(rows, autoflag_notices + inheritance_notices + exclusion_notices) + + +def _target_label(spec: Mapping) -> str: + table = spec.get("table_name") or "?" + column = spec.get("column_name") + return f"{table}.{column}" if column else table + + +def _apply_update(spec, perms, tg_cache, flag_writes, inheritance_notices, exclusion_notices) -> tuple[str, str]: + if not isinstance(spec, Mapping): + raise MCPUserError("Each update must be an object.") + + table_group_id = spec.get("table_group_id") + table_name = spec.get("table_name") + if not table_group_id: + raise MCPUserError("table_group_id is required.") + if not table_name: + raise MCPUserError("table_name is required.") + column_name = spec.get("column_name") + is_column = column_name is not None + + for arg in _BOOLEAN_ARGS: + if arg in spec and spec[arg] is not None and not isinstance(spec[arg], bool): + raise MCPUserError(f"{arg} must be true or false.") + + if not is_column: + if "xde" in spec: + raise MCPUserError("xde applies to columns only; omit it for table-level updates.") + if "pii" in spec: + raise MCPUserError("pii applies to columns only; omit it for table-level updates.") + + tg = tg_cache.get(table_group_id) + if tg is None: + tg = resolve_table_group(table_group_id) + tg_cache[table_group_id] = tg + + if "pii" in spec and not perms.has_permission("view_pii", tg.project_code): + raise MCPUserError(f"Setting pii requires permission to view PII on project `{tg.project_code}`.") + + fields = _build_fields(spec) + errors = validate_metadata_fields(fields) + if errors: + raise MCPUserError("; ".join(errors)) + + if not fields: + return "Skipped", "no metadata fields provided" + + if is_column: + target = _resolve_column(tg.id, table_name, column_name) + if target is None: + raise MCPUserError(f"Column `{column_name}` not found in table `{table_name}`.") + else: + target = _resolve_table(tg.id, table_name) + if target is None: + raise MCPUserError(f"Table `{table_name}` not found in this table group.") + + # Disable a table group's auto-detect flag only on a real change, matching the UI's confirm-to-disable + # behavior. A no-op write (e.g. cde: false on an already-false column) must not silently turn it off. + cde_changed = "cde" in spec and target.critical_data_element != spec["cde"] + pii_changed = "pii" in spec and target.pii_flag != ("MANUAL" if spec["pii"] else None) + + if is_column: + apply_column_metadata(target, fields) + else: + apply_table_metadata(target, fields) + + entry = flag_writes.setdefault(tg.id, {"tg": tg, "cde": False, "pii": False}) + if cde_changed: + entry["cde"] = True + if pii_changed: + entry["pii"] = True + + if not is_column and spec.get("cde") is True: + inheritance_notices.append( + f"Table-level CDE set; affects all columns of `{table_name}` unless explicitly overridden." + ) + if is_column and spec.get("xde") is True: + exclusion_notices.append( + f"Column `{table_name}.{column_name}` excluded; next profiling run and test generation will skip it." + ) + + return "Updated", _change_summary(fields) + + +def _build_fields(spec: Mapping) -> dict: + """Translate a spec into a {data_*_chars column: value} dict (pii bool -> MANUAL/None).""" + fields: dict = {} + if "description" in spec: + fields["description"] = spec["description"] + for tag in TAG_FIELDS: + if tag in spec: + fields[tag] = spec[tag] + if "cde" in spec: + fields["critical_data_element"] = spec["cde"] + if "xde" in spec: + fields["excluded_data_element"] = spec["xde"] + if "pii" in spec: + fields["pii_flag"] = "MANUAL" if spec["pii"] else None + return fields + + +def _resolve_table(table_groups_id, table_name: str): + matches = list(DataTable.select_where( + DataTable.table_groups_id == table_groups_id, + DataTable.table_name == table_name, + )) + return matches[0] if matches else None + + +def _resolve_column(table_groups_id, table_name: str, column_name: str): + matches = list(DataColumnChars.select_where( + DataColumnChars.table_groups_id == table_groups_id, + DataColumnChars.table_name == table_name, + DataColumnChars.column_name == column_name, + )) + return matches[0] if matches else None + + +def _change_summary(fields: dict) -> str: + parts = [] + for column, value in fields.items(): + label = _COLUMN_TO_ARG.get(column, column) + parts.append(f"{label} cleared" if value is None else f"{label} set") + return ", ".join(parts) + + +def _render(rows: list[tuple[str, str, str]], notices: list[str]) -> str: + succeeded = sum(1 for _, outcome, _ in rows if outcome == "Updated") + skipped = sum(1 for _, outcome, _ in rows if outcome == "Skipped") + failed = sum(1 for _, outcome, _ in rows if outcome == "Failed") + + doc = MdDoc() + doc.heading(1, "Catalog metadata update") + doc.field("Rows attempted", len(rows)) + doc.field("Updated", succeeded) + doc.field("Skipped", skipped) + doc.field("Failed", failed) + doc.table(["Target", "Outcome", "Details"], [list(row) for row in rows], code=[0]) + for notice in notices: + doc.text(notice) + return doc.render() diff --git a/testgen/ui/queries/profiling_queries.py b/testgen/ui/queries/profiling_queries.py index ce7ff242..3176ce75 100644 --- a/testgen/ui/queries/profiling_queries.py +++ b/testgen/ui/queries/profiling_queries.py @@ -3,19 +3,10 @@ import pandas as pd import streamlit as st +from testgen.common.data_catalog_service import TAG_FIELDS from testgen.ui.services.database_service import fetch_all_from_db, fetch_df_from_db, fetch_one_from_db from testgen.utils import is_uuid4 -TAG_FIELDS = [ - "data_source", - "source_system", - "source_process", - "business_domain", - "stakeholder_group", - "transform_level", - "aggregation_level", - "data_product", -] COLUMN_PROFILING_FIELDS = """ -- Value Counts profile_results.record_ct, diff --git a/testgen/ui/views/data_catalog.py b/testgen/ui/views/data_catalog.py index f5db0893..3b0b848f 100644 --- a/testgen/ui/views/data_catalog.py +++ b/testgen/ui/views/data_catalog.py @@ -4,16 +4,25 @@ from collections import defaultdict from datetime import datetime from functools import partial +from uuid import UUID import pandas as pd import streamlit as st from sqlalchemy.sql.expression import func as sa_func from streamlit.delta_generator import DeltaGenerator -from testgen.common.data_catalog_service import build_create_table_script, fetch_table_sample +from testgen.common.data_catalog_service import ( + apply_column_metadata, + apply_table_metadata, + build_create_table_script, + disable_autoflags, + fetch_table_sample, +) from testgen.common.enums import JobSource from testgen.common.models import database_session, with_database_session from testgen.common.models.connection import Connection +from testgen.common.models.data_column import DataColumnChars +from testgen.common.models.data_table import DataTable from testgen.common.models.job_execution import JobExecution from testgen.common.models.profiling_run import ProfilingRun from testgen.common.models.table_group import TableGroup, TableGroupMinimal @@ -603,60 +612,33 @@ def remove_table_dialog(item: dict) -> None: @with_database_session -def on_tags_changed(spinner_container: DeltaGenerator, payload: dict) -> FILE_DATA_TYPE: - attributes = ["description"] - attributes.extend(TAG_FIELDS) - +def on_tags_changed(spinner_container: DeltaGenerator, payload: dict) -> None: tags = payload["tags"] - set_attributes = [ f"{key} = NULLIF(:{key}, '')" for key in attributes if key in tags ] - params = { key: tags.get(key) or "" for key in attributes if key in tags } + + # Empty string clears the field (matches the prior NULLIF semantics). + shared_fields: dict = {key: (tags.get(key) or None) for key in ["description", *TAG_FIELDS] if key in tags} if "critical_data_element" in tags: - set_attributes.append("critical_data_element = :critical_data_element") - params["critical_data_element"] = tags.get("critical_data_element") + shared_fields["critical_data_element"] = tags.get("critical_data_element") # pii_flag and excluded_data_element are column-only fields (not in data_table_chars) - column_set_attributes = list(set_attributes) + column_fields = dict(shared_fields) if "pii_flag" in tags: - column_set_attributes.append("pii_flag = :pii_flag") - params["pii_flag"] = tags.get("pii_flag") + column_fields["pii_flag"] = tags.get("pii_flag") if "excluded_data_element" in tags: - column_set_attributes.append("excluded_data_element = :excluded_data_element") - params["excluded_data_element"] = tags.get("excluded_data_element") + column_fields["excluded_data_element"] = tags.get("excluded_data_element") - params["table_ids"] = [ item["id"] for item in payload["items"] if item["type"] == "table" ] - params["column_ids"] = [ item["id"] for item in payload["items"] if item["type"] == "column" ] + table_ids = [ UUID(item["id"]) for item in payload["items"] if item["type"] == "table" ] + column_ids = [ UUID(item["id"]) for item in payload["items"] if item["type"] == "column" ] with spinner_container: with st.spinner("Saving tags"): - if params["table_ids"] and set_attributes: - execute_db_query( - f""" - WITH selected as ( - SELECT UNNEST(ARRAY [:table_ids]) AS table_id - ) - UPDATE data_table_chars - SET {', '.join(set_attributes)} - FROM data_table_chars dtc - INNER JOIN selected ON (dtc.table_id = selected.table_id::UUID) - WHERE dtc.table_id = data_table_chars.table_id; - """, - params, - ) + if table_ids and shared_fields: + for table in DataTable.select_where(DataTable.id.in_(table_ids)): + apply_table_metadata(table, shared_fields) - if params["column_ids"] and column_set_attributes: - execute_db_query( - f""" - WITH selected as ( - SELECT UNNEST(ARRAY [:column_ids]) AS column_id - ) - UPDATE data_column_chars - SET {', '.join(column_set_attributes)} - FROM data_column_chars dcc - INNER JOIN selected ON (dcc.column_id = selected.column_id::UUID) - WHERE dcc.column_id = data_column_chars.column_id; - """, - params, - ) + if column_ids and column_fields: + for column in DataColumnChars.select_where(DataColumnChars.id.in_(column_ids)): + apply_column_metadata(column, column_fields) # Disable autodetection flags on table group if requested disable_flags = payload.get("disable_flags", []) @@ -664,14 +646,12 @@ def on_tags_changed(spinner_container: DeltaGenerator, payload: dict) -> FILE_DA table_group_id = st.query_params.get("table_group_id") if table_group_id: table_group = get_table_group(table_group_id) - changed = False - if "profile_flag_cdes" in disable_flags and table_group.profile_flag_cdes: - table_group.profile_flag_cdes = False - changed = True - if "profile_flag_pii" in disable_flags and table_group.profile_flag_pii: - table_group.profile_flag_pii = False - changed = True - if changed: + disabled = disable_autoflags( + table_group, + wrote_cde="profile_flag_cdes" in disable_flags, + wrote_pii="profile_flag_pii" in disable_flags, + ) + if disabled: table_group.save() for func in [ get_table_group_columns, get_table_by_id, get_column_by_id, get_tag_values, select_table_groups_minimal_where ]: diff --git a/testgen/ui/views/dialogs/import_metadata_dialog.py b/testgen/ui/views/dialogs/import_metadata_dialog.py index b6b1f4c8..09b0b91b 100644 --- a/testgen/ui/views/dialogs/import_metadata_dialog.py +++ b/testgen/ui/views/dialogs/import_metadata_dialog.py @@ -4,8 +4,11 @@ import pandas as pd +from testgen.common.data_catalog_service import apply_column_metadata, apply_table_metadata, disable_autoflags +from testgen.common.models.data_column import DataColumnChars +from testgen.common.models.data_table import DataTable from testgen.ui.queries.profiling_queries import TAG_FIELDS -from testgen.ui.services.database_service import execute_db_query, fetch_all_from_db +from testgen.ui.services.database_service import fetch_all_from_db from testgen.ui.services.query_cache import get_table_group from testgen.ui.session import session @@ -296,53 +299,50 @@ def _set_row_status(preview_row: dict, bad_cde: bool, bad_xde: bool, bad_pii: bo def apply_metadata_import(preview: dict, table_group_id: str | None = None) -> dict: + metadata_columns = preview["metadata_columns"] table_count = 0 column_count = 0 for row in preview.get("table_rows", []): - set_clauses, params = _build_update_params(row, preview["metadata_columns"], is_column=False) - if not set_clauses: + fields = _build_metadata_fields(row, metadata_columns, is_column=False) + if not fields: continue - params["table_id"] = row["table_id"] - execute_db_query( - f"UPDATE data_table_chars SET {', '.join(set_clauses)} WHERE table_id = CAST(:table_id AS UUID)", - params, - ) + table = DataTable.get(row["table_id"]) + if table is None: + continue + apply_table_metadata(table, fields) table_count += 1 for row in preview.get("column_rows", []): - set_clauses, params = _build_update_params(row, preview["metadata_columns"], is_column=True) - if not set_clauses: + fields = _build_metadata_fields(row, metadata_columns, is_column=True) + if not fields: + continue + column = DataColumnChars.get(row["column_id"]) + if column is None: continue - params["column_id"] = row["column_id"] - execute_db_query( - f"UPDATE data_column_chars SET {', '.join(set_clauses)} WHERE column_id = CAST(:column_id AS UUID)", - params, - ) + apply_column_metadata(column, fields) column_count += 1 if table_group_id: - _disable_autoflags(table_group_id, preview.get("metadata_columns", [])) + _disable_autoflags(table_group_id, metadata_columns) return {"table_count": table_count, "column_count": column_count} def _disable_autoflags(table_group_id: str, metadata_columns: list[str]) -> None: table_group = get_table_group(table_group_id) - changed = False - if "critical_data_element" in metadata_columns and table_group.profile_flag_cdes: - table_group.profile_flag_cdes = False - changed = True - if "pii_flag" in metadata_columns and table_group.profile_flag_pii: - table_group.profile_flag_pii = False - changed = True - if changed: + disabled = disable_autoflags( + table_group, + wrote_cde="critical_data_element" in metadata_columns, + wrote_pii="pii_flag" in metadata_columns, + ) + if disabled: table_group.save() -def _build_update_params(row: dict, metadata_columns: list[str], is_column: bool = False) -> tuple[list[str], dict]: - set_clauses = [] - params = {} +def _build_metadata_fields(row: dict, metadata_columns: list[str], is_column: bool = False) -> dict: + """Map a preview row to a {column_name: value} update dict (empty string clears the field).""" + fields: dict = {} for col in metadata_columns: if col not in row: @@ -350,22 +350,18 @@ def _build_update_params(row: dict, metadata_columns: list[str], is_column: bool value = row[col] if col == "critical_data_element": - set_clauses.append("critical_data_element = :critical_data_element") - params["critical_data_element"] = value + fields[col] = value elif col == "excluded_data_element": if is_column: - set_clauses.append("excluded_data_element = :excluded_data_element") - params["excluded_data_element"] = value + fields[col] = value elif col == "pii_flag": # Prevent user from editing PII flag if they cannot view PII if is_column and session.auth.user_has_permission("view_pii"): - set_clauses.append("pii_flag = :pii_flag") - params["pii_flag"] = value + fields[col] = value else: - set_clauses.append(f"{col} = NULLIF(:{col}, '')") - params[col] = value if value is not None else "" + fields[col] = value or None - return set_clauses, params + return fields def build_import_preview_props(preview: dict) -> dict: diff --git a/tests/unit/common/test_data_catalog_service.py b/tests/unit/common/test_data_catalog_service.py index 335ae6ce..17445bf9 100644 --- a/tests/unit/common/test_data_catalog_service.py +++ b/tests/unit/common/test_data_catalog_service.py @@ -1,12 +1,20 @@ +from types import SimpleNamespace from unittest.mock import MagicMock, patch import pytest from testgen.common.data_catalog_service import ( + DESCRIPTION_MAX_LENGTH, + TAG_FIELDS, + TAG_MAX_LENGTH, TableSampleResult, + apply_column_metadata, + apply_table_metadata, build_create_table_script, + disable_autoflags, fetch_table_sample, render_create_table_script, + validate_metadata_fields, ) from testgen.common.models.data_column import CreateScriptColumn @@ -192,3 +200,105 @@ def test_table_sample_result_defaults(): result = TableSampleResult("ND") assert result.df is None assert result.pii_redacted is False + + +# --------------------------------------------------------------------------- +# validate_metadata_fields +# --------------------------------------------------------------------------- + +def test_validate_accepts_values_within_limits(): + fields = {"description": "x" * DESCRIPTION_MAX_LENGTH, "business_domain": "y" * TAG_MAX_LENGTH} + assert validate_metadata_fields(fields) == [] + + +def test_validate_flags_description_too_long(): + errors = validate_metadata_fields({"description": "x" * (DESCRIPTION_MAX_LENGTH + 1)}) + assert len(errors) == 1 + assert "description" in errors[0] + + +def test_validate_flags_each_tag_too_long(): + fields = {"data_source": "a" * (TAG_MAX_LENGTH + 1), "data_product": "b" * (TAG_MAX_LENGTH + 1)} + errors = validate_metadata_fields(fields) + assert len(errors) == 2 + assert any("data_source" in e for e in errors) + assert any("data_product" in e for e in errors) + + +def test_validate_ignores_none_and_absent(): + # None clears the field — never a length violation; absent keys are skipped. + assert validate_metadata_fields({"description": None, "source_system": None}) == [] + assert validate_metadata_fields({}) == [] + + +# --------------------------------------------------------------------------- +# apply_table_metadata / apply_column_metadata +# --------------------------------------------------------------------------- + +def _table(): + return SimpleNamespace(description="old", business_domain="old", critical_data_element=False) + + +def _metadata_column(): + return SimpleNamespace( + description="old", data_product="old", critical_data_element=False, + excluded_data_element=False, pii_flag=None, + ) + + +def test_apply_table_sets_present_keys_only(): + table = _table() + apply_table_metadata(table, {"description": "new", "critical_data_element": True}) + assert table.description == "new" + assert table.critical_data_element is True + assert table.business_domain == "old" # absent → unchanged + + +def test_apply_table_none_clears(): + table = _table() + apply_table_metadata(table, {"business_domain": None}) + assert table.business_domain is None + + +def test_apply_column_sets_column_only_fields(): + column = _metadata_column() + apply_column_metadata(column, {"pii_flag": "MANUAL", "excluded_data_element": True, "data_product": "CRM"}) + assert column.pii_flag == "MANUAL" + assert column.excluded_data_element is True + assert column.data_product == "CRM" + + +def test_apply_column_none_clears_pii(): + column = _metadata_column() + column.pii_flag = "MANUAL" + apply_column_metadata(column, {"pii_flag": None}) + assert column.pii_flag is None + + +# --------------------------------------------------------------------------- +# disable_autoflags +# --------------------------------------------------------------------------- + +def test_disable_autoflags_disables_only_written_and_enabled(): + tg = SimpleNamespace(profile_flag_cdes=True, profile_flag_pii=True) + disabled = disable_autoflags(tg, wrote_cde=True, wrote_pii=False) + assert disabled == ["profile_flag_cdes"] + assert tg.profile_flag_cdes is False + assert tg.profile_flag_pii is True + + +def test_disable_autoflags_noop_when_already_disabled(): + tg = SimpleNamespace(profile_flag_cdes=False, profile_flag_pii=False) + assert disable_autoflags(tg, wrote_cde=True, wrote_pii=True) == [] + + +def test_disable_autoflags_disables_both(): + tg = SimpleNamespace(profile_flag_cdes=True, profile_flag_pii=True) + disabled = disable_autoflags(tg, wrote_cde=True, wrote_pii=True) + assert disabled == ["profile_flag_cdes", "profile_flag_pii"] + + +def test_tag_fields_has_eight_entries(): + assert len(TAG_FIELDS) == 8 + assert "data_source" in TAG_FIELDS + assert "aggregation_level" in TAG_FIELDS diff --git a/tests/unit/mcp/test_model_data_column.py b/tests/unit/mcp/test_model_data_column.py index aa8b0181..dffdb21e 100644 --- a/tests/unit/mcp/test_model_data_column.py +++ b/tests/unit/mcp/test_model_data_column.py @@ -4,6 +4,31 @@ from testgen.common.models.data_column import ColumnProfileDetail, DataColumnChars +_CATALOG_METADATA_COLUMNS = { + "description", + "data_source", + "source_system", + "source_process", + "business_domain", + "stakeholder_group", + "transform_level", + "aggregation_level", + "data_product", +} + + +def test_catalog_metadata_columns_are_mapped(): + mapped = set(DataColumnChars.__table__.columns.keys()) + assert _CATALOG_METADATA_COLUMNS <= mapped + + +def test_catalog_metadata_attributes_settable(): + column = DataColumnChars() + column.description = "Primary email" + column.data_product = "CRM" + assert column.description == "Primary email" + assert column.data_product == "CRM" + def _detail_row(**overrides) -> dict: """Build a dict matching every ColumnProfileDetail field.""" diff --git a/tests/unit/mcp/test_model_data_table.py b/tests/unit/mcp/test_model_data_table.py index 0f9f10e2..f5424963 100644 --- a/tests/unit/mcp/test_model_data_table.py +++ b/tests/unit/mcp/test_model_data_table.py @@ -1,9 +1,47 @@ from unittest.mock import patch from uuid import uuid4 +from sqlalchemy import select +from sqlalchemy.dialects import postgresql + from testgen.common.models.data_table import DataTable +def test_default_order_by_compiles_for_entity_select(): + # Regression: the inherited Entity default ordered by the textual label "id", which + # fails to compile here because the PK column is "table_id". A full-entity select + # (e.g. DataTable.select_where) must compile and order by a real column. + stmt = select(DataTable).order_by(*DataTable._default_order_by) + compiled = str(stmt.compile(dialect=postgresql.dialect())) + assert "ORDER BY" in compiled + assert "table_name" in compiled + +_CATALOG_METADATA_COLUMNS = { + "description", + "data_source", + "source_system", + "source_process", + "business_domain", + "stakeholder_group", + "transform_level", + "aggregation_level", + "data_product", +} + + +def test_catalog_metadata_columns_are_mapped(): + mapped = set(DataTable.__table__.columns.keys()) + assert _CATALOG_METADATA_COLUMNS <= mapped + + +def test_catalog_metadata_attributes_settable(): + table = DataTable() + table.description = "Customer master table" + table.business_domain = "Sales" + assert table.description == "Customer master table" + assert table.business_domain == "Sales" + + @patch("testgen.common.models.data_table.get_current_session") def test_select_table_names_returns_list(session_mock): session_mock.return_value.scalars.return_value.all.return_value = ["customers", "orders", "products"] diff --git a/tests/unit/mcp/test_tools_data_catalog.py b/tests/unit/mcp/test_tools_data_catalog.py new file mode 100644 index 00000000..b622cb17 --- /dev/null +++ b/tests/unit/mcp/test_tools_data_catalog.py @@ -0,0 +1,294 @@ +from unittest.mock import MagicMock, patch +from uuid import uuid4 + +import pytest + +pytestmark = pytest.mark.unit + + +@pytest.fixture(autouse=True) +def grant_perms(): + """Grant role_a disposition + view_pii on demo (overrides the conftest matrix).""" + with patch("testgen.mcp.permissions.PluginHook") as hook: + hook.instance.return_value.rbac.get_roles_with_permission.side_effect = ( + lambda perm: ["role_a"] if perm in ("disposition", "view_pii", "catalog", "view") else [] + ) + yield hook + + +def _mock_table_group(): + tg = MagicMock() + tg.id = uuid4() + tg.project_code = "demo" + tg.table_groups_name = "sales_group" + tg.profile_flag_cdes = True + tg.profile_flag_pii = True + return tg + + +def _mock_entity(): + """A stand-in table/column row that records attribute writes.""" + return MagicMock() + + +def _patch(tg=None, table_rows=None, column_rows=None): + """Patch the resolution boundary: TableGroup.get, DataTable/DataColumnChars.select_where.""" + tg = tg or _mock_table_group() + mock_tg_cls = patch("testgen.mcp.tools.common.TableGroup").start() + mock_tg_cls.get.return_value = tg + mock_dt = patch("testgen.mcp.tools.data_catalog.DataTable").start() + mock_dt.select_where.return_value = table_rows if table_rows is not None else [] + mock_dcc = patch("testgen.mcp.tools.data_catalog.DataColumnChars").start() + mock_dcc.select_where.return_value = column_rows if column_rows is not None else [] + return tg, mock_dt, mock_dcc + + +@pytest.fixture(autouse=True) +def _cleanup_patches(): + yield + patch.stopall() + + +# ---------------------------------------------------------------------- +# Happy paths +# ---------------------------------------------------------------------- + +def test_table_description_update(db_session_mock): + table = _mock_entity() + _patch(table_rows=[table]) + + from testgen.mcp.tools.data_catalog import update_catalog_metadata + result = update_catalog_metadata([ + {"table_group_id": str(uuid4()), "table_name": "orders", "description": "Order header table"}, + ]) + + assert table.description == "Order header table" + assert "Updated" in result + assert "1" in result + + +def test_column_pii_true_stores_manual(db_session_mock): + column = _mock_entity() + _patch(column_rows=[column]) + + from testgen.mcp.tools.data_catalog import update_catalog_metadata + update_catalog_metadata([ + {"table_group_id": str(uuid4()), "table_name": "orders", "column_name": "ssn", "pii": True}, + ]) + + assert column.pii_flag == "MANUAL" + + +def test_column_pii_false_clears(db_session_mock): + column = _mock_entity() + _patch(column_rows=[column]) + + from testgen.mcp.tools.data_catalog import update_catalog_metadata + update_catalog_metadata([ + {"table_group_id": str(uuid4()), "table_name": "orders", "column_name": "ssn", "pii": False}, + ]) + + assert column.pii_flag is None + + +# ---------------------------------------------------------------------- +# Per-row isolation +# ---------------------------------------------------------------------- + +def test_one_bad_row_does_not_block_others(db_session_mock): + good = _mock_entity() + _tg, _mock_dt, mock_dcc = _patch() + + # First column resolves to the good row; the second is not found in the catalog. + calls = {"n": 0} + + def _resolve(*_a, **_k): + calls["n"] += 1 + return [good] if calls["n"] == 1 else [] + + mock_dcc.select_where.side_effect = _resolve + + from testgen.mcp.tools.data_catalog import update_catalog_metadata + tg_id = str(uuid4()) + result = update_catalog_metadata([ + {"table_group_id": tg_id, "table_name": "orders", "column_name": "known", "description": "ok"}, + {"table_group_id": tg_id, "table_name": "orders", "column_name": "missing", "description": "x"}, + ]) + + assert good.description == "ok" + assert "Updated" in result and "Failed" in result + assert "not found" in result.lower() + + +# ---------------------------------------------------------------------- +# Validation / rejection +# ---------------------------------------------------------------------- + +def test_xde_on_table_row_rejected(db_session_mock): + _patch(table_rows=[_mock_entity()]) + + from testgen.mcp.tools.data_catalog import update_catalog_metadata + result = update_catalog_metadata([ + {"table_group_id": str(uuid4()), "table_name": "orders", "xde": True}, + ]) + + assert "Failed" in result + assert "xde" in result.lower() and "column" in result.lower() + + +def test_pii_without_view_pii_rejected(db_session_mock): + _patch(column_rows=[_mock_entity()]) + + from testgen.mcp.tools.data_catalog import update_catalog_metadata + with patch("testgen.mcp.tools.data_catalog.get_project_permissions") as mock_perms: + perms = MagicMock() + perms.has_permission.return_value = False # no view_pii + mock_perms.return_value = perms + result = update_catalog_metadata([ + {"table_group_id": str(uuid4()), "table_name": "orders", "column_name": "ssn", "pii": True}, + ]) + + assert "Failed" in result + assert "view_pii" not in result + assert "permission to view PII" in result + + +def test_description_too_long_rejected(db_session_mock): + _patch(table_rows=[_mock_entity()]) + + from testgen.mcp.tools.data_catalog import update_catalog_metadata + result = update_catalog_metadata([ + {"table_group_id": str(uuid4()), "table_name": "orders", "description": "x" * 1001}, + ]) + + assert "Failed" in result + assert "description" in result.lower() + + +def test_non_boolean_pii_rejected(db_session_mock): + _patch(column_rows=[_mock_entity()]) + + from testgen.mcp.tools.data_catalog import update_catalog_metadata + result = update_catalog_metadata([ + {"table_group_id": str(uuid4()), "table_name": "orders", "column_name": "ssn", "pii": "high"}, + ]) + + assert "Failed" in result + + +def test_unknown_table_not_found(db_session_mock): + _patch(table_rows=[]) # select_where finds nothing + + from testgen.mcp.tools.data_catalog import update_catalog_metadata + result = update_catalog_metadata([ + {"table_group_id": str(uuid4()), "table_name": "ghost", "description": "x"}, + ]) + + assert "Failed" in result + assert "not found" in result.lower() + + +def test_inaccessible_table_group(db_session_mock): + tg, _dt, _dcc = _patch() + # TableGroup.get returns None → resolve_table_group raises not-found-or-not-accessible + patch.stopall() + mock_tg_cls = patch("testgen.mcp.tools.common.TableGroup").start() + mock_tg_cls.get.return_value = None + + from testgen.mcp.tools.data_catalog import update_catalog_metadata + result = update_catalog_metadata([ + {"table_group_id": str(uuid4()), "table_name": "orders", "description": "x"}, + ]) + + assert "Failed" in result + assert "not found or not accessible" in result.lower() + + +# ---------------------------------------------------------------------- +# Side-effect notices +# ---------------------------------------------------------------------- + +def test_cde_write_auto_disables_flag(db_session_mock): + tg = _mock_table_group() + column = _mock_entity() + _patch(tg=tg, column_rows=[column]) + + from testgen.mcp.tools.data_catalog import update_catalog_metadata + result = update_catalog_metadata([ + {"table_group_id": str(uuid4()), "table_name": "orders", "column_name": "amount", "cde": True}, + ]) + + assert tg.profile_flag_cdes is False + assert "Auto-disabled profile_flag_cdes" in result + + +def test_cde_noop_does_not_disable_flag(db_session_mock): + tg = _mock_table_group() + column = _mock_entity() + column.critical_data_element = False # already false; cde: false is a no-op + _patch(tg=tg, column_rows=[column]) + + from testgen.mcp.tools.data_catalog import update_catalog_metadata + result = update_catalog_metadata([ + {"table_group_id": str(uuid4()), "table_name": "orders", "column_name": "amount", "cde": False}, + ]) + + assert tg.profile_flag_cdes is True + assert "Auto-disabled profile_flag_cdes" not in result + + +def test_pii_noop_does_not_disable_flag(db_session_mock): + tg = _mock_table_group() + column = _mock_entity() + column.pii_flag = None # already cleared; pii: false is a no-op + _patch(tg=tg, column_rows=[column]) + + from testgen.mcp.tools.data_catalog import update_catalog_metadata + result = update_catalog_metadata([ + {"table_group_id": str(uuid4()), "table_name": "orders", "column_name": "ssn", "pii": False}, + ]) + + assert tg.profile_flag_pii is True + assert "Auto-disabled profile_flag_pii" not in result + + +def test_table_cde_inheritance_notice(db_session_mock): + _patch(table_rows=[_mock_entity()]) + + from testgen.mcp.tools.data_catalog import update_catalog_metadata + result = update_catalog_metadata([ + {"table_group_id": str(uuid4()), "table_name": "orders", "cde": True}, + ]) + + assert "affects all columns" in result.lower() + + +def test_xde_exclusion_notice(db_session_mock): + _patch(column_rows=[_mock_entity()]) + + from testgen.mcp.tools.data_catalog import update_catalog_metadata + result = update_catalog_metadata([ + {"table_group_id": str(uuid4()), "table_name": "orders", "column_name": "scratch", "xde": True}, + ]) + + assert "excluded" in result.lower() + assert "next profiling run" in result.lower() + assert "skip" in result.lower() + + +def test_empty_fields_skipped(db_session_mock): + _patch(table_rows=[_mock_entity()]) + + from testgen.mcp.tools.data_catalog import update_catalog_metadata + result = update_catalog_metadata([ + {"table_group_id": str(uuid4()), "table_name": "orders"}, + ]) + + assert "Skipped" in result + + +def test_empty_updates_raises(db_session_mock): + from testgen.mcp.exceptions import MCPUserError + from testgen.mcp.tools.data_catalog import update_catalog_metadata + with pytest.raises(MCPUserError): + update_catalog_metadata([]) diff --git a/tests/unit/ui/test_import_metadata.py b/tests/unit/ui/test_import_metadata.py index e5f7f718..627958f8 100644 --- a/tests/unit/ui/test_import_metadata.py +++ b/tests/unit/ui/test_import_metadata.py @@ -1,4 +1,5 @@ import base64 +from unittest.mock import patch import pandas as pd import pytest @@ -6,6 +7,7 @@ from testgen.ui.views.dialogs.import_metadata_dialog import ( DESCRIPTION_MAX_LENGTH, TAG_MAX_LENGTH, + _build_metadata_fields, _extract_metadata_fields, _parse_csv, _set_row_status, @@ -326,3 +328,43 @@ def test_preview_props_unmatched_preserved(): } result = build_import_preview_props(preview) assert result["preview_rows"][0]["_status"] == "unmatched" + + +# --- _build_metadata_fields --- + +_ALL_COLUMNS = ["description", "critical_data_element", "excluded_data_element", "pii_flag", "data_source"] + + +def test_build_fields_table_excludes_column_only_fields(): + row = { + "description": "Orders table", + "critical_data_element": True, + "excluded_data_element": True, + "pii_flag": "MANUAL", + "data_source": "Salesforce", + } + fields = _build_metadata_fields(row, _ALL_COLUMNS, is_column=False) + assert fields == {"description": "Orders table", "critical_data_element": True, "data_source": "Salesforce"} + + +def test_build_fields_blank_text_clears(): + row = {"description": "", "data_source": ""} + fields = _build_metadata_fields(row, ["description", "data_source"], is_column=False) + assert fields == {"description": None, "data_source": None} + + +def test_build_fields_column_includes_pii_when_permitted(): + row = {"pii_flag": "MANUAL", "excluded_data_element": True} + with patch("testgen.ui.views.dialogs.import_metadata_dialog.session") as mock_session: + mock_session.auth.user_has_permission.return_value = True + fields = _build_metadata_fields(row, ["pii_flag", "excluded_data_element"], is_column=True) + assert fields == {"pii_flag": "MANUAL", "excluded_data_element": True} + + +def test_build_fields_column_drops_pii_without_permission(): + row = {"pii_flag": "MANUAL", "excluded_data_element": True} + with patch("testgen.ui.views.dialogs.import_metadata_dialog.session") as mock_session: + mock_session.auth.user_has_permission.return_value = False + fields = _build_metadata_fields(row, ["pii_flag", "excluded_data_element"], is_column=True) + assert fields == {"excluded_data_element": True} + assert "pii_flag" not in fields From 3216d15b61ded240b23089f62886f1ecbe84aae0 Mon Sep 17 00:00:00 2001 From: Diogo Basto Date: Mon, 1 Jun 2026 21:24:43 +0100 Subject: [PATCH 25/78] feat(mcp): add entity read tools for projects, connections, table groups, test suites (TG-1053) Adds six read tools so an MCP client can introspect TestGen configuration before invoking write tools: get_project, list_connections, get_connection, list_table_groups, get_table_group, get_test_suite. All gate on the view permission; get_connection redacts every encrypted field; get_test_suite rejects monitor suites with the unified missing-or-inaccessible wording; sql_flavor_code is rendered as its display label everywhere. Backing model helpers: TestSuite.test_definition_stats aggregates total / locked counts plus per-test-type breakdown in one query; Connection .list_for_project and TableGroup.list_for_project / list_for_connection return paginated config-focused list rows with relevant activity timestamps. Co-Authored-By: Claude Opus 4.7 --- testgen/common/models/connection.py | 40 ++++ testgen/common/models/table_group.py | 96 ++++++++++ testgen/common/models/test_suite.py | 34 +++- testgen/mcp/server.py | 23 ++- testgen/mcp/tools/connections.py | 101 +++++++++++ testgen/mcp/tools/discovery.py | 105 ++++++++++- testgen/mcp/tools/table_groups.py | 159 +++++++++++++++- tests/unit/mcp/test_tools_connections.py | 156 ++++++++++++++++ tests/unit/mcp/test_tools_discovery.py | 183 +++++++++++++++++++ tests/unit/mcp/test_tools_table_groups.py | 211 ++++++++++++++++++++++ 10 files changed, 1101 insertions(+), 7 deletions(-) diff --git a/testgen/common/models/connection.py b/testgen/common/models/connection.py index 2004a67f..863efbb0 100644 --- a/testgen/common/models/connection.py +++ b/testgen/common/models/connection.py @@ -43,6 +43,17 @@ class ConnectionMinimal(EntityMinimal): connection_name: str +@dataclass +class ConnectionListItem(EntityMinimal): + connection_id: int + connection_name: str + project_code: str + sql_flavor_code: SQLFlavorCode + project_host: str | None + project_db: str | None + table_group_count: int + + class Connection(Entity): __tablename__ = "connections" @@ -93,6 +104,35 @@ def select_minimal_where( results = cls._select_columns_where(cls._minimal_columns, *clauses, order_by=order_by) return [ConnectionMinimal(**row) for row in results] + @classmethod + def list_for_project( + cls, project_code: str, *clauses, page: int = 1, limit: int = 20 + ) -> tuple[list[ConnectionListItem], int]: + """Paginated listing of connections in a project, with their table-group counts.""" + tg_count_subq = ( + select( + TableGroup.connection_id.label("connection_id"), + func.count().label("table_group_count"), + ) + .group_by(TableGroup.connection_id) + .subquery() + ) + query = ( + select( + cls.connection_id, + cls.connection_name, + cls.project_code, + cls.sql_flavor_code, + cls.project_host, + cls.project_db, + func.coalesce(tg_count_subq.c.table_group_count, 0).label("table_group_count"), + ) + .outerjoin(tg_count_subq, tg_count_subq.c.connection_id == cls.connection_id) + .where(cls.project_code == project_code, *clauses) + .order_by(*cls._default_order_by) + ) + return cls._paginate(query, page=page, limit=limit, data_class=ConnectionListItem) + @classmethod def is_in_use(cls, ids: list[str]) -> bool: table_groups = TableGroup.select_minimal_where(TableGroup.connection_id.in_(ids)) diff --git a/testgen/common/models/table_group.py b/testgen/common/models/table_group.py index 251dcc46..aa4f9605 100644 --- a/testgen/common/models/table_group.py +++ b/testgen/common/models/table_group.py @@ -47,6 +47,19 @@ class TableGroupStats(EntityMinimal): data_point_ct: int +@dataclass +class TableGroupListItem(EntityMinimal): + id: UUID + table_groups_name: str + table_group_schema: str + project_code: str + connection_name: str | None + table_count: int + last_profiled_date: datetime | None + last_tested_date: datetime | None + testing_score: float | None + + @dataclass class TableGroupSummary(EntityMinimal): id: UUID @@ -385,6 +398,89 @@ def select_summary( total = items[0].total_count if items else 0 return items, total + @classmethod + def list_for_project( + cls, project_code: str, *, page: int = 1, limit: int = 20 + ) -> tuple[list[TableGroupListItem], int]: + """Config-focused paginated listing for a project, with table count and activity timestamps.""" + return cls._list_with_activity( + scope_sql="groups.project_code = :project_code", + scope_params={"project_code": project_code}, + page=page, + limit=limit, + ) + + @classmethod + def list_for_connection( + cls, connection_id: int, *, page: int = 1, limit: int = 20 + ) -> tuple[list[TableGroupListItem], int]: + """Config-focused paginated listing for a connection, with table count and activity timestamps.""" + return cls._list_with_activity( + scope_sql="groups.connection_id = :connection_id", + scope_params={"connection_id": connection_id}, + page=page, + limit=limit, + ) + + @classmethod + def _list_with_activity( + cls, + *, + scope_sql: str, + scope_params: dict, + page: int, + limit: int, + ) -> tuple[list[TableGroupListItem], int]: + params: dict = {**scope_params, "limit": limit, "offset": (page - 1) * limit} + query = f""" + WITH stats AS ( + SELECT table_groups_id, COUNT(*) AS table_count + FROM data_table_chars + GROUP BY table_groups_id + ), + latest_profile AS ( + SELECT pr.table_groups_id, MAX(je.started_at) AS started_at + FROM profiling_runs pr + LEFT JOIN job_executions je ON je.id = pr.job_execution_id + GROUP BY pr.table_groups_id + ), + latest_test AS ( + SELECT ts.table_groups_id, MAX(tr.test_starttime) AS test_starttime + FROM test_runs tr + JOIN test_suites ts ON ts.id = tr.test_suite_id + WHERE ts.is_monitor IS NOT TRUE + GROUP BY ts.table_groups_id + ) + SELECT + groups.id, + groups.table_groups_name, + groups.table_group_schema, + groups.project_code, + connections.connection_name, + COALESCE(stats.table_count, 0) AS table_count, + latest_profile.started_at AS last_profiled_date, + latest_test.test_starttime AS last_tested_date, + groups.dq_score_testing AS testing_score, + COUNT(*) OVER() AS total_count + FROM table_groups AS groups + LEFT JOIN connections ON connections.connection_id = groups.connection_id + LEFT JOIN stats ON stats.table_groups_id = groups.id + LEFT JOIN latest_profile ON latest_profile.table_groups_id = groups.id + LEFT JOIN latest_test ON latest_test.table_groups_id = groups.id + WHERE {scope_sql} + ORDER BY LOWER(groups.table_groups_name) + LIMIT :limit OFFSET :offset; + """ + rows = get_current_session().execute(text(query), params).mappings().all() + if not rows: + return [], 0 + total = int(rows[0]["total_count"]) + items = [ + TableGroupListItem(**{k: v for k, v in row.items() if k != "total_count"}) + for row in rows + ] + return items, total + @classmethod def is_in_use(cls, ids: list[str]) -> bool: test_suites = TestSuite.select_minimal_where(TestSuite.table_groups_id.in_(ids)) diff --git a/testgen/common/models/test_suite.py b/testgen/common/models/test_suite.py index ac29ddce..6465f8e2 100644 --- a/testgen/common/models/test_suite.py +++ b/testgen/common/models/test_suite.py @@ -29,6 +29,15 @@ class TestSuiteMinimal(EntityMinimal): export_to_observability: str +@dataclass(frozen=True) +class TestDefinitionStats: + """Aggregate test-definition counts for a test suite.""" + + total: int + locked: int + counts_by_type: dict[str, int] + + @dataclass class TestSuiteSummary(EntityMinimal): id: UUID @@ -213,6 +222,30 @@ def select_summary(cls, project_code: str, table_group_id: str | UUID | None = N results = db_session.execute(text(query), params).mappings().all() return [TestSuiteSummary(**row) for row in results] + @classmethod + def test_definition_stats(cls, test_suite_id: str | UUID) -> "TestDefinitionStats": + """Aggregate test-definition counts for a suite: total, locked, and per-test-type.""" + query = """ + SELECT + test_type, + COUNT(*) AS total, + SUM(CASE WHEN COALESCE(lock_refresh, 'N') = 'Y' THEN 1 ELSE 0 END) AS locked + FROM test_definitions + WHERE test_suite_id = :test_suite_id + GROUP BY test_type + ORDER BY test_type; + """ + rows = ( + get_current_session() + .execute(text(query), {"test_suite_id": str(test_suite_id)}) + .mappings() + .all() + ) + counts_by_type = {row["test_type"]: int(row["total"]) for row in rows} + total = sum(counts_by_type.values()) + locked = sum(int(row["locked"]) for row in rows) + return TestDefinitionStats(total=total, locked=locked, counts_by_type=counts_by_type) + @classmethod def is_in_use(cls, ids: list[str]) -> bool: query = """ @@ -248,4 +281,3 @@ def cascade_delete(cls, ids: list[str]) -> None: db_session = get_current_session() db_session.execute(text(query), {"test_suite_ids": tuple(ids)}) cls.delete_where(cls.id.in_(ids)) - diff --git a/testgen/mcp/server.py b/testgen/mcp/server.py index c628036a..83682afb 100644 --- a/testgen/mcp/server.py +++ b/testgen/mcp/server.py @@ -140,9 +140,20 @@ def build_mcp_server( profiling_overview, table_health, ) - from testgen.mcp.tools.connections import test_connection + from testgen.mcp.tools.connections import ( + 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, list_projects, list_tables, list_test_suites + from testgen.mcp.tools.discovery import ( + get_data_inventory, + get_project, + get_test_suite, + list_projects, + list_tables, + list_test_suites, + ) from testgen.mcp.tools.execution import ( cancel_profiling_run, cancel_test_run, @@ -208,6 +219,8 @@ def build_mcp_server( from testgen.mcp.tools.source_data import get_source_data, get_source_data_query, get_table_sample from testgen.mcp.tools.table_groups import ( create_table_group, + get_table_group, + list_table_groups, preview_table_group, update_table_group, ) @@ -263,8 +276,10 @@ def safe_prompt(fn): # Tools safe_tool(get_data_inventory) safe_tool(list_projects) + safe_tool(get_project) safe_tool(list_tables) safe_tool(list_test_suites) + safe_tool(get_test_suite) safe_tool(list_test_runs) safe_tool(get_test_run) safe_tool(list_test_results) @@ -329,7 +344,11 @@ def safe_prompt(fn): safe_tool(create_notification) safe_tool(update_notification) safe_tool(delete_notification) + safe_tool(list_connections) + safe_tool(get_connection) safe_tool(test_connection) + safe_tool(list_table_groups) + safe_tool(get_table_group) safe_tool(create_table_group) safe_tool(update_table_group) safe_tool(preview_table_group) diff --git a/testgen/mcp/tools/connections.py b/testgen/mcp/tools/connections.py index 7558ae55..434caf57 100644 --- a/testgen/mcp/tools/connections.py +++ b/testgen/mcp/tools/connections.py @@ -26,14 +26,115 @@ apply_connection_params, connection_display_fields, connection_field_labels, + format_page_footer, + format_page_info, infer_mode, parse_sql_flavor, resolve_connection, validate_connection_fields, + validate_limit, + validate_page, ) from testgen.mcp.tools.markdown import MdDoc +@with_database_session +@mcp_permission("view") +def list_connections(project_code: str, page: int = 1, limit: int = 20) -> str: + """List database connections in a project with their database type and table-group counts. + + Use this before changing or referencing a connection to confirm its ID, name, and host. + Credentials are never returned. + + Args: + project_code: The project code to list connections for. + page: Page number starting at 1 (default 1). + limit: Page size (default 20, max 100). + """ + validate_page(page) + validate_limit(limit, 100) + + perms = get_project_permissions() + perms.verify_access(project_code, not_found=MCPResourceNotAccessible("Project", project_code)) + + rows, total = Connection.list_for_project(project_code, page=page, limit=limit) + + if not rows: + if page > 1: + return f"No connections on page {page} (total: {total})." + return f"No connections found for project `{project_code}`." + + doc = MdDoc() + doc.heading(1, f"Connections for `{project_code}`") + doc.text(format_page_info(total, page, limit)) + table_rows: list[list[object]] = [] + for row in rows: + label = SQL_FLAVOR_CODE_TO_LABEL.get(row.sql_flavor_code) + database_type = label.value if label else row.sql_flavor_code + table_rows.append( + [ + row.connection_id, + row.connection_name, + database_type, + row.project_host, + row.project_db, + row.table_group_count, + ] + ) + doc.table( + ["ID", "Name", "Type", "Host", "Database", "Table groups"], + table_rows, + code=[0, 3, 4], + ) + if footer := format_page_footer(total, page, limit): + doc.text(footer) + return doc.render() + + +@with_database_session +@mcp_permission("view") +def get_connection(connection_id: int) -> str: + """Get a connection's configuration, including database type, host, and authentication mode. + + Credentials (password, private key, service-account key) are never returned. + Use this before editing a connection or creating a table group on it. + + Args: + connection_id: Bigint connection ID returned by `list_connections` or shown on the connections page. + """ + connection = resolve_connection(connection_id) + + doc = MdDoc() + doc.heading(1, f"Connection `{connection.connection_name}`") + doc.field("ID", connection.connection_id, code=True) + doc.field("Project", connection.project_code, code=True) + label = SQL_FLAVOR_CODE_TO_LABEL.get(connection.sql_flavor_code) + doc.field("Type", label.value if label else connection.sql_flavor_code) + + if connection.project_host: + doc.field("Host", connection.project_host, code=True) + if connection.project_port: + doc.field("Port", connection.project_port) + if connection.project_db: + doc.field("Database", connection.project_db, code=True) + if connection.project_user: + doc.field("User", connection.project_user, code=True) + if connection.connect_by_url and connection.url: + doc.field("URL", connection.url, code=True) + if connection.warehouse: + doc.field("Warehouse", connection.warehouse, code=True) + if connection.http_path: + doc.field("HTTP Path", connection.http_path, code=True) + 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) + + return doc.render() + + @with_database_session @mcp_permission("administer") def test_connection( diff --git a/testgen/mcp/tools/discovery.py b/testgen/mcp/tools/discovery.py index 978cc858..bd72f0e9 100644 --- a/testgen/mcp/tools/discovery.py +++ b/testgen/mcp/tools/discovery.py @@ -1,12 +1,21 @@ from testgen.common.mixpanel_service import MixpanelService from testgen.common.models import with_database_session +from testgen.common.models.connection import Connection from testgen.common.models.data_table import DataTable from testgen.common.models.project import Project +from testgen.common.models.table_group import TableGroup from testgen.common.models.test_run import TestRun from testgen.common.models.test_suite import TestSuite from testgen.mcp.exceptions import MCPResourceNotAccessible from testgen.mcp.permissions import get_project_permissions, mcp_permission -from testgen.mcp.tools.common import DocGroup, resolve_table_group, validate_limit, validate_page +from testgen.mcp.tools.common import ( + SQL_FLAVOR_CODE_TO_LABEL, + DocGroup, + resolve_table_group, + resolve_test_suite, + validate_limit, + validate_page, +) from testgen.mcp.tools.markdown import MdDoc _DOC_GROUP = DocGroup.DISCOVER @@ -52,6 +61,49 @@ def list_projects() -> str: return doc.render() +@with_database_session +@mcp_permission("view") +def get_project(project_code: str) -> str: + """Get a project's configuration and configuration counts. + + Returns the project name, observability and data-retention settings, and counts + of connections, table groups, test suites, test definitions, profiling runs, and + test runs scoped to the project. Use this before configuration changes to confirm + a project's current shape. + + Args: + project_code: The project code, e.g. from `list_projects`. + """ + perms = get_project_permissions() + perms.verify_access(project_code, not_found=MCPResourceNotAccessible("Project", project_code)) + + project = Project.get(project_code) + summary = Project.get_summary(project_code) + if project is None or summary is None: + raise MCPResourceNotAccessible("Project", project_code) + + doc = MdDoc() + doc.heading(1, f"Project `{project_code}`") + doc.field("Name", project.project_name) + doc.field("Connections", summary.connection_count) + doc.field("Table groups", summary.table_group_count) + doc.field("Test suites", summary.test_suite_count) + doc.field("Test definitions", summary.test_definition_count) + doc.field("Profiling runs", summary.profiling_run_count) + doc.field("Test runs", summary.test_run_count) + + doc.heading(2, "Configuration") + doc.field("Quality-score weighting", project.use_dq_score_weights) + doc.field("Data retention", project.data_retention_enabled) + if project.data_retention_enabled and project.data_retention_days is not None: + doc.field("Retention period (days)", project.data_retention_days) + doc.field("Observability export configured", summary.can_export_to_observability) + if project.observability_api_url: + doc.field("Observability API URL", project.observability_api_url, code=True) + + return doc.render() + + @with_database_session @mcp_permission("view") def list_test_suites(project_code: str) -> str: @@ -104,6 +156,57 @@ def list_test_suites(project_code: str) -> str: return doc.render() +@with_database_session +@mcp_permission("view") +def get_test_suite(test_suite_id: str) -> str: + """Get a test suite's configuration: default severity, connection, table group, and per-test-type counts. + + Returns the test suite's identity and configuration along with a breakdown of how many test + definitions it contains by type and how many are lock-refreshed (excluded from regeneration). + Use this before changing a suite's tests to understand what will be affected. + + Args: + test_suite_id: The test suite UUID, e.g. from `list_test_suites`. + """ + suite = resolve_test_suite(test_suite_id) + connection = Connection.get(suite.connection_id) if suite.connection_id else None + table_group = TableGroup.get(suite.table_groups_id) if suite.table_groups_id else None + stats = TestSuite.test_definition_stats(suite.id) + + doc = MdDoc() + doc.heading(1, f"Test Suite `{suite.test_suite}`") + doc.field("ID", str(suite.id), code=True) + doc.field("Project", suite.project_code, code=True) + + if connection is not None: + label = SQL_FLAVOR_CODE_TO_LABEL.get(connection.sql_flavor_code) + flavor_label = label.value if label else connection.sql_flavor_code + doc.field( + "Connection", + f"{connection.connection_name} (`{connection.connection_id}`, {flavor_label})", + ) + if table_group is not None: + doc.field( + "Table group", + f"{table_group.table_groups_name} (`{table_group.id}`)", + ) + + if suite.test_suite_description: + doc.field("Description", suite.test_suite_description) + doc.field("Default severity", suite.severity or "Inherit from test type") + doc.field("Export to observability", suite.export_to_observability) + + doc.field("Total tests", stats.total) + doc.field("Locked tests", stats.locked) + + if stats.counts_by_type: + doc.heading(2, "Tests by type") + rows = [[test_type, count] for test_type, count in stats.counts_by_type.items()] + doc.table(["Test type", "Count"], rows, code=[0]) + + return doc.render() + + @with_database_session @mcp_permission("catalog") def list_tables(table_group_id: str, limit: int = 200, page: int = 1) -> str: diff --git a/testgen/mcp/tools/table_groups.py b/testgen/mcp/tools/table_groups.py index 1485055c..dba9ea91 100644 --- a/testgen/mcp/tools/table_groups.py +++ b/testgen/mcp/tools/table_groups.py @@ -18,9 +18,17 @@ from testgen.common.models import with_database_session from testgen.common.models.connection import Connection from testgen.common.models.table_group import TableGroup -from testgen.mcp.exceptions import MCPUserError -from testgen.mcp.permissions import mcp_permission -from testgen.mcp.tools.common import resolve_connection, resolve_table_group +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, + format_page_footer, + format_page_info, + resolve_connection, + resolve_table_group, + validate_limit, + validate_page, +) from testgen.mcp.tools.markdown import MdDoc _DUPLICATE_NAME_MESSAGE = "A Table Group with the same name already exists." @@ -30,6 +38,151 @@ ) +@with_database_session +@mcp_permission("view") +def list_table_groups( + project_code: str | None = None, + connection_id: int | None = None, + page: int = 1, + limit: int = 20, +) -> str: + """List table groups in a project or on a specific connection. + + Pass exactly one of `project_code` or `connection_id`. Returns each group's + table count, last profile / test timestamps, and current Testing Score. + + Args: + project_code: List groups in a project. + connection_id: List groups on a specific connection. + page: Page number starting at 1 (default 1). + limit: Page size (default 20, max 100). + """ + if (project_code is None) == (connection_id is None): + raise MCPUserError("Pass either `project_code` or `connection_id`, not both.") + validate_page(page) + validate_limit(limit, 100) + + perms = get_project_permissions() + if project_code is not None: + perms.verify_access(project_code, not_found=MCPResourceNotAccessible("Project", project_code)) + rows, total = TableGroup.list_for_project(project_code, page=page, limit=limit) + heading = f"Table groups for project `{project_code}`" + else: + connection = resolve_connection(connection_id) + rows, total = TableGroup.list_for_connection(connection.connection_id, page=page, limit=limit) + heading = f"Table groups on connection `{connection.connection_name}` (`{connection.connection_id}`)" + + if not rows: + if page > 1: + return f"No table groups on page {page} (total: {total})." + return f"{heading} — none found." + + doc = MdDoc() + doc.heading(1, heading) + doc.text(format_page_info(total, page, limit)) + table_rows: list[list[object]] = [] + for row in rows: + score = f"{row.testing_score:.2f}" if row.testing_score is not None else None + table_rows.append( + [ + str(row.id), + row.table_groups_name, + row.connection_name, + row.table_group_schema, + row.table_count, + row.last_profiled_date, + row.last_tested_date, + score, + ] + ) + doc.table( + ["ID", "Name", "Connection", "Schema", "Tables", "Last profiled", "Last tested", "Testing Score"], + table_rows, + code=[0, 3], + ) + if footer := format_page_footer(total, page, limit): + doc.text(footer) + return doc.render() + + +@with_database_session +@mcp_permission("view") +def get_table_group(table_group_id: str) -> str: + """Get a table group's full configuration: filters, sampling, profiling flags, catalog tags, and recent activity. + + Mirrors the field set of the table-group edit dialog so an LLM can introspect or modify + every setting available in the UI. Use this before editing a table group or generating tests. + + Args: + table_group_id: The table group UUID, e.g. from `list_table_groups` or `get_data_inventory`. + """ + table_group = resolve_table_group(table_group_id) + connection = Connection.get(table_group.connection_id) if table_group.connection_id else None + + doc = MdDoc() + doc.heading(1, f"Table group `{table_group.table_groups_name}`") + doc.field("ID", str(table_group.id), code=True) + doc.field("Project", table_group.project_code, code=True) + if connection is not None: + label = SQL_FLAVOR_CODE_TO_LABEL.get(connection.sql_flavor_code) + flavor_label = label.value if label else connection.sql_flavor_code + doc.field( + "Connection", + f"{connection.connection_name} (`{connection.connection_id}`, {flavor_label})", + ) + doc.field("Schema", table_group.table_group_schema, code=True) + if table_group.description: + doc.field("Description", table_group.description) + + doc.heading(2, "Criteria") + if table_group.profiling_table_set: + doc.field("Explicit table list", table_group.profiling_table_set, code=True) + if table_group.profiling_include_mask: + doc.field("Tables to include mask", table_group.profiling_include_mask, code=True) + if table_group.profiling_exclude_mask: + doc.field("Tables to exclude mask", table_group.profiling_exclude_mask, code=True) + doc.field("Profiling ID column mask", table_group.profile_id_column_mask, code=True) + doc.field("Profiling surrogate-key column mask", table_group.profile_sk_column_mask, code=True) + + doc.heading(2, "Settings") + doc.field("Detect CDE during profiling", table_group.profile_flag_cdes) + doc.field("Detect PII during profiling", table_group.profile_flag_pii) + doc.field("Exclude XDE columns from profiling", table_group.profile_exclude_xde) + doc.field("Include in project dashboard", table_group.include_in_dashboard) + doc.field("Min profiling age (days)", table_group.profiling_delay_days) + + doc.heading(2, "Sampling parameters") + doc.field("Use profile sampling", table_group.profile_use_sampling) + if table_group.profile_use_sampling: + doc.field("Sample percent", table_group.profile_sample_percent) + doc.field("Min sample record count", table_group.profile_sample_min_count) + + catalog_fields: list[tuple[str, str | None]] = [ + ("Data source", table_group.data_source), + ("Source system", table_group.source_system), + ("Source process", table_group.source_process), + ("Business domain", table_group.business_domain), + ("Transform level", table_group.transform_level), + ("Data location", table_group.data_location), + ("Stakeholder group", table_group.stakeholder_group), + ("Data product", table_group.data_product), + ] + if any(value for _, value in catalog_fields): + doc.heading(2, "Catalog tags") + for label, value in catalog_fields: + if value: + doc.field(label, value) + + if table_group.dq_score_testing is not None or table_group.dq_score_profiling is not None: + doc.heading(2, "Latest activity") + if table_group.dq_score_testing is not None: + doc.field("Testing Score", f"{table_group.dq_score_testing:.2f}") + if table_group.dq_score_profiling is not None: + doc.field("Profiling Score", f"{table_group.dq_score_profiling:.2f}") + + return doc.render() + + @with_database_session @mcp_permission("edit") def create_table_group( diff --git a/tests/unit/mcp/test_tools_connections.py b/tests/unit/mcp/test_tools_connections.py index 3533023e..e6b77642 100644 --- a/tests/unit/mcp/test_tools_connections.py +++ b/tests/unit/mcp/test_tools_connections.py @@ -231,3 +231,159 @@ def test_test_connection_inline_does_not_save(mock_conn_cls, mock_runner, db_ses ) inline.save.assert_not_called() + + +# --------------------------------------------------------------------------- +# list_connections +# --------------------------------------------------------------------------- + + +def _list_item(**overrides): + from testgen.common.models.connection import ConnectionListItem + + return ConnectionListItem( + connection_id=overrides.get("connection_id", 1), + connection_name=overrides.get("connection_name", "warehouse_prod"), + project_code=overrides.get("project_code", "demo"), + sql_flavor_code=overrides.get("sql_flavor_code", "snowflake"), + project_host=overrides.get("project_host", "abc.snowflakecomputing.com"), + project_db=overrides.get("project_db", "ANALYTICS"), + table_group_count=overrides.get("table_group_count", 4), + ) + + +@patch(f"{MODULE}.Connection") +def test_list_connections_returns_rows_with_display_label(mock_conn_cls, db_session_mock): + mock_conn_cls.list_for_project.return_value = ( + [_list_item(project_host="warehouse-host.example.com")], + 1, + ) + + from testgen.mcp.tools.connections import list_connections + + with _patch_perms(permission="view"): + out = list_connections("demo") + + assert "Connections for `demo`" in out + assert "warehouse_prod" in out + # Display label is rendered, NOT the raw sql_flavor_code: + assert "Snowflake" in out + assert "snowflake" not in out + assert "warehouse-host.example.com" in out + assert "ANALYTICS" in out + assert "4" in out # table_group_count + + +@patch(f"{MODULE}.Connection") +def test_list_connections_pagination_footer(mock_conn_cls, db_session_mock): + rows = [_list_item(connection_id=i, connection_name=f"c{i}") for i in range(1, 3)] + mock_conn_cls.list_for_project.return_value = (rows, 25) + + from testgen.mcp.tools.connections import list_connections + + with _patch_perms(permission="view"): + out = list_connections("demo", page=1, limit=2) + + assert "Showing 1\u20132 of 25" in out + assert "Use `page=2`" in out + + +@patch(f"{MODULE}.Connection") +def test_list_connections_empty(mock_conn_cls, db_session_mock): + mock_conn_cls.list_for_project.return_value = ([], 0) + + from testgen.mcp.tools.connections import list_connections + + with _patch_perms(permission="view"): + out = list_connections("demo") + + assert "No connections" in out + + +@patch("testgen.mcp.permissions._compute_project_permissions") +def test_list_connections_rejects_inaccessible_project(mock_compute, db_session_mock): + mock_compute.return_value = ProjectPermissions( + memberships={"other": "role_a"}, permission="view", username="test_user", + ) + + from testgen.mcp.tools.connections import list_connections + + with pytest.raises(MCPResourceNotAccessible, match="Project `secret` not found or not accessible"): + list_connections("secret") + + +def test_list_connections_invalid_limit_raises(db_session_mock): + from testgen.mcp.tools.connections import list_connections + + with _patch_perms(permission="view"), pytest.raises(MCPUserError, match="Invalid limit"): + list_connections("demo", limit=500) + + +# --------------------------------------------------------------------------- +# get_connection +# --------------------------------------------------------------------------- + + +@patch(f"{MODULE}.resolve_connection") +def test_get_connection_returns_config_with_display_label(mock_resolve, db_session_mock): + mock_resolve.return_value = _mock_connection( + sql_flavor_code="snowflake", + connection_name="warehouse_prod", + connection_id=12, + ) + + from testgen.mcp.tools.connections import get_connection + + with _patch_perms(permission="view"): + out = get_connection(12) + + assert "Connection `warehouse_prod`" in out + assert "**ID:** `12`" in out + assert "**Type:** Snowflake" in out + # The raw flavor code MUST NOT appear in the output + assert "snowflake" not in out + + +@patch(f"{MODULE}.resolve_connection") +def test_get_connection_never_returns_secrets(mock_resolve, db_session_mock): + """Password, private key, passphrase, service-account key must not leak.""" + mock_resolve.return_value = _mock_connection( + project_pw_encrypted="hunter2", + private_key="-----BEGIN PRIVATE KEY-----secret", + private_key_passphrase="passphrase", # noqa: S106 + service_account_key={"key": "secret-account-json"}, + ) + + from testgen.mcp.tools.connections import get_connection + + with _patch_perms(permission="view"): + out = get_connection(42) + + assert "hunter2" not in out + assert "BEGIN PRIVATE KEY" not in out + assert "passphrase" not in out + assert "secret-account-json" not in out + + +@patch(f"{MODULE}.resolve_connection") +def test_get_connection_authentication_password(mock_resolve, db_session_mock): + mock_resolve.return_value = _mock_connection( + connect_by_key=False, connect_by_url=False, connect_with_identity=False, + ) + + from testgen.mcp.tools.connections import get_connection + + with _patch_perms(permission="view"): + out = get_connection(42) + + assert "**Authentication:** Password" in out + + +@patch("testgen.mcp.tools.common.Connection") +def test_get_connection_raises_not_found_for_inaccessible(mock_conn_cls, db_session_mock): + mock_conn_cls.get.return_value = None + + from testgen.mcp.tools.connections import get_connection + + with pytest.raises(MCPResourceNotAccessible, match="Connection .* not found or not accessible"): + get_connection(999) diff --git a/tests/unit/mcp/test_tools_discovery.py b/tests/unit/mcp/test_tools_discovery.py index 409b20f2..b1a09b1d 100644 --- a/tests/unit/mcp/test_tools_discovery.py +++ b/tests/unit/mcp/test_tools_discovery.py @@ -229,3 +229,186 @@ def test_list_tables_scopes_data_lookup_to_resolved_tg_project(mock_tg_cls, mock call_kwargs = mock_dt.select_table_names.call_args assert call_kwargs.kwargs["project_codes"] == ["proj_a"] + + +# --------------------------------------------------------------------------- +# get_project +# --------------------------------------------------------------------------- + + +def _mock_project(**overrides): + project = MagicMock() + project.project_code = overrides.get("project_code", "demo") + project.project_name = overrides.get("project_name", "Demo Project") + project.use_dq_score_weights = overrides.get("use_dq_score_weights", True) + project.data_retention_enabled = overrides.get("data_retention_enabled", True) + project.data_retention_days = overrides.get("data_retention_days", 180) + project.observability_api_url = overrides.get("observability_api_url", None) + return project + + +def _mock_project_summary(**overrides): + summary = MagicMock() + summary.connection_count = overrides.get("connection_count", 3) + summary.table_group_count = overrides.get("table_group_count", 5) + summary.test_suite_count = overrides.get("test_suite_count", 7) + summary.test_definition_count = overrides.get("test_definition_count", 142) + summary.profiling_run_count = overrides.get("profiling_run_count", 12) + summary.test_run_count = overrides.get("test_run_count", 38) + summary.can_export_to_observability = overrides.get("can_export_to_observability", False) + return summary + + +@patch("testgen.mcp.tools.discovery.Project") +def test_get_project_returns_counts_and_config(mock_project_cls, db_session_mock): + mock_project_cls.get.return_value = _mock_project() + mock_project_cls.get_summary.return_value = _mock_project_summary() + + from testgen.mcp.tools.discovery import get_project + + out = get_project("demo") + + assert "Project `demo`" in out + assert "Demo Project" in out + assert "**Connections:** 3" in out + assert "**Table groups:** 5" in out + assert "**Test suites:** 7" in out + assert "**Test definitions:** 142" in out + assert "**Profiling runs:** 12" in out + assert "**Test runs:** 38" in out + assert "**Quality-score weighting:** Yes" in out + assert "**Data retention:** Yes" in out + assert "**Retention period (days):** 180" in out + + +@patch("testgen.mcp.permissions._compute_project_permissions") +def test_get_project_raises_not_found_for_inaccessible(mock_compute, db_session_mock): + mock_compute.return_value = ProjectPermissions( + memberships={"other": "role_a"}, permission="view", username="test_user", + ) + + from testgen.mcp.tools.discovery import get_project + + with pytest.raises(MCPResourceNotAccessible, match="Project `secret` not found or not accessible"): + get_project("secret") + + +@patch("testgen.mcp.tools.discovery.Project") +def test_get_project_raises_not_found_when_missing(mock_project_cls, db_session_mock): + mock_project_cls.get.return_value = None + mock_project_cls.get_summary.return_value = None + + from testgen.mcp.tools.discovery import get_project + + with pytest.raises(MCPResourceNotAccessible, match="Project `demo` not found or not accessible"): + get_project("demo") + + +# --------------------------------------------------------------------------- +# get_test_suite +# --------------------------------------------------------------------------- + + +def _mock_test_suite(**overrides): + suite = MagicMock() + suite.id = overrides.get("id", uuid4()) + suite.test_suite = overrides.get("test_suite", "qa_checks") + suite.project_code = overrides.get("project_code", "demo") + suite.connection_id = overrides.get("connection_id", 42) + suite.table_groups_id = overrides.get("table_groups_id", uuid4()) + suite.test_suite_description = overrides.get("test_suite_description", "Daily QA") + suite.severity = overrides.get("severity", "Warning") + suite.export_to_observability = overrides.get("export_to_observability", True) + suite.is_monitor = overrides.get("is_monitor", False) + return suite + + +_DEFAULT_TYPE_COUNTS = {"Aggregate_Balance": 5, "Row_Count": 12} + + +def _stats(total=47, locked=3, counts_by_type=None): + from testgen.common.models.test_suite import TestDefinitionStats + + return TestDefinitionStats( + total=total, + locked=locked, + counts_by_type=_DEFAULT_TYPE_COUNTS if counts_by_type is None else counts_by_type, + ) + + +@patch("testgen.mcp.tools.discovery.TestSuite") +@patch("testgen.mcp.tools.discovery.TableGroup") +@patch("testgen.mcp.tools.discovery.Connection") +@patch("testgen.mcp.tools.common.TestSuite") +def test_get_test_suite_returns_full_config( + mock_common_suite, mock_conn_cls, mock_tg_cls, mock_suite_cls, db_session_mock, +): + suite = _mock_test_suite() + mock_common_suite.get.return_value = suite + conn = MagicMock(connection_id=42, connection_name="warehouse_prod", sql_flavor_code="snowflake") + mock_conn_cls.get.return_value = conn + tg = MagicMock(table_groups_name="curated_payments") + tg.id = suite.table_groups_id + mock_tg_cls.get.return_value = tg + mock_suite_cls.test_definition_stats.return_value = _stats() + + from testgen.mcp.tools.discovery import get_test_suite + + out = get_test_suite(str(suite.id)) + + assert "Test Suite `qa_checks`" in out + assert f"**ID:** `{suite.id}`" in out + assert "**Project:** `demo`" in out + assert "warehouse_prod" in out + assert "Snowflake" in out + assert "curated_payments" in out + assert "**Default severity:** Warning" in out + assert "**Export to observability:** Yes" in out + assert "**Total tests:** 47" in out + assert "**Locked tests:** 3" in out + assert "Tests by type" in out + assert "Aggregate_Balance" in out + assert "Row_Count" in out + + +@patch("testgen.mcp.tools.common.TestSuite") +def test_get_test_suite_rejects_monitor_suite(mock_common_suite, db_session_mock): + """resolve_test_suite returns None when suite is_monitor — uses unified wording.""" + mock_common_suite.get.return_value = None + + from testgen.mcp.tools.discovery import get_test_suite + + bogus_id = str(uuid4()) + with pytest.raises(MCPResourceNotAccessible, match="Test suite .* not found or not accessible"): + get_test_suite(bogus_id) + + +@patch("testgen.mcp.tools.common.TestSuite") +def test_get_test_suite_rejects_invalid_uuid(mock_common_suite, db_session_mock): + from testgen.mcp.exceptions import MCPUserError as _MCPUserError + from testgen.mcp.tools.discovery import get_test_suite + + with pytest.raises(_MCPUserError, match="not a valid UUID"): + get_test_suite("not-a-uuid") + + +@patch("testgen.mcp.tools.discovery.TestSuite") +@patch("testgen.mcp.tools.discovery.TableGroup") +@patch("testgen.mcp.tools.discovery.Connection") +@patch("testgen.mcp.tools.common.TestSuite") +def test_get_test_suite_no_severity_renders_inherit( + mock_common_suite, mock_conn_cls, mock_tg_cls, mock_suite_cls, db_session_mock, +): + suite = _mock_test_suite(severity=None) + mock_common_suite.get.return_value = suite + mock_conn_cls.get.return_value = None + mock_tg_cls.get.return_value = None + mock_suite_cls.test_definition_stats.return_value = _stats(total=0, locked=0, counts_by_type={}) + + from testgen.mcp.tools.discovery import get_test_suite + + out = get_test_suite(str(suite.id)) + + assert "Inherit from test type" in out + # No "Tests by type" table when there are no test definitions + assert "Tests by type" not in out diff --git a/tests/unit/mcp/test_tools_table_groups.py b/tests/unit/mcp/test_tools_table_groups.py index 9c78c95a..ead9576b 100644 --- a/tests/unit/mcp/test_tools_table_groups.py +++ b/tests/unit/mcp/test_tools_table_groups.py @@ -755,3 +755,214 @@ def test_preview_requires_edit(db_session_mock): with _patch_perms(memberships={"demo": "role_c"}), pytest.raises(MCPPermissionDenied): preview_table_group(table_group_id=str(uuid4())) + + +# --------------------------------------------------------------------------- +# list_table_groups +# --------------------------------------------------------------------------- + + +def _list_item(**overrides): + from testgen.common.models.table_group import TableGroupListItem + + return TableGroupListItem( + id=overrides.get("id", uuid4()), + table_groups_name=overrides.get("table_groups_name", "core_tables"), + table_group_schema=overrides.get("table_group_schema", "public"), + project_code=overrides.get("project_code", "demo"), + connection_name=overrides.get("connection_name", "warehouse_prod"), + table_count=overrides.get("table_count", 12), + last_profiled_date=overrides.get("last_profiled_date", None), + last_tested_date=overrides.get("last_tested_date", None), + testing_score=overrides.get("testing_score", 0.97), + ) + + +def test_list_table_groups_requires_one_arg(db_session_mock): + from testgen.mcp.tools.table_groups import list_table_groups + + with _patch_perms(permission="view"), pytest.raises( + MCPUserError, match="Pass either `project_code` or `connection_id`", + ): + list_table_groups() + + +def test_list_table_groups_rejects_both_args(db_session_mock): + from testgen.mcp.tools.table_groups import list_table_groups + + with _patch_perms(permission="view"), pytest.raises( + MCPUserError, match="Pass either `project_code` or `connection_id`", + ): + list_table_groups(project_code="demo", connection_id=12) + + +@patch(f"{MODULE}.TableGroup") +def test_list_table_groups_by_project_renders_rows(mock_tg_cls, db_session_mock): + mock_tg_cls.list_for_project.return_value = ([_list_item()], 1) + + from testgen.mcp.tools.table_groups import list_table_groups + + with _patch_perms(permission="view"): + out = list_table_groups(project_code="demo") + + assert "Table groups for project `demo`" in out + assert "core_tables" in out + assert "warehouse_prod" in out + assert "public" in out + assert "12" in out # table count + assert "0.97" in out # testing score formatted + mock_tg_cls.list_for_project.assert_called_once_with("demo", page=1, limit=20) + + +@patch(f"{MODULE}.resolve_connection") +@patch(f"{MODULE}.TableGroup") +def test_list_table_groups_by_connection_renders_rows(mock_tg_cls, mock_resolve, db_session_mock): + conn = _mock_connection() + conn.connection_id = 42 + conn.connection_name = "warehouse_prod" + mock_resolve.return_value = conn + mock_tg_cls.list_for_connection.return_value = ([_list_item()], 1) + + from testgen.mcp.tools.table_groups import list_table_groups + + with _patch_perms(permission="view"): + out = list_table_groups(connection_id=42) + + assert "Table groups on connection `warehouse_prod` (`42`)" in out + mock_tg_cls.list_for_connection.assert_called_once_with(42, page=1, limit=20) + + +@patch(f"{MODULE}.TableGroup") +def test_list_table_groups_empty(mock_tg_cls, db_session_mock): + mock_tg_cls.list_for_project.return_value = ([], 0) + + from testgen.mcp.tools.table_groups import list_table_groups + + with _patch_perms(permission="view"): + out = list_table_groups(project_code="demo") + + assert "none found" in out + + +@patch("testgen.mcp.permissions._compute_project_permissions") +def test_list_table_groups_rejects_inaccessible_project(mock_compute, db_session_mock): + mock_compute.return_value = ProjectPermissions( + memberships={"other": "role_a"}, permission="view", username="test_user", + ) + + from testgen.mcp.tools.table_groups import list_table_groups + + with pytest.raises(MCPResourceNotAccessible, match="Project `secret` not found or not accessible"): + list_table_groups(project_code="secret") + + +# --------------------------------------------------------------------------- +# get_table_group +# --------------------------------------------------------------------------- + + +def _read_mock_table_group(**overrides): + """Variant of _mock_table_group that also sets dq score fields used by get_table_group.""" + tg = _mock_table_group(**overrides) + tg.dq_score_testing = overrides.get("dq_score_testing", None) + tg.dq_score_profiling = overrides.get("dq_score_profiling", None) + return tg + + +@patch(f"{MODULE}.Connection") +@patch(f"{MODULE}.resolve_table_group") +def test_get_table_group_renders_all_dialog_sections(mock_resolve, mock_conn_cls, db_session_mock): + tg = _read_mock_table_group( + description="Curated payments tables", + profiling_include_mask="payments_%", + profiling_exclude_mask="tmp_%", + profiling_table_set="payments,refunds", + profile_use_sampling=True, + profile_sample_percent="50", + data_source="DataKitchen", + business_domain="Finance", + dq_score_testing=0.91, + dq_score_profiling=0.95, + ) + mock_resolve.return_value = tg + conn = _mock_connection(connection_name="warehouse_prod", sql_flavor_code="snowflake") + mock_conn_cls.get.return_value = conn + + from testgen.mcp.tools.table_groups import get_table_group + + with _patch_perms(permission="view"): + out = get_table_group(str(tg.id)) + + # Identity + assert "Table group `Sample TG`" in out + assert "warehouse_prod" in out + assert "Snowflake" in out + assert "`public`" in out + assert "Curated payments tables" in out + # Criteria — both table-name and column-name masks rendered + assert "## Criteria" in out + assert "payments_%" in out + assert "tmp_%" in out + assert "payments,refunds" in out + assert "Profiling ID column mask" in out + assert "Profiling surrogate-key column mask" in out + # Settings + assert "## Settings" in out + assert "Detect CDE during profiling" in out + # Sampling enabled → percent + min count rendered + assert "## Sampling parameters" in out + assert "**Sample percent:** 50" in out + assert "Min sample record count" in out + # Catalog tags only render when set + assert "## Catalog tags" in out + assert "DataKitchen" in out + assert "Finance" in out + # Latest activity + assert "## Latest activity" in out + assert "0.91" in out + assert "0.95" in out + + +@patch(f"{MODULE}.Connection") +@patch(f"{MODULE}.resolve_table_group") +def test_get_table_group_skips_catalog_when_no_tags(mock_resolve, mock_conn_cls, db_session_mock): + tg = _read_mock_table_group() # all catalog tags None + mock_resolve.return_value = tg + mock_conn_cls.get.return_value = _mock_connection() + + from testgen.mcp.tools.table_groups import get_table_group + + with _patch_perms(permission="view"): + out = get_table_group(str(tg.id)) + + assert "## Catalog tags" not in out + + +@patch(f"{MODULE}.Connection") +@patch(f"{MODULE}.resolve_table_group") +def test_get_table_group_skips_sample_details_when_sampling_off( + mock_resolve, mock_conn_cls, db_session_mock, +): + tg = _read_mock_table_group(profile_use_sampling=False) + mock_resolve.return_value = tg + mock_conn_cls.get.return_value = _mock_connection() + + from testgen.mcp.tools.table_groups import get_table_group + + with _patch_perms(permission="view"): + out = get_table_group(str(tg.id)) + + # Sampling section header still present (with the toggle), but percent/min not shown + assert "Use profile sampling" in out + assert "**Sample percent:**" not in out + assert "Min sample record count" not in out + + +@patch("testgen.mcp.tools.common.TableGroup") +def test_get_table_group_raises_not_found_for_inaccessible(mock_tg_cls, db_session_mock): + mock_tg_cls.get.return_value = None + + from testgen.mcp.tools.table_groups import get_table_group + + with pytest.raises(MCPResourceNotAccessible, match="Table group .* not found or not accessible"): + get_table_group(str(uuid4())) From 13eff2bec1e6ad707819d670f485deaca841f5c5 Mon Sep 17 00:00:00 2001 From: Diogo Basto Date: Tue, 9 Jun 2026 18:18:15 +0100 Subject: [PATCH 26/78] refactor(mcp): apply TG-1053 review feedback Correctness: - Project.get_summary: exclude is_monitor suites from test_suite_count / test_definition_count / test_run_count subqueries (the rule applied per CLAUDE.md monitor-exclusion). - TableGroup._list_with_activity: separate COUNT(*) so total is correct on out-of-range pages (window function over an empty page returned 0). - TestSuite.test_definition_stats: render user-facing test_name_short labels via LEFT JOIN to test_types instead of leaking raw test_type codes. Output shape: - list_table_groups: expose Columns + Rows + all 3 quality scores (profiling, testing, overall), column header is "Quality Score" (overall). - get_table_group: drop dev-facing first sentence; field labels via the _DIFF_LABELS map for parity with create/update tools; Latest activity now includes the calculated overall Quality Score via a new TableGroup.quality_score property. - get_connection: rendering consolidated through a new _render_connection_body helper shared with _render_created_connection. - get_project: restructured into UI-mirroring sections (Configuration / Observability Integration / Data Retention) with UI-matching labels. Defense in depth: - get_test_suite / get_table_group route Connection / TableGroup lookups through resolve_connection / resolve_table_group instead of Model.get(). - format_flavor_label helper centralizes the SQL_FLAVOR_CODE_TO_LABEL + fallback pattern used across the read tools. Tests: - get_connection redaction test uses a real Connection with EncryptedBytea fields populated; previously vacuous (mocks bypassed the columns entirely). - get_test_suite monitor-rejection test now drives an actual is_monitor=True suite through resolve_test_suite's filter (was only exercising the missing- entity path). - Updated mocks and assertions for the new label/field set. Docs: - mcp-patterns.md followable-IDs table updated with the new producers / consumers introduced by this MR. Co-Authored-By: Claude Opus 4.7 --- testgen/common/models/project.py | 6 +- testgen/common/models/table_group.py | 54 ++++++++++++---- testgen/common/models/test_suite.py | 22 ++++--- testgen/mcp/tools/common.py | 12 ++++ testgen/mcp/tools/connections.py | 53 +++++----------- testgen/mcp/tools/discovery.py | 40 ++++++------ testgen/mcp/tools/table_groups.py | 77 +++++++++++------------ tests/unit/mcp/test_tools_connections.py | 50 ++++++++++++--- tests/unit/mcp/test_tools_discovery.py | 71 +++++++++++++++------ tests/unit/mcp/test_tools_table_groups.py | 67 ++++++++++++-------- 10 files changed, 282 insertions(+), 170 deletions(-) diff --git a/testgen/common/models/project.py b/testgen/common/models/project.py index 40d04974..bb132802 100644 --- a/testgen/common/models/project.py +++ b/testgen/common/models/project.py @@ -66,19 +66,23 @@ def get_summary(cls, project_code: str) -> ProjectSummary | None: WHERE table_groups.project_code = :project_code ) AS profiling_run_count, ( - SELECT COUNT(*) FROM test_suites WHERE test_suites.project_code = :project_code + SELECT COUNT(*) FROM test_suites + WHERE test_suites.project_code = :project_code + AND test_suites.is_monitor IS NOT TRUE ) AS test_suite_count, ( SELECT COUNT(*) FROM test_definitions LEFT JOIN test_suites ON test_definitions.test_suite_id = test_suites.id WHERE test_suites.project_code = :project_code + AND test_suites.is_monitor IS NOT TRUE ) AS test_definition_count, ( SELECT COUNT(*) FROM test_runs LEFT JOIN test_suites ON test_runs.test_suite_id = test_suites.id WHERE test_suites.project_code = :project_code + AND test_suites.is_monitor IS NOT TRUE ) AS test_run_count, ( SELECT COALESCE(observability_api_key, '') <> '' diff --git a/testgen/common/models/table_group.py b/testgen/common/models/table_group.py index aa4f9605..f30ab780 100644 --- a/testgen/common/models/table_group.py +++ b/testgen/common/models/table_group.py @@ -55,9 +55,13 @@ class TableGroupListItem(EntityMinimal): project_code: str connection_name: str | None table_count: int + column_count: int | None + row_count: int | None last_profiled_date: datetime | None last_tested_date: datetime | None + profiling_score: float | None testing_score: float | None + quality_score: float | None @dataclass @@ -162,6 +166,20 @@ class TableGroup(Entity): dq_score_testing, ) + @property + def quality_score(self) -> float: + """Overall quality score per ``utils.score``: profiling * testing when both + exist; the non-null one when only one ran; ``0.0`` when neither has run. + + Mirrors what Project Dashboard, data_catalog, and inventory_service display. + Render through ``utils.friendly_score`` to get the user-facing percentage form + (``"95.0"`` rather than ``"0.95"``). The list-side equivalent is computed in + SQL inside ``_list_with_activity`` — keep both in sync if the formula changes. + """ + from testgen.utils import score + + return score(self.dq_score_profiling, self.dq_score_testing) + @classmethod def get_minimal(cls, id_: str | UUID) -> TableGroupMinimal | None: result = cls._get_columns(id_, cls._minimal_columns) @@ -431,10 +449,22 @@ def _list_with_activity( page: int, limit: int, ) -> tuple[list[TableGroupListItem], int]: + session = get_current_session() + + # Separate COUNT(*) query — keeps `total` correct on out-of-range pages where + # the page rows are empty (a window function over zero rows would return 0). + total_query = f"SELECT COUNT(*) FROM table_groups AS groups WHERE {scope_sql};" + total = session.execute(text(total_query), scope_params).scalar() or 0 + if total == 0: + return [], 0 + params: dict = {**scope_params, "limit": limit, "offset": (page - 1) * limit} - query = f""" + rows_query = f""" WITH stats AS ( - SELECT table_groups_id, COUNT(*) AS table_count + SELECT table_groups_id, + COUNT(*) AS table_count, + SUM(column_ct) AS column_count, + SUM(COALESCE(record_ct, approx_record_ct)) AS row_count FROM data_table_chars GROUP BY table_groups_id ), @@ -458,10 +488,18 @@ def _list_with_activity( groups.project_code, connections.connection_name, COALESCE(stats.table_count, 0) AS table_count, + stats.column_count, + stats.row_count, latest_profile.started_at AS last_profiled_date, latest_test.test_starttime AS last_tested_date, + groups.dq_score_profiling AS profiling_score, groups.dq_score_testing AS testing_score, - COUNT(*) OVER() AS total_count + -- Mirrors utils.score: product when both exist, else the non-null one, else NULL. + CASE + WHEN groups.dq_score_profiling IS NOT NULL AND groups.dq_score_testing IS NOT NULL + THEN groups.dq_score_profiling * groups.dq_score_testing + ELSE COALESCE(groups.dq_score_profiling, groups.dq_score_testing) + END AS quality_score FROM table_groups AS groups LEFT JOIN connections ON connections.connection_id = groups.connection_id LEFT JOIN stats ON stats.table_groups_id = groups.id @@ -471,14 +509,8 @@ def _list_with_activity( ORDER BY LOWER(groups.table_groups_name) LIMIT :limit OFFSET :offset; """ - rows = get_current_session().execute(text(query), params).mappings().all() - if not rows: - return [], 0 - total = int(rows[0]["total_count"]) - items = [ - TableGroupListItem(**{k: v for k, v in row.items() if k != "total_count"}) - for row in rows - ] + rows = session.execute(text(rows_query), params).mappings().all() + items = [TableGroupListItem(**row) for row in rows] return items, total @classmethod diff --git a/testgen/common/models/test_suite.py b/testgen/common/models/test_suite.py index 6465f8e2..57284dbe 100644 --- a/testgen/common/models/test_suite.py +++ b/testgen/common/models/test_suite.py @@ -224,16 +224,22 @@ def select_summary(cls, project_code: str, table_group_id: str | UUID | None = N @classmethod def test_definition_stats(cls, test_suite_id: str | UUID) -> "TestDefinitionStats": - """Aggregate test-definition counts for a suite: total, locked, and per-test-type.""" + """Aggregate test-definition counts for a suite: total, locked, and per-test-type. + + The per-type bucket uses the user-facing ``test_name_short`` label from + ``test_types``, falling back to the raw ``test_type`` code when the lookup + row is missing (defensive — every shipping test type has a short name). + """ query = """ SELECT - test_type, + COALESCE(tt.test_name_short, td.test_type) AS type_label, COUNT(*) AS total, - SUM(CASE WHEN COALESCE(lock_refresh, 'N') = 'Y' THEN 1 ELSE 0 END) AS locked - FROM test_definitions - WHERE test_suite_id = :test_suite_id - GROUP BY test_type - ORDER BY test_type; + SUM(CASE WHEN COALESCE(td.lock_refresh, 'N') = 'Y' THEN 1 ELSE 0 END) AS locked + FROM test_definitions td + LEFT JOIN test_types tt ON tt.test_type = td.test_type + WHERE td.test_suite_id = :test_suite_id + GROUP BY type_label + ORDER BY type_label; """ rows = ( get_current_session() @@ -241,7 +247,7 @@ def test_definition_stats(cls, test_suite_id: str | UUID) -> "TestDefinitionStat .mappings() .all() ) - counts_by_type = {row["test_type"]: int(row["total"]) for row in rows} + counts_by_type = {row["type_label"]: int(row["total"]) for row in rows} total = sum(counts_by_type.values()) locked = sum(int(row["locked"]) for row in rows) return TestDefinitionStats(total=total, locked=locked, counts_by_type=counts_by_type) diff --git a/testgen/mcp/tools/common.py b/testgen/mcp/tools/common.py index d05511ce..c65a91e7 100644 --- a/testgen/mcp/tools/common.py +++ b/testgen/mcp/tools/common.py @@ -874,6 +874,18 @@ def parse_sql_flavor(value: str) -> tuple[SqlFlavorLabel, str, str]: return label, code, SQL_FLAVOR_CODE_TO_FAMILY[code] +def format_flavor_label(sql_flavor_code: str | None) -> str: + """Map a stored ``sql_flavor_code`` to its user-facing display label. + + Returns the raw code as a fallback when the code is not in the registry — defensive + against a never-shipping-but-still-stored value rather than letting an LLM see ``None``. + """ + if sql_flavor_code is None: + return "" + label = SQL_FLAVOR_CODE_TO_LABEL.get(sql_flavor_code) + return label.value if label else sql_flavor_code + + # =========================================================================== # Connection-parameter contract (MCP input vocabulary) # diff --git a/testgen/mcp/tools/connections.py b/testgen/mcp/tools/connections.py index 434caf57..9e4cf55f 100644 --- a/testgen/mcp/tools/connections.py +++ b/testgen/mcp/tools/connections.py @@ -26,6 +26,7 @@ apply_connection_params, connection_display_fields, connection_field_labels, + format_flavor_label, format_page_footer, format_page_info, infer_mode, @@ -69,13 +70,11 @@ def list_connections(project_code: str, page: int = 1, limit: int = 20) -> str: doc.text(format_page_info(total, page, limit)) table_rows: list[list[object]] = [] for row in rows: - label = SQL_FLAVOR_CODE_TO_LABEL.get(row.sql_flavor_code) - database_type = label.value if label else row.sql_flavor_code table_rows.append( [ row.connection_id, row.connection_name, - database_type, + format_flavor_label(row.sql_flavor_code), row.project_host, row.project_db, row.table_group_count, @@ -100,38 +99,12 @@ def get_connection(connection_id: int) -> str: Use this before editing a connection or creating a table group on it. Args: - connection_id: Bigint connection ID returned by `list_connections` or shown on the connections page. + connection_id: Bigint connection ID returned by `list_connections`. """ connection = resolve_connection(connection_id) - doc = MdDoc() doc.heading(1, f"Connection `{connection.connection_name}`") - doc.field("ID", connection.connection_id, code=True) - doc.field("Project", connection.project_code, code=True) - label = SQL_FLAVOR_CODE_TO_LABEL.get(connection.sql_flavor_code) - doc.field("Type", label.value if label else connection.sql_flavor_code) - - if connection.project_host: - doc.field("Host", connection.project_host, code=True) - if connection.project_port: - doc.field("Port", connection.project_port) - if connection.project_db: - doc.field("Database", connection.project_db, code=True) - if connection.project_user: - doc.field("User", connection.project_user, code=True) - if connection.connect_by_url and connection.url: - doc.field("URL", connection.url, code=True) - if connection.warehouse: - doc.field("Warehouse", connection.warehouse, code=True) - if connection.http_path: - doc.field("HTTP Path", connection.http_path, code=True) - 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) - + _render_connection_body(doc, connection) return doc.render() @@ -236,12 +209,22 @@ def _raise_validation_error(errors: list[str], header: str) -> None: 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) - label = SQL_FLAVOR_CODE_TO_LABEL.get(connection.sql_flavor_code) - doc.field("Type", label.value if label else connection.sql_flavor_code) + doc.field("Type", format_flavor_label(connection.sql_flavor_code)) - # Render each populated, non-secret field under its flavor-specific label + # 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: @@ -257,8 +240,6 @@ def _render_created_connection(connection: Connection) -> str: if connection.max_query_chars is not None: doc.field("Max Expression Length", connection.max_query_chars) - return doc.render() - def _authentication_label(connection: Connection) -> str: """The connection's auth method: the active connection mode for multi-mode diff --git a/testgen/mcp/tools/discovery.py b/testgen/mcp/tools/discovery.py index bd72f0e9..e3194ce0 100644 --- a/testgen/mcp/tools/discovery.py +++ b/testgen/mcp/tools/discovery.py @@ -1,16 +1,15 @@ from testgen.common.mixpanel_service import MixpanelService from testgen.common.models import with_database_session -from testgen.common.models.connection import Connection from testgen.common.models.data_table import DataTable from testgen.common.models.project import Project -from testgen.common.models.table_group import TableGroup from testgen.common.models.test_run import TestRun from testgen.common.models.test_suite import TestSuite from testgen.mcp.exceptions import MCPResourceNotAccessible from testgen.mcp.permissions import get_project_permissions, mcp_permission from testgen.mcp.tools.common import ( - SQL_FLAVOR_CODE_TO_LABEL, DocGroup, + format_flavor_label, + resolve_connection, resolve_table_group, resolve_test_suite, validate_limit, @@ -93,13 +92,17 @@ def get_project(project_code: str) -> str: doc.field("Test runs", summary.test_run_count) doc.heading(2, "Configuration") - doc.field("Quality-score weighting", project.use_dq_score_weights) - doc.field("Data retention", project.data_retention_enabled) - if project.data_retention_enabled and project.data_retention_days is not None: - doc.field("Retention period (days)", project.data_retention_days) - doc.field("Observability export configured", summary.can_export_to_observability) + doc.field("Weighted data quality scoring", project.use_dq_score_weights) + + doc.heading(2, "Observability Integration") + doc.field("Configured", summary.can_export_to_observability) if project.observability_api_url: - doc.field("Observability API URL", project.observability_api_url, code=True) + doc.field("API URL", project.observability_api_url, code=True) + + doc.heading(2, "Data Retention") + doc.field("Automatically delete old profiling and test history", project.data_retention_enabled) + if project.data_retention_enabled and project.data_retention_days is not None: + doc.field("Delete history older than (days)", project.data_retention_days) return doc.render() @@ -159,18 +162,21 @@ def list_test_suites(project_code: str) -> str: @with_database_session @mcp_permission("view") def get_test_suite(test_suite_id: str) -> str: - """Get a test suite's configuration: default severity, connection, table group, and per-test-type counts. + """Get a test suite's configuration: connection, table group, default severity, and per-test-type counts. Returns the test suite's identity and configuration along with a breakdown of how many test - definitions it contains by type and how many are lock-refreshed (excluded from regeneration). + definitions it contains by type and how many are locked (excluded from regeneration). Use this before changing a suite's tests to understand what will be affected. Args: test_suite_id: The test suite UUID, e.g. from `list_test_suites`. """ suite = resolve_test_suite(test_suite_id) - connection = Connection.get(suite.connection_id) if suite.connection_id else None - table_group = TableGroup.get(suite.table_groups_id) if suite.table_groups_id else None + # Defense in depth: resolve via perm-filtered helpers rather than `Model.get(...)`. + # FK constraints guarantee same-project today; the resolvers are the established wrapper + # for project-scoped lookups and keep us aligned if those guarantees ever change. + connection = resolve_connection(suite.connection_id) if suite.connection_id else None + table_group = resolve_table_group(str(suite.table_groups_id)) if suite.table_groups_id else None stats = TestSuite.test_definition_stats(suite.id) doc = MdDoc() @@ -179,11 +185,9 @@ def get_test_suite(test_suite_id: str) -> str: doc.field("Project", suite.project_code, code=True) if connection is not None: - label = SQL_FLAVOR_CODE_TO_LABEL.get(connection.sql_flavor_code) - flavor_label = label.value if label else connection.sql_flavor_code doc.field( "Connection", - f"{connection.connection_name} (`{connection.connection_id}`, {flavor_label})", + f"{connection.connection_name} (`{connection.connection_id}`, {format_flavor_label(connection.sql_flavor_code)})", ) if table_group is not None: doc.field( @@ -201,8 +205,8 @@ def get_test_suite(test_suite_id: str) -> str: if stats.counts_by_type: doc.heading(2, "Tests by type") - rows = [[test_type, count] for test_type, count in stats.counts_by_type.items()] - doc.table(["Test type", "Count"], rows, code=[0]) + rows = [[type_label, count] for type_label, count in stats.counts_by_type.items()] + doc.table(["Test", "Count"], rows) return doc.render() diff --git a/testgen/mcp/tools/table_groups.py b/testgen/mcp/tools/table_groups.py index dba9ea91..6c947437 100644 --- a/testgen/mcp/tools/table_groups.py +++ b/testgen/mcp/tools/table_groups.py @@ -21,7 +21,7 @@ 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, + format_flavor_label, format_page_footer, format_page_info, resolve_connection, @@ -30,6 +30,7 @@ validate_page, ) from testgen.mcp.tools.markdown import MdDoc +from testgen.utils import friendly_score _DUPLICATE_NAME_MESSAGE = "A Table Group with the same name already exists." _SCHEMA_LOCKED_MESSAGE = ( @@ -49,7 +50,7 @@ def list_table_groups( """List table groups in a project or on a specific connection. Pass exactly one of `project_code` or `connection_id`. Returns each group's - table count, last profile / test timestamps, and current Testing Score. + table count, last profile / test timestamps, and current quality score. Args: project_code: List groups in a project. @@ -82,7 +83,6 @@ def list_table_groups( doc.text(format_page_info(total, page, limit)) table_rows: list[list[object]] = [] for row in rows: - score = f"{row.testing_score:.2f}" if row.testing_score is not None else None table_rows.append( [ str(row.id), @@ -90,13 +90,15 @@ def list_table_groups( row.connection_name, row.table_group_schema, row.table_count, + row.column_count, + row.row_count, row.last_profiled_date, row.last_tested_date, - score, + friendly_score(row.quality_score), ] ) doc.table( - ["ID", "Name", "Connection", "Schema", "Tables", "Last profiled", "Last tested", "Testing Score"], + ["ID", "Name", "Connection", "Schema", "Tables", "Columns", "Rows", "Last profiled", "Last tested", "Quality Score"], table_rows, code=[0, 3], ) @@ -110,25 +112,23 @@ def list_table_groups( def get_table_group(table_group_id: str) -> str: """Get a table group's full configuration: filters, sampling, profiling flags, catalog tags, and recent activity. - Mirrors the field set of the table-group edit dialog so an LLM can introspect or modify - every setting available in the UI. Use this before editing a table group or generating tests. + Use this before editing a table group or generating tests. Args: table_group_id: The table group UUID, e.g. from `list_table_groups` or `get_data_inventory`. """ table_group = resolve_table_group(table_group_id) - connection = Connection.get(table_group.connection_id) if table_group.connection_id else None + # Defense in depth: route through resolve_connection (perm-scoped) rather than Connection.get. + connection = resolve_connection(table_group.connection_id) if table_group.connection_id else None doc = MdDoc() doc.heading(1, f"Table group `{table_group.table_groups_name}`") doc.field("ID", str(table_group.id), code=True) doc.field("Project", table_group.project_code, code=True) if connection is not None: - label = SQL_FLAVOR_CODE_TO_LABEL.get(connection.sql_flavor_code) - flavor_label = label.value if label else connection.sql_flavor_code doc.field( "Connection", - f"{connection.connection_name} (`{connection.connection_id}`, {flavor_label})", + f"{connection.connection_name} (`{connection.connection_id}`, {format_flavor_label(connection.sql_flavor_code)})", ) doc.field("Schema", table_group.table_group_schema, code=True) if table_group.description: @@ -136,49 +136,42 @@ def get_table_group(table_group_id: str) -> str: doc.heading(2, "Criteria") if table_group.profiling_table_set: - doc.field("Explicit table list", table_group.profiling_table_set, code=True) + doc.field(_DIFF_LABELS["profiling_table_set"], table_group.profiling_table_set, code=True) if table_group.profiling_include_mask: - doc.field("Tables to include mask", table_group.profiling_include_mask, code=True) + doc.field(_DIFF_LABELS["profiling_include_mask"], table_group.profiling_include_mask, code=True) if table_group.profiling_exclude_mask: - doc.field("Tables to exclude mask", table_group.profiling_exclude_mask, code=True) - doc.field("Profiling ID column mask", table_group.profile_id_column_mask, code=True) - doc.field("Profiling surrogate-key column mask", table_group.profile_sk_column_mask, code=True) + doc.field(_DIFF_LABELS["profiling_exclude_mask"], table_group.profiling_exclude_mask, code=True) + doc.field(_DIFF_LABELS["profile_id_column_mask"], table_group.profile_id_column_mask, code=True) + doc.field(_DIFF_LABELS["profile_sk_column_mask"], table_group.profile_sk_column_mask, code=True) doc.heading(2, "Settings") - doc.field("Detect CDE during profiling", table_group.profile_flag_cdes) - doc.field("Detect PII during profiling", table_group.profile_flag_pii) - doc.field("Exclude XDE columns from profiling", table_group.profile_exclude_xde) - doc.field("Include in project dashboard", table_group.include_in_dashboard) - doc.field("Min profiling age (days)", table_group.profiling_delay_days) + doc.field(_DIFF_LABELS["profile_flag_cdes"], table_group.profile_flag_cdes) + doc.field(_DIFF_LABELS["profile_flag_pii"], table_group.profile_flag_pii) + doc.field(_DIFF_LABELS["profile_exclude_xde"], table_group.profile_exclude_xde) + doc.field(_DIFF_LABELS["include_in_dashboard"], table_group.include_in_dashboard) + doc.field(_DIFF_LABELS["profiling_delay_days"], table_group.profiling_delay_days) doc.heading(2, "Sampling parameters") - doc.field("Use profile sampling", table_group.profile_use_sampling) + doc.field(_DIFF_LABELS["profile_use_sampling"], table_group.profile_use_sampling) if table_group.profile_use_sampling: - doc.field("Sample percent", table_group.profile_sample_percent) - doc.field("Min sample record count", table_group.profile_sample_min_count) - - catalog_fields: list[tuple[str, str | None]] = [ - ("Data source", table_group.data_source), - ("Source system", table_group.source_system), - ("Source process", table_group.source_process), - ("Business domain", table_group.business_domain), - ("Transform level", table_group.transform_level), - ("Data location", table_group.data_location), - ("Stakeholder group", table_group.stakeholder_group), - ("Data product", table_group.data_product), - ] - if any(value for _, value in catalog_fields): + doc.field(_DIFF_LABELS["profile_sample_percent"], table_group.profile_sample_percent) + doc.field(_DIFF_LABELS["profile_sample_min_count"], table_group.profile_sample_min_count) + + if any(getattr(table_group, attr, None) for attr in _CATALOG_ATTRS): doc.heading(2, "Catalog tags") - for label, value in catalog_fields: + for attr in _CATALOG_ATTRS: + value = getattr(table_group, attr, None) if value: - doc.field(label, value) + doc.field(_DIFF_LABELS[attr], value) if table_group.dq_score_testing is not None or table_group.dq_score_profiling is not None: doc.heading(2, "Latest activity") - if table_group.dq_score_testing is not None: - doc.field("Testing Score", f"{table_group.dq_score_testing:.2f}") - if table_group.dq_score_profiling is not None: - doc.field("Profiling Score", f"{table_group.dq_score_profiling:.2f}") + if (profiling := friendly_score(table_group.dq_score_profiling)) is not None: + doc.field("Profiling Score", profiling) + if (testing := friendly_score(table_group.dq_score_testing)) is not None: + doc.field("Testing Score", testing) + if (quality := friendly_score(table_group.quality_score)) is not None: + doc.field("Quality Score", quality) return doc.render() diff --git a/tests/unit/mcp/test_tools_connections.py b/tests/unit/mcp/test_tools_connections.py index e6b77642..59c9d77d 100644 --- a/tests/unit/mcp/test_tools_connections.py +++ b/tests/unit/mcp/test_tools_connections.py @@ -6,6 +6,7 @@ """ from unittest.mock import MagicMock, patch +from uuid import uuid4 import pytest @@ -346,23 +347,52 @@ def test_get_connection_returns_config_with_display_label(mock_resolve, db_sessi @patch(f"{MODULE}.resolve_connection") def test_get_connection_never_returns_secrets(mock_resolve, db_session_mock): - """Password, private key, passphrase, service-account key must not leak.""" - mock_resolve.return_value = _mock_connection( - project_pw_encrypted="hunter2", - private_key="-----BEGIN PRIVATE KEY-----secret", - private_key_passphrase="passphrase", # noqa: S106 - service_account_key={"key": "secret-account-json"}, + """Password, private key, passphrase, service-account key must not leak. + + Uses a real ``Connection`` ORM instance so the test would fail if the rendering + started reading the encrypted attributes (a MagicMock would silently return + a ``MagicMock`` for any attribute, masking the regression). + """ + from testgen.common.models.connection import Connection + + conn = Connection( + id=uuid4(), + connection_id=42, + project_code="demo", + connection_name="warehouse_prod", + sql_flavor="postgresql", + sql_flavor_code="postgresql", + project_host="localhost", + project_port="5432", + project_db="testgen", + project_user="dq_user", + project_pw_encrypted="sentinel-password-hunter2", + private_key="-----BEGIN PRIVATE KEY-----\nsentinel-private-key-body\n-----END PRIVATE KEY-----", + private_key_passphrase="sentinel-passphrase", # noqa: S106 + service_account_key={"client_email": "sentinel-account@example.com", "private_key": "sentinel-sak-key"}, ) + mock_resolve.return_value = conn from testgen.mcp.tools.connections import get_connection with _patch_perms(permission="view"): out = get_connection(42) - assert "hunter2" not in out - assert "BEGIN PRIVATE KEY" not in out - assert "passphrase" not in out - assert "secret-account-json" not in out + # Every populated secret value must NOT appear in output. + for needle in ( + "sentinel-password-hunter2", + "sentinel-private-key-body", + "BEGIN PRIVATE KEY", + "sentinel-passphrase", + "sentinel-account@example.com", + "sentinel-sak-key", + ): + assert needle not in out, f"secret material `{needle}` leaked into get_connection output" + + # The non-secret content we DO render should be present — sanity check that the + # test is exercising the real rendering path (not silently bypassing it). + assert "warehouse_prod" in out + assert "localhost" in out @patch(f"{MODULE}.resolve_connection") diff --git a/tests/unit/mcp/test_tools_discovery.py b/tests/unit/mcp/test_tools_discovery.py index b1a09b1d..44052d8a 100644 --- a/tests/unit/mcp/test_tools_discovery.py +++ b/tests/unit/mcp/test_tools_discovery.py @@ -276,9 +276,12 @@ def test_get_project_returns_counts_and_config(mock_project_cls, db_session_mock assert "**Test definitions:** 142" in out assert "**Profiling runs:** 12" in out assert "**Test runs:** 38" in out - assert "**Quality-score weighting:** Yes" in out - assert "**Data retention:** Yes" in out - assert "**Retention period (days):** 180" in out + assert "**Weighted data quality scoring:** Yes" in out + assert "## Observability Integration" in out + assert "**Configured:** No" in out # default summary has can_export_to_observability=False + assert "## Data Retention" in out + assert "**Automatically delete old profiling and test history:** Yes" in out + assert "**Delete history older than (days):** 180" in out @patch("testgen.mcp.permissions._compute_project_permissions") @@ -323,7 +326,7 @@ def _mock_test_suite(**overrides): return suite -_DEFAULT_TYPE_COUNTS = {"Aggregate_Balance": 5, "Row_Count": 12} +_DEFAULT_TYPE_COUNTS = {"Aggregate Balance": 5, "Row Count": 12} def _stats(total=47, locked=3, counts_by_type=None): @@ -337,19 +340,19 @@ def _stats(total=47, locked=3, counts_by_type=None): @patch("testgen.mcp.tools.discovery.TestSuite") -@patch("testgen.mcp.tools.discovery.TableGroup") -@patch("testgen.mcp.tools.discovery.Connection") +@patch("testgen.mcp.tools.discovery.resolve_table_group") +@patch("testgen.mcp.tools.discovery.resolve_connection") @patch("testgen.mcp.tools.common.TestSuite") def test_get_test_suite_returns_full_config( - mock_common_suite, mock_conn_cls, mock_tg_cls, mock_suite_cls, db_session_mock, + mock_common_suite, mock_resolve_conn, mock_resolve_tg, mock_suite_cls, db_session_mock, ): suite = _mock_test_suite() mock_common_suite.get.return_value = suite conn = MagicMock(connection_id=42, connection_name="warehouse_prod", sql_flavor_code="snowflake") - mock_conn_cls.get.return_value = conn + mock_resolve_conn.return_value = conn tg = MagicMock(table_groups_name="curated_payments") tg.id = suite.table_groups_id - mock_tg_cls.get.return_value = tg + mock_resolve_tg.return_value = tg mock_suite_cls.test_definition_stats.return_value = _stats() from testgen.mcp.tools.discovery import get_test_suite @@ -367,13 +370,15 @@ def test_get_test_suite_returns_full_config( assert "**Total tests:** 47" in out assert "**Locked tests:** 3" in out assert "Tests by type" in out - assert "Aggregate_Balance" in out - assert "Row_Count" in out + # Renders short_name labels (e.g. "Aggregate Balance"), NOT raw test_type codes (e.g. "Aggregate_Balance") + assert "Aggregate Balance" in out + assert "Row Count" in out + assert "Aggregate_Balance" not in out # the raw test_type code must NOT appear @patch("testgen.mcp.tools.common.TestSuite") -def test_get_test_suite_rejects_monitor_suite(mock_common_suite, db_session_mock): - """resolve_test_suite returns None when suite is_monitor — uses unified wording.""" +def test_get_test_suite_rejects_inaccessible_id(mock_common_suite, db_session_mock): + """A genuinely missing / inaccessible id (TestSuite.get returns None) raises the unified error.""" mock_common_suite.get.return_value = None from testgen.mcp.tools.discovery import get_test_suite @@ -383,6 +388,33 @@ def test_get_test_suite_rejects_monitor_suite(mock_common_suite, db_session_mock get_test_suite(bogus_id) +@patch("testgen.mcp.tools.common.TestSuite") +def test_get_test_suite_rejects_actual_monitor_suite(mock_common_suite, db_session_mock): + """An existing ``is_monitor=True`` suite is rejected because resolve_test_suite + passes ``TestSuite.is_monitor.isnot(True)`` as a filter clause. + + Simulates the real DB filter behaviour: ``TestSuite.get`` returns the monitor + suite when called without the filter clause, ``None`` when the clause is present. + """ + monitor_suite = _mock_test_suite(is_monitor=True) + + def filtered_get(_uuid, *clauses): + # The resolver's contract: pass an `is_monitor.isnot(True)` clause to TestSuite.get. + # If the clause is present, a DB query against it would not match this monitor row. + for clause in clauses: + clause_str = str(clause).lower() + if "is_monitor" in clause_str and "not" in clause_str: + return None + return monitor_suite + + mock_common_suite.get.side_effect = filtered_get + + from testgen.mcp.tools.discovery import get_test_suite + + with pytest.raises(MCPResourceNotAccessible, match="Test suite .* not found or not accessible"): + get_test_suite(str(monitor_suite.id)) + + @patch("testgen.mcp.tools.common.TestSuite") def test_get_test_suite_rejects_invalid_uuid(mock_common_suite, db_session_mock): from testgen.mcp.exceptions import MCPUserError as _MCPUserError @@ -393,16 +425,15 @@ def test_get_test_suite_rejects_invalid_uuid(mock_common_suite, db_session_mock) @patch("testgen.mcp.tools.discovery.TestSuite") -@patch("testgen.mcp.tools.discovery.TableGroup") -@patch("testgen.mcp.tools.discovery.Connection") +@patch("testgen.mcp.tools.discovery.resolve_table_group") +@patch("testgen.mcp.tools.discovery.resolve_connection") @patch("testgen.mcp.tools.common.TestSuite") def test_get_test_suite_no_severity_renders_inherit( - mock_common_suite, mock_conn_cls, mock_tg_cls, mock_suite_cls, db_session_mock, + mock_common_suite, mock_resolve_conn, mock_resolve_tg, mock_suite_cls, db_session_mock, ): - suite = _mock_test_suite(severity=None) + suite = _mock_test_suite(severity=None, connection_id=None, table_groups_id=None) mock_common_suite.get.return_value = suite - mock_conn_cls.get.return_value = None - mock_tg_cls.get.return_value = None + # connection_id / table_groups_id are None, so resolvers should not be called mock_suite_cls.test_definition_stats.return_value = _stats(total=0, locked=0, counts_by_type={}) from testgen.mcp.tools.discovery import get_test_suite @@ -412,3 +443,5 @@ def test_get_test_suite_no_severity_renders_inherit( assert "Inherit from test type" in out # No "Tests by type" table when there are no test definitions assert "Tests by type" not in out + mock_resolve_conn.assert_not_called() + mock_resolve_tg.assert_not_called() diff --git a/tests/unit/mcp/test_tools_table_groups.py b/tests/unit/mcp/test_tools_table_groups.py index ead9576b..ddbed94b 100644 --- a/tests/unit/mcp/test_tools_table_groups.py +++ b/tests/unit/mcp/test_tools_table_groups.py @@ -772,9 +772,13 @@ def _list_item(**overrides): project_code=overrides.get("project_code", "demo"), connection_name=overrides.get("connection_name", "warehouse_prod"), table_count=overrides.get("table_count", 12), + column_count=overrides.get("column_count", 84), + row_count=overrides.get("row_count", 100_000), last_profiled_date=overrides.get("last_profiled_date", None), last_tested_date=overrides.get("last_tested_date", None), + profiling_score=overrides.get("profiling_score", 0.95), testing_score=overrides.get("testing_score", 0.97), + quality_score=overrides.get("quality_score", 0.92), ) @@ -809,8 +813,12 @@ def test_list_table_groups_by_project_renders_rows(mock_tg_cls, db_session_mock) assert "core_tables" in out assert "warehouse_prod" in out assert "public" in out + assert "Quality Score" in out # column header (#8) + assert "Columns" in out # column header (#10) + assert "Rows" in out # column header (#10) assert "12" in out # table count - assert "0.97" in out # testing score formatted + # quality_score 0.92 rendered via friendly_score → "92.0" (percentage form, mirrors UI). + assert "92.0" in out mock_tg_cls.list_for_project.assert_called_once_with("demo", page=1, limit=20) @@ -863,15 +871,22 @@ def test_list_table_groups_rejects_inaccessible_project(mock_compute, db_session def _read_mock_table_group(**overrides): """Variant of _mock_table_group that also sets dq score fields used by get_table_group.""" + from testgen.utils import score + tg = _mock_table_group(**overrides) - tg.dq_score_testing = overrides.get("dq_score_testing", None) - tg.dq_score_profiling = overrides.get("dq_score_profiling", None) + testing = overrides.get("dq_score_testing", None) + profiling = overrides.get("dq_score_profiling", None) + tg.dq_score_testing = testing + tg.dq_score_profiling = profiling + # Mirror TableGroup.quality_score property — MagicMock would return a MagicMock for + # the attribute, so wire it explicitly through the same `score` helper the property uses. + tg.quality_score = score(profiling, testing) return tg -@patch(f"{MODULE}.Connection") +@patch(f"{MODULE}.resolve_connection") @patch(f"{MODULE}.resolve_table_group") -def test_get_table_group_renders_all_dialog_sections(mock_resolve, mock_conn_cls, db_session_mock): +def test_get_table_group_renders_all_dialog_sections(mock_resolve, mock_resolve_conn, db_session_mock): tg = _read_mock_table_group( description="Curated payments tables", profiling_include_mask="payments_%", @@ -886,7 +901,7 @@ def test_get_table_group_renders_all_dialog_sections(mock_resolve, mock_conn_cls ) mock_resolve.return_value = tg conn = _mock_connection(connection_name="warehouse_prod", sql_flavor_code="snowflake") - mock_conn_cls.get.return_value = conn + mock_resolve_conn.return_value = conn from testgen.mcp.tools.table_groups import get_table_group @@ -899,36 +914,38 @@ def test_get_table_group_renders_all_dialog_sections(mock_resolve, mock_conn_cls assert "Snowflake" in out assert "`public`" in out assert "Curated payments tables" in out - # Criteria — both table-name and column-name masks rendered + # Criteria — both table-name and column-name masks rendered (labels via _DIFF_LABELS) assert "## Criteria" in out assert "payments_%" in out assert "tmp_%" in out assert "payments,refunds" in out - assert "Profiling ID column mask" in out - assert "Profiling surrogate-key column mask" in out + assert "ID column mask" in out + assert "SK column mask" in out # Settings assert "## Settings" in out - assert "Detect CDE during profiling" in out - # Sampling enabled → percent + min count rendered + assert "Flag CDEs" in out # _DIFF_LABELS["profile_flag_cdes"] + # Sampling enabled → percent + min count rendered (labels via _DIFF_LABELS) assert "## Sampling parameters" in out - assert "**Sample percent:** 50" in out - assert "Min sample record count" in out + assert "**Sample %:** 50" in out + assert "Sample min rows" in out # Catalog tags only render when set assert "## Catalog tags" in out assert "DataKitchen" in out assert "Finance" in out - # Latest activity + # Latest activity — scores rendered via friendly_score (percentage form, mirrors UI). assert "## Latest activity" in out - assert "0.91" in out - assert "0.95" in out + assert "**Profiling Score:** 95.0" in out + assert "**Testing Score:** 91.0" in out + # 0.95 * 0.91 = 0.8645 → friendly_score → "86.4" (Python banker's rounding of 86.45). + assert "**Quality Score:** 86.4" in out -@patch(f"{MODULE}.Connection") +@patch(f"{MODULE}.resolve_connection") @patch(f"{MODULE}.resolve_table_group") -def test_get_table_group_skips_catalog_when_no_tags(mock_resolve, mock_conn_cls, db_session_mock): +def test_get_table_group_skips_catalog_when_no_tags(mock_resolve, mock_resolve_conn, db_session_mock): tg = _read_mock_table_group() # all catalog tags None mock_resolve.return_value = tg - mock_conn_cls.get.return_value = _mock_connection() + mock_resolve_conn.return_value = _mock_connection() from testgen.mcp.tools.table_groups import get_table_group @@ -938,14 +955,14 @@ def test_get_table_group_skips_catalog_when_no_tags(mock_resolve, mock_conn_cls, assert "## Catalog tags" not in out -@patch(f"{MODULE}.Connection") +@patch(f"{MODULE}.resolve_connection") @patch(f"{MODULE}.resolve_table_group") def test_get_table_group_skips_sample_details_when_sampling_off( - mock_resolve, mock_conn_cls, db_session_mock, + mock_resolve, mock_resolve_conn, db_session_mock, ): tg = _read_mock_table_group(profile_use_sampling=False) mock_resolve.return_value = tg - mock_conn_cls.get.return_value = _mock_connection() + mock_resolve_conn.return_value = _mock_connection() from testgen.mcp.tools.table_groups import get_table_group @@ -953,9 +970,9 @@ def test_get_table_group_skips_sample_details_when_sampling_off( out = get_table_group(str(tg.id)) # Sampling section header still present (with the toggle), but percent/min not shown - assert "Use profile sampling" in out - assert "**Sample percent:**" not in out - assert "Min sample record count" not in out + assert "**Sampling:** No" in out # _DIFF_LABELS["profile_use_sampling"] + assert "**Sample %:**" not in out + assert "Sample min rows" not in out @patch("testgen.mcp.tools.common.TableGroup") From d37ae12fd83d4aba9263ed444f897350881a47ea Mon Sep 17 00:00:00 2001 From: Luis Date: Wed, 10 Jun 2026 16:31:22 -0400 Subject: [PATCH 27/78] fix(table-groups): gate profile_flag_pii edits on view_pii permission MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The "Detect PII during profiling" checkbox was no longer disabled for users without view_pii, and nothing enforced it server-side — so a less-privileged user (or an MCP client) could change the flag on an existing table group, and the next profiling run could overwrite manually marked PII columns, exposing that data to users who shouldn't see it. Without view_pii a user can no longer change profile_flag_pii on an existing table group. Creation is unaffected: a new group has no manual PII marks to overwrite, so the model default (PII on) applies for everyone. - MCP update_table_group rejects changing the flag without view_pii (change-detection — an unchanged re-send still passes) - UI edit dialog disables the checkbox, and the save path skips the field for users without view_pii so the stored value is preserved - thread can_view_pii into the edit dialog and table group list page Tests run the deny/allow pair against a real ProjectPermissions using a role with edit-but-not-view_pii, proving the gate checks view_pii specifically (a mutation to any other permission fails the deny tests). Co-Authored-By: Claude Opus 4.8 (1M context) --- testgen/mcp/tools/table_groups.py | 13 +- .../frontend/js/pages/table_group_list.js | 1 + .../js/components/table_group_edit_dialog.js | 2 + testgen/ui/views/table_groups.py | 4 + tests/unit/mcp/conftest.py | 8 +- tests/unit/mcp/test_tools_table_groups.py | 112 +++++++++++++++++- 6 files changed, 135 insertions(+), 5 deletions(-) diff --git a/testgen/mcp/tools/table_groups.py b/testgen/mcp/tools/table_groups.py index 6c947437..ef9d9046 100644 --- a/testgen/mcp/tools/table_groups.py +++ b/testgen/mcp/tools/table_groups.py @@ -18,7 +18,7 @@ from testgen.common.models import with_database_session from testgen.common.models.connection import Connection from testgen.common.models.table_group import TableGroup -from testgen.mcp.exceptions import MCPResourceNotAccessible, MCPUserError +from testgen.mcp.exceptions import MCPPermissionDenied, MCPResourceNotAccessible, MCPUserError from testgen.mcp.permissions import get_project_permissions, mcp_permission from testgen.mcp.tools.common import ( format_flavor_label, @@ -33,6 +33,10 @@ from testgen.utils import friendly_score _DUPLICATE_NAME_MESSAGE = "A Table Group with the same name already exists." +_PII_FLAG_DENIED_MESSAGE = ( + "Changing PII detection requires permission to view PII. " + "Leave this setting unchanged or contact your administrator." +) _SCHEMA_LOCKED_MESSAGE = ( "Schema cannot be changed once the table group has been used. " "Delete and recreate the table group to use a different schema." @@ -402,6 +406,13 @@ def update_table_group( ): raise MCPUserError(_SCHEMA_LOCKED_MESSAGE) + if ( + profile_flag_pii is not None + and profile_flag_pii != table_group.profile_flag_pii + and not get_project_permissions().has_permission("view_pii", table_group.project_code) + ): + raise MCPPermissionDenied(_PII_FLAG_DENIED_MESSAGE) + before = _snapshot(table_group) _apply_args_to_table_group(table_group, **supplied) diff --git a/testgen/ui/components/frontend/js/pages/table_group_list.js b/testgen/ui/components/frontend/js/pages/table_group_list.js index 3f52a24a..154010dd 100644 --- a/testgen/ui/components/frontend/js/pages/table_group_list.js +++ b/testgen/ui/components/frontend/js/pages/table_group_list.js @@ -120,6 +120,7 @@ const TableGroupList = (props) => { connections: van.derive(() => getValue(props.edit_dialog)?.connections), table_group: van.derive(() => getValue(props.edit_dialog)?.table_group), is_in_use: van.derive(() => getValue(props.edit_dialog)?.is_in_use), + can_view_pii: van.derive(() => getValue(props.permissions)?.can_view_pii), table_group_preview: van.derive(() => getValue(props.edit_dialog)?.table_group_preview), result: van.derive(() => getValue(props.edit_dialog)?.result), })); diff --git a/testgen/ui/static/js/components/table_group_edit_dialog.js b/testgen/ui/static/js/components/table_group_edit_dialog.js index a5af0c5d..ac00f375 100644 --- a/testgen/ui/static/js/components/table_group_edit_dialog.js +++ b/testgen/ui/static/js/components/table_group_edit_dialog.js @@ -14,6 +14,7 @@ * @property {Connection[]} connections * @property {TableGroup} table_group * @property {boolean} is_in_use + * @property {boolean?} can_view_pii * @property {TableGroupPreview?} table_group_preview * @property {EditResult?} result */ @@ -61,6 +62,7 @@ const TableGroupEditDialog = (props) => { showConnectionSelector: connections.length > 1, disableConnectionSelector: false, disableSchemaField: getValue(props.is_in_use) ?? false, + disablePiiFlag: !(getValue(props.can_view_pii) ?? false), onChange: (updatedTableGroup, state) => { tableGroupState.val = updatedTableGroup; formValid.val = state.valid; diff --git a/testgen/ui/views/table_groups.py b/testgen/ui/views/table_groups.py index bb9e186c..408b86c5 100644 --- a/testgen/ui/views/table_groups.py +++ b/testgen/ui/views/table_groups.py @@ -539,10 +539,14 @@ def on_close_edit(_params: dict) -> None: is_in_use = TableGroup.is_in_use([table_group_id]) edit_tg_data = get_edit_tg() + can_view_pii = session.auth.user_has_permission("view_pii") add_scorecard_definition = False for key, value in edit_tg_data.items(): if key == "add_scorecard_definition": add_scorecard_definition = value + elif key == "profile_flag_pii" and not can_view_pii: + # Users without view_pii cannot change the PII flag — keep the stored value. + continue else: setattr(table_group, key, value) diff --git a/tests/unit/mcp/conftest.py b/tests/unit/mcp/conftest.py index 53e23a5d..fc50fef2 100644 --- a/tests/unit/mcp/conftest.py +++ b/tests/unit/mcp/conftest.py @@ -5,12 +5,16 @@ from testgen.mcp.permissions import set_mcp_token, set_mcp_username -# Fictional role matrix for tests. role_a has full access, role_c is restricted. +# Fictional role matrix for tests. role_a has full access (but NOT view_pii — several +# tests rely on that to exercise the no-view_pii path), role_c is restricted, and +# role_d holds edit + view_pii so deny/allow pairs can be distinguished against a real +# ProjectPermissions without role_a accidentally granting view_pii. TEST_PERM_MATRIX = { "view": ["role_a", "role_b"], "catalog": ["role_a", "role_b", "role_c"], - "edit": ["role_a"], + "edit": ["role_a", "role_d"], "administer": ["role_a"], + "view_pii": ["role_d"], } diff --git a/tests/unit/mcp/test_tools_table_groups.py b/tests/unit/mcp/test_tools_table_groups.py index ddbed94b..c6c63415 100644 --- a/tests/unit/mcp/test_tools_table_groups.py +++ b/tests/unit/mcp/test_tools_table_groups.py @@ -19,8 +19,9 @@ # --------------------------------------------------------------------------- -def _patch_perms(allowed=("demo",), memberships=None, permission="edit"): - memberships = memberships or dict.fromkeys(allowed, "role_a") +def _patch_perms(allowed=("demo",), memberships=None, permission="edit", role="role_a"): + # role_a has edit but NOT view_pii; role_d has edit + view_pii (see conftest matrix). + memberships = memberships or dict.fromkeys(allowed, role) return patch( "testgen.mcp.permissions._compute_project_permissions", return_value=ProjectPermissions( @@ -363,6 +364,34 @@ def test_create_table_group_requires_edit(db_session_mock): ) +# --------------------------------------------------------------------------- +# create_table_group — PII flag is not gated on create +# --------------------------------------------------------------------------- + + +@patch(f"{MODULE}.TableGroup") +@patch(f"{MODULE}.resolve_connection") +def test_create_table_group_pii_on_allowed_without_view_pii(mock_resolve, mock_tg_cls, db_session_mock): + """A new table group has no manually-marked PII to overwrite, so creating with the + flag on is allowed even without view_pii — the gate applies only to editing it.""" + mock_resolve.return_value = _mock_connection() + instance = _mock_table_group(profile_flag_pii=False) + mock_tg_cls.return_value = instance + + from testgen.mcp.tools.table_groups import create_table_group + + with _patch_perms(): # role_a: edit, no view_pii + create_table_group( + connection_id=42, + table_group_name="Sample TG", + schema="public", + profile_flag_pii=True, + ) + + instance.save.assert_called_once() + assert instance.profile_flag_pii is True + + # --------------------------------------------------------------------------- # update_table_group # --------------------------------------------------------------------------- @@ -525,6 +554,85 @@ def test_update_table_group_delay_days_int_cast_to_str(mock_resolve, mock_tg_cls assert tg.profiling_delay_days == "3" +# --------------------------------------------------------------------------- +# update_table_group — PII flag gating (view_pii permission) +# --------------------------------------------------------------------------- + + +@patch(f"{MODULE}.TableGroup") +@patch(f"{MODULE}.resolve_table_group") +def test_update_table_group_enable_pii_denied_without_view_pii(mock_resolve, mock_tg_cls, db_session_mock): + """role_a has edit but not view_pii (real ProjectPermissions) — enabling PII is denied. + + role_a *does* hold administer, so this also proves the gate checks view_pii + specifically, not some broader permission. + """ + tg = _mock_table_group(profile_flag_pii=False) + mock_resolve.return_value = tg + mock_tg_cls.is_in_use.return_value = False + + from testgen.mcp.tools.table_groups import update_table_group + + with _patch_perms(role="role_a"), pytest.raises(MCPPermissionDenied): + update_table_group(table_group_id=str(tg.id), profile_flag_pii=True) + tg.save.assert_not_called() + + +@patch(f"{MODULE}.TableGroup") +@patch(f"{MODULE}.resolve_table_group") +def test_update_table_group_disable_pii_denied_without_view_pii(mock_resolve, mock_tg_cls, db_session_mock): + """Change-detection mirrors the disabled checkbox — the value can't be touched either way.""" + tg = _mock_table_group(profile_flag_pii=True) + mock_resolve.return_value = tg + mock_tg_cls.is_in_use.return_value = False + + from testgen.mcp.tools.table_groups import update_table_group + + with _patch_perms(role="role_a"), pytest.raises(MCPPermissionDenied): + update_table_group(table_group_id=str(tg.id), profile_flag_pii=False) + tg.save.assert_not_called() + + +@patch(f"{MODULE}.TableGroup") +@patch(f"{MODULE}.resolve_table_group") +def test_update_table_group_unchanged_pii_allowed_without_view_pii(mock_resolve, mock_tg_cls, db_session_mock): + """Re-sending the current PII value (as the disabled UI checkbox does) is not a change.""" + tg = _mock_table_group(profile_flag_pii=True, description=None) + mock_resolve.return_value = tg + mock_tg_cls.is_in_use.return_value = False + + from testgen.mcp.tools.table_groups import update_table_group + + with _patch_perms(role="role_a"): + out = update_table_group( + table_group_id=str(tg.id), + profile_flag_pii=True, + description="Edited elsewhere", + ) + + tg.save.assert_called_once() + assert tg.profile_flag_pii is True + assert "Description" in out + + +@patch(f"{MODULE}.TableGroup") +@patch(f"{MODULE}.resolve_table_group") +def test_update_table_group_enable_pii_allowed_with_view_pii(mock_resolve, mock_tg_cls, db_session_mock): + """role_d holds edit + view_pii (real ProjectPermissions) — enabling PII is allowed.""" + tg = _mock_table_group(profile_flag_pii=False) + mock_resolve.return_value = tg + mock_tg_cls.is_in_use.return_value = False + + from testgen.mcp.tools.table_groups import update_table_group + + with _patch_perms(role="role_d"): + out = update_table_group(table_group_id=str(tg.id), profile_flag_pii=True) + + tg.save.assert_called_once() + assert tg.profile_flag_pii is True + assert "Flag PII" in out + + # --------------------------------------------------------------------------- # preview_table_group # --------------------------------------------------------------------------- From 644a152e1fac8274a0ced9ccc6994f3ebaea88ab Mon Sep 17 00:00:00 2001 From: Ricardo Boni Date: Thu, 11 Jun 2026 11:00:27 -0400 Subject: [PATCH 28/78] feat(api): normalize run-summary count shapes (TG-1121) GET /profiling-runs/{job_id}: replace 'issues' with nested 'issue_counts' (hygiene_issues by likelihood, potential_pii by risk, dismissed total). GET /test-runs/{job_id}: rename 'tests' to 'result_counts' (same fields). New shared HygieneIssue.count_for_run() is the canonical breakdown source. Co-Authored-By: Claude Fable 5 --- testgen/api/runs.py | 22 +++-------- testgen/api/schemas.py | 26 ++++++++++--- testgen/common/models/hygiene_issue.py | 52 +++++++++++++++++++------- tests/unit/api/test_runs.py | 23 +++++++----- 4 files changed, 78 insertions(+), 45 deletions(-) diff --git a/testgen/api/runs.py b/testgen/api/runs.py index be2f62f1..ac3e645d 100644 --- a/testgen/api/runs.py +++ b/testgen/api/runs.py @@ -6,10 +6,10 @@ from testgen.api.deps import db_session, resolve_job from testgen.api.schemas import ( ErrorResponse, - IssueBreakdown, + IssueCounts, ProfilingRunResponse, ProfilingRunResult, - TestBreakdown, + ResultCounts, TestRunResponse, TestRunResult, ) @@ -42,14 +42,7 @@ def get_test_run(job: JobExecution = resolve_job("view", JobExecution.job_key == counts = TestResult.count_by_status(test_run.id) result = TestRunResult( score=test_run.dq_score_test_run, - tests=TestBreakdown( - passed=counts.passed, - failed=counts.failed, - warning=counts.warning, - error=counts.error, - log=counts.log, - dismissed=counts.dismissed, - ), + result_counts=ResultCounts.model_validate(counts, from_attributes=True), ) test_suite_id = test_run.test_suite_id if test_run else None @@ -80,18 +73,13 @@ def get_profiling_run(job: JobExecution = resolve_job("view", JobExecution.job_k result = None if profiling_run: - counts = HygieneIssue.count_by_likelihood(profiling_run.id) + counts = HygieneIssue.count_for_run(profiling_run.id) result = ProfilingRunResult( score=profiling_run.dq_score_profiling, table_ct=profiling_run.table_ct, column_ct=profiling_run.column_ct, record_ct=profiling_run.record_ct, - issues=IssueBreakdown( - definite=counts.definite, - likely=counts.likely, - possible=counts.possible, - dismissed=counts.dismissed, - ), + issue_counts=IssueCounts.model_validate(counts, from_attributes=True), ) return ProfilingRunResponse( diff --git a/testgen/api/schemas.py b/testgen/api/schemas.py index fbe56584..61a75766 100644 --- a/testgen/api/schemas.py +++ b/testgen/api/schemas.py @@ -48,8 +48,8 @@ class JobListResponse(BaseModel): # --- Test Runs --- -class TestBreakdown(BaseModel): - """Counts of test results by outcome status.""" +class ResultCounts(BaseModel): + """Counts of test results by outcome status, with dismissed results separated.""" passed: int = 0 failed: int = 0 @@ -63,7 +63,7 @@ class TestRunResult(BaseModel): """Run-specific data populated when execution completes.""" score: float | None = None - tests: TestBreakdown + result_counts: ResultCounts class TestRunResponse(BaseModel): @@ -81,12 +81,26 @@ class TestRunResponse(BaseModel): # --- Profiling Runs --- -class IssueBreakdown(BaseModel): - """Counts of hygiene issues by likelihood category.""" +class HygieneIssueCounts(BaseModel): + """Counts of active data-quality hygiene issues by likelihood category.""" definite: int = 0 likely: int = 0 possible: int = 0 + + +class PotentialPiiCounts(BaseModel): + """Counts of active Potential PII findings by risk level.""" + + high: int = 0 + moderate: int = 0 + + +class IssueCounts(BaseModel): + """Profiling-finding breakdown: active counts by kind, plus a single dismissed total.""" + + hygiene_issues: HygieneIssueCounts + potential_pii: PotentialPiiCounts dismissed: int = 0 @@ -97,7 +111,7 @@ class ProfilingRunResult(BaseModel): table_ct: int | None = None column_ct: int | None = None record_ct: int | None = None - issues: IssueBreakdown + issue_counts: IssueCounts class ProfilingRunResponse(BaseModel): diff --git a/testgen/common/models/hygiene_issue.py b/testgen/common/models/hygiene_issue.py index c7683479..24103ffe 100644 --- a/testgen/common/models/hygiene_issue.py +++ b/testgen/common/models/hygiene_issue.py @@ -11,7 +11,7 @@ from sqlalchemy.orm import aliased, relationship from sqlalchemy.sql.functions import func -from testgen.common.enums import Disposition +from testgen.common.enums import Disposition, IssueLikelihood, PiiRisk from testgen.common.models import Base, get_current_session from testgen.common.models.entity import Entity from testgen.common.models.job_execution import JobExecution @@ -23,12 +23,28 @@ @dataclass -class IssueLikelihoodCounts: - """Counts of hygiene issues by likelihood category, with dismissed/inactive separated.""" +class HygieneIssueCounts: + """Counts of active data-quality hygiene issues by likelihood category.""" definite: int = 0 likely: int = 0 possible: int = 0 + + +@dataclass +class PotentialPiiCounts: + """Counts of active Potential PII findings by risk level.""" + + high: int = 0 + moderate: int = 0 + + +@dataclass +class IssueCounts: + """Profiling-finding breakdown: active counts by kind, plus a single dismissed total.""" + + hygiene_issues: HygieneIssueCounts + potential_pii: PotentialPiiCounts dismissed: int = 0 @@ -202,19 +218,25 @@ def select_count_by_priority(cls, profiling_run_id: UUID) -> dict[str, IssueCoun return result @classmethod - def count_by_likelihood(cls, profile_run_id: UUID) -> IssueLikelihoodCounts: - """Count hygiene issues by likelihood category for a single profiling run.""" - dismissed = func.coalesce(cls.disposition, "Confirmed").in_(("Dismissed", "Inactive")) + def count_for_run(cls, profile_run_id: UUID) -> IssueCounts: + """Count profiling findings for a single run: active hygiene issues by likelihood, + active Potential PII findings by risk level, and a single dismissed total.""" + dismissed = func.coalesce(cls.disposition, Disposition.CONFIRMED).in_( + (Disposition.DISMISSED, Disposition.INACTIVE) + ) + is_pii = HygieneIssueType.likelihood == IssueLikelihood.POTENTIAL_PII - def _count_active(likelihood_values: tuple[str, ...]): - return func.sum(case((~dismissed & HygieneIssueType.likelihood.in_(likelihood_values), 1), else_=0)) + def _count(condition): + return func.coalesce(func.sum(case((condition, 1), else_=0)), 0) query = ( select( - _count_active(("Definite",)).label("definite"), - _count_active(("Likely",)).label("likely"), - _count_active(("Possible", "Potential PII")).label("possible"), - func.sum(case((dismissed, 1), else_=0)).label("dismissed"), + _count(~dismissed & (HygieneIssueType.likelihood == IssueLikelihood.DEFINITE)).label("definite"), + _count(~dismissed & (HygieneIssueType.likelihood == IssueLikelihood.LIKELY)).label("likely"), + _count(~dismissed & (HygieneIssueType.likelihood == IssueLikelihood.POSSIBLE)).label("possible"), + _count(~dismissed & is_pii & (cls.priority == PiiRisk.HIGH)).label("high"), + _count(~dismissed & is_pii & (cls.priority == PiiRisk.MODERATE)).label("moderate"), + _count(dismissed).label("dismissed"), ) .select_from(cls) .join(HygieneIssueType, HygieneIssueType.id == cls.type_id) @@ -222,7 +244,11 @@ def _count_active(likelihood_values: tuple[str, ...]): ) row = get_current_session().execute(query).first() - return IssueLikelihoodCounts(**{k: v for k, v in row._mapping.items() if v is not None}) + return IssueCounts( + hygiene_issues=HygieneIssueCounts(definite=row.definite, likely=row.likely, possible=row.possible), + potential_pii=PotentialPiiCounts(high=row.high, moderate=row.moderate), + dismissed=row.dismissed, + ) @classmethod def _priority_order(cls): diff --git a/tests/unit/api/test_runs.py b/tests/unit/api/test_runs.py index 0fe8143c..f1dc1b0d 100644 --- a/tests/unit/api/test_runs.py +++ b/tests/unit/api/test_runs.py @@ -7,7 +7,7 @@ import pytest from testgen.api.runs import get_profiling_run, get_test_run -from testgen.common.models.hygiene_issue import IssueLikelihoodCounts +from testgen.common.models.hygiene_issue import HygieneIssueCounts, IssueCounts, PotentialPiiCounts from testgen.common.models.test_result import ResultStatusCounts pytestmark = pytest.mark.unit @@ -84,9 +84,9 @@ def test_get_test_run_completed(mock_tr_cls, mock_result_cls, mock_session): assert result.table_group_id == TABLE_GROUP_ID assert result.result is not None assert result.result.score == 0.95 - assert result.result.tests.passed == 90 - assert result.result.tests.failed == 5 - assert result.result.tests.dismissed == 12 + assert result.result.result_counts.passed == 90 + assert result.result.result_counts.failed == 5 + assert result.result.result_counts.dismissed == 12 @patch(f"{MODULE}.TestRun") @@ -111,8 +111,10 @@ def test_get_test_run_pending_no_run(mock_tr_cls): def test_get_profiling_run_completed(mock_pr_cls, mock_issue_cls): job = _mock_job() mock_pr_cls.get_by_id_or_job.return_value = _mock_profiling_run() - mock_issue_cls.count_by_likelihood.return_value = IssueLikelihoodCounts( - definite=5, likely=3, possible=8, dismissed=2, + mock_issue_cls.count_for_run.return_value = IssueCounts( + hygiene_issues=HygieneIssueCounts(definite=5, likely=3, possible=8), + potential_pii=PotentialPiiCounts(high=4, moderate=6), + dismissed=2, ) result = get_profiling_run(job) @@ -123,9 +125,12 @@ def test_get_profiling_run_completed(mock_pr_cls, mock_issue_cls): assert result.result is not None assert result.result.score == 0.88 assert result.result.table_ct == 10 - assert result.result.issues.definite == 5 - assert result.result.issues.likely == 3 - assert result.result.issues.dismissed == 2 + assert result.result.issue_counts.hygiene_issues.definite == 5 + assert result.result.issue_counts.hygiene_issues.likely == 3 + assert result.result.issue_counts.hygiene_issues.possible == 8 + assert result.result.issue_counts.potential_pii.high == 4 + assert result.result.issue_counts.potential_pii.moderate == 6 + assert result.result.issue_counts.dismissed == 2 @patch(f"{MODULE}.ProfilingRun") From c5a8feed72276fe129d6f9a690848a1999a3591a Mon Sep 17 00:00:00 2001 From: Diogo Basto Date: Thu, 11 Jun 2026 16:55:43 +0100 Subject: [PATCH 29/78] =?UTF-8?q?feat(mcp):=20add=20monitor=20read=20tools?= =?UTF-8?q?=20=E2=80=94=20group=20summary=20and=20per-table=20inventory=20?= =?UTF-8?q?(TG-1090)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Adds get_monitor_summary and list_monitored_tables so an MCP client can ask "how healthy are this group's monitors?" and "which tables have anomalies?" without falling back to the dashboard. Both gate on view; both resolve the table group's linked monitor_test_suite_id and return the explicit string "This table group is not monitored." when the link is absent, distinct from the unified missing-or-inaccessible error. Backing model methods on TableGroup: list_monitor_table_summaries and get_monitor_group_summary, sharing one CTE keyed off the monitor suite's lookback (overridable per call). The monitors dashboard's raw-SQL helpers delegate to these methods instead of duplicating the CTE. Monitor type values cross the user/DB boundary via a new MonitorType StrEnum and parse_monitor_type / format_monitor_type — callers and output use the short labels (freshness / volume / schema / metric); internal test_type codes never leak. ``MonitorTableSort`` carries the four sort_by values (table_name / anomaly_count_desc / latest_update_desc / row_count_change_desc); the renamed row_count_change_desc backs a TableGroup._MONITOR_SORT_COLUMN entry that sorts by the signed delta (monitor_tables.row_count - baseline_tables.previous_row_count). mcp-patterns.md documents the dedicated monitor read surface and the user-label ↔ DB-code mapping; replaces the prior "roadmap, out of scope" note in the monitor-suites section. Output details: - get_monitor_summary status cells use "no results yet or not configured" for the pending state — reads as either "no signal yet" or "monitor not configured" rather than implying user error. - get_monitor_summary converts next_scheduled_run to UTC before MdDoc rendering. JobSchedule yields tz-aware datetimes in the schedule's cron_tz; MdDoc blindly appends a "UTC" suffix, so without the conversion a local-tz timestamp gets mislabeled. - get_monitor_summary renders the lookback window as two MdDoc fields ("Window start", "Window end"), routed through the formatter so microseconds are stripped and the UTC suffix is added uniformly. - format_page_info returns empty when the requested page is past the last page — beyond-last responses show only the caller's "No rows on page N (total: M)" message instead of a nonsense `Showing 6–5 of 5 (page 6).` range above it. Applies to every paginated MCP tool, not just monitors. - list_monitored_tables Schema column shows just the anomaly count (or pending / error); the verbose detail moves to a sibling "Schema change" column rendering "Table added with N columns.", "Table dropped with N columns.", or a per-kind breakdown ("1 column added. 2 columns dropped."). - list_monitored_tables Row count change column shows the signed delta (current − previous_row_count) instead of the raw current count, reflecting net change across the window. Co-Authored-By: Claude Opus 4.7 --- testgen/common/enums.py | 9 + testgen/common/models/table_group.py | 429 +++++++++++++++++++++++ testgen/mcp/server.py | 3 + testgen/mcp/tools/common.py | 89 ++++- testgen/mcp/tools/monitors.py | 309 +++++++++++++++++ testgen/ui/views/monitors_dashboard.py | 277 +++------------ tests/unit/mcp/test_tools_common.py | 139 ++++++++ tests/unit/mcp/test_tools_monitors.py | 458 +++++++++++++++++++++++++ 8 files changed, 1486 insertions(+), 227 deletions(-) create mode 100644 testgen/mcp/tools/monitors.py create mode 100644 tests/unit/mcp/test_tools_monitors.py diff --git a/testgen/common/enums.py b/testgen/common/enums.py index f315023d..84fffb65 100644 --- a/testgen/common/enums.py +++ b/testgen/common/enums.py @@ -102,3 +102,12 @@ class PiiRisk(StrEnum): """Risk level extracted from PII issue ``detail`` strings via ``priority`` hybrid.""" HIGH = "High" MODERATE = "Moderate" + + +class MonitorType(StrEnum): + """Stored ``test_type`` values for the four monitor test types. Surfaced to users + as the lowercase short labels (freshness / volume / schema / metric).""" + FRESHNESS = "Freshness_Trend" + VOLUME = "Volume_Trend" + SCHEMA = "Schema_Drift" + METRIC = "Metric_Trend" diff --git a/testgen/common/models/table_group.py b/testgen/common/models/table_group.py index f30ab780..147b1ef6 100644 --- a/testgen/common/models/table_group.py +++ b/testgen/common/models/table_group.py @@ -107,6 +107,118 @@ class TableGroupSummary(EntityMinimal): total_count: int = 0 +@dataclass +class MonitorTableSummary: + """One row per monitored table in a table group's lookback window. + + Per-type ``*_anomalies`` count result_code = 0 results; ``*_is_training`` reflects the + latest run's training-mode state (result_code = -1); ``*_is_pending`` is True when no + result of that type exists in the window (monitor not yet configured / executed). + Schema is special-cased: ``*_is_pending`` only when no events at all, and ``table_state`` + captures whether the table was added / dropped / column-modified in the window. + """ + table_name: str + lookback: int + lookback_start: datetime | None + lookback_end: datetime | None + freshness_anomalies: int + volume_anomalies: int + schema_anomalies: int + metric_anomalies: int + freshness_is_training: bool | None + volume_is_training: bool | None + metric_is_training: bool | None + freshness_is_pending: bool + volume_is_pending: bool + schema_is_pending: bool + metric_is_pending: bool + freshness_error_message: str | None + volume_error_message: str | None + schema_error_message: str | None + metric_error_message: str | None + latest_update: datetime | None + row_count: int | None + previous_row_count: int | None + column_adds: int + column_drops: int + column_mods: int + table_state: str | None + + +@dataclass +class MonitorGroupSummary: + """Aggregated monitor health for a single table group across its lookback window. + + Booleans are group-wide: ``*_has_errors`` is True iff at least one monitored table errored; + ``*_is_training`` is True iff every monitored table is in training (and at least one is); + ``*_is_pending`` is True iff every monitored table is pending. + """ + lookback: int + lookback_start: datetime | None + lookback_end: datetime | None + total_monitored_tables: int + freshness_anomalies: int + volume_anomalies: int + schema_anomalies: int + metric_anomalies: int + freshness_has_errors: bool + volume_has_errors: bool + schema_has_errors: bool + metric_has_errors: bool + freshness_is_training: bool + volume_is_training: bool + metric_is_training: bool + freshness_is_pending: bool + volume_is_pending: bool + schema_is_pending: bool + metric_is_pending: bool + + +_ANOMALY_TYPE_TO_COLUMN: dict[str, str] = { + "Freshness_Trend": "freshness_anomalies", + "Volume_Trend": "volume_anomalies", + "Schema_Drift": "schema_anomalies", + "Metric_Trend": "metric_anomalies", +} + +_MONITOR_SORT_COLUMN: dict[str, str] = { + "table_name": "LOWER(monitor_tables.table_name)", + "freshness_anomalies": "monitor_tables.freshness_anomalies", + "volume_anomalies": "monitor_tables.volume_anomalies", + "schema_anomalies": "monitor_tables.schema_anomalies", + "metric_anomalies": "monitor_tables.metric_anomalies", + "total_anomalies": ( + "monitor_tables.freshness_anomalies + monitor_tables.volume_anomalies" + " + monitor_tables.schema_anomalies + monitor_tables.metric_anomalies" + ), + "latest_update": "monitor_tables.latest_update", + "row_count": "monitor_tables.row_count", + "row_count_change": "(monitor_tables.row_count - baseline_tables.previous_row_count)", +} + + +def _build_monitor_order_clause(sort_by: str | None, anomaly_type: str | None) -> str: + """Build an ORDER BY clause for the monitor-changes-by-tables query. + + ``sort_by`` is a field name with optional ``_desc`` suffix. When ``sort_by`` is + ``total_anomalies_desc`` and ``anomaly_type`` is set, the sort collapses to that + type's column so callers see "most anomalies of the filtered type first." Falls + back to ``LOWER(table_name) ASC`` for unknown / missing values. + """ + descending = False + field = sort_by + if field and field.endswith("_desc"): + descending = True + field = field[: -len("_desc")] + + if field == "total_anomalies" and anomaly_type is not None: + field = _ANOMALY_TYPE_TO_COLUMN[anomaly_type] + + column = _MONITOR_SORT_COLUMN.get(field or "table_name", _MONITOR_SORT_COLUMN["table_name"]) + direction = "DESC" if descending else "ASC" + return f"ORDER BY {column} {direction} NULLS LAST" + + class TableGroup(Entity): __tablename__ = "table_groups" @@ -513,6 +625,323 @@ def _list_with_activity( items = [TableGroupListItem(**row) for row in rows] return items, total + @classmethod + def list_monitor_table_summaries( + cls, + table_group_id: str | UUID, + *, + anomaly_types: list[str] | None = None, + sort_by: str | None = None, + lookback_override: int | None = None, + table_name_filter: str | None = None, + page: int = 1, + limit: int = 20, + ) -> tuple[list[MonitorTableSummary], int]: + """Per-monitored-table summary, paginated within the group's lookback window. + + ``anomaly_types`` are internal ``test_type`` values (``Freshness_Trend`` / + ``Volume_Trend`` / ``Schema_Drift`` / ``Metric_Trend``) — filters to tables with + at least one anomaly of any listed type in the window. ``lookback_override`` + replaces the suite-configured lookback for ad-hoc views. ``sort_by`` accepts + the dataclass field names (e.g. ``freshness_anomalies``, ``latest_update``, + ``row_count``, ``total_anomalies``) suffixed with ``_desc`` for descending + order; ``table_name`` is the default and sorts case-insensitively. + """ + # Run the CTE twice (rows + COUNT) so the total reflects every matching table, + # not just rows on the requested page — a window function would short-cut to 0 + # on out-of-range pages. + query, params = cls._monitor_changes_by_tables_query( + table_group_id, + anomaly_types=anomaly_types, + sort_by=sort_by, + lookback_override=lookback_override, + table_name_filter=table_name_filter, + limit=limit, + offset=(page - 1) * limit, + ) + count_query, count_params = cls._monitor_changes_by_tables_query( + table_group_id, + anomaly_types=anomaly_types, + lookback_override=lookback_override, + table_name_filter=table_name_filter, + ) + session = get_current_session() + rows = session.execute(text(query), params).mappings().all() + total = session.scalar( + text(f"SELECT COUNT(*) FROM ({count_query}) AS subquery"), count_params + ) or 0 + return [MonitorTableSummary(**row) for row in rows], int(total) + + @classmethod + def get_monitor_group_summary( + cls, + table_group_id: str | UUID, + *, + lookback_override: int | None = None, + ) -> MonitorGroupSummary: + """Group-level monitor health across the lookback window. + + Aggregates per-table results from ``list_monitor_table_summaries`` into a single + row. Returns zeroed counts and all-pending booleans when no tables are + monitored — callers detect "not monitored" upstream via the linked suite. + """ + inner_query, params = cls._monitor_changes_by_tables_query( + table_group_id, lookback_override=lookback_override, + ) + query = f""" + SELECT + COALESCE(MAX(lookback), :default_lookback)::INTEGER AS lookback, + MIN(lookback_start) AS lookback_start, + MAX(lookback_end) AS lookback_end, + COUNT(*)::INTEGER AS total_monitored_tables, + COALESCE(SUM(freshness_anomalies), 0)::INTEGER AS freshness_anomalies, + COALESCE(SUM(volume_anomalies), 0)::INTEGER AS volume_anomalies, + COALESCE(SUM(schema_anomalies), 0)::INTEGER AS schema_anomalies, + COALESCE(SUM(metric_anomalies), 0)::INTEGER AS metric_anomalies, + COALESCE(BOOL_OR(freshness_error_message IS NOT NULL), FALSE) AS freshness_has_errors, + COALESCE(BOOL_OR(volume_error_message IS NOT NULL), FALSE) AS volume_has_errors, + COALESCE(BOOL_OR(schema_error_message IS NOT NULL), FALSE) AS schema_has_errors, + COALESCE(BOOL_OR(metric_error_message IS NOT NULL), FALSE) AS metric_has_errors, + COALESCE( + BOOL_OR(freshness_is_training) AND BOOL_AND(freshness_is_training OR freshness_is_pending), + FALSE + ) AS freshness_is_training, + COALESCE( + BOOL_OR(volume_is_training) AND BOOL_AND(volume_is_training OR volume_is_pending), + FALSE + ) AS volume_is_training, + COALESCE( + BOOL_OR(metric_is_training) AND BOOL_AND(metric_is_training OR metric_is_pending), + FALSE + ) AS metric_is_training, + COALESCE(BOOL_AND(freshness_is_pending), TRUE) AS freshness_is_pending, + COALESCE(BOOL_AND(volume_is_pending), TRUE) AS volume_is_pending, + COALESCE(BOOL_AND(schema_is_pending), TRUE) AS schema_is_pending, + COALESCE(BOOL_AND(metric_is_pending), TRUE) AS metric_is_pending + FROM ({inner_query}) AS subquery + """ + params = {**params, "default_lookback": lookback_override or 1} + # Outer query has no GROUP BY — aggregates over zero rows still yield one + # COALESCE'd row, so .first() never returns None here. + row = get_current_session().execute(text(query), params).mappings().first() + return MonitorGroupSummary(**row) + + @classmethod + def _monitor_changes_by_tables_query( + cls, + table_group_id: str | UUID, + *, + anomaly_types: list[str] | None = None, + sort_by: str | None = None, + lookback_override: int | None = None, + table_name_filter: str | None = None, + limit: int | None = None, + offset: int | None = None, + ) -> tuple[str, dict]: + """Build the CTE that produces one ``MonitorTableSummary``-shaped row per table. + + Shared by ``list_monitor_table_summaries`` and ``get_monitor_group_summary`` + and by the monitors dashboard. ``anomaly_types`` filters the outer SELECT (so + the group summary can omit it and get all tables). ``sort_by`` is the + dataclass field name with optional ``_desc`` suffix — the caller (MCP / + dashboard) is responsible for validating it against a known set. + """ + lookback_expr = ( + ":lookback_override" if lookback_override is not None + else "COALESCE(test_suites.monitor_lookback, 1)" + ) + + anomaly_filter_clause = "" + if anomaly_types: + columns = [_ANOMALY_TYPE_TO_COLUMN[t] for t in anomaly_types] + anomaly_filter_clause = ( + "WHERE (" + + " OR ".join(f"monitor_tables.{col} > 0" for col in columns) + + ")" + ) + + sort_anomaly_type = anomaly_types[0] if anomaly_types and len(anomaly_types) == 1 else None + order_clause = _build_monitor_order_clause(sort_by, sort_anomaly_type) + + query = f""" + WITH ranked_test_runs AS ( + SELECT + test_runs.id, + test_runs.test_starttime, + {lookback_expr} AS lookback, + ROW_NUMBER() OVER (PARTITION BY test_runs.test_suite_id ORDER BY test_runs.test_starttime DESC) AS position + FROM table_groups + INNER JOIN test_runs + ON (test_runs.test_suite_id = table_groups.monitor_test_suite_id) + INNER JOIN test_suites + ON (table_groups.monitor_test_suite_id = test_suites.id) + WHERE table_groups.id = :table_group_id + ), + lookback_window AS ( + SELECT MIN(test_starttime) AS lookback_start + FROM ranked_test_runs + WHERE position <= lookback + ), + latest_tables AS ( + SELECT DISTINCT + table_chars.schema_name, + table_chars.table_name + FROM data_table_chars table_chars + CROSS JOIN lookback_window + WHERE table_chars.table_groups_id = :table_group_id + -- Include current tables and tables dropped within lookback window + AND (table_chars.drop_date IS NULL OR table_chars.drop_date >= lookback_window.lookback_start) + {"AND table_chars.table_name ILIKE :table_name_filter" if table_name_filter else ''} + ), + monitor_results AS ( + SELECT + latest_tables.table_name, + results.test_time, + results.test_type, + results.result_code, + ranked_test_runs.lookback, + ranked_test_runs.position, + ranked_test_runs.test_starttime, + -- result_code = -1 indicates training mode + CASE WHEN results.result_code = -1 THEN 1 ELSE 0 END AS is_training, + CASE WHEN results.test_type = 'Freshness_Trend' AND results.result_code = 0 THEN 1 ELSE 0 END AS freshness_anomaly, + CASE WHEN results.test_type = 'Volume_Trend' AND results.result_code = 0 THEN 1 ELSE 0 END AS volume_anomaly, + CASE WHEN results.test_type = 'Schema_Drift' AND results.result_code = 0 THEN 1 ELSE 0 END AS schema_anomaly, + CASE WHEN results.test_type = 'Metric_Trend' AND results.result_code = 0 THEN 1 ELSE 0 END AS metric_anomaly, + CASE WHEN results.test_type = 'Freshness_Trend' THEN results.result_signal ELSE NULL END AS freshness_interval, + CASE WHEN results.test_type = 'Volume_Trend' THEN results.result_signal::BIGINT ELSE NULL END AS row_count, + CASE WHEN results.test_type = 'Schema_Drift' THEN SPLIT_PART(results.result_signal, '|', 1) ELSE NULL END AS table_change, + CASE WHEN results.test_type = 'Schema_Drift' THEN NULLIF(SPLIT_PART(results.result_signal, '|', 2), '')::INT ELSE 0 END AS col_adds, + CASE WHEN results.test_type = 'Schema_Drift' THEN NULLIF(SPLIT_PART(results.result_signal, '|', 3), '')::INT ELSE 0 END AS col_drops, + CASE WHEN results.test_type = 'Schema_Drift' THEN NULLIF(SPLIT_PART(results.result_signal, '|', 4), '')::INT ELSE 0 END AS col_mods, + CASE WHEN results.result_status = 'Error' THEN results.result_message ELSE NULL END AS error_message + FROM latest_tables + LEFT JOIN ranked_test_runs ON TRUE + LEFT JOIN test_results AS results + ON results.test_run_id = ranked_test_runs.id + AND results.table_name = latest_tables.table_name + WHERE ranked_test_runs.position IS NULL + -- Also capture 1 run before the lookback to get baseline results + OR ranked_test_runs.position <= ranked_test_runs.lookback + 1 + ), + monitor_tables AS ( + SELECT + table_name, + MAX(lookback)::INTEGER AS lookback, + COALESCE(SUM(freshness_anomaly), 0)::INTEGER AS freshness_anomalies, + COALESCE(SUM(volume_anomaly), 0)::INTEGER AS volume_anomalies, + COALESCE(SUM(schema_anomaly), 0)::INTEGER AS schema_anomalies, + COALESCE(SUM(metric_anomaly), 0)::INTEGER AS metric_anomalies, + MAX(test_time - (COALESCE(NULLIF(freshness_interval, 'Unknown')::INTEGER, 0) * INTERVAL '1 minute')) + FILTER (WHERE test_type = 'Freshness_Trend' AND position = 1) AS latest_update, + MAX(row_count) FILTER (WHERE position = 1) AS row_count, + COALESCE(SUM(col_adds), 0)::INTEGER AS column_adds, + COALESCE(SUM(col_drops), 0)::INTEGER AS column_drops, + COALESCE(SUM(col_mods), 0)::INTEGER AS column_mods, + MAX(error_message) FILTER (WHERE test_type = 'Freshness_Trend' AND position = 1) AS freshness_error_message, + MAX(error_message) FILTER (WHERE test_type = 'Volume_Trend' AND position = 1) AS volume_error_message, + MAX(error_message) FILTER (WHERE test_type = 'Schema_Drift' AND position = 1) AS schema_error_message, + MAX(error_message) FILTER (WHERE test_type = 'Metric_Trend' AND position = 1) AS metric_error_message, + BOOL_OR(is_training = 1) FILTER (WHERE test_type = 'Freshness_Trend' AND position = 1) AS freshness_is_training, + BOOL_OR(is_training = 1) FILTER (WHERE test_type = 'Volume_Trend' AND position = 1) AS volume_is_training, + BOOL_OR(is_training = 1) FILTER (WHERE test_type = 'Metric_Trend' AND position = 1) AS metric_is_training, + BOOL_OR(test_type = 'Freshness_Trend') IS NOT TRUE AS freshness_is_pending, + BOOL_OR(test_type = 'Volume_Trend') IS NOT TRUE AS volume_is_pending, + -- Schema monitor only creates results on schema changes (Failed) + -- Mark it as pending only if there are no results of any test type + BOOL_OR(test_time IS NOT NULL) IS NOT TRUE AS schema_is_pending, + BOOL_OR(test_type = 'Metric_Trend') IS NOT TRUE AS metric_is_pending, + CASE + -- Mark as Dropped if latest Schema Drift result for the table indicates it was dropped + WHEN (ARRAY_AGG(table_change ORDER BY test_time DESC) FILTER (WHERE table_change IS NOT NULL))[1] = 'D' + THEN 'dropped' + -- Only mark as Added if latest change does not indicate a drop + WHEN MAX(CASE WHEN table_change = 'A' THEN 1 ELSE 0 END) = 1 + THEN 'added' + WHEN SUM(schema_anomaly) > 0 + THEN 'modified' + ELSE NULL + END AS table_state + FROM monitor_results + -- Only aggregate within lookback runs + WHERE position IS NULL OR position <= COALESCE(lookback, 1) + GROUP BY table_name + ), + table_bounds AS ( + SELECT + table_name, + MIN(position) AS min_position, + MAX(position) AS max_position + FROM monitor_results + WHERE position IS NOT NULL + GROUP BY table_name + ), + baseline_tables AS ( + SELECT + monitor_results.table_name, + MIN(monitor_results.test_starttime) FILTER ( + WHERE monitor_results.position = LEAST(monitor_results.lookback + 1, table_bounds.max_position) + ) AS lookback_start, + MAX(monitor_results.test_starttime) FILTER ( + WHERE monitor_results.position = GREATEST(1, table_bounds.min_position) + ) AS lookback_end, + MAX(monitor_results.row_count) FILTER ( + WHERE monitor_results.test_type = 'Volume_Trend' + AND monitor_results.position = LEAST(monitor_results.lookback + 1, table_bounds.max_position) + ) AS previous_row_count + FROM monitor_results + JOIN table_bounds ON monitor_results.table_name = table_bounds.table_name + GROUP BY monitor_results.table_name + ) + SELECT + monitor_tables.table_name, + monitor_tables.lookback, + baseline_tables.lookback_start, + baseline_tables.lookback_end, + monitor_tables.freshness_anomalies, + monitor_tables.volume_anomalies, + monitor_tables.schema_anomalies, + monitor_tables.metric_anomalies, + monitor_tables.freshness_is_training, + monitor_tables.volume_is_training, + monitor_tables.metric_is_training, + monitor_tables.freshness_is_pending, + monitor_tables.volume_is_pending, + monitor_tables.schema_is_pending, + monitor_tables.metric_is_pending, + monitor_tables.freshness_error_message, + monitor_tables.volume_error_message, + monitor_tables.schema_error_message, + monitor_tables.metric_error_message, + monitor_tables.latest_update, + monitor_tables.row_count, + baseline_tables.previous_row_count, + monitor_tables.column_adds, + monitor_tables.column_drops, + monitor_tables.column_mods, + monitor_tables.table_state + FROM monitor_tables + LEFT JOIN baseline_tables ON monitor_tables.table_name = baseline_tables.table_name + {anomaly_filter_clause} + {order_clause} + {"LIMIT :limit" if limit is not None else ""} + {"OFFSET :offset" if offset is not None else ""} + """ + + escaped_name_filter = ( + table_name_filter.replace("_", "\\_") if table_name_filter else None + ) + params: dict = {"table_group_id": str(table_group_id)} + if escaped_name_filter is not None: + params["table_name_filter"] = f"%{escaped_name_filter}%" + if lookback_override is not None: + params["lookback_override"] = lookback_override + if limit is not None: + params["limit"] = limit + if offset is not None: + params["offset"] = offset + return query, params + @classmethod def is_in_use(cls, ids: list[str]) -> bool: test_suites = TestSuite.select_minimal_where(TestSuite.table_groups_id.in_(ids)) diff --git a/testgen/mcp/server.py b/testgen/mcp/server.py index aaf42471..0be45764 100644 --- a/testgen/mcp/server.py +++ b/testgen/mcp/server.py @@ -172,6 +172,7 @@ def build_mcp_server( search_hygiene_issues, update_hygiene_issue, ) + from testgen.mcp.tools.monitors import get_monitor_summary, list_monitored_tables from testgen.mcp.tools.notifications import ( create_notification, delete_notification, @@ -316,6 +317,8 @@ def safe_prompt(fn): safe_tool(compare_profiling_runs) safe_tool(get_profiling_trends) safe_tool(get_schema_history) + safe_tool(get_monitor_summary) + safe_tool(list_monitored_tables) safe_tool(run_tests) safe_tool(run_profiling) safe_tool(cancel_test_run) diff --git a/testgen/mcp/tools/common.py b/testgen/mcp/tools/common.py index c65a91e7..366537de 100644 --- a/testgen/mcp/tools/common.py +++ b/testgen/mcp/tools/common.py @@ -6,7 +6,15 @@ from sqlalchemy import select from testgen.common.date_service import parse_since -from testgen.common.enums import Disposition, ImpactDimension, IssueLikelihood, JobStatus, PiiRisk, QualityDimension +from testgen.common.enums import ( + Disposition, + ImpactDimension, + IssueLikelihood, + JobStatus, + MonitorType, + PiiRisk, + QualityDimension, +) from testgen.common.flavors import FLAVOR_CODE_TO_FAMILY, FLAVOR_CODE_TO_LABEL, SqlFlavorLabel from testgen.common.models import get_current_session from testgen.common.models.connection import Connection @@ -57,6 +65,7 @@ class DocGroup(StrEnum): DISCOVER = "Discover what TestGen knows about" INVESTIGATE = "Investigate quality issues" BROWSE_PROFILING = "Browse profiling results" + MONITORS = "Browse monitor health and events" TRIGGER = "Trigger profiling, tests, and test generation" SCORING = "Track data quality scores" MANAGE = "Manage TestGen configuration" @@ -340,6 +349,58 @@ def next_scheduled_run( return min(s.get_sample_triggering_timestamps(1)[0] for s in schedules) +# Monitor type — internal ``test_type`` value ↔ user-facing short label (the form +# callers pass and the form rendered in output). +_MONITOR_TYPE_USER_TO_DB: dict[str, MonitorType] = { + "freshness": MonitorType.FRESHNESS, + "volume": MonitorType.VOLUME, + "schema": MonitorType.SCHEMA, + "metric": MonitorType.METRIC, +} +_MONITOR_TYPE_DB_TO_USER: dict[MonitorType, str] = {v: k for k, v in _MONITOR_TYPE_USER_TO_DB.items()} + + +def parse_monitor_type(value: str) -> MonitorType: + """Validate a user-facing monitor type label and return the stored ``MonitorType``. + + Accepts ``freshness`` / ``volume`` / ``schema`` / ``metric``. + """ + db_value = _MONITOR_TYPE_USER_TO_DB.get(value) + if db_value is None: + valid = ", ".join(_MONITOR_TYPE_USER_TO_DB) + raise MCPUserError(f"Invalid monitor_type `{value}`. Valid values: {valid}") + return db_value + + +def format_monitor_type(value: MonitorType | str) -> str: + """Map a stored monitor ``test_type`` to its user-facing short label.""" + try: + return _MONITOR_TYPE_DB_TO_USER[MonitorType(value)] + except ValueError: + return str(value) + + +class MonitorTableSort(StrEnum): + """User-facing values accepted for the ``sort_by`` argument on ``list_monitored_tables``. + + When an ``anomaly_type`` filter is set, ``anomaly_count_desc`` sorts by that type's + count; with no filter, it sorts by total anomalies across all types. + """ + + TABLE_NAME = "table_name" + ANOMALY_COUNT_DESC = "anomaly_count_desc" + LATEST_UPDATE_DESC = "latest_update_desc" + ROW_COUNT_CHANGE_DESC = "row_count_change_desc" + + +def parse_monitor_table_sort(value: str) -> MonitorTableSort: + try: + return MonitorTableSort(value) + except ValueError as err: + valid = ", ".join(s.value for s in MonitorTableSort) + raise MCPUserError(f"Invalid sort_by `{value}`. Valid values: {valid}") from err + + def parse_disposition(value: str) -> Disposition: """Validate a user-facing disposition label and return the stored ``Disposition``. @@ -529,10 +590,17 @@ def resolve_issue_type(name: str) -> str: def format_page_info(total: int, page: int, limit: int) -> str: - """Shared pagination summary line for MCP tool output.""" + """Shared pagination summary line for MCP tool output. + + Returns empty for a zero total *or* when ``page`` is past the last page \u2014 + the caller's own "no rows on page N" message is more useful than a + ``Showing 6\u20135 of 5`` nonsense range. + """ if total == 0: return "" start = (page - 1) * limit + 1 + if start > total: + return "" end = min(start + limit - 1, total) return f"Showing {start}\u2013{end} of {total} (page {page})." @@ -595,6 +663,23 @@ def resolve_test_suite(test_suite_id: str) -> TestSuite: return suite +def resolve_monitored_table_group(table_group_id: str) -> tuple[TableGroup, TestSuite | None]: + """Resolve a table group ID and look up its linked monitor suite. + + Returns ``(table_group, monitor_suite)``. ``monitor_suite`` is ``None`` when the + table group has no ``monitor_test_suite_id`` set, or that pointer doesn't resolve + to an ``is_monitor=True`` suite — callers render the "not monitored" output in + that case rather than raising an inaccessible error. Raises + ``MCPResourceNotAccessible`` when the table group itself is missing or out of + scope. + """ + tg = resolve_table_group(table_group_id) + if tg.monitor_test_suite_id is None: + return tg, None + suite = TestSuite.get(tg.monitor_test_suite_id, TestSuite.is_monitor.is_(True)) + return tg, suite + + def resolve_profiling_run(job_execution_id: str) -> ProfilingRun: """Resolve a profiling run by id-or-JE-id, scoped to allowed projects. diff --git a/testgen/mcp/tools/monitors.py b/testgen/mcp/tools/monitors.py new file mode 100644 index 00000000..b3837eea --- /dev/null +++ b/testgen/mcp/tools/monitors.py @@ -0,0 +1,309 @@ +from datetime import UTC + +from testgen.common.enums import MonitorType +from testgen.common.models import with_database_session +from testgen.common.models.scheduler import RUN_MONITORS_JOB_KEY +from testgen.common.models.table_group import MonitorTableSummary, TableGroup +from testgen.mcp.exceptions import MCPUserError +from testgen.mcp.permissions import mcp_permission +from testgen.mcp.tools.common import ( + DocGroup, + MonitorTableSort, + format_page_footer, + format_page_info, + next_scheduled_run, + parse_monitor_table_sort, + parse_monitor_type, + resolve_monitored_table_group, + validate_limit, + validate_page, +) +from testgen.mcp.tools.markdown import MdDoc + +_DOC_GROUP = DocGroup.MONITORS + +_NOT_MONITORED_OUTPUT = "This table group is not monitored." + +_MONITOR_LABEL: dict[MonitorType, str] = { + MonitorType.FRESHNESS: "Freshness", + MonitorType.VOLUME: "Volume", + MonitorType.SCHEMA: "Schema", + MonitorType.METRIC: "Metric", +} + +# Maps a ``MonitorTableSort`` value to the model-method ``sort_by`` argument. The +# anomaly-count sort uses the ``total_anomalies`` field — the model layer collapses +# it to the filtered type's column when exactly one ``anomaly_type`` is set. +_SORT_TO_MODEL_FIELD: dict[MonitorTableSort, str] = { + MonitorTableSort.TABLE_NAME: "table_name", + MonitorTableSort.ANOMALY_COUNT_DESC: "total_anomalies_desc", + MonitorTableSort.LATEST_UPDATE_DESC: "latest_update_desc", + MonitorTableSort.ROW_COUNT_CHANGE_DESC: "row_count_change_desc", +} + + +@with_database_session +@mcp_permission("view") +def get_monitor_summary(table_group_id: str, lookback: int | None = None) -> str: + """Get monitor health for a table group: per-type anomaly counts, error / training / pending status, and the active lookback window. + + Args: + table_group_id: UUID of the table group, e.g. from ``list_table_groups``. + lookback: Number of monitor runs to summarize. Omit to use the lookback runs configured for the table group. + """ + if lookback is not None and not 1 <= lookback <= 365: + raise MCPUserError(f"Invalid lookback `{lookback}`: must be between 1 and 365.") + + tg, monitor_suite = resolve_monitored_table_group(table_group_id) + if monitor_suite is None: + return _NOT_MONITORED_OUTPUT + + summary = TableGroup.get_monitor_group_summary(tg.id, lookback_override=lookback) + next_run = next_scheduled_run( + RUN_MONITORS_JOB_KEY, {"test_suite_id": str(monitor_suite.id)}, tg.project_code + ) + + doc = MdDoc() + doc.heading(1, f"Monitor summary for `{tg.table_groups_name}`") + doc.field("Project", tg.project_code, code=True) + doc.field("Monitored tables", summary.total_monitored_tables) + lookback_label = f"{summary.lookback} runs" + if lookback is not None: + lookback_label += " (override)" + doc.field("Lookback", lookback_label) + if summary.lookback_start: + doc.field("Window start", summary.lookback_start) + if summary.lookback_end: + doc.field("Window end", summary.lookback_end) + if next_run: + # Next-firing timestamps from JobSchedule are tz-aware in the schedule's + # configured cron_tz. MdDoc renders any datetime with a blind "UTC" suffix, + # so convert here to avoid claiming a local-tz timestamp is UTC. + doc.field("Next scheduled run", next_run.astimezone(UTC)) + else: + doc.field("Next scheduled run", "not scheduled") + + doc.heading(2, "Anomalies by type") + doc.table( + ["Monitor", "Anomalies", "Status"], + [ + [ + _MONITOR_LABEL[MonitorType.FRESHNESS], + summary.freshness_anomalies, + _summary_status( + is_pending=summary.freshness_is_pending, + is_training=summary.freshness_is_training, + has_errors=summary.freshness_has_errors, + ), + ], + [ + _MONITOR_LABEL[MonitorType.VOLUME], + summary.volume_anomalies, + _summary_status( + is_pending=summary.volume_is_pending, + is_training=summary.volume_is_training, + has_errors=summary.volume_has_errors, + ), + ], + [ + _MONITOR_LABEL[MonitorType.SCHEMA], + summary.schema_anomalies, + _summary_status( + is_pending=summary.schema_is_pending, + is_training=False, + has_errors=summary.schema_has_errors, + ), + ], + [ + _MONITOR_LABEL[MonitorType.METRIC], + summary.metric_anomalies, + _summary_status( + is_pending=summary.metric_is_pending, + is_training=summary.metric_is_training, + has_errors=summary.metric_has_errors, + ), + ], + ], + ) + + return doc.render() + + +@with_database_session +@mcp_permission("view") +def list_monitored_tables( + table_group_id: str, + anomaly_type: str | None = None, + sort_by: str | None = None, + limit: int = 20, + page: int = 1, +) -> str: + """List monitored tables in a table group with per-type anomaly counts, training / pending / error status, latest update timestamp, and row count change. + + Args: + table_group_id: UUID of the table group, e.g. from ``list_table_groups``. + anomaly_type: Filter to tables with at least one anomaly of this type. One of ``freshness`` / ``volume`` / ``schema`` / ``metric``. + sort_by: Sort order. One of ``table_name`` (default, case-insensitive ascending), ``anomaly_count_desc`` (sorts by the filtered type when ``anomaly_type`` is set, total anomalies otherwise), ``latest_update_desc``, ``row_count_change_desc``. + limit: Page size (default 20, max 100). + page: Page number starting at 1 (default 1). + """ + validate_page(page) + validate_limit(limit, 100) + + monitor_type = parse_monitor_type(anomaly_type) if anomaly_type is not None else None + sort = parse_monitor_table_sort(sort_by) if sort_by is not None else None + + tg, monitor_suite = resolve_monitored_table_group(table_group_id) + if monitor_suite is None: + return _NOT_MONITORED_OUTPUT + + rows, total = TableGroup.list_monitor_table_summaries( + tg.id, + anomaly_types=[monitor_type.value] if monitor_type is not None else None, + sort_by=_SORT_TO_MODEL_FIELD[sort] if sort is not None else None, + page=page, + limit=limit, + ) + + doc = MdDoc() + scope = f" — anomaly type `{anomaly_type}`" if anomaly_type else "" + doc.heading(1, f"Monitored tables in `{tg.table_groups_name}`{scope}") + page_info = format_page_info(total, page, limit) + if page_info: + doc.text(page_info) + + if not rows: + if page > 1: + doc.text(f"No tables on page {page} (total: {total}).") + else: + doc.text("_No monitored tables match this filter._") + return doc.render() + + doc.table( + [ + "Table", + "Freshness", + "Volume", + "Schema", + "Schema change", + "Metric", + "Latest update", + "Row count change", + ], + [ + [ + row.table_name, + _format_monitor_cell( + count=row.freshness_anomalies, + is_pending=row.freshness_is_pending, + is_training=row.freshness_is_training, + has_error=row.freshness_error_message is not None, + ), + _format_monitor_cell( + count=row.volume_anomalies, + is_pending=row.volume_is_pending, + is_training=row.volume_is_training, + has_error=row.volume_error_message is not None, + ), + _format_schema_cell(row), + _format_schema_change(row), + _format_monitor_cell( + count=row.metric_anomalies, + is_pending=row.metric_is_pending, + is_training=row.metric_is_training, + has_error=row.metric_error_message is not None, + ), + row.latest_update, + _format_row_count_change(row), + ] + for row in rows + ], + code=[0], + ) + + if footer := format_page_footer(total, page, limit): + doc.text(footer) + + return doc.render() + + +def _summary_status(*, is_pending: bool, is_training: bool, has_errors: bool) -> str: + """Single status word for the group-level summary cell.""" + if has_errors: + return "error" + if is_pending: + return "no results yet or not configured" + if is_training: + return "training" + return "ok" + + +def _format_monitor_cell( + *, + count: int, + is_pending: bool, + is_training: bool | None, + has_error: bool, +) -> str: + """Render a per-table, per-monitor cell (non-schema).""" + if has_error: + return "error" + if is_pending: + return "pending" + if is_training: + return "training" + return str(count) + + +def _format_schema_cell(row: MonitorTableSummary) -> str: + """Schema anomaly count, or pending / error status. Verbose detail is in the + sibling ``Schema change`` column.""" + if row.schema_error_message is not None: + return "error" + if row.schema_is_pending: + return "pending" + return str(row.schema_anomalies) + + +def _format_schema_change(row: MonitorTableSummary) -> str | None: + """Verbose description of schema events in the lookback window. + + Mirrors the dashboard tooltip wording: "Table added with N columns.", + "Table dropped with N columns.", or a per-kind breakdown when the table was + modified ("1 column added. 2 columns dropped."). Empty when there are no + schema events or the monitor is in an unfinished state. + """ + if row.schema_is_pending or row.schema_error_message is not None: + return None + state = row.table_state + if state == "added": + return f"Table added with {row.column_adds} columns." + if state == "dropped": + return f"Table dropped with {row.column_drops} columns." + if state == "modified": + parts: list[str] = [] + if row.column_adds: + parts.append(f"{row.column_adds} column{'' if row.column_adds == 1 else 's'} added") + if row.column_drops: + parts.append(f"{row.column_drops} column{'' if row.column_drops == 1 else 's'} dropped") + if row.column_mods: + parts.append(f"{row.column_mods} column{'' if row.column_mods == 1 else 's'} modified") + return ". ".join(parts) + "." if parts else None + return None + + +def _format_row_count_change(row: MonitorTableSummary) -> str | None: + """Signed delta between the latest and pre-window row count. + + ``+1,234`` / ``-1,234`` / ``0`` when both endpoints are known; ``None`` (em-dash) + when either is missing (e.g. first run with no baseline). Sign reflects net + change across the window, not run-to-run variance. + """ + current = row.row_count + previous = row.previous_row_count + if current is None or previous is None: + return None + delta = current - previous + if delta == 0: + return "0" + return f"{delta:+,}" diff --git a/testgen/ui/views/monitors_dashboard.py b/testgen/ui/views/monitors_dashboard.py index e4fd3719..2719109a 100644 --- a/testgen/ui/views/monitors_dashboard.py +++ b/testgen/ui/views/monitors_dashboard.py @@ -1,3 +1,4 @@ +import dataclasses import logging from datetime import UTC, date, datetime from math import ceil @@ -24,7 +25,7 @@ from testgen.ui.navigation.page import Page from testgen.ui.navigation.router import Router from testgen.ui.queries.profiling_queries import get_tables_by_table_group -from testgen.ui.services.database_service import execute_db_query, fetch_all_from_db, fetch_one_from_db +from testgen.ui.services.database_service import fetch_all_from_db from testgen.ui.services.query_cache import ( get_monitor_schedule, get_project_summary, @@ -41,6 +42,15 @@ from testgen.ui.views.dialogs.manage_notifications import NotificationSettingsDialogBase from testgen.utils import make_json_safe +# Maps the user-facing ``anomaly_type_filter`` values supplied by the dashboard +# frontend to the internal ``test_type`` codes the model method expects. +_DASHBOARD_ANOMALY_TYPE_TO_DB: dict[str, str] = { + "freshness": "Freshness_Trend", + "volume": "Volume_Trend", + "schema": "Schema_Drift", + "metrics": "Metric_Trend", +} + PAGE_ICON = "apps_outage" PAGE_TITLE = "Monitors" LOG = logging.getLogger("testgen") @@ -106,7 +116,7 @@ def render( all_monitored_tables_count = 0 monitor_changes_summary = None auto_open_table = None - + current_page = int(current_page) items_per_page = int(items_per_page) page_start = current_page * items_per_page @@ -349,18 +359,22 @@ def get_monitor_changes_by_tables( limit: int | None = None, offset: int | None = None, ) -> list[dict]: - query, params = _monitor_changes_by_tables_query( + """Per-monitored-table summaries shaped for the dashboard's JSON payload. + + Returns plain dicts (rather than ``MonitorTableSummary`` dataclasses) because the + monitor-dashboard widget consumes the payload via ``make_json_safe``. Each row is + augmented with ``table_group_id`` to match the historical payload shape. + """ + page = 1 + (offset // limit) if limit and offset else 1 + summaries, _total = TableGroup.list_monitor_table_summaries( table_group_id, + anomaly_types=_dashboard_anomaly_types(anomaly_type_filter), + sort_by=_dashboard_sort_to_model(sort_field, sort_order), table_name_filter=table_name_filter, - anomaly_type_filter=anomaly_type_filter, - sort_field=sort_field, - sort_order=sort_order, - limit=limit, - offset=offset, + page=page, + limit=limit or 1000, ) - - results = fetch_all_from_db(query, params) - return [ dict(row) for row in results ] + return [{**dataclasses.asdict(s), "table_group_id": table_group_id} for s in summaries] @st.cache_data(show_spinner=False) @@ -369,228 +383,41 @@ def count_monitor_changes_by_tables( table_name_filter: str | None = None, anomaly_type_filter: list[str] | None = None, ) -> int: - query, params = _monitor_changes_by_tables_query( + _items, total = TableGroup.list_monitor_table_summaries( table_group_id, + anomaly_types=_dashboard_anomaly_types(anomaly_type_filter), table_name_filter=table_name_filter, - anomaly_type_filter=anomaly_type_filter, + page=1, + limit=1, ) - count_query = f"SELECT COUNT(*) AS count FROM ({query}) AS subquery" - result = execute_db_query(count_query, params) - return result or 0 + return total @st.cache_data(show_spinner=False) def summarize_monitor_changes(table_group_id: str) -> dict: - query, params = _monitor_changes_by_tables_query(table_group_id) - count_query = f""" - SELECT - lookback, - MIN(lookback_start) AS lookback_start, - MAX(lookback_end) AS lookback_end, - SUM(freshness_anomalies)::INTEGER AS freshness_anomalies, - SUM(volume_anomalies)::INTEGER AS volume_anomalies, - SUM(schema_anomalies)::INTEGER AS schema_anomalies, - SUM(metric_anomalies)::INTEGER AS metric_anomalies, - BOOL_OR(freshness_error_message IS NOT NULL) AS freshness_has_errors, - BOOL_OR(volume_error_message IS NOT NULL) AS volume_has_errors, - BOOL_OR(schema_error_message IS NOT NULL) AS schema_has_errors, - BOOL_OR(metric_error_message IS NOT NULL) AS metric_has_errors, - BOOL_OR(freshness_is_training) AND BOOL_AND(freshness_is_training OR freshness_is_pending) AS freshness_is_training, - BOOL_OR(volume_is_training) AND BOOL_AND(volume_is_training OR volume_is_pending) AS volume_is_training, - BOOL_OR(metric_is_training) AND BOOL_AND(metric_is_training OR metric_is_pending) AS metric_is_training, - BOOL_AND(freshness_is_pending) AS freshness_is_pending, - BOOL_AND(volume_is_pending) AS volume_is_pending, - BOOL_AND(schema_is_pending) AS schema_is_pending, - BOOL_AND(metric_is_pending) AS metric_is_pending - FROM ({query}) AS subquery - GROUP BY lookback - """ - - result = fetch_one_from_db(count_query, params) - return {**result} if result else { - "lookback": 0, - "freshness_anomalies": 0, - "volume_anomalies": 0, - "schema_anomalies": 0, - "metric_anomalies": 0, - "freshness_is_training": False, - "volume_is_training": False, - "metric_is_training": False, - "freshness_is_pending": False, - "volume_is_pending": False, - "schema_is_pending": False, - "metric_is_pending": False, - "freshness_has_errors": False, - "volume_has_errors": False, - "schema_has_errors": False, - "metric_has_errors": False, - } + return dataclasses.asdict(TableGroup.get_monitor_group_summary(table_group_id)) -def _monitor_changes_by_tables_query( - table_group_id: str, - table_name_filter: str | None = None, - anomaly_type_filter: list[str] | None = None, - sort_field: str | None = None, - sort_order: Literal["asc"] | Literal["desc"] | None = None, - limit: int | None = None, - offset: int | None = None, -) -> tuple[str, dict]: - query = f""" - WITH ranked_test_runs AS ( - SELECT - test_runs.id, - test_runs.test_starttime, - COALESCE(test_suites.monitor_lookback, 1) AS lookback, - ROW_NUMBER() OVER (PARTITION BY test_runs.test_suite_id ORDER BY test_runs.test_starttime DESC) AS position - FROM table_groups - INNER JOIN test_runs - ON (test_runs.test_suite_id = table_groups.monitor_test_suite_id) - INNER JOIN test_suites - ON (table_groups.monitor_test_suite_id = test_suites.id) - WHERE table_groups.id = :table_group_id - ), - lookback_window AS ( - SELECT MIN(test_starttime) AS lookback_start - FROM ranked_test_runs - WHERE position <= lookback - ), - latest_tables AS ( - SELECT DISTINCT - table_chars.schema_name, - table_chars.table_name - FROM data_table_chars table_chars - CROSS JOIN lookback_window - WHERE table_chars.table_groups_id = :table_group_id - -- Include current tables and tables dropped within lookback window - AND (table_chars.drop_date IS NULL OR table_chars.drop_date >= lookback_window.lookback_start) - {"AND table_chars.table_name ILIKE :table_name_filter" if table_name_filter else ''} - ), - monitor_results AS ( - SELECT - latest_tables.table_name, - results.test_time, - results.test_type, - results.result_code, - ranked_test_runs.lookback, - ranked_test_runs.position, - ranked_test_runs.test_starttime, - -- result_code = -1 indicates training mode - CASE WHEN results.result_code = -1 THEN 1 ELSE 0 END AS is_training, - CASE WHEN results.test_type = 'Freshness_Trend' AND results.result_code = 0 THEN 1 ELSE 0 END AS freshness_anomaly, - CASE WHEN results.test_type = 'Volume_Trend' AND results.result_code = 0 THEN 1 ELSE 0 END AS volume_anomaly, - CASE WHEN results.test_type = 'Schema_Drift' AND results.result_code = 0 THEN 1 ELSE 0 END AS schema_anomaly, - CASE WHEN results.test_type = 'Metric_Trend' AND results.result_code = 0 THEN 1 ELSE 0 END AS metric_anomaly, - CASE WHEN results.test_type = 'Freshness_Trend' THEN results.result_signal ELSE NULL END AS freshness_interval, - CASE WHEN results.test_type = 'Volume_Trend' THEN results.result_signal::BIGINT ELSE NULL END AS row_count, - CASE WHEN results.test_type = 'Schema_Drift' THEN SPLIT_PART(results.result_signal, '|', 1) ELSE NULL END AS table_change, - CASE WHEN results.test_type = 'Schema_Drift' THEN NULLIF(SPLIT_PART(results.result_signal, '|', 2), '')::INT ELSE 0 END AS col_adds, - CASE WHEN results.test_type = 'Schema_Drift' THEN NULLIF(SPLIT_PART(results.result_signal, '|', 3), '')::INT ELSE 0 END AS col_drops, - CASE WHEN results.test_type = 'Schema_Drift' THEN NULLIF(SPLIT_PART(results.result_signal, '|', 4), '')::INT ELSE 0 END AS col_mods, - CASE WHEN results.result_status = 'Error' THEN results.result_message ELSE NULL END AS error_message - FROM latest_tables - LEFT JOIN ranked_test_runs ON TRUE - LEFT JOIN test_results AS results - ON results.test_run_id = ranked_test_runs.id - AND results.table_name = latest_tables.table_name - WHERE ranked_test_runs.position IS NULL - -- Also capture 1 run before the lookback to get baseline results - OR ranked_test_runs.position <= ranked_test_runs.lookback + 1 - ), - monitor_tables AS ( - SELECT - :table_group_id AS table_group_id, - table_name, - MAX(lookback) AS lookback, - SUM(freshness_anomaly) AS freshness_anomalies, - SUM(volume_anomaly) AS volume_anomalies, - SUM(schema_anomaly) AS schema_anomalies, - SUM(metric_anomaly) AS metric_anomalies, - MAX(test_time - (COALESCE(NULLIF(freshness_interval, 'Unknown')::INTEGER, 0) * INTERVAL '1 minute')) - FILTER (WHERE test_type = 'Freshness_Trend' AND position = 1) AS latest_update, - MAX(row_count) FILTER (WHERE position = 1) AS row_count, - SUM(col_adds) AS column_adds, - SUM(col_drops) AS column_drops, - SUM(col_mods) AS column_mods, - MAX(error_message) FILTER (WHERE test_type = 'Freshness_Trend' AND position = 1) AS freshness_error_message, - MAX(error_message) FILTER (WHERE test_type = 'Volume_Trend' AND position = 1) AS volume_error_message, - MAX(error_message) FILTER (WHERE test_type = 'Schema_Drift' AND position = 1) AS schema_error_message, - MAX(error_message) FILTER (WHERE test_type = 'Metric_Trend' AND position = 1) AS metric_error_message, - BOOL_OR(is_training = 1) FILTER (WHERE test_type = 'Freshness_Trend' AND position = 1) AS freshness_is_training, - BOOL_OR(is_training = 1) FILTER (WHERE test_type = 'Volume_Trend' AND position = 1) AS volume_is_training, - BOOL_OR(is_training = 1) FILTER (WHERE test_type = 'Metric_Trend' AND position = 1) AS metric_is_training, - BOOL_OR(test_type = 'Freshness_Trend') IS NOT TRUE AS freshness_is_pending, - BOOL_OR(test_type = 'Volume_Trend') IS NOT TRUE AS volume_is_pending, - -- Schema monitor only creates results on schema changes (Failed) - -- Mark it as pending only if there are no results of any test type - BOOL_OR(test_time IS NOT NULL) IS NOT TRUE AS schema_is_pending, - BOOL_OR(test_type = 'Metric_Trend') IS NOT TRUE AS metric_is_pending, - CASE - -- Mark as Dropped if latest Schema Drift result for the table indicates it was dropped - WHEN (ARRAY_AGG(table_change ORDER BY test_time DESC) FILTER (WHERE table_change IS NOT NULL))[1] = 'D' - THEN 'dropped' - -- Only mark as Added if latest change does not indicate a drop - WHEN MAX(CASE WHEN table_change = 'A' THEN 1 ELSE 0 END) = 1 - THEN 'added' - WHEN SUM(schema_anomaly) > 0 - THEN 'modified' - ELSE NULL - END AS table_state - FROM monitor_results - -- Only aggregate within lookback runs - WHERE position IS NULL OR position <= COALESCE(lookback, 1) - GROUP BY table_name - ), - table_bounds AS ( - SELECT - table_name, - MIN(position) AS min_position, - MAX(position) AS max_position - FROM monitor_results - WHERE position IS NOT NULL - GROUP BY table_name - ), - baseline_tables AS ( - SELECT - monitor_results.table_name, - MIN(monitor_results.test_starttime) FILTER ( - WHERE monitor_results.position = LEAST(monitor_results.lookback + 1, table_bounds.max_position) - ) AS lookback_start, - MAX(monitor_results.test_starttime) FILTER ( - WHERE monitor_results.position = GREATEST(1, table_bounds.min_position) - ) AS lookback_end, - MAX(monitor_results.row_count) FILTER ( - WHERE monitor_results.test_type = 'Volume_Trend' - AND monitor_results.position = LEAST(monitor_results.lookback + 1, table_bounds.max_position) - ) AS previous_row_count - FROM monitor_results - JOIN table_bounds ON monitor_results.table_name = table_bounds.table_name - GROUP BY monitor_results.table_name - ) - SELECT - monitor_tables.*, - baseline_tables.lookback_start, - baseline_tables.lookback_end, - baseline_tables.previous_row_count - FROM monitor_tables - LEFT JOIN baseline_tables ON monitor_tables.table_name = baseline_tables.table_name - {f"WHERE ({' OR '.join(f'{ANOMALY_TYPE_FILTERS[t]} > 0' for t in anomaly_type_filter)})" if anomaly_type_filter else ""} - ORDER BY {"LOWER(monitor_tables.table_name)" if not sort_field or sort_field == "table_name" else f"monitor_tables.{sort_field}"} - {"DESC" if sort_order == "desc" else "ASC"} NULLS LAST - {"LIMIT :limit" if limit else ""} - {"OFFSET :offset" if offset else ""} - """ - - escaped_table_name_filter = table_name_filter.replace("_", "\\_") if table_name_filter else None - params = { - "table_group_id": table_group_id, - "table_name_filter": f"%{escaped_table_name_filter}%" if escaped_table_name_filter else None, - "sort_field": sort_field, - "limit": limit, - "offset": offset, - } - - return query, params +def _dashboard_anomaly_types(anomaly_type_filter: list[str] | None) -> list[str] | None: + """Map dashboard-form anomaly type labels to internal ``test_type`` codes.""" + if not anomaly_type_filter: + return None + return [ + _DASHBOARD_ANOMALY_TYPE_TO_DB[t] + for t in anomaly_type_filter + if t in _DASHBOARD_ANOMALY_TYPE_TO_DB + ] or None + + +def _dashboard_sort_to_model( + sort_field: str | None, + sort_order: Literal["asc"] | Literal["desc"] | None, +) -> str | None: + """Translate the dashboard's (sort_field, sort_order) pair into the model's + ``sort_by`` form (the field name, optionally suffixed with ``_desc``).""" + if not sort_field: + return None + return f"{sort_field}_desc" if sort_order == "desc" else sort_field def set_param_values(payload: dict) -> None: @@ -913,7 +740,7 @@ def get_monitor_events_for_table(test_suite_id: str, table_name: str, lookback_m CROSS JOIN target_tests tt LEFT JOIN test_results AS results ON ( - results.test_run_id = active_runs.id + results.test_run_id = active_runs.id AND results.test_type = tt.test_type AND results.table_name = :table_name ) diff --git a/tests/unit/mcp/test_tools_common.py b/tests/unit/mcp/test_tools_common.py index 068d1096..31ee061a 100644 --- a/tests/unit/mcp/test_tools_common.py +++ b/tests/unit/mcp/test_tools_common.py @@ -834,3 +834,142 @@ def test_resolve_test_result_invalid_uuid(): with pytest.raises(MCPUserError, match="Invalid test_result_id"): resolve_test_result("not-a-uuid") + + +# --- Monitor helpers --- + + +@pytest.mark.parametrize( + "label,expected_value", + [ + ("freshness", "Freshness_Trend"), + ("volume", "Volume_Trend"), + ("schema", "Schema_Drift"), + ("metric", "Metric_Trend"), + ], +) +def test_parse_monitor_type_user_labels(label, expected_value): + from testgen.mcp.tools.common import parse_monitor_type + + assert parse_monitor_type(label).value == expected_value + + +def test_parse_monitor_type_rejects_db_codes(): + """Internal codes like ``Freshness_Trend`` are not accepted on the input boundary — + only the lowercase user-facing short labels are.""" + from testgen.mcp.tools.common import parse_monitor_type + + with pytest.raises(MCPUserError, match="Invalid monitor_type"): + parse_monitor_type("Freshness_Trend") + + +def test_parse_monitor_type_lists_valid_values_on_error(): + from testgen.mcp.tools.common import parse_monitor_type + + with pytest.raises(MCPUserError, match="Valid values:") as exc: + parse_monitor_type("metrics") + msg = str(exc.value) + for label in ("freshness", "volume", "schema", "metric"): + assert label in msg + + +@pytest.mark.parametrize( + "value", + ["table_name", "anomaly_count_desc", "latest_update_desc", "row_count_change_desc"], +) +def test_parse_monitor_table_sort_accepts_documented_values(value): + from testgen.mcp.tools.common import parse_monitor_table_sort + + assert parse_monitor_table_sort(value).value == value + + +def test_parse_monitor_table_sort_rejects_unknown(): + from testgen.mcp.tools.common import parse_monitor_table_sort + + with pytest.raises(MCPUserError, match="Invalid sort_by") as exc: + parse_monitor_table_sort("alphabetical") + msg = str(exc.value) + for valid in ("table_name", "anomaly_count_desc", "latest_update_desc", "row_count_change_desc"): + assert valid in msg + + +def test_format_monitor_type_round_trips(): + from testgen.common.enums import MonitorType + from testgen.mcp.tools.common import format_monitor_type + + for member in MonitorType: + formatted = format_monitor_type(member) + assert format_monitor_type(member.value) == formatted + + +def test_resolve_monitored_table_group_returns_suite(): + from testgen.common.models.table_group import TableGroup + from testgen.common.models.test_suite import TestSuite + from testgen.mcp.tools.common import resolve_monitored_table_group + + tg = MagicMock(spec=TableGroup) + tg.id = uuid4() + tg.monitor_test_suite_id = uuid4() + suite = MagicMock(spec=TestSuite) + suite.is_monitor = True + + with ( + patch("testgen.mcp.tools.common.resolve_table_group", return_value=tg), + patch("testgen.mcp.tools.common.TestSuite.get", return_value=suite) as mock_get, + ): + out_tg, out_suite = resolve_monitored_table_group(str(tg.id)) + + assert out_tg is tg + assert out_suite is suite + assert mock_get.call_args.args[0] == tg.monitor_test_suite_id + + +def test_resolve_monitored_table_group_returns_none_when_unlinked(): + """Table group exists but has no monitor suite pointer.""" + from testgen.common.models.table_group import TableGroup + from testgen.mcp.tools.common import resolve_monitored_table_group + + tg = MagicMock(spec=TableGroup) + tg.id = uuid4() + tg.monitor_test_suite_id = None + + with patch("testgen.mcp.tools.common.resolve_table_group", return_value=tg): + out_tg, suite = resolve_monitored_table_group(str(tg.id)) + + assert out_tg is tg + assert suite is None + + +def test_resolve_monitored_table_group_returns_none_when_pointer_stale(): + """Pointer set, but suite no longer exists or no longer ``is_monitor=True``.""" + from testgen.common.models.table_group import TableGroup + from testgen.mcp.tools.common import resolve_monitored_table_group + + tg = MagicMock(spec=TableGroup) + tg.id = uuid4() + tg.monitor_test_suite_id = uuid4() + + with ( + patch("testgen.mcp.tools.common.resolve_table_group", return_value=tg), + patch("testgen.mcp.tools.common.TestSuite.get", return_value=None), + ): + out_tg, suite = resolve_monitored_table_group(str(tg.id)) + + assert out_tg is tg + assert suite is None + + +def test_resolve_monitored_table_group_raises_when_tg_inaccessible(): + """Inaccessible TG propagates ``MCPResourceNotAccessible`` from ``resolve_table_group`` + — the "not monitored" path must not mask an access failure.""" + from testgen.mcp.tools.common import resolve_monitored_table_group + + bad_id = str(uuid4()) + with ( + patch( + "testgen.mcp.tools.common.resolve_table_group", + side_effect=MCPResourceNotAccessible("Table group", bad_id), + ), + pytest.raises(MCPResourceNotAccessible), + ): + resolve_monitored_table_group(bad_id) diff --git a/tests/unit/mcp/test_tools_monitors.py b/tests/unit/mcp/test_tools_monitors.py new file mode 100644 index 00000000..581ff099 --- /dev/null +++ b/tests/unit/mcp/test_tools_monitors.py @@ -0,0 +1,458 @@ +"""Tests for the MCP monitor read tools — ``get_monitor_summary`` and ``list_monitored_tables``.""" + +from datetime import UTC, datetime +from unittest.mock import MagicMock, patch +from uuid import uuid4 + +import pytest + +from testgen.common.models.table_group import MonitorGroupSummary, MonitorTableSummary +from testgen.mcp.exceptions import MCPResourceNotAccessible, MCPUserError +from testgen.mcp.permissions import ProjectPermissions + +pytestmark = pytest.mark.unit + +MODULE = "testgen.mcp.tools.monitors" + + +# --------------------------------------------------------------------------- +# Fixtures / helpers +# --------------------------------------------------------------------------- + + +def _patch_perms(allowed=("demo",), memberships=None, permission="view"): + memberships = memberships or dict.fromkeys(allowed, "role_a") + return patch( + "testgen.mcp.permissions._compute_project_permissions", + return_value=ProjectPermissions( + memberships=memberships, permission=permission, username="test_user", + ), + ) + + +def _mock_table_group(**overrides) -> MagicMock: + tg = MagicMock() + tg.id = overrides.get("id", uuid4()) + tg.project_code = overrides.get("project_code", "demo") + tg.table_groups_name = overrides.get("table_groups_name", "Sales") + tg.monitor_test_suite_id = overrides.get("monitor_test_suite_id", uuid4()) + return tg + + +def _mock_monitor_suite(**overrides) -> MagicMock: + suite = MagicMock() + suite.id = overrides.get("id", uuid4()) + suite.is_monitor = True + suite.monitor_lookback = overrides.get("monitor_lookback", 7) + return suite + + +def _group_summary(**overrides) -> MonitorGroupSummary: + defaults: dict = { + "lookback": 7, + "lookback_start": datetime(2026, 5, 25, 12, 0, tzinfo=UTC), + "lookback_end": datetime(2026, 6, 1, 12, 0, tzinfo=UTC), + "total_monitored_tables": 12, + "freshness_anomalies": 2, + "volume_anomalies": 0, + "schema_anomalies": 1, + "metric_anomalies": 0, + "freshness_has_errors": False, + "volume_has_errors": False, + "schema_has_errors": False, + "metric_has_errors": False, + "freshness_is_training": False, + "volume_is_training": False, + "metric_is_training": True, + "freshness_is_pending": False, + "volume_is_pending": False, + "schema_is_pending": False, + "metric_is_pending": False, + } + defaults.update(overrides) + return MonitorGroupSummary(**defaults) + + +def _table_summary(**overrides) -> MonitorTableSummary: + defaults: dict = { + "table_name": "orders", + "lookback": 7, + "lookback_start": datetime(2026, 5, 25, 12, 0, tzinfo=UTC), + "lookback_end": datetime(2026, 6, 1, 12, 0, tzinfo=UTC), + "freshness_anomalies": 2, + "volume_anomalies": 0, + "schema_anomalies": 0, + "metric_anomalies": 0, + "freshness_is_training": False, + "volume_is_training": False, + "metric_is_training": None, + "freshness_is_pending": False, + "volume_is_pending": False, + "schema_is_pending": False, + "metric_is_pending": True, + "freshness_error_message": None, + "volume_error_message": None, + "schema_error_message": None, + "metric_error_message": None, + "latest_update": datetime(2026, 6, 1, 8, 30, tzinfo=UTC), + "row_count": 1_234_567, + "previous_row_count": 1_200_000, + "column_adds": 0, + "column_drops": 0, + "column_mods": 0, + "table_state": None, + } + defaults.update(overrides) + return MonitorTableSummary(**defaults) + + +# --------------------------------------------------------------------------- +# get_monitor_summary +# --------------------------------------------------------------------------- + + +@patch(f"{MODULE}.next_scheduled_run", return_value=None) +@patch(f"{MODULE}.TableGroup") +@patch(f"{MODULE}.resolve_monitored_table_group") +def test_get_monitor_summary_happy_path(mock_resolve, mock_tg_cls, mock_next, db_session_mock): + tg = _mock_table_group() + suite = _mock_monitor_suite() + mock_resolve.return_value = (tg, suite) + mock_tg_cls.get_monitor_group_summary.return_value = _group_summary() + + from testgen.mcp.tools.monitors import get_monitor_summary + + with _patch_perms(): + out = get_monitor_summary(str(tg.id)) + + assert "# Monitor summary for `Sales`" in out + assert "**Project:** `demo`" in out + assert "**Monitored tables:** 12" in out + assert "**Lookback:** 7 runs" in out + assert "(override)" not in out + assert "**Next scheduled run:** not scheduled" in out + # Per-type rows + assert "| Freshness | 2 | ok |" in out + assert "| Volume | 0 | ok |" in out + assert "| Schema | 1 | ok |" in out + assert "| Metric | 0 | training |" in out + + +@patch(f"{MODULE}.next_scheduled_run", return_value=datetime(2026, 6, 2, 18, 0, tzinfo=UTC)) +@patch(f"{MODULE}.TableGroup") +@patch(f"{MODULE}.resolve_monitored_table_group") +def test_get_monitor_summary_renders_next_scheduled(mock_resolve, mock_tg_cls, mock_next, db_session_mock): + tg = _mock_table_group() + mock_resolve.return_value = (tg, _mock_monitor_suite()) + mock_tg_cls.get_monitor_group_summary.return_value = _group_summary() + + from testgen.mcp.tools.monitors import get_monitor_summary + + with _patch_perms(): + out = get_monitor_summary(str(tg.id)) + + assert "Next scheduled run" in out + assert "2026-06-02" in out + + +@patch(f"{MODULE}.resolve_monitored_table_group") +def test_get_monitor_summary_not_monitored(mock_resolve, db_session_mock): + tg = _mock_table_group(monitor_test_suite_id=None) + mock_resolve.return_value = (tg, None) + + from testgen.mcp.tools.monitors import get_monitor_summary + + with _patch_perms(): + out = get_monitor_summary(str(tg.id)) + + assert out == "This table group is not monitored." + + +@patch(f"{MODULE}.resolve_monitored_table_group") +def test_get_monitor_summary_inaccessible(mock_resolve, db_session_mock): + bad_id = str(uuid4()) + mock_resolve.side_effect = MCPResourceNotAccessible("Table group", bad_id) + + from testgen.mcp.tools.monitors import get_monitor_summary + + with _patch_perms(), pytest.raises(MCPResourceNotAccessible): + get_monitor_summary(bad_id) + + +@patch(f"{MODULE}.next_scheduled_run", return_value=None) +@patch(f"{MODULE}.TableGroup") +@patch(f"{MODULE}.resolve_monitored_table_group") +def test_get_monitor_summary_lookback_override_applied(mock_resolve, mock_tg_cls, mock_next, db_session_mock): + tg = _mock_table_group() + mock_resolve.return_value = (tg, _mock_monitor_suite()) + mock_tg_cls.get_monitor_group_summary.return_value = _group_summary(lookback=14) + + from testgen.mcp.tools.monitors import get_monitor_summary + + with _patch_perms(): + out = get_monitor_summary(str(tg.id), lookback=14) + + mock_tg_cls.get_monitor_group_summary.assert_called_once_with(tg.id, lookback_override=14) + assert "**Lookback:** 14 runs (override)" in out + + +def test_get_monitor_summary_lookback_out_of_range(db_session_mock): + from testgen.mcp.tools.monitors import get_monitor_summary + + with _patch_perms(), pytest.raises(MCPUserError) as exc: + get_monitor_summary(str(uuid4()), lookback=0) + assert "between 1 and 365" in str(exc.value) + + +@patch(f"{MODULE}.next_scheduled_run", return_value=None) +@patch(f"{MODULE}.TableGroup") +@patch(f"{MODULE}.resolve_monitored_table_group") +def test_get_monitor_summary_renders_error_and_pending_states( + mock_resolve, mock_tg_cls, mock_next, db_session_mock, +): + tg = _mock_table_group() + mock_resolve.return_value = (tg, _mock_monitor_suite()) + mock_tg_cls.get_monitor_group_summary.return_value = _group_summary( + freshness_has_errors=True, + volume_is_pending=True, + schema_is_pending=True, + metric_is_pending=True, + ) + + from testgen.mcp.tools.monitors import get_monitor_summary + + with _patch_perms(): + out = get_monitor_summary(str(tg.id)) + + assert "| Freshness | 2 | error |" in out + assert "| Volume | 0 | no results yet or not configured |" in out + assert "| Schema | 1 | no results yet or not configured |" in out + assert "| Metric | 0 | no results yet or not configured |" in out + + +# --------------------------------------------------------------------------- +# list_monitored_tables +# --------------------------------------------------------------------------- + + +@patch(f"{MODULE}.TableGroup") +@patch(f"{MODULE}.resolve_monitored_table_group") +def test_list_monitored_tables_happy_path(mock_resolve, mock_tg_cls, db_session_mock): + tg = _mock_table_group() + mock_resolve.return_value = (tg, _mock_monitor_suite()) + mock_tg_cls.list_monitor_table_summaries.return_value = ( + [ + _table_summary(table_name="orders", freshness_anomalies=2), + _table_summary(table_name="customers", freshness_anomalies=0, row_count=42_000), + ], + 2, + ) + + from testgen.mcp.tools.monitors import list_monitored_tables + + with _patch_perms(): + out = list_monitored_tables(str(tg.id)) + + assert "# Monitored tables in `Sales`" in out + assert "Showing 1–2 of 2" in out # noqa: RUF001 — page-info formatter uses EN DASH + assert "`orders`" in out + assert "`customers`" in out + # "Row count change" column renders the signed delta (row_count - previous_row_count), + # not the raw current count. Defaults are row_count=1_234_567, previous=1_200_000. + assert "+34,567" in out + # customers overrides row_count to 42_000 — delta vs default previous of 1_200_000 is -1,158,000. + assert "-1,158,000" in out + # Absolute current count should not appear in the column. + assert "1,234,567" not in out + # Default sort_by is None → table_name asc + mock_tg_cls.list_monitor_table_summaries.assert_called_once_with( + tg.id, anomaly_types=None, sort_by=None, page=1, limit=20, + ) + + +@patch(f"{MODULE}.TableGroup") +@patch(f"{MODULE}.resolve_monitored_table_group") +def test_list_monitored_tables_anomaly_type_filter_translated( + mock_resolve, mock_tg_cls, db_session_mock, +): + tg = _mock_table_group() + mock_resolve.return_value = (tg, _mock_monitor_suite()) + mock_tg_cls.list_monitor_table_summaries.return_value = ([], 0) + + from testgen.mcp.tools.monitors import list_monitored_tables + + with _patch_perms(): + list_monitored_tables(str(tg.id), anomaly_type="freshness") + + mock_tg_cls.list_monitor_table_summaries.assert_called_once_with( + tg.id, anomaly_types=["Freshness_Trend"], sort_by=None, page=1, limit=20, + ) + + +@patch(f"{MODULE}.TableGroup") +@patch(f"{MODULE}.resolve_monitored_table_group") +def test_list_monitored_tables_sort_by_translated(mock_resolve, mock_tg_cls, db_session_mock): + tg = _mock_table_group() + mock_resolve.return_value = (tg, _mock_monitor_suite()) + mock_tg_cls.list_monitor_table_summaries.return_value = ([], 0) + + from testgen.mcp.tools.monitors import list_monitored_tables + + with _patch_perms(): + list_monitored_tables(str(tg.id), sort_by="anomaly_count_desc") + + mock_tg_cls.list_monitor_table_summaries.assert_called_once_with( + tg.id, anomaly_types=None, sort_by="total_anomalies_desc", page=1, limit=20, + ) + + +@patch(f"{MODULE}.resolve_monitored_table_group") +def test_list_monitored_tables_not_monitored(mock_resolve, db_session_mock): + tg = _mock_table_group(monitor_test_suite_id=None) + mock_resolve.return_value = (tg, None) + + from testgen.mcp.tools.monitors import list_monitored_tables + + with _patch_perms(): + out = list_monitored_tables(str(tg.id)) + + assert out == "This table group is not monitored." + + +def test_list_monitored_tables_invalid_anomaly_type(db_session_mock): + from testgen.mcp.tools.monitors import list_monitored_tables + + with _patch_perms(), pytest.raises(MCPUserError) as exc: + list_monitored_tables(str(uuid4()), anomaly_type="bogus") + assert "Invalid monitor_type" in str(exc.value) + + +def test_list_monitored_tables_invalid_sort_by(db_session_mock): + from testgen.mcp.tools.monitors import list_monitored_tables + + with _patch_perms(), pytest.raises(MCPUserError) as exc: + list_monitored_tables(str(uuid4()), sort_by="wat") + assert "Invalid sort_by" in str(exc.value) + + +def test_list_monitored_tables_invalid_page(db_session_mock): + from testgen.mcp.tools.monitors import list_monitored_tables + + with _patch_perms(), pytest.raises(MCPUserError): + list_monitored_tables(str(uuid4()), page=0) + + +def test_list_monitored_tables_limit_out_of_range(db_session_mock): + from testgen.mcp.tools.monitors import list_monitored_tables + + with _patch_perms(), pytest.raises(MCPUserError): + list_monitored_tables(str(uuid4()), limit=500) + + +@patch(f"{MODULE}.TableGroup") +@patch(f"{MODULE}.resolve_monitored_table_group") +def test_list_monitored_tables_schema_change_column( + mock_resolve, mock_tg_cls, db_session_mock, +): + """The Schema column shows just the anomaly count; the new Schema change column + renders the verbose description (added with N columns / modified breakdown / + dropped with N columns).""" + tg = _mock_table_group() + mock_resolve.return_value = (tg, _mock_monitor_suite()) + mock_tg_cls.list_monitor_table_summaries.return_value = ( + [ + _table_summary( + table_name="t_mod", schema_anomalies=2, table_state="modified", + column_adds=1, column_drops=2, column_mods=0, + ), + _table_summary( + table_name="t_add", schema_anomalies=1, table_state="added", + column_adds=5, column_drops=0, column_mods=0, + ), + _table_summary( + table_name="t_drop", schema_anomalies=1, table_state="dropped", + column_adds=0, column_drops=10, column_mods=0, + ), + _table_summary(table_name="t_quiet", schema_anomalies=0, table_state=None), + ], + 4, + ) + + from testgen.mcp.tools.monitors import list_monitored_tables + + with _patch_perms(): + out = list_monitored_tables(str(tg.id)) + + # New column appears in the header + assert "Schema change" in out + # Verbose strings appear in their respective rows + assert "Table added with 5 columns." in out + assert "Table dropped with 10 columns." in out + assert "1 column added. 2 columns dropped." in out + # No more parenthetical states on the Schema column + assert "(columns)" not in out + assert "(added)" not in out + assert "(dropped)" not in out + # Quiet row's Schema column is the raw count, Schema change is em-dash + quiet_row = next(line for line in out.splitlines() if "`t_quiet`" in line) + assert quiet_row.count(" 0 ") >= 1 # Schema = 0 for quiet + + +@patch(f"{MODULE}.TableGroup") +@patch(f"{MODULE}.resolve_monitored_table_group") +def test_list_monitored_tables_training_and_pending_cells( + mock_resolve, mock_tg_cls, db_session_mock, +): + tg = _mock_table_group() + mock_resolve.return_value = (tg, _mock_monitor_suite()) + mock_tg_cls.list_monitor_table_summaries.return_value = ( + [ + _table_summary( + table_name="t1", + freshness_is_training=True, + volume_is_pending=True, + metric_error_message="boom", + ), + ], + 1, + ) + + from testgen.mcp.tools.monitors import list_monitored_tables + + with _patch_perms(): + out = list_monitored_tables(str(tg.id)) + + row = next(line for line in out.splitlines() if "`t1`" in line) + assert "training" in row + assert "pending" in row + assert "error" in row + + +@patch(f"{MODULE}.TableGroup") +@patch(f"{MODULE}.resolve_monitored_table_group") +def test_list_monitored_tables_empty(mock_resolve, mock_tg_cls, db_session_mock): + tg = _mock_table_group() + mock_resolve.return_value = (tg, _mock_monitor_suite()) + mock_tg_cls.list_monitor_table_summaries.return_value = ([], 0) + + from testgen.mcp.tools.monitors import list_monitored_tables + + with _patch_perms(): + out = list_monitored_tables(str(tg.id)) + + assert "_No monitored tables match this filter._" in out + + +@patch(f"{MODULE}.TableGroup") +@patch(f"{MODULE}.resolve_monitored_table_group") +def test_list_monitored_tables_empty_beyond_last_page(mock_resolve, mock_tg_cls, db_session_mock): + tg = _mock_table_group() + mock_resolve.return_value = (tg, _mock_monitor_suite()) + mock_tg_cls.list_monitor_table_summaries.return_value = ([], 7) + + from testgen.mcp.tools.monitors import list_monitored_tables + + with _patch_perms(): + out = list_monitored_tables(str(tg.id), page=99) + + assert "No tables on page 99 (total: 7)." in out From 100a40a8e936800d80590a360a2f90b630b2b3b1 Mon Sep 17 00:00:00 2001 From: Ricardo Boni Date: Thu, 11 Jun 2026 16:06:57 -0400 Subject: [PATCH 30/78] refactor(api): TD export/import domain exceptions and version marker The export/import service raises InvalidImportPayload (duplicate natural keys, unsupported version) and ImportStrictViolation (apply_strict with skips, carrying the projected ImportResponse) instead of FastAPI HTTPException / silently withholding the apply phase; the REST endpoint adapts both to the existing wire shapes. ExportDocument.version is required so default-omitting serialization keeps the schema-version marker, and imports reject unsupported versions. The service and its domain schemas move to common/test_definition_export_import_service.py (no REST dependencies remain); only the wire types stay in api/schemas.py. Co-Authored-By: Claude Fable 5 --- testgen/api/schemas.py | 175 +------------ testgen/api/test_definitions.py | 42 +-- .../test_definition_export_import_service.py} | 240 ++++++++++++++++-- tests/unit/api/test_td_export_import.py | 186 ++++++++++---- 4 files changed, 376 insertions(+), 267 deletions(-) rename testgen/{api/test_definition_service.py => common/test_definition_export_import_service.py} (66%) diff --git a/testgen/api/schemas.py b/testgen/api/schemas.py index fbe56584..60f4a10d 100644 --- a/testgen/api/schemas.py +++ b/testgen/api/schemas.py @@ -1,12 +1,12 @@ """Pydantic request/response models for API v1 endpoints.""" from datetime import datetime -from enum import StrEnum from uuid import UUID -from pydantic import BaseModel, field_validator +from pydantic import BaseModel from testgen.common.enums import JobSource, JobStatus, PublicJobKey +from testgen.common.test_definition_export_import_service import ImportConfig, ImportPayload, ImportResponse # --- Jobs --- @@ -127,153 +127,7 @@ class ErrorResponse(BaseModel): errors: list[ErrorDetail] -# --- Test Definition Export/Import --- - - -class Origin(StrEnum): - manual = "manual" - auto = "auto" - both = "both" - - -class ImportMode(StrEnum): - preview = "preview" - apply = "apply" - apply_strict = "apply_strict" - - -class OnMatch(StrEnum): - overwrite_all = "overwrite_all" - overwrite_unlocked = "overwrite_unlocked" - skip = "skip" - - -class OnNew(StrEnum): - skip = "skip" - create = "create" - create_and_lock = "create_and_lock" - - -class OnAbsence(StrEnum): - do_nothing = "do_nothing" - delete_all = "delete_all" - delete_unlocked = "delete_unlocked" - - -class ImportAction(StrEnum): - create = "create" - update = "update" - skip = "skip" - delete = "delete" - - -class ImportReason(StrEnum): - matched = "matched" - no_match = "no_match" - policy = "policy" - locked = "locked" - invalid_test_type = "invalid_test_type" - invalid_table = "invalid_table" - missing_external_id = "missing_external_id" - absent = "absent" - - -# Non-None defaults must match the ORM column defaults in TestDefinition: -# test_active=True (YNString default="Y"), lock_refresh=False (YNString default="N"), -# skip_errors=0 (ZeroIfEmptyInteger), window_days=0 (ZeroIfEmptyInteger), -# history_lookback=0 (Column default=0). -# On export, the model_serializer omits fields matching these defaults to keep the file compact. -# On import, model_fields_set distinguishes explicit from defaulted. -class TestDefinitionExport(BaseModel): - """Test definition fields included in the export/import file.""" - - model_config = {"from_attributes": True} - - # Matching / identity - test_type: str - external_id: UUID | None = None - last_auto_gen_date: datetime | None = None - - # Definition fields - table_name: str | None = None - column_name: str | None = None - test_description: str | None = None - test_active: bool = True - severity: str | None = None - lock_refresh: bool = False - export_to_observability: bool | None = None - skip_errors: int = 0 - - # Calibration fields - baseline_ct: str | None = None - baseline_unique_ct: str | None = None - baseline_value: str | None = None - baseline_value_ct: str | None = None - threshold_value: str | None = None - baseline_sum: str | None = None - baseline_avg: str | None = None - baseline_sd: str | None = None - lower_tolerance: str | None = None - upper_tolerance: str | None = None - - # Subset / grouping - subset_condition: str | None = None - groupby_names: str | None = None - having_condition: str | None = None - window_date_column: str | None = None - window_days: int = 0 - - # Referential - match_schema_name: str | None = None - match_table_name: str | None = None - match_column_names: str | None = None - match_subset_condition: str | None = None - match_groupby_names: str | None = None - match_having_condition: str | None = None - - # Query / history - custom_query: str | None = None - history_calculation: str | None = None - history_calculation_upper: str | None = None - history_lookback: int = 0 - - @field_validator("skip_errors", "window_days", "history_lookback", mode="before") - @classmethod - def _coerce_none_to_zero(cls, v: int | None) -> int: - return v if v is not None else 0 - - -class ExportSource(BaseModel): - project_code: str - test_suite: str - table_group: str - table_group_schema: str - exported_at: datetime - testgen_version: str | None = None - - -class ExportDocument(BaseModel): - version: int = 1 - source: ExportSource - definitions: list[TestDefinitionExport] - - -# --- Import --- - - -class ImportConfig(BaseModel): - mode: ImportMode - on_match: OnMatch - on_new: OnNew - on_absence: OnAbsence - - -class ImportPayload(BaseModel): - """Import payload — same structure as an export document, but definitions are typed.""" - - version: int = 1 - source: ExportSource | None = None - definitions: list[TestDefinitionExport] +# --- Test Definition Export/Import (wire types) --- class ImportRequest(BaseModel): @@ -281,29 +135,6 @@ class ImportRequest(BaseModel): payload: ImportPayload -class ImportItemTD(BaseModel): - idx: int | None = None - target_id: UUID | None = None - - -class ImportItem(BaseModel): - action: ImportAction - reason: ImportReason - tds: list[ImportItemTD] - - -class ImportSummary(BaseModel): - created: int = 0 - updated: int = 0 - skipped: int = 0 - deleted: int = 0 - - -class ImportResponse(BaseModel): - summary: ImportSummary - items: list[ImportItem] - - class ImportStrictError(ErrorResponse): """400 response for apply_strict when entries would be skipped.""" diff --git a/testgen/api/test_definitions.py b/testgen/api/test_definitions.py index 6d2d0e41..eadb0b92 100644 --- a/testgen/api/test_definitions.py +++ b/testgen/api/test_definitions.py @@ -2,19 +2,23 @@ from fastapi import APIRouter, Depends, HTTPException, Query -from testgen.api import test_definition_service -from testgen.api.deps import db_session, resolve_test_suite +from testgen.api.deps import api_error, db_session, resolve_test_suite from testgen.api.schemas import ( ErrorDetail, ErrorResponse, - ExportDocument, - ImportMode, ImportRequest, - ImportResponse, ImportStrictError, - Origin, ) from testgen.common.models.test_suite import TestSuite +from testgen.common.test_definition_export_import_service import ( + ExportDocument, + ImportResponse, + ImportStrictViolation, + InvalidImportPayload, + Origin, + export_definitions, + import_definitions, +) _error_responses = { 404: {"model": ErrorResponse, "description": "Not found"}, @@ -39,14 +43,17 @@ def export_test_definitions( test_type: str | None = Query(default=None), ) -> ExportDocument: """Export test definitions from a test suite as a portable JSON document.""" - return test_definition_service.export_definitions(test_suite, origin, table_name, test_type) + return export_definitions(test_suite, origin, table_name, test_type) @router.post( "/test-suites/{test_suite_id}/test-definition-import", response_model=ImportResponse, responses={ - 400: {"model": ImportStrictError, "description": "Invalid request or strict validation failed"}, + 400: { + "model": ImportStrictError | ErrorResponse, + "description": "Strict validation failed (includes the projected import result) or invalid payload", + }, }, ) def import_test_definitions( @@ -54,18 +61,15 @@ def import_test_definitions( test_suite: TestSuite = resolve_test_suite("edit"), # noqa: B008 ) -> ImportResponse: """Import test definitions into a test suite from a portable JSON document.""" - result = test_definition_service.import_definitions(test_suite, body.config, body.payload) - - if body.config.mode == ImportMode.apply_strict and result.summary.skipped > 0: + try: + return import_definitions(test_suite, body.config, body.payload) + except InvalidImportPayload as err: + raise api_error(400, err.code, str(err)) from err + except ImportStrictViolation as err: raise HTTPException( status_code=400, detail=ImportStrictError( - errors=[ErrorDetail( - code="strict_validation_failed", - detail=f"{result.summary.skipped} test definition(s) would be skipped", - )], - import_result=result, + errors=[ErrorDetail(code="strict_validation_failed", detail=str(err))], + import_result=err.result, ).model_dump(mode="json"), - ) - - return result + ) from err diff --git a/testgen/api/test_definition_service.py b/testgen/common/test_definition_export_import_service.py similarity index 66% rename from testgen/api/test_definition_service.py rename to testgen/common/test_definition_export_import_service.py index e8b2199f..1e831d6b 100644 --- a/testgen/api/test_definition_service.py +++ b/testgen/common/test_definition_export_import_service.py @@ -1,39 +1,216 @@ -"""Business logic for test definition export/import.""" +"""Test definition export/import — portable document schemas and import/export logic.""" from dataclasses import dataclass from datetime import UTC, datetime +from enum import StrEnum from typing import Any from uuid import UUID +from pydantic import BaseModel, field_validator from sqlalchemy import delete, func, select, update from sqlalchemy.engine import Row from testgen import settings -from testgen.api.deps import api_error -from testgen.api.schemas import ( - ExportDocument, - ExportSource, - ImportAction, - ImportConfig, - ImportItem, - ImportItemTD, - ImportMode, - ImportPayload, - ImportReason, - ImportResponse, - ImportSummary, - OnAbsence, - OnMatch, - OnNew, - Origin, - TestDefinitionExport, -) from testgen.common.models import get_current_session from testgen.common.models.data_table import DataTable from testgen.common.models.table_group import TableGroup from testgen.common.models.test_definition import TestDefinition, TestType from testgen.common.models.test_suite import TestSuite +EXPORT_FORMAT_VERSION = 1 + + +class Origin(StrEnum): + manual = "manual" + auto = "auto" + both = "both" + + +class ImportMode(StrEnum): + preview = "preview" + apply = "apply" + apply_strict = "apply_strict" + + +class OnMatch(StrEnum): + overwrite_all = "overwrite_all" + overwrite_unlocked = "overwrite_unlocked" + skip = "skip" + + +class OnNew(StrEnum): + skip = "skip" + create = "create" + create_and_lock = "create_and_lock" + + +class OnAbsence(StrEnum): + do_nothing = "do_nothing" + delete_all = "delete_all" + delete_unlocked = "delete_unlocked" + + +class ImportAction(StrEnum): + create = "create" + update = "update" + skip = "skip" + delete = "delete" + + +class ImportReason(StrEnum): + matched = "matched" + no_match = "no_match" + policy = "policy" + locked = "locked" + invalid_test_type = "invalid_test_type" + invalid_table = "invalid_table" + missing_external_id = "missing_external_id" + absent = "absent" + + +# Non-None defaults must match the ORM column defaults in TestDefinition: +# test_active=True (YNString default="Y"), lock_refresh=False (YNString default="N"), +# skip_errors=0 (ZeroIfEmptyInteger), window_days=0 (ZeroIfEmptyInteger), +# history_lookback=0 (Column default=0). +# On export, fields equal to these defaults are omitted to keep the file compact — any +# serialization of the document must pass ``exclude_defaults=True`` (the REST route does this +# via ``response_model_exclude_defaults=True``). On import, model_fields_set distinguishes +# explicit from defaulted, so a non-compact file would force-write every defaulted field. +class TestDefinitionExport(BaseModel): + """Test definition fields included in the export/import file.""" + + model_config = {"from_attributes": True} + + # Matching / identity + test_type: str + external_id: UUID | None = None + last_auto_gen_date: datetime | None = None + + # Definition fields + table_name: str | None = None + column_name: str | None = None + test_description: str | None = None + test_active: bool = True + severity: str | None = None + lock_refresh: bool = False + export_to_observability: bool | None = None + skip_errors: int = 0 + + # Calibration fields + baseline_ct: str | None = None + baseline_unique_ct: str | None = None + baseline_value: str | None = None + baseline_value_ct: str | None = None + threshold_value: str | None = None + baseline_sum: str | None = None + baseline_avg: str | None = None + baseline_sd: str | None = None + lower_tolerance: str | None = None + upper_tolerance: str | None = None + + # Subset / grouping + subset_condition: str | None = None + groupby_names: str | None = None + having_condition: str | None = None + window_date_column: str | None = None + window_days: int = 0 + + # Referential + match_schema_name: str | None = None + match_table_name: str | None = None + match_column_names: str | None = None + match_subset_condition: str | None = None + match_groupby_names: str | None = None + match_having_condition: str | None = None + + # Query / history + custom_query: str | None = None + history_calculation: str | None = None + history_calculation_upper: str | None = None + history_lookback: int = 0 + + @field_validator("skip_errors", "window_days", "history_lookback", mode="before") + @classmethod + def _coerce_none_to_zero(cls, v: int | None) -> int: + return v if v is not None else 0 + + +class ExportSource(BaseModel): + project_code: str + test_suite: str + table_group: str + table_group_schema: str + exported_at: datetime + testgen_version: str | None = None + + +class ExportDocument(BaseModel): + # No default: a defaulted version is stripped by default-omitting serialization, + # leaving exported documents without their schema-version marker. + version: int + source: ExportSource + definitions: list[TestDefinitionExport] + + +class ImportConfig(BaseModel): + mode: ImportMode + on_match: OnMatch + on_new: OnNew + on_absence: OnAbsence + + +class ImportPayload(BaseModel): + """Import payload — same structure as an export document, but definitions are typed.""" + + version: int = EXPORT_FORMAT_VERSION + source: ExportSource | None = None + definitions: list[TestDefinitionExport] + + +class ImportItemTD(BaseModel): + idx: int | None = None + target_id: UUID | None = None + + +class ImportItem(BaseModel): + action: ImportAction + reason: ImportReason + tds: list[ImportItemTD] + + +class ImportSummary(BaseModel): + created: int = 0 + updated: int = 0 + skipped: int = 0 + deleted: int = 0 + + +class ImportResponse(BaseModel): + summary: ImportSummary + items: list[ImportItem] + + +class ImportErrorCode(StrEnum): + duplicate_natural_key = "duplicate_natural_key" + unsupported_version = "unsupported_version" + + +class InvalidImportPayload(ValueError): + """Import payload rejected before any matching or persistence.""" + + def __init__(self, code: ImportErrorCode, detail: str): + super().__init__(detail) + self.code = code + + +class ImportStrictViolation(Exception): + """``apply_strict`` import with test definitions that would be skipped; nothing was applied.""" + + def __init__(self, result: ImportResponse): + super().__init__(f"{result.summary.skipped} test definitions would be skipped") + self.result = result + + # Fields that must never be written from the import payload on update. # These are either identity fields (set once on create) or determined by matching logic. _UPDATE_EXCLUDE_FIELDS = frozenset({"test_type", "last_auto_gen_date", "external_id"}) @@ -87,6 +264,7 @@ def export_definitions( definitions = [TestDefinitionExport.model_validate(td, from_attributes=True) for td in tds] return ExportDocument( + version=EXPORT_FORMAT_VERSION, source=ExportSource( project_code=test_suite.project_code, test_suite=test_suite.test_suite, @@ -110,6 +288,11 @@ def import_definitions( profiled_tables = set(DataTable.select_table_names(test_suite.table_groups_id, limit=None)) # --- Phase 1: Upfront validation --- + if payload.version != EXPORT_FORMAT_VERSION: + raise InvalidImportPayload( + ImportErrorCode.unsupported_version, + f"Unsupported export document version {payload.version}; supported version: {EXPORT_FORMAT_VERSION}", + ) _check_duplicate_keys(incoming) # --- Phase 2: Matching --- @@ -175,10 +358,13 @@ def import_definitions( # Locked TDs surviving delete_unlocked are omitted entirely (per design) # --- Phase 3: Apply --- - should_apply = config.mode in (ImportMode.apply, ImportMode.apply_strict) has_skips = any(a.action == ImportAction.skip for a in actions) + if config.mode == ImportMode.apply_strict and has_skips: + # Built before apply, so created TDs have no target_id — correct, nothing was persisted. + # The post-apply build below is NOT redundant: _apply_actions fills planned.target for creates. + raise ImportStrictViolation(_build_response(actions)) - if should_apply and not (config.mode == ImportMode.apply_strict and has_skips): + if config.mode in (ImportMode.apply, ImportMode.apply_strict): _apply_actions(actions, test_suite, table_group, config) return _build_response(actions) @@ -218,9 +404,8 @@ def _check_duplicate_keys(incoming: list[TestDefinitionExport]) -> None: if td.last_auto_gen_date is not None: key = (td.test_type, td.table_name, td.column_name) if key in auto_keys: - raise api_error( - 400, - "duplicate_natural_key", + raise InvalidImportPayload( + ImportErrorCode.duplicate_natural_key, f"Duplicate auto-gen key at index {idx}: ({td.test_type}, {td.table_name}, {td.column_name})", ) auto_keys.add(key) @@ -228,9 +413,8 @@ def _check_duplicate_keys(incoming: list[TestDefinitionExport]) -> None: if td.external_id is None: continue if td.external_id in manual_keys: - raise api_error( - 400, - "duplicate_natural_key", + raise InvalidImportPayload( + ImportErrorCode.duplicate_natural_key, f"Duplicate external_id at index {idx}: {td.external_id}", ) manual_keys.add(td.external_id) diff --git a/tests/unit/api/test_td_export_import.py b/tests/unit/api/test_td_export_import.py index ca3a1b0b..5a875136 100644 --- a/tests/unit/api/test_td_export_import.py +++ b/tests/unit/api/test_td_export_import.py @@ -7,14 +7,20 @@ import pytest from fastapi import HTTPException -from testgen.api.schemas import ( +from testgen.api.schemas import ImportRequest +from testgen.common.test_definition_export_import_service import ( ExportDocument, + ExportSource, ImportAction, ImportConfig, + ImportErrorCode, ImportMode, ImportPayload, ImportReason, ImportResponse, + ImportStrictViolation, + ImportSummary, + InvalidImportPayload, OnAbsence, OnMatch, OnNew, @@ -24,7 +30,7 @@ pytestmark = pytest.mark.unit -SERVICE_MODULE = "testgen.api.test_definition_service" +SERVICE_MODULE = "testgen.common.test_definition_export_import_service" ENDPOINT_MODULE = "testgen.api.test_definitions" @@ -162,7 +168,7 @@ def test_export_builds_document(self, mock_session_fn, mock_tg_cls, mock_setting session.scalars.return_value.all.return_value = [td_obj] - from testgen.api.test_definition_service import export_definitions + from testgen.common.test_definition_export_import_service import export_definitions result = export_definitions(ts, Origin.both, None, None) @@ -186,7 +192,7 @@ def test_export_assigns_external_id_to_manual_tds(self, mock_session_fn, mock_tg ts = _make_test_suite() session.scalars.return_value.all.return_value = [] - from testgen.api.test_definition_service import export_definitions + from testgen.common.test_definition_export_import_service import export_definitions export_definitions(ts, Origin.manual, None, None) @@ -207,7 +213,7 @@ def test_export_skips_external_id_assignment_for_auto_only(self, mock_session_fn ts = _make_test_suite() session.scalars.return_value.all.return_value = [] - from testgen.api.test_definition_service import export_definitions + from testgen.common.test_definition_export_import_service import export_definitions export_definitions(ts, Origin.auto, None, None) @@ -237,7 +243,7 @@ def test_auto_td_matches_by_natural_key(self, mock_session_fn, mock_tg_cls, mock td = _make_import_td(test_type="Alpha", table_name="t1", column_name="c1") config = _make_config(on_match=OnMatch.overwrite_unlocked) - from testgen.api.test_definition_service import import_definitions + from testgen.common.test_definition_export_import_service import import_definitions with patch(f"{SERVICE_MODULE}._load_valid_test_types", return_value={"Alpha"}): result = import_definitions(ts, config, _make_payload(td)) @@ -265,7 +271,7 @@ def test_manual_td_matches_by_external_id(self, mock_session_fn, mock_tg_cls, mo td = _make_import_td(last_auto_gen_date=None, external_id=ext_id, table_name="t1") config = _make_config(on_match=OnMatch.overwrite_unlocked) - from testgen.api.test_definition_service import import_definitions + from testgen.common.test_definition_export_import_service import import_definitions with patch(f"{SERVICE_MODULE}._load_valid_test_types", return_value={"Alpha"}): result = import_definitions(ts, config, _make_payload(td)) @@ -287,7 +293,7 @@ def test_no_match_creates(self, mock_session_fn, mock_tg_cls, mock_dt_cls): td = _make_import_td(test_type="Alpha", table_name="t1") config = _make_config(on_new=OnNew.create) - from testgen.api.test_definition_service import import_definitions + from testgen.common.test_definition_export_import_service import import_definitions with patch(f"{SERVICE_MODULE}._load_valid_test_types", return_value={"Alpha"}): result = import_definitions(ts, config, _make_payload(td)) @@ -317,7 +323,7 @@ def test_on_match_skip(self, mock_session_fn, mock_tg_cls, mock_dt_cls): td = _make_import_td() config = _make_config(on_match=OnMatch.skip) - from testgen.api.test_definition_service import import_definitions + from testgen.common.test_definition_export_import_service import import_definitions with patch(f"{SERVICE_MODULE}._load_valid_test_types", return_value={"Alpha"}): result = import_definitions(ts, config, _make_payload(td)) @@ -341,7 +347,7 @@ def test_on_match_overwrite_unlocked_skips_locked(self, mock_session_fn, mock_tg td = _make_import_td() config = _make_config(on_match=OnMatch.overwrite_unlocked) - from testgen.api.test_definition_service import import_definitions + from testgen.common.test_definition_export_import_service import import_definitions with patch(f"{SERVICE_MODULE}._load_valid_test_types", return_value={"Alpha"}): result = import_definitions(ts, config, _make_payload(td)) @@ -365,7 +371,7 @@ def test_on_match_overwrite_all_ignores_lock(self, mock_session_fn, mock_tg_cls, td = _make_import_td() config = _make_config(on_match=OnMatch.overwrite_all) - from testgen.api.test_definition_service import import_definitions + from testgen.common.test_definition_export_import_service import import_definitions with patch(f"{SERVICE_MODULE}._load_valid_test_types", return_value={"Alpha"}): result = import_definitions(ts, config, _make_payload(td)) @@ -386,7 +392,7 @@ def test_on_new_skip(self, mock_session_fn, mock_tg_cls, mock_dt_cls): td = _make_import_td() config = _make_config(on_new=OnNew.skip) - from testgen.api.test_definition_service import import_definitions + from testgen.common.test_definition_export_import_service import import_definitions with patch(f"{SERVICE_MODULE}._load_valid_test_types", return_value={"Alpha"}): result = import_definitions(ts, config, _make_payload(td)) @@ -412,7 +418,7 @@ def test_on_absence_delete_all(self, mock_session_fn, mock_tg_cls, mock_dt_cls): td = _make_import_td(test_type="Alpha", table_name="t1", column_name="c1") config = _make_config(on_new=OnNew.create, on_absence=OnAbsence.delete_all) - from testgen.api.test_definition_service import import_definitions + from testgen.common.test_definition_export_import_service import import_definitions with patch(f"{SERVICE_MODULE}._load_valid_test_types", return_value={"Alpha"}): result = import_definitions(ts, config, _make_payload(td)) @@ -436,7 +442,7 @@ def test_on_absence_delete_unlocked_spares_locked(self, mock_session_fn, mock_tg td = _make_import_td(test_type="Alpha", table_name="t1") config = _make_config(on_new=OnNew.create, on_absence=OnAbsence.delete_unlocked) - from testgen.api.test_definition_service import import_definitions + from testgen.common.test_definition_export_import_service import import_definitions with patch(f"{SERVICE_MODULE}._load_valid_test_types", return_value={"Alpha"}): result = import_definitions(ts, config, _make_payload(td)) @@ -465,7 +471,7 @@ def test_invalid_test_type_skipped(self, mock_session_fn, mock_tg_cls, mock_dt_c td = _make_import_td(test_type="NonExistent", table_name="t1") config = _make_config() - from testgen.api.test_definition_service import import_definitions + from testgen.common.test_definition_export_import_service import import_definitions with patch(f"{SERVICE_MODULE}._load_valid_test_types", return_value={"Alpha"}): result = import_definitions(ts, config, _make_payload(td)) @@ -487,7 +493,7 @@ def test_invalid_table_skipped(self, mock_session_fn, mock_tg_cls, mock_dt_cls): td = _make_import_td(test_type="Alpha", table_name="t1") config = _make_config() - from testgen.api.test_definition_service import import_definitions + from testgen.common.test_definition_export_import_service import import_definitions with patch(f"{SERVICE_MODULE}._load_valid_test_types", return_value={"Alpha"}): result = import_definitions(ts, config, _make_payload(td)) @@ -510,7 +516,7 @@ def test_manual_td_without_external_id_skipped(self, mock_session_fn, mock_tg_cl td = _make_import_td(last_auto_gen_date=None, external_id=None, table_name="t1") config = _make_config() - from testgen.api.test_definition_service import import_definitions + from testgen.common.test_definition_export_import_service import import_definitions with patch(f"{SERVICE_MODULE}._load_valid_test_types", return_value={"Alpha"}): result = import_definitions(ts, config, _make_payload(td)) @@ -533,7 +539,7 @@ def test_td_with_null_table_name_passes_validation(self, mock_session_fn, mock_t td = _make_import_td(test_type="Alpha", table_name=None) config = _make_config() - from testgen.api.test_definition_service import import_definitions + from testgen.common.test_definition_export_import_service import import_definitions with patch(f"{SERVICE_MODULE}._load_valid_test_types", return_value={"Alpha"}): result = import_definitions(ts, config, _make_payload(td)) @@ -555,13 +561,13 @@ def test_duplicate_auto_keys_rejected(self, mock_session_fn, mock_tg_cls, mock_d td2 = _make_import_td(test_type="Alpha", table_name="t1", column_name="c1") config = _make_config() - from testgen.api.test_definition_service import import_definitions + from testgen.common.test_definition_export_import_service import import_definitions with patch(f"{SERVICE_MODULE}._load_valid_test_types", return_value={"Alpha"}): - with pytest.raises(HTTPException) as exc_info: + with pytest.raises(InvalidImportPayload) as exc_info: import_definitions(ts, config, _make_payload(td1, td2)) - assert exc_info.value.status_code == 400 - assert "duplicate_natural_key" in str(exc_info.value.detail) + assert exc_info.value.code == "duplicate_natural_key" + assert "Duplicate auto-gen key" in str(exc_info.value) @patch(f"{SERVICE_MODULE}.DataTable") @patch(f"{SERVICE_MODULE}.TableGroup") @@ -579,12 +585,12 @@ def test_duplicate_external_ids_rejected(self, mock_session_fn, mock_tg_cls, moc td2 = _make_import_td(last_auto_gen_date=None, external_id=ext_id, table_name="t1") config = _make_config() - from testgen.api.test_definition_service import import_definitions + from testgen.common.test_definition_export_import_service import import_definitions with patch(f"{SERVICE_MODULE}._load_valid_test_types", return_value={"Alpha"}): - with pytest.raises(HTTPException) as exc_info: + with pytest.raises(InvalidImportPayload) as exc_info: import_definitions(ts, config, _make_payload(td1, td2)) - assert exc_info.value.status_code == 400 + assert exc_info.value.code == "duplicate_natural_key" # --- Import: apply_strict mode --- @@ -607,14 +613,17 @@ def test_apply_strict_does_not_apply_when_skips_exist(self, mock_session_fn, moc bad_td = _make_import_td(test_type="NonExistent", table_name="t1") config = _make_config(mode=ImportMode.apply_strict) - from testgen.api.test_definition_service import import_definitions + from testgen.common.test_definition_export_import_service import import_definitions with patch(f"{SERVICE_MODULE}._load_valid_test_types", return_value={"Alpha"}): - result = import_definitions(ts, config, _make_payload(good_td, bad_td)) + with pytest.raises(ImportStrictViolation) as exc_info: + import_definitions(ts, config, _make_payload(good_td, bad_td)) - # Result reports what would happen, but nothing was applied + # The exception carries the projected result, but nothing was applied + result = exc_info.value.result assert result.summary.created == 1 assert result.summary.skipped == 1 + assert "1 test definitions would be skipped" in str(exc_info.value) # No session.add calls (create not applied) session.add.assert_not_called() @@ -632,7 +641,7 @@ def test_apply_strict_applies_when_no_skips(self, mock_session_fn, mock_tg_cls, td = _make_import_td(test_type="Alpha", table_name="t1") config = _make_config(mode=ImportMode.apply_strict) - from testgen.api.test_definition_service import import_definitions + from testgen.common.test_definition_export_import_service import import_definitions with patch(f"{SERVICE_MODULE}._load_valid_test_types", return_value={"Alpha"}): with patch(f"{SERVICE_MODULE}.TestDefinition") as mock_td_cls: @@ -663,7 +672,7 @@ def test_create_auto_td_sets_last_auto_gen_date_to_now(self, mock_session_fn, mo td = _make_import_td(test_type="Alpha", table_name="t1", last_auto_gen_date=datetime(2024, 1, 1, tzinfo=UTC)) config = _make_config(mode=ImportMode.apply, on_new=OnNew.create) - from testgen.api.test_definition_service import import_definitions + from testgen.common.test_definition_export_import_service import import_definitions before = datetime.now(UTC) with patch(f"{SERVICE_MODULE}._load_valid_test_types", return_value={"Alpha"}): @@ -690,7 +699,7 @@ def test_create_manual_td_sets_null_last_auto_gen_date(self, mock_session_fn, mo td = _make_import_td(last_auto_gen_date=None, external_id=ext_id, table_name="t1") config = _make_config(mode=ImportMode.apply, on_new=OnNew.create) - from testgen.api.test_definition_service import import_definitions + from testgen.common.test_definition_export_import_service import import_definitions with patch(f"{SERVICE_MODULE}._load_valid_test_types", return_value={"Alpha"}): with patch(f"{SERVICE_MODULE}.TestDefinition") as mock_td_cls: @@ -714,7 +723,7 @@ def test_create_and_lock_forces_lock_for_auto_td(self, mock_session_fn, mock_tg_ td = _make_import_td(test_type="Alpha", table_name="t1", lock_refresh=False) config = _make_config(mode=ImportMode.apply, on_new=OnNew.create_and_lock) - from testgen.api.test_definition_service import import_definitions + from testgen.common.test_definition_export_import_service import import_definitions with patch(f"{SERVICE_MODULE}._load_valid_test_types", return_value={"Alpha"}): with patch(f"{SERVICE_MODULE}.TestDefinition") as mock_td_cls: @@ -738,7 +747,7 @@ def test_create_sets_last_manual_update(self, mock_session_fn, mock_tg_cls, mock td = _make_import_td(test_type="Alpha", table_name="t1") config = _make_config(mode=ImportMode.apply, on_new=OnNew.create) - from testgen.api.test_definition_service import import_definitions + from testgen.common.test_definition_export_import_service import import_definitions before = datetime.now(UTC) with patch(f"{SERVICE_MODULE}._load_valid_test_types", return_value={"Alpha"}): @@ -775,7 +784,7 @@ def test_update_excludes_identity_fields(self, mock_session_fn, mock_tg_cls, moc ) config = _make_config(mode=ImportMode.apply, on_match=OnMatch.overwrite_unlocked) - from testgen.api.test_definition_service import import_definitions + from testgen.common.test_definition_export_import_service import import_definitions with patch(f"{SERVICE_MODULE}._load_valid_test_types", return_value={"Alpha"}): import_definitions(ts, config, _make_payload(td)) @@ -810,7 +819,7 @@ def test_update_inherits_external_id_if_target_has_none(self, mock_session_fn, m ) config = _make_config(mode=ImportMode.apply, on_match=OnMatch.overwrite_unlocked) - from testgen.api.test_definition_service import import_definitions + from testgen.common.test_definition_export_import_service import import_definitions with patch(f"{SERVICE_MODULE}._load_valid_test_types", return_value={"Alpha"}): import_definitions(ts, config, _make_payload(td)) @@ -848,7 +857,7 @@ def test_invalid_td_still_protects_match_from_absence_delete(self, mock_session_ td = _make_import_td(test_type="InvalidType", table_name="t1", column_name="c1") config = _make_config(on_absence=OnAbsence.delete_all) - from testgen.api.test_definition_service import import_definitions + from testgen.common.test_definition_export_import_service import import_definitions with patch(f"{SERVICE_MODULE}._load_valid_test_types", return_value={"Alpha"}): result = import_definitions(ts, config, _make_payload(td)) @@ -877,7 +886,7 @@ def test_preview_does_not_apply(self, mock_session_fn, mock_tg_cls, mock_dt_cls) td = _make_import_td(test_type="Alpha", table_name="t1") config = _make_config(mode=ImportMode.preview) - from testgen.api.test_definition_service import import_definitions + from testgen.common.test_definition_export_import_service import import_definitions with patch(f"{SERVICE_MODULE}._load_valid_test_types", return_value={"Alpha"}): result = import_definitions(ts, config, _make_payload(td)) @@ -900,7 +909,7 @@ def test_preview_target_id_is_none_for_creates(self, mock_session_fn, mock_tg_cl td = _make_import_td(test_type="Alpha", table_name="t1") config = _make_config(mode=ImportMode.preview) - from testgen.api.test_definition_service import import_definitions + from testgen.common.test_definition_export_import_service import import_definitions with patch(f"{SERVICE_MODULE}._load_valid_test_types", return_value={"Alpha"}): result = import_definitions(ts, config, _make_payload(td)) @@ -915,7 +924,6 @@ def test_preview_target_id_is_none_for_creates(self, mock_session_fn, mock_tg_cl class Test_import_endpoint_strict: def test_strict_raises_400_on_skips(self): - from testgen.api.schemas import ImportRequest from testgen.api.test_definitions import import_test_definitions ts = _make_test_suite() @@ -923,14 +931,14 @@ def test_strict_raises_400_on_skips(self): payload = _make_payload(_make_import_td(test_type="Alpha", table_name="t1")) request = ImportRequest(config=config, payload=payload) - # Mock the service to return a response with skips + # The service signals strict violations with an exception carrying the projected result mock_result = ImportResponse( summary=ImportSummary(created=1, skipped=1), items=[], ) - with patch(f"{ENDPOINT_MODULE}.test_definition_service") as mock_service: - mock_service.import_definitions.return_value = mock_result + with patch(f"{ENDPOINT_MODULE}.import_definitions") as mock_import: + mock_import.side_effect = ImportStrictViolation(mock_result) with pytest.raises(HTTPException) as exc_info: import_test_definitions(body=request, test_suite=ts) @@ -939,7 +947,6 @@ def test_strict_raises_400_on_skips(self): assert "import_result" in exc_info.value.detail def test_strict_returns_200_when_no_skips(self): - from testgen.api.schemas import ImportRequest from testgen.api.test_definitions import import_test_definitions ts = _make_test_suite() @@ -952,13 +959,99 @@ def test_strict_returns_200_when_no_skips(self): items=[], ) - with patch(f"{ENDPOINT_MODULE}.test_definition_service") as mock_service: - mock_service.import_definitions.return_value = mock_result + with patch(f"{ENDPOINT_MODULE}.import_definitions") as mock_import: + mock_import.return_value = mock_result result = import_test_definitions(body=request, test_suite=ts) assert result.summary.created == 1 +def test_invalid_payload_adapted_to_400(): + from testgen.api.test_definitions import import_test_definitions + + ts = _make_test_suite() + request = ImportRequest( + config=_make_config(), + payload=_make_payload(_make_import_td(test_type="Alpha", table_name="t1")), + ) + + with patch(f"{ENDPOINT_MODULE}.import_definitions") as mock_import: + mock_import.side_effect = InvalidImportPayload( + ImportErrorCode.duplicate_natural_key, "Duplicate external_id at index 1" + ) + with pytest.raises(HTTPException) as exc_info: + import_test_definitions(body=request, test_suite=ts) + + assert exc_info.value.status_code == 400 + assert "duplicate_natural_key" in str(exc_info.value.detail) + assert "Duplicate external_id at index 1" in str(exc_info.value.detail) + + +@patch(f"{SERVICE_MODULE}.DataTable") +@patch(f"{SERVICE_MODULE}.TableGroup") +@patch(f"{SERVICE_MODULE}.get_current_session") +def test_unsupported_payload_version_rejected(mock_session_fn, mock_tg_cls, mock_dt_cls): + mock_session_fn.return_value = MagicMock() + mock_tg_cls.get.return_value = _make_table_group() + mock_dt_cls.select_table_names.return_value = ["t1"] + + ts = _make_test_suite() + payload = _make_payload(_make_import_td(test_type="Alpha", table_name="t1")) + payload.version = 2 + + from testgen.common.test_definition_export_import_service import import_definitions + + with pytest.raises(InvalidImportPayload) as exc_info: + import_definitions(ts, _make_config(), payload) + + assert exc_info.value.code == ImportErrorCode.unsupported_version + assert "version 2" in str(exc_info.value) + + +@patch(f"{SERVICE_MODULE}.DataTable") +@patch(f"{SERVICE_MODULE}.TableGroup") +@patch(f"{SERVICE_MODULE}.get_current_session") +def test_apply_with_skips_still_applies_valid_actions(mock_session_fn, mock_tg_cls, mock_dt_cls): + session = MagicMock() + mock_session_fn.return_value = session + mock_tg_cls.get.return_value = _make_table_group() + mock_dt_cls.select_table_names.return_value = ["t1"] + session.execute.return_value.all.return_value = [] + + ts = _make_test_suite() + good_td = _make_import_td(test_type="Alpha", table_name="t1") + bad_td = _make_import_td(test_type="NonExistent", table_name="t1") + config = _make_config(mode=ImportMode.apply) + + from testgen.common.test_definition_export_import_service import import_definitions + + with patch(f"{SERVICE_MODULE}._load_valid_test_types", return_value={"Alpha"}): + with patch(f"{SERVICE_MODULE}.TestDefinition") as mock_td_cls: + mock_td_cls.return_value = MagicMock(id=uuid4()) + result = import_definitions(ts, config, _make_payload(good_td, bad_td)) + + # Plain apply tolerates skips: the valid create is applied, no strict violation raised + assert result.summary.created == 1 + assert result.summary.skipped == 1 + session.add.assert_called_once() + + +def test_export_document_version_survives_default_omitting_serialization(): + doc = ExportDocument( + version=1, + source=ExportSource( + project_code="DEFAULT", + test_suite="suite", + table_group="tg", + table_group_schema="public", + exported_at=datetime(2026, 6, 11, tzinfo=UTC), + ), + definitions=[], + ) + + assert '"version":1' in doc.model_dump_json(exclude_defaults=True) + + # --- Schema: TestDefinitionExport --- @@ -988,6 +1081,3 @@ def test_defaults_match_expected_values(self): assert td.window_days == 0 assert td.history_lookback == 0 - -# Need to import this here for the endpoint test -from testgen.api.schemas import ImportSummary From e715c5571dbbe6fdecfd56225724ff185cbc0fda Mon Sep 17 00:00:00 2001 From: Aarthy Adityan Date: Fri, 12 Jun 2026 17:27:38 -0400 Subject: [PATCH 31/78] fix(monitors): store freshness staleness threshold for full-week active schedules The staleness threshold (median gap x sensitivity factor) was only assigned inside the schedule-aware recompute branch, which requires excluded days or a sub-daily window. Tables with an active schedule covering all 7 days at a non-sub-daily frequency (e.g., a 12h cadence) never stored it, so threshold_value fell back to the deadline-based upper tolerance and the first 2-3 missed updates were not flagged. Store the staleness threshold for every active-stage schedule: full-week schedules use the first-pass gap thresholds, which are already in business minutes since no days are excluded. Non-active stages keep the conservative fallback to upper_tolerance to avoid false weekend anomalies. Co-Authored-By: Claude --- .../commands/test_thresholds_prediction.py | 14 ++++-- tests/unit/common/conftest.py | 18 ++++++++ tests/unit/common/test_freshness_scenarios.py | 46 +++++++++++++++++++ 3 files changed, 74 insertions(+), 4 deletions(-) diff --git a/testgen/commands/test_thresholds_prediction.py b/testgen/commands/test_thresholds_prediction.py index 7c501e98..d357693b 100644 --- a/testgen/commands/test_thresholds_prediction.py +++ b/testgen/commands/test_thresholds_prediction.py @@ -89,7 +89,7 @@ def run(self) -> None: # Freshness update events are fetched as secondary data only when the suite # is a monitor — Volume_Trend / Metric_Trend in monitor suites couple to the - # Freshness_Trend signal to avoid stairstep false positives. + # Freshness_Trend signal to avoid stairstep false positives. freshness_updates_by_table: dict[tuple[str, str], list[str]] = ( self._fetch_freshness_updates_by_table() if self.test_suite.is_monitor else {} ) @@ -268,10 +268,16 @@ def compute_freshness_threshold( window_end=schedule.window_end if has_window else None, ) lower, upper = result.lower, result.upper - staleness = result.staleness except NotEnoughData: pass # Keep first-pass thresholds + # An active schedule certifies the cadence is regular, so the tight staleness + # threshold (median * factor) applies. Full-week schedules skip the recompute + # above and use the first-pass result — with no excluded days, those gaps are + # already in business minutes. Non-active stages keep staleness = None so the + # prediction SQL falls back to upper_tolerance (avoids false weekend anomalies). + staleness = result.staleness + # Override upper threshold with schedule-based deadline (daily/weekly only) if schedule.frequency != "sub_daily": holiday_dates = resolve_holiday_dates(holiday_codes, history.index) if holiday_codes else None @@ -344,7 +350,7 @@ def compute_volume_or_metric_threshold( This avoids the "stairstep" false-positive shape where inter-change plateaus collapse the SE estimate. The returned prediction JSON is augmented with `freshness_gated` and `baseline_value` so that test execution can apply dual-branch evaluation. - + If the filtered fit fails for any reason, falls back to fit SARIMAX on the raw value series and emits a prediction JSON without the freshness-gating markers. @@ -381,5 +387,5 @@ def compute_volume_or_metric_threshold( exclude_weekends=exclude_weekends, holiday_codes=holiday_codes, schedule_tz=schedule_tz, - ) + ) return lower, upper, None, prediction diff --git a/tests/unit/common/conftest.py b/tests/unit/common/conftest.py index 875d147b..9b895ce9 100644 --- a/tests/unit/common/conftest.py +++ b/tests/unit/common/conftest.py @@ -294,6 +294,24 @@ def _gen_subdaily_gap_schedule_phase() -> list[tuple[pd.Timestamp, float]]: return _to_csv_rows(_make_observations(start, end, 2, updates)) +def _gen_twice_daily_outage() -> list[tuple[pd.Timestamp, float]]: + """Updates every 12h (00:00/12:00 UTC), 7 days a week, 5 weeks, then updates stop. + + The 12h cadence classifies as "daily" frequency with a full-week active schedule, + so no excluded days or sub-daily window apply — staleness must come from the + first-pass gap thresholds. + """ + start = datetime(2025, 10, 6, 0, 0) + end = datetime(2025, 11, 11, 12, 0) + outage_start = datetime(2025, 11, 10, 0, 0) + updates: set[datetime] = set() + d = start + while d < outage_start: + updates.add(d) + d += timedelta(hours=12) + return _to_csv_rows(_make_observations(start, end, 12, updates)) + + def _gen_weekly_early() -> list[tuple[pd.Timestamp, float]]: start = datetime(2025, 8, 7, 10, 0) end = datetime(2025, 11, 6, 22, 0) diff --git a/tests/unit/common/test_freshness_scenarios.py b/tests/unit/common/test_freshness_scenarios.py index a20d2a52..a9e1f289 100644 --- a/tests/unit/common/test_freshness_scenarios.py +++ b/tests/unit/common/test_freshness_scenarios.py @@ -25,6 +25,7 @@ _gen_subdaily_gap_schedule_phase, _gen_subdaily_regular, _gen_training_only, + _gen_twice_daily_outage, _gen_weekly_early, _run_scenario, ) @@ -472,3 +473,48 @@ def test_recovery_passes(self, results: list[ScenarioPoint]) -> None: assert post_recovery[0].result_code == 0 # Second update after recovery should pass assert post_recovery[1].result_code == 1 + + +# ─── Scenario 9: Twice-Daily Outage (full-week schedule) ──────────── + + +class Test_TwiceDailyOutage: + """Updates every 12h around the clock for 5 weeks, then updates stop. + + Full-week "daily"-frequency schedule: no excluded days and no sub-daily + window, so the staleness threshold must come from the first-pass gap + computation. With a 720-min median gap, staleness = 720 * 0.85 = 612 min, + so the first missed check (gap 720 min) flags Late. + """ + + @pytest.fixture(scope="class") + def results(self) -> list[ScenarioPoint]: + rows = _gen_twice_daily_outage() + return _run_scenario(rows, PredictSensitivity.medium, exclude_weekends=False, tz="UTC") + + def test_schedule_goes_active_all_days(self, results: list[ScenarioPoint]) -> None: + sched = _schedule(_updates(results)[-1]) + assert sched is not None + assert sched.get("schedule_stage") == "active" + assert sched.get("active_days") == [0, 1, 2, 3, 4, 5, 6] + + def test_staleness_stored_once_active(self, results: list[ScenarioPoint]) -> None: + last_update = _updates(results)[-1] + assert last_update.staleness == pytest.approx(720 * 0.85) + + def test_no_anomalies_before_outage(self, results: list[ScenarioPoint]) -> None: + outage_start = pd.Timestamp("2025-11-10") + assert all(p.timestamp >= outage_start for p in _anomalies(results)) + + def test_first_missed_check_flags_late(self, results: list[ScenarioPoint]) -> None: + anomalies = _anomalies(results) + assert len(anomalies) > 0 + first = anomalies[0] + assert first.timestamp == pd.Timestamp("2025-11-10 00:00") + assert first.value == 720 + + def test_all_missed_checks_flag(self, results: list[ScenarioPoint]) -> None: + outage_start = pd.Timestamp("2025-11-10") + missed_checks = [p for p in results if p.timestamp >= outage_start] + assert len(missed_checks) == 4 + assert all(p.result_code == 0 for p in missed_checks) From cc905ab12d4a5485a9a95620a07dc5b8a81181d6 Mon Sep 17 00:00:00 2001 From: Ricardo Boni Date: Mon, 15 Jun 2026 12:26:57 -0400 Subject: [PATCH 32/78] refactor(runs): apply TG-1047 review feedback - drop fabricated job_execution_id from the profile-history test mock; the run primary key is the job execution id, so assertions use run.id - on_delete_runs: only delete the job execution explicitly when no run row exists; cascade_delete already removes it via the FK chain otherwise Co-Authored-By: Claude Fable 5 --- testgen/ui/views/profiling_runs.py | 3 ++- testgen/ui/views/test_runs.py | 3 ++- tests/unit/mcp/test_tools_profile_history.py | 16 +++++++--------- 3 files changed, 11 insertions(+), 11 deletions(-) diff --git a/testgen/ui/views/profiling_runs.py b/testgen/ui/views/profiling_runs.py index 9f034615..78d24df2 100644 --- a/testgen/ui/views/profiling_runs.py +++ b/testgen/ui/views/profiling_runs.py @@ -307,7 +307,8 @@ def on_delete_runs(job_execution_ids: list[str]) -> None: profiling_run = next(iter(select_profiling_runs_where(ProfilingRun.id == je_id)), None) if profiling_run: ProfilingRun.cascade_delete([str(profiling_run.id)]) - get_current_session().delete(job_exec) + else: + get_current_session().delete(job_exec) get_profiling_run_summaries.clear() Router().set_query_params({"page": 1}) except Exception: diff --git a/testgen/ui/views/test_runs.py b/testgen/ui/views/test_runs.py index 84bc515b..fe52d051 100644 --- a/testgen/ui/views/test_runs.py +++ b/testgen/ui/views/test_runs.py @@ -319,7 +319,8 @@ def on_delete_runs(job_execution_ids: list[str]) -> None: test_run = next(iter(select_test_runs_where(TestRun.id == je_id)), None) if test_run: TestRun.cascade_delete([str(test_run.id)]) - get_current_session().delete(job_exec) + else: + get_current_session().delete(job_exec) get_test_run_summaries.clear() Router().set_query_params({"page": 1}) except Exception: diff --git a/tests/unit/mcp/test_tools_profile_history.py b/tests/unit/mcp/test_tools_profile_history.py index e82b5bf3..910b0764 100644 --- a/tests/unit/mcp/test_tools_profile_history.py +++ b/tests/unit/mcp/test_tools_profile_history.py @@ -96,7 +96,6 @@ def _profile_row( def _profiling_run( id_=None, - job_execution_id=None, table_groups_id=None, status="Complete", profiling_starttime=None, @@ -105,7 +104,6 @@ def _profiling_run( ): run = MagicMock() run.id = id_ or uuid4() - run.job_execution_id = job_execution_id or uuid4() run.table_groups_id = table_groups_id or uuid4() run.status = status run.profiling_starttime = profiling_starttime or datetime(2026, 5, 10, 12, 0) @@ -283,7 +281,7 @@ def test_compare_profiling_runs_auto_baseline( mock_iss_type.select_where.return_value = [] with _patch_session([_je(), _je()]): - result = compare_profiling_runs(str(target_run.job_execution_id)) + result = compare_profiling_runs(str(target_run.id)) assert "Profiling Run Comparison" in result assert "Target" in result and "Baseline" in result @@ -298,7 +296,7 @@ def test_compare_profiling_runs_rejects_non_completed_target(mock_resolve, db_se with _patch_session([_je(status=JobStatus.RUNNING)]): with pytest.raises(MCPUserError, match="Target run is in `Running` state"): - compare_profiling_runs(str(target_run.job_execution_id)) + compare_profiling_runs(str(target_run.id)) @patch("testgen.mcp.tools.profile_history.resolve_profiling_run") @@ -308,7 +306,7 @@ def test_compare_profiling_runs_rejects_canceled_target(mock_resolve, db_session with _patch_session([_je(status=JobStatus.CANCELED)]): with pytest.raises(MCPUserError, match="`Canceled`"): - compare_profiling_runs(str(target_run.job_execution_id)) + compare_profiling_runs(str(target_run.id)) @patch("testgen.mcp.tools.profile_history.resolve_profiling_run") @@ -320,8 +318,8 @@ def test_compare_profiling_runs_rejects_cross_table_group(mock_resolve, db_sessi with _patch_session([_je()]): with pytest.raises(MCPUserError, match="same table group"): compare_profiling_runs( - str(target_run.job_execution_id), - str(baseline_run.job_execution_id), + str(target_run.id), + str(baseline_run.id), ) @@ -338,7 +336,7 @@ def test_compare_profiling_runs_auto_baseline_first_run(mock_resolve, db_session with _patch_session([_je()]): with pytest.raises(MCPUserError, match="no earlier completed profiling run"): - compare_profiling_runs(str(target_run.job_execution_id)) + compare_profiling_runs(str(target_run.id)) @patch("testgen.mcp.tools.profile_history.HygieneIssue") @@ -361,7 +359,7 @@ def test_compare_profiling_runs_identical_runs_renders_no_changes( mock_iss_type.select_where.return_value = [] with _patch_session([_je(), _je()]): - result = compare_profiling_runs(str(target_run.job_execution_id)) + result = compare_profiling_runs(str(target_run.id)) assert "No changes between target and baseline" in result From 9f600203f33586fcd1c680922ec635266d4481d2 Mon Sep 17 00:00:00 2001 From: Ricardo Boni Date: Mon, 15 Jun 2026 12:27:20 -0400 Subject: [PATCH 33/78] refactor(runs): renumber migration 0194 -> 0195 Avoids a collision with enterprise's 0194 (mssql schema-name realign) that landed on enterprise after this branch was last synced. TG-1047's dedup migration moves to the next free slot per the migration-numbering convention. Co-Authored-By: Claude Fable 5 --- ...{0194_incremental_upgrade.sql => 0195_incremental_upgrade.sql} | 0 1 file changed, 0 insertions(+), 0 deletions(-) rename testgen/template/dbupgrade/{0194_incremental_upgrade.sql => 0195_incremental_upgrade.sql} (100%) diff --git a/testgen/template/dbupgrade/0194_incremental_upgrade.sql b/testgen/template/dbupgrade/0195_incremental_upgrade.sql similarity index 100% rename from testgen/template/dbupgrade/0194_incremental_upgrade.sql rename to testgen/template/dbupgrade/0195_incremental_upgrade.sql From e5ed1fd764252d27e704e2d92ce5ad75663574a5 Mon Sep 17 00:00:00 2001 From: Aarthy Adityan Date: Mon, 15 Jun 2026 12:49:15 -0400 Subject: [PATCH 34/78] fix(profiling): prevent integer overflow in Invalid_Zip3_USA prevalence formula The dq_score_prevalence_formula cast record_ct (BIGINT) to INT, overflowing on tables with more than 2.1B rows and aborting hygiene issue detection for the whole profiling run. Drop the cast so the numerator stays BIGINT, matching the other operand and the Column_Pattern_Mismatch formula. Reaches existing deployments via the static-metadata refresh on upgrade; no migration needed. TG-1120 Co-Authored-By: Claude Opus 4.8 (1M context) --- .../profile_anomaly_types_Invalid_Zip3_USA.yaml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/testgen/template/dbsetup_anomaly_types/profile_anomaly_types_Invalid_Zip3_USA.yaml b/testgen/template/dbsetup_anomaly_types/profile_anomaly_types_Invalid_Zip3_USA.yaml index 723b3862..a20bc8df 100644 --- a/testgen/template/dbsetup_anomaly_types/profile_anomaly_types_Invalid_Zip3_USA.yaml +++ b/testgen/template/dbsetup_anomaly_types/profile_anomaly_types_Invalid_Zip3_USA.yaml @@ -16,7 +16,7 @@ profile_anomaly_types: suggested_action: |- Review your source data, ingestion process, and any processing steps that update this column. dq_score_prevalence_formula: |- - (NULLIF(p.record_ct, 0)::INT - NULLIF(SPLIT_PART(p.top_patterns, ' | ', 1), '')::BIGINT)::FLOAT/NULLIF(p.record_ct, 0)::FLOAT + (NULLIF(p.record_ct, 0) - NULLIF(SPLIT_PART(p.top_patterns, ' | ', 1), '')::BIGINT)::FLOAT/NULLIF(p.record_ct, 0)::FLOAT dq_score_risk_factor: '1' dq_dimension: Validity impact_dimension: Conformance From da00f50664bdbc09797241bbf0c2060e9d3119c5 Mon Sep 17 00:00:00 2001 From: Aarthy Adityan Date: Mon, 15 Jun 2026 12:49:15 -0400 Subject: [PATCH 35/78] fix(security): patch HIGH/CRITICAL image CVEs - pyarrow 21.0.0 -> 23.0.1 (CVE-2026-25087, HIGH) - PyJWT 2.12.0 -> 2.13.0 (CVE-2026-48526 HIGH + 3 MED + 1 LOW) - cryptography 46.0.6 -> 46.0.7 (CVE-2026-39892, MED) - base image: remove unused system Arrow packages (libarrow, apache-arrow-dev). pyarrow ships a self-contained musllinux wheel, so the apk Arrow libs were dead weight pulling in grpc (CVE-2026-33186, CRITICAL) and protobuf (CVE-2026-0994, HIGH) as runtime deps. Removing them also drops ~110 MB. - pip 26.0 -> 26.1.2 (3 MED) Co-Authored-By: Claude Opus 4.8 (1M context) --- deploy/testgen-base.dockerfile | 9 +++------ pyproject.toml | 6 +++--- 2 files changed, 6 insertions(+), 9 deletions(-) diff --git a/deploy/testgen-base.dockerfile b/deploy/testgen-base.dockerfile index 2fafe213..9b32f85a 100644 --- a/deploy/testgen-base.dockerfile +++ b/deploy/testgen-base.dockerfile @@ -27,9 +27,7 @@ RUN apk update && apk upgrade && apk add --no-cache \ openblas=0.3.30-r2 \ openblas-dev=0.3.30-r2 \ unixodbc=2.3.14-r0 \ - unixodbc-dev=2.3.14-r0 \ - libarrow=21.0.0-r4 \ - apache-arrow-dev=21.0.0-r4 + unixodbc-dev=2.3.14-r0 COPY --chmod=775 ./deploy/install_linuxodbc.sh /tmp/dk/install_linuxodbc.sh RUN /tmp/dk/install_linuxodbc.sh @@ -39,7 +37,7 @@ COPY ./pyproject.toml /tmp/dk/pyproject.toml RUN mkdir /dk # Upgrading pip for security -RUN python3 -m pip install --no-cache-dir --upgrade pip==26.0 +RUN python3 -m pip install --no-cache-dir --upgrade pip==26.1.2 # hdbcli only ships manylinux wheels (no musl). pip 26+ correctly rejects these on Alpine. # We download the wheel for the correct arch, then extract it directly into site-packages @@ -73,8 +71,7 @@ RUN apk del \ openssl \ linux-headers \ openblas-dev \ - unixodbc-dev \ - apache-arrow-dev + unixodbc-dev # Remove interactive ODBC tools — not needed at runtime, and iusql triggers # false-positive secret detection in security scanners (SECRET-3010) diff --git a/pyproject.toml b/pyproject.toml index 7af55506..3e7ca2df 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -61,7 +61,7 @@ dependencies = [ "xlsxwriter==3.2.0", "psutil==5.9.8", "concurrent_log_handler==0.9.25", - "cryptography==46.0.6", + "cryptography==46.0.7", "validators==0.33.0", "reportlab==4.2.2", "cron-converter==1.2.1", @@ -72,7 +72,7 @@ dependencies = [ "holidays~=0.89", # Pinned to match the manually compiled libs or for security - "pyarrow==21.0.0", + "pyarrow==23.0.1", "matplotlib==3.9.2", "scipy==1.14.1", "jinja2==3.1.6", @@ -82,7 +82,7 @@ dependencies = [ # MCP server "mcp[cli]==1.26.0", "uvicorn==0.41.0", - "PyJWT==2.12.0", + "PyJWT==2.13.0", "bcrypt==5.0.0", # API & OAuth server From 6360df553f840713e2a887f0f98093edf4476a77 Mon Sep 17 00:00:00 2001 From: Aarthy Adityan Date: Mon, 15 Jun 2026 12:49:15 -0400 Subject: [PATCH 36/78] fix(db): grant write access to project_memberships table project_memberships is mutated at runtime but was missing from the testgen_execute_role write-grant list, so membership changes failed with a permission error. Applied on existing deployments via the grant refresh on upgrade; no migration needed. --- testgen/template/dbsetup/075_grant_role_rights.sql | 1 + 1 file changed, 1 insertion(+) diff --git a/testgen/template/dbsetup/075_grant_role_rights.sql b/testgen/template/dbsetup/075_grant_role_rights.sql index a18f20b0..287663a8 100644 --- a/testgen/template/dbsetup/075_grant_role_rights.sql +++ b/testgen/template/dbsetup/075_grant_role_rights.sql @@ -44,6 +44,7 @@ GRANT SELECT, INSERT, DELETE, UPDATE ON {SCHEMA_NAME}.job_executions, {SCHEMA_NAME}.settings, {SCHEMA_NAME}.notification_settings, + {SCHEMA_NAME}.project_memberships, {SCHEMA_NAME}.test_definition_notes, {SCHEMA_NAME}.oauth2_clients, {SCHEMA_NAME}.oauth2_authorization_codes, From 32ffc20c1e1147daab34ecc1ddbe0683a530cbff Mon Sep 17 00:00:00 2001 From: Diogo Basto Date: Tue, 16 Jun 2026 13:11:58 +0100 Subject: [PATCH 37/78] refactor(mcp): apply TG-1090 round-2 review feedback - _format_monitor_cell: precedence reordered so a positive anomaly count beats training / pending. Previously a row that surfaced via anomaly_type=X could show "training" or "pending" in the X column, hiding the value that put it in the filtered set. - get_monitor_group_summary: empty-state lookback signal restored. The outer query now COALESCEs MAX(lookback) to 0 (was 1), preserving the pre-refactor dashboard signal for "No monitor runs yet". The fabricated default_lookback param is dropped. - parse_monitor_type: accepts an optional label so the error message names the caller's public arg. list_monitored_tables exposes the arg as anomaly_type, so the error now says "Invalid anomaly_type ..." instead of leaking the helper's internal monitor_type name. - common.py: drop unused format_monitor_type helper (no callers on this branch; re-added in TG-1092 where list_monitors will use it). Tests: - list_monitored_tables: new test pinning count > 0 precedence over training / pending, and verifying error still wins over count. - get_monitor_summary: new empty-state test (lookback=0; Window start/end fields suppressed; all four cells render the full "no results yet or not configured" phrase). - parse_monitor_type: new label-override test for the anomaly_type caller path. Co-Authored-By: Claude Opus 4.7 --- testgen/common/models/table_group.py | 7 +- testgen/mcp/tools/common.py | 17 ++--- testgen/mcp/tools/monitors.py | 14 +++- tests/unit/mcp/test_tools_common.py | 18 ++--- tests/unit/mcp/test_tools_monitors.py | 100 ++++++++++++++++++++++++-- 5 files changed, 125 insertions(+), 31 deletions(-) diff --git a/testgen/common/models/table_group.py b/testgen/common/models/table_group.py index b7aba489..b107e1d3 100644 --- a/testgen/common/models/table_group.py +++ b/testgen/common/models/table_group.py @@ -689,7 +689,7 @@ def get_monitor_group_summary( ) query = f""" SELECT - COALESCE(MAX(lookback), :default_lookback)::INTEGER AS lookback, + COALESCE(MAX(lookback), 0)::INTEGER AS lookback, MIN(lookback_start) AS lookback_start, MAX(lookback_end) AS lookback_end, COUNT(*)::INTEGER AS total_monitored_tables, @@ -719,9 +719,10 @@ def get_monitor_group_summary( COALESCE(BOOL_AND(metric_is_pending), TRUE) AS metric_is_pending FROM ({inner_query}) AS subquery """ - params = {**params, "default_lookback": lookback_override or 1} # Outer query has no GROUP BY — aggregates over zero rows still yield one - # COALESCE'd row, so .first() never returns None here. + # COALESCE'd row, so .first() never returns None here. ``lookback`` is 0 + # when the per-table CTE has no rows OR no runs against the monitor suite, + # so the dashboard / MCP can render the "no monitor runs yet" state. row = get_current_session().execute(text(query), params).mappings().first() return MonitorGroupSummary(**row) diff --git a/testgen/mcp/tools/common.py b/testgen/mcp/tools/common.py index 1a823f35..0aa555a6 100644 --- a/testgen/mcp/tools/common.py +++ b/testgen/mcp/tools/common.py @@ -357,29 +357,22 @@ def next_scheduled_run( "schema": MonitorType.SCHEMA, "metric": MonitorType.METRIC, } -_MONITOR_TYPE_DB_TO_USER: dict[MonitorType, str] = {v: k for k, v in _MONITOR_TYPE_USER_TO_DB.items()} -def parse_monitor_type(value: str) -> MonitorType: +def parse_monitor_type(value: str, label: str = "monitor_type") -> MonitorType: """Validate a user-facing monitor type label and return the stored ``MonitorType``. - Accepts ``freshness`` / ``volume`` / ``schema`` / ``metric``. + Accepts ``freshness`` / ``volume`` / ``schema`` / ``metric``. ``label`` names the + caller's argument in the error message — pass ``"anomaly_type"`` when the + public arg is named differently from ``monitor_type``. """ db_value = _MONITOR_TYPE_USER_TO_DB.get(value) if db_value is None: valid = ", ".join(_MONITOR_TYPE_USER_TO_DB) - raise MCPUserError(f"Invalid monitor_type `{value}`. Valid values: {valid}") + raise MCPUserError(f"Invalid {label} `{value}`. Valid values: {valid}") return db_value -def format_monitor_type(value: MonitorType | str) -> str: - """Map a stored monitor ``test_type`` to its user-facing short label.""" - try: - return _MONITOR_TYPE_DB_TO_USER[MonitorType(value)] - except ValueError: - return str(value) - - class MonitorTableSort(StrEnum): """User-facing values accepted for the ``sort_by`` argument on ``list_monitored_tables``. diff --git a/testgen/mcp/tools/monitors.py b/testgen/mcp/tools/monitors.py index b3837eea..3e91a98a 100644 --- a/testgen/mcp/tools/monitors.py +++ b/testgen/mcp/tools/monitors.py @@ -150,7 +150,7 @@ def list_monitored_tables( validate_page(page) validate_limit(limit, 100) - monitor_type = parse_monitor_type(anomaly_type) if anomaly_type is not None else None + monitor_type = parse_monitor_type(anomaly_type, "anomaly_type") if anomaly_type is not None else None sort = parse_monitor_table_sort(sort_by) if sort_by is not None else None tg, monitor_suite = resolve_monitored_table_group(table_group_id) @@ -245,14 +245,22 @@ def _format_monitor_cell( is_training: bool | None, has_error: bool, ) -> str: - """Render a per-table, per-monitor cell (non-schema).""" + """Render a per-table, per-monitor cell (non-schema). + + Precedence: error > positive count > pending > training > zero. A positive + count wins over training/pending so anomalies stay visible when a monitor is + still learning, and so a row that surfaces via ``anomaly_type=X`` actually + shows the X count in its column. + """ if has_error: return "error" + if count > 0: + return str(count) if is_pending: return "pending" if is_training: return "training" - return str(count) + return "0" def _format_schema_cell(row: MonitorTableSummary) -> str: diff --git a/tests/unit/mcp/test_tools_common.py b/tests/unit/mcp/test_tools_common.py index 7cd7c187..a2014d4f 100644 --- a/tests/unit/mcp/test_tools_common.py +++ b/tests/unit/mcp/test_tools_common.py @@ -873,6 +873,15 @@ def test_parse_monitor_type_lists_valid_values_on_error(): assert label in msg +def test_parse_monitor_type_label_override(): + """``label`` argument lets callers tailor the error to their public arg name + (e.g. ``list_monitored_tables`` exposes it as ``anomaly_type``).""" + from testgen.mcp.tools.common import parse_monitor_type + + with pytest.raises(MCPUserError, match=r"Invalid anomaly_type `bogus`"): + parse_monitor_type("bogus", "anomaly_type") + + @pytest.mark.parametrize( "value", ["table_name", "anomaly_count_desc", "latest_update_desc", "row_count_change_desc"], @@ -893,15 +902,6 @@ def test_parse_monitor_table_sort_rejects_unknown(): assert valid in msg -def test_format_monitor_type_round_trips(): - from testgen.common.enums import MonitorType - from testgen.mcp.tools.common import format_monitor_type - - for member in MonitorType: - formatted = format_monitor_type(member) - assert format_monitor_type(member.value) == formatted - - def test_resolve_monitored_table_group_returns_suite(): from testgen.common.models.table_group import TableGroup from testgen.common.models.test_suite import TestSuite diff --git a/tests/unit/mcp/test_tools_monitors.py b/tests/unit/mcp/test_tools_monitors.py index 581ff099..a7da2f7a 100644 --- a/tests/unit/mcp/test_tools_monitors.py +++ b/tests/unit/mcp/test_tools_monitors.py @@ -204,6 +204,45 @@ def test_get_monitor_summary_lookback_out_of_range(db_session_mock): assert "between 1 and 365" in str(exc.value) +@patch(f"{MODULE}.next_scheduled_run", return_value=None) +@patch(f"{MODULE}.TableGroup") +@patch(f"{MODULE}.resolve_monitored_table_group") +def test_get_monitor_summary_empty_state_lookback_zero( + mock_resolve, mock_tg_cls, mock_next, db_session_mock, +): + """When a monitor suite is configured but no runs have happened yet, the model + method returns ``lookback=0`` (preserves the pre-refactor signal the dashboard + uses to render "No monitor runs yet"). The MCP output reflects the empty state + rather than fabricating a one-run window.""" + tg = _mock_table_group() + mock_resolve.return_value = (tg, _mock_monitor_suite()) + mock_tg_cls.get_monitor_group_summary.return_value = _group_summary( + lookback=0, + lookback_start=None, + lookback_end=None, + total_monitored_tables=0, + freshness_anomalies=0, volume_anomalies=0, + schema_anomalies=0, metric_anomalies=0, + freshness_is_pending=True, volume_is_pending=True, + schema_is_pending=True, metric_is_pending=True, + freshness_is_training=False, volume_is_training=False, + metric_is_training=False, + ) + + from testgen.mcp.tools.monitors import get_monitor_summary + + with _patch_perms(): + out = get_monitor_summary(str(tg.id)) + + assert "**Lookback:** 0 runs" in out, "must show 0, not a fabricated default like 1" + # Window start / end fields are absent in the empty case (None values render as em-dash + # at most, but the tool omits them when the value is falsy) + assert "**Window start:**" not in out + assert "**Window end:**" not in out + # All per-type status cells reflect "no results" + assert out.count("no results yet or not configured") == 4 + + @patch(f"{MODULE}.next_scheduled_run", return_value=None) @patch(f"{MODULE}.TableGroup") @patch(f"{MODULE}.resolve_monitored_table_group") @@ -320,11 +359,15 @@ def test_list_monitored_tables_not_monitored(mock_resolve, db_session_mock): def test_list_monitored_tables_invalid_anomaly_type(db_session_mock): + """Error message must reference the caller's public arg name (``anomaly_type``) + rather than the helper's internal arg name (``monitor_type``).""" from testgen.mcp.tools.monitors import list_monitored_tables with _patch_perms(), pytest.raises(MCPUserError) as exc: list_monitored_tables(str(uuid4()), anomaly_type="bogus") - assert "Invalid monitor_type" in str(exc.value) + msg = str(exc.value) + assert "Invalid anomaly_type" in msg + assert "Invalid monitor_type" not in msg def test_list_monitored_tables_invalid_sort_by(db_session_mock): @@ -403,15 +446,16 @@ def test_list_monitored_tables_schema_change_column( def test_list_monitored_tables_training_and_pending_cells( mock_resolve, mock_tg_cls, db_session_mock, ): + """Status words render when the per-type count is zero.""" tg = _mock_table_group() mock_resolve.return_value = (tg, _mock_monitor_suite()) mock_tg_cls.list_monitor_table_summaries.return_value = ( [ _table_summary( table_name="t1", - freshness_is_training=True, - volume_is_pending=True, - metric_error_message="boom", + freshness_anomalies=0, freshness_is_training=True, + volume_anomalies=0, volume_is_pending=True, + metric_anomalies=0, metric_error_message="boom", ), ], 1, @@ -428,6 +472,54 @@ def test_list_monitored_tables_training_and_pending_cells( assert "error" in row +@patch(f"{MODULE}.TableGroup") +@patch(f"{MODULE}.resolve_monitored_table_group") +def test_list_monitored_tables_count_wins_over_training_and_pending( + mock_resolve, mock_tg_cls, db_session_mock, +): + """A positive anomaly count must surface even when the monitor is in training + or pending state — otherwise the cell hides the value that made the row match + an ``anomaly_type`` filter, and the table doesn't agree with itself. + + Precedence is error > positive count > pending > training > zero. Errors still + win (the latest measurement is suspect, so the historic count is misleading). + """ + tg = _mock_table_group() + mock_resolve.return_value = (tg, _mock_monitor_suite()) + mock_tg_cls.list_monitor_table_summaries.return_value = ( + [ + _table_summary( + table_name="t_busy_during_learn", + freshness_anomalies=5, freshness_is_training=True, + volume_anomalies=3, volume_is_pending=True, + metric_anomalies=2, metric_is_training=True, + ), + _table_summary( + table_name="t_error_with_count", + freshness_anomalies=4, freshness_error_message="db down", + ), + ], + 2, + ) + + from testgen.mcp.tools.monitors import list_monitored_tables + + with _patch_perms(): + out = list_monitored_tables(str(tg.id)) + + # Row 1: counts visible despite training/pending + row_count = next(line for line in out.splitlines() if "`t_busy_during_learn`" in line) + assert " 5 " in row_count, "freshness count should render despite is_training" + assert " 3 " in row_count, "volume count should render despite is_pending" + assert " 2 " in row_count, "metric count should render despite is_training" + assert "training" not in row_count + assert "pending" not in row_count + # Row 2: error still wins over count + row_error = next(line for line in out.splitlines() if "`t_error_with_count`" in line) + assert "error" in row_error + assert " 4 " not in row_error, "error must win over count (measurement is suspect)" + + @patch(f"{MODULE}.TableGroup") @patch(f"{MODULE}.resolve_monitored_table_group") def test_list_monitored_tables_empty(mock_resolve, mock_tg_cls, db_session_mock): From 2de9b0ba29ea1e539712252858621898c7ccc0bc Mon Sep 17 00:00:00 2001 From: Diogo Basto Date: Tue, 16 Jun 2026 13:22:30 +0100 Subject: [PATCH 38/78] refactor(mcp): collapse monitor dashboard helpers and close two unit-test gaps MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Perf — uncached dashboard render now runs the per-table monitor CTE 3 times instead of 5 (review feedback #4). ``get_monitor_changes_by_tables`` already had the total in hand from ``list_monitor_table_summaries`` and was throwing it away; ``count_monitor_changes_by_tables`` then re-ran the full CTE (rows + COUNT) just to recover that total. Helper now returns ``(rows, total)`` directly and the count helper is gone. Render code unpacks the tuple. Tests: - ``test_get_monitor_summary_lookback_out_of_range`` is parametrized over the lower bound, upper bound, beyond-upper, and a negative — was only pinning ``lookback=0`` before, so the 1-365 ceiling was untested. - New ``test_parse_monitor_table_sort_rejects_legacy_row_count_desc`` asserts the legacy ``row_count_desc`` name is rejected and the error surfaces the renamed ``row_count_change_desc`` — guards against an accidental revert of the column-rename signal that the dashboard cell shows a delta, not the raw current count. Co-Authored-By: Claude Opus 4.7 --- testgen/ui/views/monitors_dashboard.py | 39 ++++++++------------------ tests/unit/mcp/test_tools_common.py | 12 ++++++++ tests/unit/mcp/test_tools_monitors.py | 7 +++-- 3 files changed, 28 insertions(+), 30 deletions(-) diff --git a/testgen/ui/views/monitors_dashboard.py b/testgen/ui/views/monitors_dashboard.py index 2719109a..1a4043c5 100644 --- a/testgen/ui/views/monitors_dashboard.py +++ b/testgen/ui/views/monitors_dashboard.py @@ -133,7 +133,7 @@ def render( if sort_field and sort_field not in ALLOWED_SORT_FIELDS: sort_field = None - monitored_tables_page = get_monitor_changes_by_tables( + monitored_tables_page, all_monitored_tables_count = get_monitor_changes_by_tables( table_group_id, table_name_filter=table_name_filter, anomaly_type_filter=anomaly_type_filter, @@ -142,11 +142,6 @@ def render( limit=int(items_per_page), offset=page_start, ) - all_monitored_tables_count = count_monitor_changes_by_tables( - table_group_id, - table_name_filter=table_name_filter, - anomaly_type_filter=anomaly_type_filter, - ) monitor_changes_summary = summarize_monitor_changes(table_group_id) monitored_table_names = {table["table_name"] for table in monitored_tables_page} @@ -358,15 +353,18 @@ def get_monitor_changes_by_tables( sort_order: Literal["asc"] | Literal["desc"] | None = None, limit: int | None = None, offset: int | None = None, -) -> list[dict]: +) -> tuple[list[dict], int]: """Per-monitored-table summaries shaped for the dashboard's JSON payload. - Returns plain dicts (rather than ``MonitorTableSummary`` dataclasses) because the - monitor-dashboard widget consumes the payload via ``make_json_safe``. Each row is - augmented with ``table_group_id`` to match the historical payload shape. + Returns ``(rows, total)`` so the dashboard can fill its pager without an extra + round-trip to the model (which would re-run the heavy CTE twice — once for the + rows, once for the count — just to throw the rows away). Rows are dicts (rather + than ``MonitorTableSummary`` dataclasses) because the monitor-dashboard widget + consumes the payload via ``make_json_safe``. Each row is augmented with + ``table_group_id`` to match the historical payload shape. """ page = 1 + (offset // limit) if limit and offset else 1 - summaries, _total = TableGroup.list_monitor_table_summaries( + summaries, total = TableGroup.list_monitor_table_summaries( table_group_id, anomaly_types=_dashboard_anomaly_types(anomaly_type_filter), sort_by=_dashboard_sort_to_model(sort_field, sort_order), @@ -374,23 +372,8 @@ def get_monitor_changes_by_tables( page=page, limit=limit or 1000, ) - return [{**dataclasses.asdict(s), "table_group_id": table_group_id} for s in summaries] - - -@st.cache_data(show_spinner=False) -def count_monitor_changes_by_tables( - table_group_id: str, - table_name_filter: str | None = None, - anomaly_type_filter: list[str] | None = None, -) -> int: - _items, total = TableGroup.list_monitor_table_summaries( - table_group_id, - anomaly_types=_dashboard_anomaly_types(anomaly_type_filter), - table_name_filter=table_name_filter, - page=1, - limit=1, - ) - return total + rows = [{**dataclasses.asdict(s), "table_group_id": table_group_id} for s in summaries] + return rows, total @st.cache_data(show_spinner=False) diff --git a/tests/unit/mcp/test_tools_common.py b/tests/unit/mcp/test_tools_common.py index a2014d4f..a1d81f9a 100644 --- a/tests/unit/mcp/test_tools_common.py +++ b/tests/unit/mcp/test_tools_common.py @@ -902,6 +902,18 @@ def test_parse_monitor_table_sort_rejects_unknown(): assert valid in msg +def test_parse_monitor_table_sort_rejects_legacy_row_count_desc(): + """Guard against the pre-review-feedback ``row_count_desc`` name accidentally + coming back: the rename to ``row_count_change_desc`` is the canonical signal + that the column shows a delta, not the raw current count. Drop this only when + introducing a deliberate replacement.""" + from testgen.mcp.tools.common import parse_monitor_table_sort + + with pytest.raises(MCPUserError, match="Invalid sort_by") as exc: + parse_monitor_table_sort("row_count_desc") + assert "row_count_change_desc" in str(exc.value) + + def test_resolve_monitored_table_group_returns_suite(): from testgen.common.models.table_group import TableGroup from testgen.common.models.test_suite import TestSuite diff --git a/tests/unit/mcp/test_tools_monitors.py b/tests/unit/mcp/test_tools_monitors.py index a7da2f7a..327ce204 100644 --- a/tests/unit/mcp/test_tools_monitors.py +++ b/tests/unit/mcp/test_tools_monitors.py @@ -196,12 +196,15 @@ def test_get_monitor_summary_lookback_override_applied(mock_resolve, mock_tg_cls assert "**Lookback:** 14 runs (override)" in out -def test_get_monitor_summary_lookback_out_of_range(db_session_mock): +@pytest.mark.parametrize("bad_lookback", [0, 366, 500, -1]) +def test_get_monitor_summary_lookback_out_of_range(bad_lookback, db_session_mock): + """Both bounds (and beyond) are pinned — the model accepts 1..365 inclusive.""" from testgen.mcp.tools.monitors import get_monitor_summary with _patch_perms(), pytest.raises(MCPUserError) as exc: - get_monitor_summary(str(uuid4()), lookback=0) + get_monitor_summary(str(uuid4()), lookback=bad_lookback) assert "between 1 and 365" in str(exc.value) + assert f"`{bad_lookback}`" in str(exc.value) @patch(f"{MODULE}.next_scheduled_run", return_value=None) From 4dbea0efb933748a811afba7fc9bcab9bdb732cc Mon Sep 17 00:00:00 2001 From: testgen-ci-bot Date: Tue, 16 Jun 2026 18:32:05 +0000 Subject: [PATCH 39/78] ci: bump base image to v17 --- deploy/testgen.dockerfile | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/deploy/testgen.dockerfile b/deploy/testgen.dockerfile index 800cfa4f..d60654aa 100644 --- a/deploy/testgen.dockerfile +++ b/deploy/testgen.dockerfile @@ -1,4 +1,4 @@ -ARG TESTGEN_BASE_LABEL=v16 +ARG TESTGEN_BASE_LABEL=v17 FROM datakitchen/dataops-testgen-base:${TESTGEN_BASE_LABEL} AS release-image From 6bbfd7587cb70e6f6e4d43b3230876d346c2fc42 Mon Sep 17 00:00:00 2001 From: Ricardo Boni Date: Tue, 16 Jun 2026 16:16:41 -0400 Subject: [PATCH 40/78] refactor(mcp): remove connection write tools from the core connection surface (TG-1124) Remove create_connection and update_connection (with their private helpers and tool tests) from the core MCP connection tools. The core connection surface is now list / get / test. Co-Authored-By: Claude Opus 4.8 (1M context) --- testgen/mcp/server.py | 4 +- testgen/mcp/tools/connections.py | 92 ++---------------------- tests/unit/mcp/test_tools_connections.py | 4 +- 3 files changed, 11 insertions(+), 89 deletions(-) diff --git a/testgen/mcp/server.py b/testgen/mcp/server.py index 0be45764..88f69854 100644 --- a/testgen/mcp/server.py +++ b/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, diff --git a/testgen/mcp/tools/connections.py b/testgen/mcp/tools/connections.py index 9e4cf55f..ea6ee2fc 100644 --- a/testgen/mcp/tools/connections.py +++ b/testgen/mcp/tools/connections.py @@ -1,23 +1,19 @@ -"""MCP tools for database connection — create, update, 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. +"""MCP tools for database connections — list, get, and test database connections. + +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 @@ -25,7 +21,6 @@ SQL_FLAVOR_CODE_TO_LABEL, apply_connection_params, connection_display_fields, - connection_field_labels, format_flavor_label, format_page_footer, format_page_info, @@ -206,18 +201,9 @@ def _raise_validation_error(errors: list[str], header: str) -> None: 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) @@ -251,67 +237,3 @@ def _authentication_label(connection: Connection) -> str: 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, -} diff --git a/tests/unit/mcp/test_tools_connections.py b/tests/unit/mcp/test_tools_connections.py index 59c9d77d..b324eb42 100644 --- a/tests/unit/mcp/test_tools_connections.py +++ b/tests/unit/mcp/test_tools_connections.py @@ -1,4 +1,4 @@ -"""Tests for the MCP connection CRUD tools — create / update / test. +"""Tests for the core MCP connection tools — list / get / test. The tools take a flavor-shaped ``connection_params`` dict (keyed by UI label) plus an explicit ``connection_mode``; mapping + validation are delegated to @@ -11,7 +11,7 @@ import pytest from testgen.common.database.connection_service import ConnectionStatus -from testgen.mcp.exceptions import MCPPermissionDenied, MCPResourceNotAccessible, MCPUserError +from testgen.mcp.exceptions import MCPResourceNotAccessible, MCPUserError from testgen.mcp.permissions import ProjectPermissions pytestmark = pytest.mark.unit From 32b1d82a65eebd64c6084adb84932ec81d3da40d Mon Sep 17 00:00:00 2001 From: Ricardo Boni Date: Tue, 16 Jun 2026 16:16:54 -0400 Subject: [PATCH 41/78] feat(mcp): add plugin MCP-tools hook and promote shared connection helpers (TG-1124) Add PluginSpec.get_mcp_tools() (default empty) and register plugin-provided tools in build_mcp_server through the same safe_tool wrapper as core tools. Promote the shared connection render/validation helpers (effective_mode, raise_validation_error, render_connection_body, authentication_label) to mcp/tools/common.py so the read/test tools and any plugin tools share them. Tag the connection and table-group tool modules with _DOC_GROUP=MANAGE so they group under 'Manage TestGen configuration'. Co-Authored-By: Claude Opus 4.8 (1M context) --- testgen/mcp/server.py | 11 +++++ testgen/mcp/tools/common.py | 52 +++++++++++++++++++++++ testgen/mcp/tools/connections.py | 70 ++++--------------------------- testgen/mcp/tools/table_groups.py | 3 ++ testgen/utils/plugins.py | 12 +++++- 5 files changed, 86 insertions(+), 62 deletions(-) diff --git a/testgen/mcp/server.py b/testgen/mcp/server.py index 88f69854..cf64af9a 100644 --- a/testgen/mcp/server.py +++ b/testgen/mcp/server.py @@ -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 diff --git a/testgen/mcp/tools/common.py b/testgen/mcp/tools/common.py index 0aa555a6..2543bf3f 100644 --- a/testgen/mcp/tools/common.py +++ b/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" diff --git a/testgen/mcp/tools/connections.py b/testgen/mcp/tools/connections.py index ea6ee2fc..55423144 100644 --- a/testgen/mcp/tools/connections.py +++ b/testgen/mcp/tools/connections.py @@ -19,13 +19,15 @@ 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, + 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, @@ -33,6 +35,8 @@ ) from testgen.mcp.tools.markdown import MdDoc +_DOC_GROUP = DocGroup.MANAGE + @with_database_session @mcp_permission("view") @@ -99,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() @@ -151,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) @@ -160,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) @@ -181,59 +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_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" diff --git a/testgen/mcp/tools/table_groups.py b/testgen/mcp/tools/table_groups.py index ef9d9046..109a670f 100644 --- a/testgen/mcp/tools/table_groups.py +++ b/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. " diff --git a/testgen/utils/plugins.py b/testgen/utils/plugins.py index f4bfbe7f..07a87c8b 100644 --- a/testgen/utils/plugins.py +++ b/testgen/utils/plugins.py @@ -3,7 +3,7 @@ import importlib import importlib.metadata import inspect -from collections.abc import Generator +from collections.abc import Callable, Generator from typing import ClassVar, get_args from testgen.ui.assets import get_asset_path @@ -61,6 +61,16 @@ def configure_ui(cls) -> None: is actually running. Called by ``Plugin.load_streamlit()``, never by ``Plugin.load()``. """ + @classmethod + def get_mcp_tools(cls) -> list[Callable]: + """Return MCP tool callables to register into the MCP server. + + Override in plugins to register additional MCP tools. Implementations import their + tool modules inside the method (like ``configure_ui``) so plugin import stays light + for processes that never build the MCP server. + """ + return [] + class PluginHook: """Singleton holding resolved plugin values, pre-loaded with defaults.""" From 885053fe9671cd07ccf8abc26fc094c20539a0da Mon Sep 17 00:00:00 2001 From: Ricardo Boni Date: Tue, 16 Jun 2026 16:17:10 -0400 Subject: [PATCH 42/78] test: extract shared unit-test fixtures to testgen.testing (TG-1124) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Move db_session_mock/patched_settings/mcp_user (+ the perm matrix) into an importable testgen.testing.fixtures module so the core test suite and any plugin test suite register them from one source. Conftest discovery does not span separate test trees, so an importable module is the only clean way to share — the core conftests now import from it. Co-Authored-By: Claude Opus 4.8 (1M context) --- testgen/testing/__init__.py | 7 ++++ testgen/testing/fixtures.py | 76 +++++++++++++++++++++++++++++++++++++ tests/unit/conftest.py | 17 +-------- tests/unit/mcp/conftest.py | 57 +--------------------------- 4 files changed, 87 insertions(+), 70 deletions(-) create mode 100644 testgen/testing/__init__.py create mode 100644 testgen/testing/fixtures.py diff --git a/testgen/testing/__init__.py b/testgen/testing/__init__.py new file mode 100644 index 00000000..36919a40 --- /dev/null +++ b/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. +""" diff --git a/testgen/testing/fixtures.py b/testgen/testing/fixtures.py new file mode 100644 index 00000000..0ae02560 --- /dev/null +++ b/testgen/testing/fixtures.py @@ -0,0 +1,76 @@ +"""Reusable pytest fixtures for unit-testing TestGen. + +Importable by the core test suite and any plugin test suite. To use, import the fixtures +you need into a ``conftest.py`` — pytest registers fixtures that are imported into a +conftest, preserving their ``autouse`` flag for that subtree:: + + from testgen.testing.fixtures import db_session_mock, mcp_user # noqa: F401 +""" + +from unittest.mock import MagicMock, patch +from uuid import uuid4 + +import pytest + +from testgen.mcp.permissions import set_mcp_token, set_mcp_username + +# Fictional role matrix for tests. role_a has full access (but NOT view_pii — several +# tests rely on that to exercise the no-view_pii path), role_c is restricted, and +# role_d holds edit + view_pii so deny/allow pairs can be distinguished against a real +# ProjectPermissions without role_a accidentally granting view_pii. +TEST_PERM_MATRIX = { + "view": ["role_a", "role_b"], + "catalog": ["role_a", "role_b", "role_c"], + "edit": ["role_a", "role_d"], + "administer": ["role_a"], + "view_pii": ["role_d"], +} + + +def _test_roles_with_permission(permission): + return TEST_PERM_MATRIX.get(permission, []) + + +@pytest.fixture(autouse=True) +def patched_settings(): + with patch("testgen.settings.UI_BASE_URL", "http://tg-base-url"): + yield + + +@pytest.fixture +def db_session_mock(): + with patch("testgen.common.models.Session") as factory_mock: + yield factory_mock().__enter__() + + +@pytest.fixture(autouse=True) +def mcp_user(): + """Set up an authenticated MCP user for all tool tests. + + Default: user has 'role_a' on 'demo' project (full access). + The @mcp_permission decorator passes for any permission. + + Tests needing scoped access patch _compute_project_permissions directly. + """ + set_mcp_username("test_user") + set_mcp_token("test_bearer_token") + user = MagicMock() + user.id = uuid4() + + membership = MagicMock() + membership.project_code = "demo" + membership.role = "role_a" + + with ( + patch("testgen.common.auth.authorize_token", return_value=user), + patch("testgen.common.models.get_current_session", return_value=MagicMock()), + patch("testgen.mcp.permissions.ProjectMembership") as mock_membership, + patch("testgen.mcp.permissions.PluginHook") as mock_hook, + ): + mock_membership.get_memberships_for_user.return_value = [membership] + mock_hook.instance.return_value.rbac.get_roles_with_permission.side_effect = ( + _test_roles_with_permission + ) + yield user + set_mcp_username(None) + set_mcp_token(None) diff --git a/tests/unit/conftest.py b/tests/unit/conftest.py index 57794a74..fc7a48d3 100644 --- a/tests/unit/conftest.py +++ b/tests/unit/conftest.py @@ -1,15 +1,2 @@ -from unittest.mock import patch - -import pytest - - -@pytest.fixture(autouse=True) -def patched_settings(): - with patch("testgen.settings.UI_BASE_URL", "http://tg-base-url"): - yield - - -@pytest.fixture -def db_session_mock(): - with patch("testgen.common.models.Session") as factory_mock: - yield factory_mock().__enter__() +# Importing the fixtures registers them for this test subtree (autouse flags preserved). +from testgen.testing.fixtures import db_session_mock, patched_settings # noqa: F401 diff --git a/tests/unit/mcp/conftest.py b/tests/unit/mcp/conftest.py index fc50fef2..a2906ff5 100644 --- a/tests/unit/mcp/conftest.py +++ b/tests/unit/mcp/conftest.py @@ -1,55 +1,2 @@ -from unittest.mock import MagicMock, patch -from uuid import uuid4 - -import pytest - -from testgen.mcp.permissions import set_mcp_token, set_mcp_username - -# Fictional role matrix for tests. role_a has full access (but NOT view_pii — several -# tests rely on that to exercise the no-view_pii path), role_c is restricted, and -# role_d holds edit + view_pii so deny/allow pairs can be distinguished against a real -# ProjectPermissions without role_a accidentally granting view_pii. -TEST_PERM_MATRIX = { - "view": ["role_a", "role_b"], - "catalog": ["role_a", "role_b", "role_c"], - "edit": ["role_a", "role_d"], - "administer": ["role_a"], - "view_pii": ["role_d"], -} - - -def _test_roles_with_permission(permission): - return TEST_PERM_MATRIX.get(permission, []) - - -@pytest.fixture(autouse=True) -def mcp_user(): - """Set up an authenticated MCP user for all tool tests. - - Default: user has 'role_a' on 'demo' project (full access). - The @mcp_permission decorator passes for any permission. - - Tests needing scoped access patch _compute_project_permissions directly. - """ - set_mcp_username("test_user") - set_mcp_token("test_bearer_token") - user = MagicMock() - user.id = uuid4() - - membership = MagicMock() - membership.project_code = "demo" - membership.role = "role_a" - - with ( - patch("testgen.common.auth.authorize_token", return_value=user), - patch("testgen.common.models.get_current_session", return_value=MagicMock()), - patch("testgen.mcp.permissions.ProjectMembership") as mock_membership, - patch("testgen.mcp.permissions.PluginHook") as mock_hook, - ): - mock_membership.get_memberships_for_user.return_value = [membership] - mock_hook.instance.return_value.rbac.get_roles_with_permission.side_effect = ( - _test_roles_with_permission - ) - yield user - set_mcp_username(None) - set_mcp_token(None) +# Importing the fixture registers the autouse authenticated-MCP-user setup for these tests. +from testgen.testing.fixtures import mcp_user # noqa: F401 From 3b00bbaf43f3ccecd94e4653fe4c6a1be5dfa15e Mon Sep 17 00:00:00 2001 From: Ricardo Boni Date: Wed, 17 Jun 2026 11:56:27 -0400 Subject: [PATCH 43/78] fix(mcp): keep testgen.testing out of the built wheel; tidy connection docs (TG-1124) Exclude testgen.testing* from packages.find so the shared test-fixtures module (which imports pytest at module top) is not packaged into the runtime wheel; editable dev/CI installs still resolve it from source. Drop the dead testlib* exclude (no such package). Describe the testgen://connection-parameters resource by its param shapes rather than enumerating specific tools. Co-Authored-By: Claude Opus 4.8 (1M context) --- pyproject.toml | 2 +- testgen/mcp/tools/reference.py | 3 +-- 2 files changed, 2 insertions(+), 3 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 3e7ca2df..3b48accc 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -145,7 +145,7 @@ include-package-data = true include = [ "testgen*", ] -exclude = [ "*.tests", "tests*", "deploy*", "invocations*", "testlib*"] +exclude = [ "*.tests", "tests*", "deploy*", "invocations*", "testgen.testing*"] [tool.pytest.ini_options] minversion = "7.0" diff --git a/testgen/mcp/tools/reference.py b/testgen/mcp/tools/reference.py index a41704f6..94ada81c 100644 --- a/testgen/mcp/tools/reference.py +++ b/testgen/mcp/tools/reference.py @@ -410,8 +410,7 @@ def _append_mode(doc: MdDoc, mode: FlavorMode, schema: FlavorSchema, *, url_offe def connection_parameters_resource(flavor: str) -> str: """Per-flavor connection parameter shapes: the auth modes and the exact - ``connection_params`` keys (with required/optional + secret notes) used by - ``create_connection`` / ``update_connection`` / ``test_connection``. + ``connection_params`` keys (with required/optional + secret notes) for the flavor. Args: flavor: Flavor code, e.g. ``snowflake``, ``azure_mssql``, ``salesforce_data360``. From 25210bea56115e925539a982bb412a534dd65fd1 Mon Sep 17 00:00:00 2001 From: Ricardo Boni Date: Tue, 9 Jun 2026 15:42:33 -0400 Subject: [PATCH 44/78] feat(mcp): monitor lifecycle & settings tools (TG-1094) Add four MCP tools to manage a table group's monitoring end-to-end: enable_monitors, get_monitor_settings, update_monitor_settings, and disable_monitors. Extract the enable/update/disable logic out of the UI into a shared common/monitor_service.py, so the MCP tools, the monitors dashboard, and the table-group creation wizard all drive monitoring through one path. Move the holiday-calendar helpers (get_holiday_dates plus a new is_supported_holiday_code validator) into common/holiday_service.py. - enable/disable reject when the target state already matches; cron is validated by the scheduler's parser with its message surfaced verbatim - holiday_codes accepts a list, validated against the holidays package and serialized server-side; unknown codes are rejected with guidance - next-run is rendered in real UTC; schedule edits are applied in place, so the schedule id and its run history are preserved - JobSchedule.get_for_monitor_suite is shared by the MCP tools and the dashboard Co-Authored-By: Claude Opus 4.8 (1M context) --- testgen/common/freshness_service.py | 3 +- testgen/common/holiday_service.py | 64 +++++ testgen/common/models/scheduler.py | 8 + testgen/common/monitor_service.py | 161 ++++++++++++ testgen/common/time_series_service.py | 32 +-- testgen/mcp/server.py | 13 +- testgen/mcp/tools/monitors.py | 260 +++++++++++++++++- testgen/ui/services/query_cache.py | 7 +- testgen/ui/views/monitors_dashboard.py | 60 ++--- testgen/ui/views/table_groups.py | 42 +-- tests/unit/common/test_holiday_service.py | 28 ++ tests/unit/common/test_monitor_service.py | 118 +++++++++ tests/unit/mcp/test_tools_monitors.py | 305 +++++++++++++++++++++- 13 files changed, 995 insertions(+), 106 deletions(-) create mode 100644 testgen/common/holiday_service.py create mode 100644 testgen/common/monitor_service.py create mode 100644 tests/unit/common/test_holiday_service.py create mode 100644 tests/unit/common/test_monitor_service.py diff --git a/testgen/common/freshness_service.py b/testgen/common/freshness_service.py index b87e5a93..2e74abe8 100644 --- a/testgen/common/freshness_service.py +++ b/testgen/common/freshness_service.py @@ -8,7 +8,8 @@ import numpy as np import pandas as pd -from testgen.common.time_series_service import NotEnoughData, get_holiday_dates +from testgen.common.holiday_service import get_holiday_dates +from testgen.common.time_series_service import NotEnoughData LOG = logging.getLogger("testgen") diff --git a/testgen/common/holiday_service.py b/testgen/common/holiday_service.py new file mode 100644 index 00000000..e72e0fdd --- /dev/null +++ b/testgen/common/holiday_service.py @@ -0,0 +1,64 @@ +"""Holiday calendar utilities: resolve and validate country / financial-market holiday codes. + +Holiday codes name the calendars excluded from prediction baselines — used as SARIMAX +exogenous regressors and in freshness business-time gap calculations. Each code resolves +via the ``holidays`` package as either a country (ISO code, e.g. ``US``) or a financial +market (e.g. ``NYSE``). +""" + +import logging +from datetime import datetime + +import holidays +import pandas as pd + +LOG = logging.getLogger("testgen") + + +def get_holiday_dates(holiday_codes: list[str], datetime_index: pd.DatetimeIndex) -> set[datetime]: + years = list(range(datetime_index.year.min(), datetime_index.year.max() + 1)) + + holiday_dates = set() + if holiday_codes: + for code in holiday_codes: + code = code.strip().upper() + found = False + + try: + country_holidays = holidays.country_holidays(code, years=years) + holiday_dates.update(country_holidays.keys()) + found = True + except NotImplementedError: + pass # Not a valid country code + + if not found: + try: + financial_holidays = holidays.financial_holidays(code, years=years) + holiday_dates.update(financial_holidays.keys()) + found = True + except NotImplementedError: + pass # Not a valid financial code + + if not found: + LOG.warning(f"Holiday code '{code}' could not be resolved as a country or financial market") + + return holiday_dates + + +def is_supported_holiday_code(code: str) -> bool: + """Whether a holiday code resolves to a country or financial-market calendar. + + Applies the same normalization as :func:`get_holiday_dates` (strip + upper), so a code + that passes here is one that resolver will actually honor. + """ + normalized = code.strip().upper() + if not normalized: + return False + for resolver in (holidays.country_holidays, holidays.financial_holidays): + try: + resolver(normalized) + except NotImplementedError: + continue + else: + return True + return False diff --git a/testgen/common/models/scheduler.py b/testgen/common/models/scheduler.py index d1c03657..8b9e1e63 100644 --- a/testgen/common/models/scheduler.py +++ b/testgen/common/models/scheduler.py @@ -135,6 +135,14 @@ def select_active_by_kwargs( query = query.where(cls.kwargs[k].astext == str(v)) return list(get_current_session().scalars(query).all()) + @classmethod + def get_for_monitor_suite(cls, monitor_suite_id: str | UUID) -> Self | None: + """The run-monitors schedule for a monitor suite, active or paused.""" + return cls.get( + cls.key == RUN_MONITORS_JOB_KEY, + cls.kwargs["test_suite_id"].astext == str(monitor_suite_id), + ) + @classmethod def upsert_for_retention( cls, diff --git a/testgen/common/monitor_service.py b/testgen/common/monitor_service.py new file mode 100644 index 00000000..ebacf645 --- /dev/null +++ b/testgen/common/monitor_service.py @@ -0,0 +1,161 @@ +"""Monitor lifecycle: enable, update, and disable monitoring for a table group. + +Shared by the MCP tools, the monitors dashboard, and the table-group creation +wizard so all three drive monitoring through one path. Functions mutate the +monitor ``TestSuite``, its ``JobSchedule``, and the table-group link, and raise +stdlib exceptions — the MCP and UI layers translate those into their own +user-facing errors. +""" + +import logging +from typing import Any + +from sqlalchemy import func, select + +from testgen.commands.test_generation import run_monitor_generation +from testgen.common.models import get_current_session +from testgen.common.models.scheduler import RUN_MONITORS_JOB_KEY, JobSchedule +from testgen.common.models.table_group import TableGroup +from testgen.common.models.test_definition import TestDefinition +from testgen.common.models.test_result import TestResult +from testgen.common.models.test_run import TestRun +from testgen.common.models.test_suite import TestSuite + +LOG = logging.getLogger("testgen") + +# Monitors generated when monitoring is first enabled. Freshness_Trend depends on +# per-table freshness fingerprinting and is generated separately, so it is not part +# of the initial bootstrap set. +INITIAL_MONITOR_TYPES = ["Volume_Trend", "Schema_Drift"] + +# Default monitor configuration applied at bootstrap, mirroring the UI's setup form. +_DEFAULT_SUITE_ATTRS: dict[str, Any] = { + "monitor_lookback": 14, + "monitor_regenerate_freshness": True, + "predict_min_lookback": 30, + "predict_sensitivity": "medium", + "predict_exclude_weekends": False, + "predict_holiday_codes": None, +} + +# The monitor-configuration columns callers may set via ``suite_attrs``. Used as a +# whitelist so a caller's dict (e.g. a UI form payload) can't write arbitrary suite columns. +_MONITOR_SETTING_COLUMNS = tuple(_DEFAULT_SUITE_ATTRS) + + +def enable_monitoring( + table_group: TableGroup, + cron_expr: str, + cron_tz: str = "UTC", + *, + suite_attrs: dict[str, Any] | None = None, + active: bool = True, +) -> tuple[TestSuite, int]: + """Bootstrap monitoring for a table group. + + Creates the monitor test suite, generates the initial monitors, creates the + run-monitors schedule, and links the suite to the table group. + + ``suite_attrs`` overrides the default monitor configuration. + + Returns ``(monitor_suite, monitor_count)``. Raises ``ValueError`` if the table + group already has monitoring enabled. + """ + if table_group.monitor_test_suite_id: + raise ValueError("Monitoring is already enabled for this table group.") + + provided = suite_attrs or {} + attrs = dict(_DEFAULT_SUITE_ATTRS) + for key in _MONITOR_SETTING_COLUMNS: + if (value := provided.get(key)) is not None: + attrs[key] = value + + monitor_suite = TestSuite( + project_code=table_group.project_code, + test_suite=f"{table_group.table_groups_name} Monitors", + connection_id=table_group.connection_id, + table_groups_id=table_group.id, + export_to_observability=False, + dq_score_exclude=True, + is_monitor=True, + **attrs, + ) + monitor_suite.save() + + JobSchedule( + project_code=table_group.project_code, + key=RUN_MONITORS_JOB_KEY, + kwargs={"test_suite_id": str(monitor_suite.id)}, + cron_expr=cron_expr, + cron_tz=cron_tz, + active=active, + ).save() + + table_group.monitor_test_suite_id = monitor_suite.id + table_group.save() + + # Commit needed to make the test suite visible to run_monitor_generation's separate DB connection. + get_current_session().commit() + run_monitor_generation(monitor_suite.id, INITIAL_MONITOR_TYPES) + + count = get_current_session().scalar( + select(func.count()).select_from(TestDefinition).where(TestDefinition.test_suite_id == monitor_suite.id) + ) + return monitor_suite, count or 0 + + +def update_monitoring( + monitor_suite: TestSuite, + schedule: JobSchedule, + *, + suite_attrs: dict[str, Any] | None = None, + cron_expr: str | None = None, + cron_tz: str | None = None, + active: bool | None = None, +) -> None: + """Apply a partial update to monitor settings and/or schedule. + + ``suite_attrs`` maps monitor ``TestSuite`` columns to new values; only the keys in + ``_MONITOR_SETTING_COLUMNS`` that are present are applied (a present ``None`` clears + the column). Supplied schedule fields are updated in place. Does not commit — the + caller's session lifecycle does. + """ + provided = suite_attrs or {} + for key in _MONITOR_SETTING_COLUMNS: + if key in provided: + setattr(monitor_suite, key, provided[key]) + monitor_suite.save() + + if cron_expr is not None: + schedule.cron_expr = cron_expr + if cron_tz is not None: + schedule.cron_tz = cron_tz + if active is not None: + schedule.active = active + schedule.save() + + +def disable_monitoring(monitor_suite: TestSuite) -> dict[str, int]: + """Remove a monitor suite and all its monitors, runs, and history. + + Counts what will be removed before cascading the delete. Returns counts keyed + ``monitors`` (test definitions), ``events`` (test results), and ``runs``. + """ + session = get_current_session() + suite_id = monitor_suite.id + counts = { + "monitors": session.scalar( + select(func.count()).select_from(TestDefinition).where(TestDefinition.test_suite_id == suite_id) + ) + or 0, + "events": session.scalar( + select(func.count()).select_from(TestResult).where(TestResult.test_suite_id == suite_id) + ) + or 0, + "runs": session.scalar( + select(func.count()).select_from(TestRun).where(TestRun.test_suite_id == suite_id) + ) + or 0, + } + TestSuite.cascade_delete([suite_id]) + return counts diff --git a/testgen/common/time_series_service.py b/testgen/common/time_series_service.py index aeabb180..7a68fb35 100644 --- a/testgen/common/time_series_service.py +++ b/testgen/common/time_series_service.py @@ -1,11 +1,11 @@ import logging -from datetime import datetime -import holidays import numpy as np import pandas as pd from statsmodels.tsa.statespace.sarimax import SARIMAX +from testgen.common.holiday_service import get_holiday_dates + LOG = logging.getLogger("testgen") # This is a heuristic minimum to get a reasonable prediction @@ -137,31 +137,3 @@ def infer_frequency(datetime_series: pd.Series) -> str: return frequency if frequency != "0min" else f"{int(total_seconds)}S" -def get_holiday_dates(holiday_codes: list[str], datetime_index: pd.DatetimeIndex) -> set[datetime]: - years = list(range(datetime_index.year.min(), datetime_index.year.max() + 1)) - - holiday_dates = set() - if holiday_codes: - for code in holiday_codes: - code = code.strip().upper() - found = False - - try: - country_holidays = holidays.country_holidays(code, years=years) - holiday_dates.update(country_holidays.keys()) - found = True - except NotImplementedError: - pass # Not a valid country code - - if not found: - try: - financial_holidays = holidays.financial_holidays(code, years=years) - holiday_dates.update(financial_holidays.keys()) - found = True - except NotImplementedError: - pass # Not a valid financial code - - if not found: - LOG.warning(f"Holiday code '{code}' could not be resolved as a country or financial market") - - return holiday_dates diff --git a/testgen/mcp/server.py b/testgen/mcp/server.py index cf64af9a..cb35b33f 100644 --- a/testgen/mcp/server.py +++ b/testgen/mcp/server.py @@ -172,7 +172,14 @@ def build_mcp_server( search_hygiene_issues, update_hygiene_issue, ) - from testgen.mcp.tools.monitors import get_monitor_summary, list_monitored_tables + from testgen.mcp.tools.monitors import ( + disable_monitors, + enable_monitors, + get_monitor_settings, + get_monitor_summary, + list_monitored_tables, + update_monitor_settings, + ) from testgen.mcp.tools.notifications import ( create_notification, delete_notification, @@ -319,6 +326,10 @@ def safe_prompt(fn): safe_tool(get_schema_history) safe_tool(get_monitor_summary) safe_tool(list_monitored_tables) + safe_tool(enable_monitors) + safe_tool(get_monitor_settings) + safe_tool(update_monitor_settings) + safe_tool(disable_monitors) safe_tool(run_tests) safe_tool(run_profiling) safe_tool(cancel_test_run) diff --git a/testgen/mcp/tools/monitors.py b/testgen/mcp/tools/monitors.py index 3e91a98a..9aeef2f6 100644 --- a/testgen/mcp/tools/monitors.py +++ b/testgen/mcp/tools/monitors.py @@ -1,9 +1,18 @@ -from datetime import UTC +from datetime import UTC, datetime +from typing import Any, cast +from uuid import UUID +from sqlalchemy import select + +from testgen.common.cron_service import describe_cron, get_cron_sample from testgen.common.enums import MonitorType -from testgen.common.models import with_database_session -from testgen.common.models.scheduler import RUN_MONITORS_JOB_KEY +from testgen.common.holiday_service import is_supported_holiday_code +from testgen.common.models import get_current_session, with_database_session +from testgen.common.models.job_execution import JobExecution +from testgen.common.models.scheduler import RUN_MONITORS_JOB_KEY, JobSchedule from testgen.common.models.table_group import MonitorTableSummary, TableGroup +from testgen.common.models.test_suite import PredictSensitivity, TestSuite +from testgen.common.monitor_service import disable_monitoring, enable_monitoring, update_monitoring from testgen.mcp.exceptions import MCPUserError from testgen.mcp.permissions import mcp_permission from testgen.mcp.tools.common import ( @@ -315,3 +324,248 @@ def _format_row_count_change(row: MonitorTableSummary) -> str | None: if delta == 0: return "0" return f"{delta:+,}" + + +# --------------------------------------------------------------------------- +# Lifecycle & settings tools +# --------------------------------------------------------------------------- + +_LOOKBACK_RUNS_RANGE = (1, 200) +_MIN_TRAINING_LOOKBACK_RANGE = (20, 1000) + + +@with_database_session +@mcp_permission("edit") +def enable_monitors(table_group_id: str, cron_expression: str, cron_tz: str = "UTC") -> str: + """Turn on monitoring for a table group: create its monitors and put them on a schedule. + + Sets up the initial Volume and Schema monitors with default settings; adjust them afterward + with ``update_monitor_settings``. Fails if monitoring is already on for the group. + + Args: + table_group_id: UUID of the table group, e.g. from ``list_table_groups``. + cron_expression: Cron expression for the monitor schedule (required — monitors only run on a schedule), e.g. ``0 6 * * *`` for 6 AM daily. + cron_tz: IANA timezone for the schedule. Defaults to ``UTC``. + """ + tg, monitor_suite = resolve_monitored_table_group(table_group_id) + if monitor_suite is not None: + raise MCPUserError("Monitoring is already enabled for this table group.") + _validate_cron_verbatim(cron_expression, cron_tz) + + _, count = enable_monitoring(tg, cron_expression, cron_tz) + + doc = MdDoc() + doc.heading(1, f"Monitoring enabled for `{tg.table_groups_name}`") + doc.field("Initial monitors created", count) + doc.field("Cron expression", cron_expression, code=True) + if (readable := describe_cron(cron_expression)) is not None: + doc.field("Cron description", readable) + doc.field("Timezone", cron_tz) + return doc.render() + + +@with_database_session +@mcp_permission("view") +def get_monitor_settings(table_group_id: str) -> str: + """Get a table group's monitor configuration and schedule. + + Args: + table_group_id: UUID of the table group, e.g. from ``list_table_groups``. + """ + tg, monitor_suite = resolve_monitored_table_group(table_group_id) + if monitor_suite is None: + return _NOT_MONITORED_OUTPUT + + # A monitored table group always has a run-monitors schedule (enable_monitors creates it). + schedule = cast(JobSchedule, JobSchedule.get_for_monitor_suite(monitor_suite.id)) + + doc = MdDoc() + doc.heading(1, f"Monitor settings for `{tg.table_groups_name}`") + doc.field("Project", tg.project_code, code=True) + _render_monitor_settings(doc, monitor_suite, schedule) + return doc.render() + + +@with_database_session +@mcp_permission("edit") +def update_monitor_settings( + table_group_id: str, + sensitivity: str | None = None, + lookback_runs: int | None = None, + min_training_lookback: int | None = None, + exclude_weekends: bool | None = None, + holiday_codes: list[str] | None = None, + regenerate_freshness: bool | None = None, + cron_expression: str | None = None, + cron_tz: str | None = None, + active: bool | None = None, +) -> str: + """Update a table group's monitor configuration and/or schedule. Partial — omitted fields are left unchanged. + + Args: + table_group_id: UUID of the table group, e.g. from ``list_table_groups``. + sensitivity: Anomaly-detection sensitivity. One of ``low`` / ``medium`` / ``high``. + lookback_runs: Monitor runs aggregated for dashboard summaries (1-200). + min_training_lookback: Minimum runs before predictions activate (20-1000). + exclude_weekends: Whether to exclude weekends from prediction baselines. + holiday_codes: Holiday calendars to exclude from baselines — ISO country codes (e.g. ``US``, ``GB``) or financial-market codes (e.g. ``NYSE``, ``ECB``); see https://holidays.readthedocs.io/en/latest/#available-countries. Pass an empty list to clear. + regenerate_freshness: Whether to auto-regenerate Freshness monitors when the schema shifts. + cron_expression: New cron expression for the schedule, e.g. ``0 6 * * *``. + cron_tz: New IANA timezone for the schedule, e.g. ``America/New_York``. + active: ``True`` to resume the schedule, ``False`` to pause it. + """ + tg, monitor_suite = resolve_monitored_table_group(table_group_id) + if monitor_suite is None: + return _NOT_MONITORED_OUTPUT + + if all( + value is None + for value in ( + sensitivity, + lookback_runs, + min_training_lookback, + exclude_weekends, + holiday_codes, + regenerate_freshness, + cron_expression, + cron_tz, + active, + ) + ): + raise MCPUserError("No fields supplied to update.") + + # A monitored table group always has a run-monitors schedule (enable_monitors creates it). + schedule = cast(JobSchedule, JobSchedule.get_for_monitor_suite(monitor_suite.id)) + + suite_attrs: dict[str, Any] = {} + if sensitivity is not None: + suite_attrs["predict_sensitivity"] = _parse_sensitivity(sensitivity) + if lookback_runs is not None: + _validate_range(lookback_runs, "lookback_runs", *_LOOKBACK_RUNS_RANGE) + suite_attrs["monitor_lookback"] = lookback_runs + if min_training_lookback is not None: + _validate_range(min_training_lookback, "min_training_lookback", *_MIN_TRAINING_LOOKBACK_RANGE) + suite_attrs["predict_min_lookback"] = min_training_lookback + if exclude_weekends is not None: + suite_attrs["predict_exclude_weekends"] = exclude_weekends + if holiday_codes is not None: + cleaned = [code.strip() for code in holiday_codes if code.strip()] + invalid = [code for code in cleaned if not is_supported_holiday_code(code)] + if invalid: + raise MCPUserError( + f"Unknown holiday codes: {', '.join(invalid)}. Use ISO country codes (e.g. `US`, `GB`) " + "or financial-market codes (e.g. `NYSE`, `ECB`); see " + "https://holidays.readthedocs.io/en/latest/#available-countries." + ) + suite_attrs["predict_holiday_codes"] = ",".join(cleaned) if cleaned else None + if regenerate_freshness is not None: + suite_attrs["monitor_regenerate_freshness"] = regenerate_freshness + + if cron_expression is not None or cron_tz is not None: + effective_expr = cron_expression if cron_expression is not None else schedule.cron_expr + effective_tz = cron_tz if cron_tz is not None else schedule.cron_tz + _validate_cron_verbatim(effective_expr, effective_tz) + + update_monitoring( + monitor_suite, + schedule, + suite_attrs=suite_attrs, + cron_expr=cron_expression, + cron_tz=cron_tz, + active=active, + ) + + doc = MdDoc() + doc.heading(1, f"Monitor settings updated for `{tg.table_groups_name}`") + _render_monitor_settings(doc, monitor_suite, schedule) + return doc.render() + + +@with_database_session +@mcp_permission("edit") +def disable_monitors(table_group_id: str) -> str: + """Turn off monitoring for a table group, removing all its monitors, the schedule, and monitor history. + + This is irreversible — the monitors and their accumulated history are deleted. + + Args: + table_group_id: UUID of the table group, e.g. from ``list_table_groups``. + """ + tg, monitor_suite = resolve_monitored_table_group(table_group_id) + if monitor_suite is None: + raise MCPUserError("Monitoring is not enabled for this table group.") + + counts = disable_monitoring(monitor_suite) + + doc = MdDoc() + doc.heading(1, f"Monitoring disabled for `{tg.table_groups_name}`") + doc.text("Removed all monitors, the schedule, and monitor history.") + doc.field("Monitors removed", counts["monitors"]) + doc.field("Events removed", counts["events"]) + doc.field("Runs removed", counts["runs"]) + return doc.render() + + +def _parse_sensitivity(value: str) -> PredictSensitivity: + try: + return PredictSensitivity(value) + except ValueError as err: + valid = ", ".join(s.value for s in PredictSensitivity) + raise MCPUserError(f"Invalid sensitivity `{value}`. Valid values: {valid}") from err + + +def _validate_range(value: int, label: str, low: int, high: int) -> None: + if not low <= value <= high: + raise MCPUserError(f"Invalid {label} `{value}`: must be between {low} and {high}.") + + +def _validate_cron_verbatim(cron_expr: str, cron_tz: str) -> None: + """Validate a cron expression + timezone, surfacing the parser's message verbatim.""" + sample = get_cron_sample(cron_expr, cron_tz, sample_count=1) + if "error" in sample: + raise MCPUserError(sample["error"]) + + +def _last_monitor_run(schedule_id: UUID) -> datetime | None: + je = ( + get_current_session() + .scalars( + select(JobExecution) + .where(JobExecution.job_schedule_id == schedule_id) + .order_by(JobExecution.created_at.desc()) + .limit(1) + ) + .first() + ) + if je is None: + return None + return je.completed_at or je.started_at + + +def _render_monitor_settings(doc: MdDoc, monitor_suite: TestSuite, schedule: JobSchedule) -> None: + sensitivity = monitor_suite.predict_sensitivity.value if monitor_suite.predict_sensitivity else None + doc.field("Sensitivity", sensitivity) + doc.field("Lookback runs", monitor_suite.monitor_lookback) + doc.field("Min training lookback", monitor_suite.predict_min_lookback) + doc.field("Exclude weekends", monitor_suite.predict_exclude_weekends) + holidays = monitor_suite.holiday_codes_list + doc.field("Holiday codes", ", ".join(holidays) if holidays else None) + doc.field("Regenerate freshness", monitor_suite.monitor_regenerate_freshness) + + doc.heading(2, "Schedule") + doc.field("Cron expression", schedule.cron_expr, code=True) + if (readable := describe_cron(schedule.cron_expr)) is not None: + doc.field("Cron description", readable) + doc.field("Timezone", schedule.cron_tz) + doc.field("Status", "Active" if schedule.active else "Paused") + if schedule.active: + try: + next_runs = schedule.get_sample_triggering_timestamps(1) + except Exception: + next_runs = [] + if next_runs: + # Cron samples are tz-aware in the schedule's timezone; convert to UTC so the + # rendered " UTC" suffix is accurate. + doc.field("Next run", next_runs[0].astimezone(UTC)) + if (last_run := _last_monitor_run(schedule.id)) is not None: + doc.field("Last run", last_run) diff --git a/testgen/ui/services/query_cache.py b/testgen/ui/services/query_cache.py index 55be1e67..2e7367a7 100644 --- a/testgen/ui/services/query_cache.py +++ b/testgen/ui/services/query_cache.py @@ -17,7 +17,7 @@ from testgen.common.models.profiling_run import ProfilingRun, ProfilingRunMinimal, ProfilingRunSummary from testgen.common.models.project import Project, ProjectSummary from testgen.common.models.project_membership import ProjectMembership -from testgen.common.models.scheduler import RUN_MONITORS_JOB_KEY, JobSchedule +from testgen.common.models.scheduler import JobSchedule from testgen.common.models.table_group import ( TableGroup, TableGroupMinimal, @@ -146,10 +146,7 @@ def get_profiling_run_summaries( @st.cache_data(show_spinner=False) def get_monitor_schedule(monitor_suite_id: str | UUID) -> JobSchedule | None: - return JobSchedule.get( - JobSchedule.key == RUN_MONITORS_JOB_KEY, - JobSchedule.kwargs["test_suite_id"].astext == str(monitor_suite_id), - ) + return JobSchedule.get_for_monitor_suite(monitor_suite_id) # -- Connection --------------------------------------------------------------- diff --git a/testgen/ui/views/monitors_dashboard.py b/testgen/ui/views/monitors_dashboard.py index 1a4043c5..8e44d4ba 100644 --- a/testgen/ui/views/monitors_dashboard.py +++ b/testgen/ui/views/monitors_dashboard.py @@ -2,24 +2,24 @@ import logging from datetime import UTC, date, datetime from math import ceil -from typing import Any, ClassVar, Literal +from typing import Any, ClassVar, Literal, cast import pandas as pd import streamlit as st -from testgen.commands.test_generation import run_monitor_generation from testgen.common.cron_service import get_cron_sample from testgen.common.freshness_service import add_business_minutes, get_schedule_params, resolve_holiday_dates -from testgen.common.models import get_current_session, with_database_session +from testgen.common.models import with_database_session from testgen.common.models.notification_settings import ( MonitorNotificationSettings, MonitorNotificationTrigger, NotificationEvent, ) -from testgen.common.models.scheduler import RUN_MONITORS_JOB_KEY, JobSchedule +from testgen.common.models.scheduler import JobSchedule from testgen.common.models.table_group import TableGroup, TableGroupMinimal from testgen.common.models.test_definition import TestDefinition, TestDefinitionSummary from testgen.common.models.test_suite import PredictSensitivity, TestSuite +from testgen.common.monitor_service import disable_monitoring, enable_monitoring, update_monitoring from testgen.ui.components import widgets as testgen from testgen.ui.navigation.menu import MenuItem from testgen.ui.navigation.page import Page @@ -437,39 +437,25 @@ def on_save_settings_clicked(payload: dict) -> None: get_monitor_suite, set_monitor_suite = temp_value(f"monitors:updated_suite:{monitor_suite_id}", default={}) if should_save(): - for key, value in get_monitor_suite().items(): - setattr(monitor_suite, key, value) - - is_new = not monitor_suite.id - monitor_suite.save() - - new_schedule_config = get_schedule() - if ( # Check if schedule has to be created/recreated - not schedule - or schedule.cron_tz != new_schedule_config["cron_tz"] - or schedule.cron_expr != new_schedule_config["cron_expr"] - ): - if schedule: - JobSchedule.delete(schedule.id) - - new_schedule = JobSchedule( - project_code=table_group.project_code, - key=RUN_MONITORS_JOB_KEY, - kwargs={"test_suite_id": str(monitor_suite.id)}, - **new_schedule_config, + schedule_config = get_schedule() + if monitor_suite_id: + # An existing monitor suite always has a run-monitors schedule. + update_monitoring( + monitor_suite, + cast(JobSchedule, schedule), + suite_attrs=get_monitor_suite(), + cron_expr=schedule_config["cron_expr"], + cron_tz=schedule_config["cron_tz"], + active=schedule_config["active"], + ) + else: + enable_monitoring( + get_table_group(table_group.id), + schedule_config["cron_expr"], + schedule_config["cron_tz"], + suite_attrs=get_monitor_suite(), + active=schedule_config["active"], ) - new_schedule.save() - - elif schedule.active != new_schedule_config["active"]: # Only active status changed - JobSchedule.update_active(schedule.id, new_schedule_config["active"]) - - if is_new: - updated_table_group = get_table_group(table_group.id) - updated_table_group.monitor_test_suite_id = monitor_suite.id - updated_table_group.save() - # Commit needed to make test suite visible to run_monitor_generation's separate DB connection - get_current_session().commit() - run_monitor_generation(monitor_suite.id, ["Volume_Trend", "Schema_Drift"]) st.session_state.pop(EDIT_MONITOR_SETTINGS_DIALOG_KEY, None) safe_rerun() @@ -497,7 +483,7 @@ def on_save_settings_clicked(payload: dict) -> None: def delete_monitor_suite(table_group: TableGroupMinimal) -> None: try: monitor_suite = get_test_suite(table_group.monitor_test_suite_id) - TestSuite.cascade_delete([monitor_suite.id]) + disable_monitoring(monitor_suite) st.cache_data.clear() except Exception: LOG.exception("Failed to delete monitor suite") diff --git a/testgen/ui/views/table_groups.py b/testgen/ui/views/table_groups.py index 408b86c5..6d568198 100644 --- a/testgen/ui/views/table_groups.py +++ b/testgen/ui/views/table_groups.py @@ -6,17 +6,17 @@ import streamlit as st from sqlalchemy.exc import IntegrityError -from testgen.commands.test_generation import run_monitor_generation from testgen.common.enums import JobSource from testgen.common.models import get_current_session, with_database_session from testgen.common.models.connection import Connection from testgen.common.models.job_execution import JobExecution from testgen.common.models.notification_settings import ProfilingRunNotificationSettings from testgen.common.models.profiling_run import ProfilingRun -from testgen.common.models.scheduler import RUN_MONITORS_JOB_KEY, RUN_TESTS_JOB_KEY, JobSchedule +from testgen.common.models.scheduler import RUN_TESTS_JOB_KEY, JobSchedule from testgen.common.models.table_group import TableGroup, TableGroupMinimal from testgen.common.models.test_run import TestRun from testgen.common.models.test_suite import TestSuite +from testgen.common.monitor_service import enable_monitoring from testgen.ui.components import widgets as testgen from testgen.ui.navigation.menu import MenuItem from testgen.ui.navigation.page import Page @@ -419,33 +419,19 @@ def on_close_clicked(_params: dict) -> None: monitor_test_suite_data = get_monitor_test_suite_data() or {} if monitor_test_suite_data.get("generate"): generate_monitor_suite = True - monitor_test_suite = TestSuite( - project_code=project_code, - test_suite=f"{table_group.table_groups_name} Monitors", - connection_id=table_group.connection_id, - table_groups_id=table_group.id, - export_to_observability=False, - dq_score_exclude=True, - is_monitor=True, - monitor_lookback=monitor_test_suite_data.get("monitor_lookback") or 14, - monitor_regenerate_freshness=monitor_test_suite_data.get("monitor_regenerate_freshness") or True, - predict_min_lookback=monitor_test_suite_data.get("predict_min_lookback") or 30, - predict_sensitivity=monitor_test_suite_data.get("predict_sensitivity") or "medium", - predict_exclude_weekends=monitor_test_suite_data.get("predict_exclude_weekends") or False, - predict_holiday_codes=monitor_test_suite_data.get("predict_holiday_codes") or None, + monitor_test_suite, _ = enable_monitoring( + table_group, + monitor_test_suite_data.get("schedule"), + monitor_test_suite_data.get("timezone") or "UTC", + suite_attrs={ + "monitor_lookback": monitor_test_suite_data.get("monitor_lookback"), + "monitor_regenerate_freshness": monitor_test_suite_data.get("monitor_regenerate_freshness"), + "predict_min_lookback": monitor_test_suite_data.get("predict_min_lookback"), + "predict_sensitivity": monitor_test_suite_data.get("predict_sensitivity"), + "predict_exclude_weekends": monitor_test_suite_data.get("predict_exclude_weekends"), + "predict_holiday_codes": monitor_test_suite_data.get("predict_holiday_codes"), + }, ) - monitor_test_suite.save() - # Commit needed to make test suite visible to run_monitor_generation's separate DB connection - get_current_session().commit() - run_monitor_generation(monitor_test_suite.id, ["Volume_Trend", "Schema_Drift"]) - - JobSchedule( - project_code=project_code, - key=RUN_MONITORS_JOB_KEY, - cron_expr=monitor_test_suite_data.get("schedule"), - cron_tz=monitor_test_suite_data.get("timezone"), - kwargs={"test_suite_id": str(monitor_test_suite.id)}, - ).save() if standard_test_suite or monitor_test_suite: table_group.default_test_suite_id = standard_test_suite.id if standard_test_suite else None diff --git a/tests/unit/common/test_holiday_service.py b/tests/unit/common/test_holiday_service.py new file mode 100644 index 00000000..a2b7ddf8 --- /dev/null +++ b/tests/unit/common/test_holiday_service.py @@ -0,0 +1,28 @@ +"""Tests for holiday-code validation in ``common/holiday_service.py``.""" + +import pytest + +from testgen.common.holiday_service import is_supported_holiday_code + +pytestmark = pytest.mark.unit + + +@pytest.mark.parametrize("code", ["US", "GB", "CA", "USA", "NYSE", "ECB", " us "]) +def test_supported_codes(code): + """ISO country codes, uppercase aliases, and financial-market codes resolve.""" + assert is_supported_holiday_code(code) is True + + +@pytest.mark.parametrize( + "code", + [ + "US_FEDERAL", # not a holidays-package calendar + "CA_STAT", + "Canada", # the upper-cased "CANADA" is not a valid key — documented quirk + "NOTAREALCODE", + "", + " ", + ], +) +def test_unsupported_codes(code): + assert is_supported_holiday_code(code) is False diff --git a/tests/unit/common/test_monitor_service.py b/tests/unit/common/test_monitor_service.py new file mode 100644 index 00000000..11c17c50 --- /dev/null +++ b/tests/unit/common/test_monitor_service.py @@ -0,0 +1,118 @@ +"""Tests for ``common/monitor_service.py`` — enable / update / disable orchestration.""" + +from types import SimpleNamespace +from unittest.mock import MagicMock, patch +from uuid import uuid4 + +import pytest + +pytestmark = pytest.mark.unit + +MODULE = "testgen.common.monitor_service" + + +def _table_group(**overrides) -> MagicMock: + tg = MagicMock() + tg.id = overrides.get("id", uuid4()) + tg.project_code = overrides.get("project_code", "demo") + tg.table_groups_name = overrides.get("table_groups_name", "Sales") + tg.connection_id = overrides.get("connection_id", 1) + tg.monitor_test_suite_id = overrides.get("monitor_test_suite_id", None) + return tg + + +# --------------------------------------------------------------------------- +# enable_monitoring +# --------------------------------------------------------------------------- + + +def test_enable_monitoring_rejects_when_already_enabled(): + from testgen.common.monitor_service import enable_monitoring + + with pytest.raises(ValueError, match="already enabled"): + enable_monitoring(_table_group(monitor_test_suite_id=uuid4()), "0 6 * * *") + + +@patch(f"{MODULE}.run_monitor_generation") +@patch(f"{MODULE}.get_current_session") +@patch(f"{MODULE}.JobSchedule") +@patch(f"{MODULE}.TestSuite") +def test_enable_monitoring_merges_defaults_skipping_none(mock_ts, mock_js, mock_session, mock_gen): + tg = _table_group() + suite = mock_ts.return_value + suite.id = uuid4() + mock_session.return_value.scalar.return_value = 4 + + from testgen.common.monitor_service import enable_monitoring + + returned_suite, count = enable_monitoring( + tg, + "0 6 * * *", + "UTC", + # monitor_lookback=None must fall back to the default; False must be honored (not None). + suite_attrs={"monitor_lookback": None, "predict_sensitivity": "high", "predict_exclude_weekends": False}, + ) + + kwargs = mock_ts.call_args.kwargs + assert kwargs["is_monitor"] is True + assert kwargs["monitor_lookback"] == 14 # None override skipped → default + assert kwargs["predict_sensitivity"] == "high" # override applied + assert kwargs["predict_exclude_weekends"] is False # False is not None → applied + assert kwargs["predict_min_lookback"] == 30 # untouched default + assert returned_suite is suite + assert count == 4 + mock_gen.assert_called_once_with(suite.id, ["Volume_Trend", "Schema_Drift"]) + tg.save.assert_called() + assert tg.monitor_test_suite_id == suite.id + + +# --------------------------------------------------------------------------- +# update_monitoring +# --------------------------------------------------------------------------- + + +def test_update_monitoring_whitelists_suite_attrs_and_edits_schedule_in_place(): + suite = SimpleNamespace(save=MagicMock()) + schedule = SimpleNamespace(cron_expr="0 6 * * *", cron_tz="UTC", active=True, save=MagicMock()) + + from testgen.common.monitor_service import update_monitoring + + update_monitoring( + suite, + schedule, + suite_attrs={"predict_sensitivity": "high", "predict_holiday_codes": None, "is_monitor": False}, + cron_expr="0 */12 * * *", + active=False, + ) + + # Whitelisted columns applied (a present None clears); non-whitelisted key ignored. + assert suite.predict_sensitivity == "high" + assert suite.predict_holiday_codes is None + assert not hasattr(suite, "is_monitor") + suite.save.assert_called_once() + + # Schedule edited in place — same object, no recreate. + assert schedule.cron_expr == "0 */12 * * *" + assert schedule.cron_tz == "UTC" # not supplied → unchanged + assert schedule.active is False + schedule.save.assert_called_once() + + +# --------------------------------------------------------------------------- +# disable_monitoring +# --------------------------------------------------------------------------- + + +@patch(f"{MODULE}.TestSuite") +@patch(f"{MODULE}.get_current_session") +def test_disable_monitoring_counts_before_cascade(mock_session, mock_ts): + suite = MagicMock() + suite.id = uuid4() + mock_session.return_value.scalar.side_effect = [4, 2, 1] # monitors, events, runs + + from testgen.common.monitor_service import disable_monitoring + + counts = disable_monitoring(suite) + + assert counts == {"monitors": 4, "events": 2, "runs": 1} + mock_ts.cascade_delete.assert_called_once_with([suite.id]) diff --git a/tests/unit/mcp/test_tools_monitors.py b/tests/unit/mcp/test_tools_monitors.py index 327ce204..553d8c88 100644 --- a/tests/unit/mcp/test_tools_monitors.py +++ b/tests/unit/mcp/test_tools_monitors.py @@ -1,4 +1,6 @@ -"""Tests for the MCP monitor read tools — ``get_monitor_summary`` and ``list_monitored_tables``.""" +"""Tests for the MCP monitor tools — read (``get_monitor_summary`` / ``list_monitored_tables``) +and lifecycle/settings (``enable_monitors`` / ``get_monitor_settings`` / ``update_monitor_settings`` +/ ``disable_monitors``).""" from datetime import UTC, datetime from unittest.mock import MagicMock, patch @@ -7,6 +9,7 @@ import pytest from testgen.common.models.table_group import MonitorGroupSummary, MonitorTableSummary +from testgen.common.models.test_suite import PredictSensitivity from testgen.mcp.exceptions import MCPResourceNotAccessible, MCPUserError from testgen.mcp.permissions import ProjectPermissions @@ -551,3 +554,303 @@ def test_list_monitored_tables_empty_beyond_last_page(mock_resolve, mock_tg_cls, out = list_monitored_tables(str(tg.id), page=99) assert "No tables on page 99 (total: 7)." in out + + +# --------------------------------------------------------------------------- +# Lifecycle & settings helpers +# --------------------------------------------------------------------------- + + +def _mock_schedule(**overrides) -> MagicMock: + sched = MagicMock() + sched.id = overrides.get("id", uuid4()) + sched.cron_expr = overrides.get("cron_expr", "0 6 * * *") + sched.cron_tz = overrides.get("cron_tz", "UTC") + sched.active = overrides.get("active", True) + sched.get_sample_triggering_timestamps.return_value = [ + overrides.get("next_run", datetime(2026, 6, 10, 6, 0, tzinfo=UTC)) + ] + return sched + + +def _settings_suite(**overrides) -> MagicMock: + suite = _mock_monitor_suite(**overrides) + suite.predict_sensitivity = overrides.get("predict_sensitivity", PredictSensitivity.medium) + suite.monitor_lookback = overrides.get("monitor_lookback", 14) + suite.predict_min_lookback = overrides.get("predict_min_lookback", 30) + suite.predict_exclude_weekends = overrides.get("predict_exclude_weekends", False) + suite.monitor_regenerate_freshness = overrides.get("monitor_regenerate_freshness", True) + suite.holiday_codes_list = overrides.get("holiday_codes_list", None) + return suite + + +# --------------------------------------------------------------------------- +# enable_monitors +# --------------------------------------------------------------------------- + + +@patch(f"{MODULE}.enable_monitoring") +@patch(f"{MODULE}.resolve_monitored_table_group") +def test_enable_monitors_happy_path(mock_resolve, mock_enable, db_session_mock): + tg = _mock_table_group(monitor_test_suite_id=None) + mock_resolve.return_value = (tg, None) + mock_enable.return_value = (MagicMock(), 4) + + from testgen.mcp.tools.monitors import enable_monitors + + with _patch_perms(permission="edit"): + out = enable_monitors(str(tg.id), "0 6 * * *", "America/New_York") + + assert "# Monitoring enabled for `Sales`" in out + assert "**Initial monitors created:** 4" in out + assert "**Cron expression:** `0 6 * * *`" in out + assert "America/New_York" in out + mock_enable.assert_called_once_with(tg, "0 6 * * *", "America/New_York") + + +@patch(f"{MODULE}.enable_monitoring") +@patch(f"{MODULE}.resolve_monitored_table_group") +def test_enable_monitors_already_enabled(mock_resolve, mock_enable, db_session_mock): + tg = _mock_table_group() + mock_resolve.return_value = (tg, _mock_monitor_suite()) + + from testgen.mcp.tools.monitors import enable_monitors + + with _patch_perms(permission="edit"), pytest.raises(MCPUserError, match="already enabled"): + enable_monitors(str(tg.id), "0 6 * * *") + + mock_enable.assert_not_called() + + +@patch(f"{MODULE}.enable_monitoring") +@patch(f"{MODULE}.resolve_monitored_table_group") +def test_enable_monitors_invalid_cron_rejected_before_side_effects(mock_resolve, mock_enable, db_session_mock): + tg = _mock_table_group(monitor_test_suite_id=None) + mock_resolve.return_value = (tg, None) + + from testgen.mcp.tools.monitors import enable_monitors + + with _patch_perms(permission="edit"), pytest.raises(MCPUserError): + enable_monitors(str(tg.id), "not a cron") + + mock_enable.assert_not_called() + + +# --------------------------------------------------------------------------- +# get_monitor_settings +# --------------------------------------------------------------------------- + + +@patch(f"{MODULE}._last_monitor_run", return_value=None) +@patch(f"{MODULE}.JobSchedule") +@patch(f"{MODULE}.resolve_monitored_table_group") +def test_get_monitor_settings_happy_path(mock_resolve, mock_js, mock_last, db_session_mock): + tg = _mock_table_group() + suite = _settings_suite( + predict_sensitivity=PredictSensitivity.high, monitor_lookback=50, holiday_codes_list=["US", "NYSE"] + ) + mock_resolve.return_value = (tg, suite) + mock_js.get_for_monitor_suite.return_value = _mock_schedule(cron_expr="0 */12 * * *", cron_tz="UTC", active=True) + + from testgen.mcp.tools.monitors import get_monitor_settings + + with _patch_perms(): + out = get_monitor_settings(str(tg.id)) + + assert "# Monitor settings for `Sales`" in out + assert "**Sensitivity:** high" in out + assert "**Lookback runs:** 50" in out + assert "**Holiday codes:** US, NYSE" in out + assert "**Regenerate freshness:** Yes" in out + assert "## Schedule" in out + assert "**Cron expression:** `0 */12 * * *`" in out + assert "**Status:** Active" in out + assert "Next run" in out + + +@patch(f"{MODULE}.resolve_monitored_table_group") +def test_get_monitor_settings_not_monitored(mock_resolve, db_session_mock): + tg = _mock_table_group(monitor_test_suite_id=None) + mock_resolve.return_value = (tg, None) + + from testgen.mcp.tools.monitors import get_monitor_settings + + with _patch_perms(): + out = get_monitor_settings(str(tg.id)) + + assert out == "This table group is not monitored." + + +# --------------------------------------------------------------------------- +# update_monitor_settings +# --------------------------------------------------------------------------- + + +@patch(f"{MODULE}._last_monitor_run", return_value=None) +@patch(f"{MODULE}.JobSchedule") +@patch(f"{MODULE}.update_monitoring") +@patch(f"{MODULE}.resolve_monitored_table_group") +def test_update_monitor_settings_partial_maps_args(mock_resolve, mock_update, mock_js, mock_last, db_session_mock): + tg = _mock_table_group() + mock_resolve.return_value = (tg, _settings_suite(predict_sensitivity=PredictSensitivity.high)) + mock_js.get_for_monitor_suite.return_value = _mock_schedule() + + from testgen.mcp.tools.monitors import update_monitor_settings + + with _patch_perms(permission="edit"): + out = update_monitor_settings(str(tg.id), sensitivity="high", lookback_runs=50, exclude_weekends=True) + + assert "# Monitor settings updated for `Sales`" in out + mock_update.assert_called_once() + suite_attrs = mock_update.call_args.kwargs["suite_attrs"] + assert suite_attrs["predict_sensitivity"] == PredictSensitivity.high + assert suite_attrs["monitor_lookback"] == 50 + assert suite_attrs["predict_exclude_weekends"] is True + + +@patch(f"{MODULE}.resolve_monitored_table_group") +def test_update_monitor_settings_not_monitored(mock_resolve, db_session_mock): + tg = _mock_table_group(monitor_test_suite_id=None) + mock_resolve.return_value = (tg, None) + + from testgen.mcp.tools.monitors import update_monitor_settings + + with _patch_perms(permission="edit"): + out = update_monitor_settings(str(tg.id), sensitivity="high") + + assert out == "This table group is not monitored." + + +@patch(f"{MODULE}.update_monitoring") +@patch(f"{MODULE}.resolve_monitored_table_group") +def test_update_monitor_settings_no_fields(mock_resolve, mock_update, db_session_mock): + tg = _mock_table_group() + mock_resolve.return_value = (tg, _mock_monitor_suite()) + + from testgen.mcp.tools.monitors import update_monitor_settings + + with _patch_perms(permission="edit"), pytest.raises(MCPUserError, match="No fields supplied"): + update_monitor_settings(str(tg.id)) + + mock_update.assert_not_called() + + +@patch(f"{MODULE}.update_monitoring") +@patch(f"{MODULE}.JobSchedule") +@patch(f"{MODULE}.resolve_monitored_table_group") +def test_update_monitor_settings_invalid_sensitivity(mock_resolve, mock_js, mock_update, db_session_mock): + tg = _mock_table_group() + mock_resolve.return_value = (tg, _mock_monitor_suite()) + mock_js.get_for_monitor_suite.return_value = _mock_schedule() + + from testgen.mcp.tools.monitors import update_monitor_settings + + with _patch_perms(permission="edit"), pytest.raises(MCPUserError, match="Invalid sensitivity"): + update_monitor_settings(str(tg.id), sensitivity="extreme") + + mock_update.assert_not_called() + + +@patch(f"{MODULE}.update_monitoring") +@patch(f"{MODULE}.JobSchedule") +@patch(f"{MODULE}.resolve_monitored_table_group") +def test_update_monitor_settings_lookback_out_of_range(mock_resolve, mock_js, mock_update, db_session_mock): + tg = _mock_table_group() + mock_resolve.return_value = (tg, _mock_monitor_suite()) + mock_js.get_for_monitor_suite.return_value = _mock_schedule() + + from testgen.mcp.tools.monitors import update_monitor_settings + + with _patch_perms(permission="edit"), pytest.raises(MCPUserError, match="between 1 and 200"): + update_monitor_settings(str(tg.id), lookback_runs=5000) + + mock_update.assert_not_called() + + +@patch(f"{MODULE}.update_monitoring") +@patch(f"{MODULE}.JobSchedule") +@patch(f"{MODULE}.resolve_monitored_table_group") +def test_update_monitor_settings_unknown_holiday_code(mock_resolve, mock_js, mock_update, db_session_mock): + # ``is_supported_holiday_code`` runs for real here — ``US_FEDERAL`` is not a valid calendar. + tg = _mock_table_group() + mock_resolve.return_value = (tg, _mock_monitor_suite()) + mock_js.get_for_monitor_suite.return_value = _mock_schedule() + + from testgen.mcp.tools.monitors import update_monitor_settings + + with _patch_perms(permission="edit"), pytest.raises(MCPUserError, match="Unknown holiday codes"): + update_monitor_settings(str(tg.id), holiday_codes=["US_FEDERAL"]) + + mock_update.assert_not_called() + + +@patch(f"{MODULE}._last_monitor_run", return_value=None) +@patch(f"{MODULE}.JobSchedule") +@patch(f"{MODULE}.update_monitoring") +@patch(f"{MODULE}.resolve_monitored_table_group") +def test_update_monitor_settings_holiday_codes_serialized(mock_resolve, mock_update, mock_js, mock_last, db_session_mock): + tg = _mock_table_group() + mock_resolve.return_value = (tg, _settings_suite()) + mock_js.get_for_monitor_suite.return_value = _mock_schedule() + + from testgen.mcp.tools.monitors import update_monitor_settings + + with _patch_perms(permission="edit"): + update_monitor_settings(str(tg.id), holiday_codes=["US", "NYSE"]) + + assert mock_update.call_args.kwargs["suite_attrs"]["predict_holiday_codes"] == "US,NYSE" + + +@patch(f"{MODULE}._last_monitor_run", return_value=None) +@patch(f"{MODULE}.JobSchedule") +@patch(f"{MODULE}.update_monitoring") +@patch(f"{MODULE}.resolve_monitored_table_group") +def test_update_monitor_settings_empty_holiday_codes_clears(mock_resolve, mock_update, mock_js, mock_last, db_session_mock): + tg = _mock_table_group() + mock_resolve.return_value = (tg, _settings_suite()) + mock_js.get_for_monitor_suite.return_value = _mock_schedule() + + from testgen.mcp.tools.monitors import update_monitor_settings + + with _patch_perms(permission="edit"): + update_monitor_settings(str(tg.id), holiday_codes=[]) + + assert mock_update.call_args.kwargs["suite_attrs"]["predict_holiday_codes"] is None + + +# --------------------------------------------------------------------------- +# disable_monitors +# --------------------------------------------------------------------------- + + +@patch(f"{MODULE}.disable_monitoring", return_value={"monitors": 4, "events": 2, "runs": 1}) +@patch(f"{MODULE}.resolve_monitored_table_group") +def test_disable_monitors_happy_path(mock_resolve, mock_disable, db_session_mock): + tg = _mock_table_group() + suite = _mock_monitor_suite() + mock_resolve.return_value = (tg, suite) + + from testgen.mcp.tools.monitors import disable_monitors + + with _patch_perms(permission="edit"): + out = disable_monitors(str(tg.id)) + + assert "# Monitoring disabled for `Sales`" in out + assert "**Monitors removed:** 4" in out + assert "**Events removed:** 2" in out + assert "**Runs removed:** 1" in out + mock_disable.assert_called_once_with(suite) + + +@patch(f"{MODULE}.disable_monitoring") +@patch(f"{MODULE}.resolve_monitored_table_group") +def test_disable_monitors_not_enabled(mock_resolve, mock_disable, db_session_mock): + tg = _mock_table_group(monitor_test_suite_id=None) + mock_resolve.return_value = (tg, None) + + from testgen.mcp.tools.monitors import disable_monitors + + with _patch_perms(permission="edit"), pytest.raises(MCPUserError, match="not enabled"): + disable_monitors(str(tg.id)) + + mock_disable.assert_not_called() From ede85890a0b901e1c8a61907b08a9ed175b2c616 Mon Sep 17 00:00:00 2001 From: Aarthy Adityan Date: Tue, 16 Jun 2026 22:56:24 -0400 Subject: [PATCH 45/78] fix(security): bump cryptography to 48.0.1 for patched bundled OpenSSL cryptography wheels through 46.x bundle a vulnerable OpenSSL (GHSA-537c-gmf6-5ccf, HIGH). 48.0.1 ships patched OpenSSL 4.0.1. Satisfies all dependent constraints (tightest is <49). Co-Authored-By: Claude Opus 4.8 (1M context) --- pyproject.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index 3b48accc..92f75a7f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -61,7 +61,7 @@ dependencies = [ "xlsxwriter==3.2.0", "psutil==5.9.8", "concurrent_log_handler==0.9.25", - "cryptography==46.0.7", + "cryptography==48.0.1", "validators==0.33.0", "reportlab==4.2.2", "cron-converter==1.2.1", From 1b55a9c1dfd327714aea40ab03b74d539402bb91 Mon Sep 17 00:00:00 2001 From: Diogo Basto Date: Tue, 16 Jun 2026 18:32:53 +0100 Subject: [PATCH 46/78] fix(monitors): show error icon over anomaly count in monitor table cell MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit When a monitor's latest run errored AND prior runs in the lookback window produced anomalies, the table cell on the monitor dashboard rendered the count as near-white text on a transparent background — visually empty but still in the DOM. The warning icon and its result-message tooltip never rendered. Two contributing causes: - AnomalyTag's content branch checked `anomalies > 0` before `hasErrors`, so the count won display whenever both were true. - `has-anomalies` and `has-errors` CSS classes were both applied. The `.has-errors` rule overrides background to transparent without touching text color, leaving the count in `var(--empty-light)` on the page. Reorder the content branch so `hasErrors` wins, and make `has-errors` and `has-anomalies` mutually exclusive in the class string so the pill styling no longer bleeds into the error state. Matches the shape that monitor_anomalies_summary.js already uses for the top-of-dashboard tags. Closes TG-1126 Co-Authored-By: Claude Opus 4.7 --- .../frontend/js/pages/monitors_dashboard.js | 44 +++++++++---------- 1 file changed, 22 insertions(+), 22 deletions(-) diff --git a/testgen/ui/components/frontend/js/pages/monitors_dashboard.js b/testgen/ui/components/frontend/js/pages/monitors_dashboard.js index f3eb9612..2ac446fe 100644 --- a/testgen/ui/components/frontend/js/pages/monitors_dashboard.js +++ b/testgen/ui/components/frontend/js/pages/monitors_dashboard.js @@ -1,13 +1,13 @@ /** * @import { MonitorSummary } from '/app/static/js/components/monitor_anomalies_summary.js'; * @import { CronSample, FilterOption, ProjectSummary } from '../types.js'; - * + * * @typedef Schedule * @type {object} * @property {boolean} active * @property {string} cron_tz * @property {CronSample} cron_sample - * + * * @typedef Monitor * @type {object} * @property {string} table_group_id @@ -36,25 +36,25 @@ * @property {number?} column_adds * @property {number?} column_drops * @property {number?} column_mods - * + * * @typedef MonitorList * @type {object} * @property {Monitor[]} items * @property {number} current_page * @property {number} items_per_page * @property {number} total_count - * + * * @typedef MonitorListFilters * @type {object} * @property {string?} table_group_id * @property {string?} table_name_filter * @property {string?} anomaly_type_filter - * + * * @typedef MonitorListSort * @type {object} * @property {string?} sort_field * @property {('asc'|'desc')?} sort_order - * + * * @typedef Permissions * @type {object} * @property {boolean} can_edit @@ -213,17 +213,17 @@ const MonitorsDashboard = (/** @type Properties */ props) => { class: 'flex-row fx-gap-1 schema-changes', onclick: () => { const summary = getValue(props.summary); - emit('OpenSchemaChanges', { payload: { + emit('OpenSchemaChanges', { payload: { table_name: monitor.table_name, start_time: summary?.lookback_start, end_time: summary?.lookback_end, }}); }, }, - monitor.table_state === 'added' + monitor.table_state === 'added' ? Icon({size: 20, classes: 'schema-icon', filled: true}, 'add_box') : null, - monitor.table_state === 'dropped' + monitor.table_state === 'dropped' ? Icon({size: 20, classes: 'schema-icon', filled: true}, 'indeterminate_check_box') : null, monitor.column_adds ? div( @@ -245,7 +245,7 @@ const MonitorsDashboard = (/** @type Properties */ props) => { { text: div( {class: 'flex-column fx-align-flex-start'}, - monitor.table_state === 'added' + monitor.table_state === 'added' ? span({class: 'mb-1', style: 'font-size: 14px;'}, 'Table added.') : null, monitor.table_state === 'dropped' @@ -329,7 +329,7 @@ const MonitorsDashboard = (/** @type Properties */ props) => { tooltipPosition: 'bottom-left', color: 'basic', type: 'stroked', - style: 'background: var(--button-generic-background-color);', + style: 'background: var(--button-generic-background-color);', onclick: () => emit('EditNotifications', {}), }), Button({ @@ -338,7 +338,7 @@ const MonitorsDashboard = (/** @type Properties */ props) => { tooltipPosition: 'bottom-left', color: 'basic', type: 'stroked', - style: 'background: var(--button-generic-background-color);', + style: 'background: var(--button-generic-background-color);', onclick: () => emit('EditMonitorSettings', {}), }), Button({ @@ -517,14 +517,14 @@ const MonitorsDashboard = (/** @type Properties */ props) => { result: van.derive(() => getValue(props.notifications_dialog)?.result), onClose: () => emit('NotificationsDialogClosed', {}), }), - EditMonitorSettings({ emit, + EditMonitorSettings({ emit, table_group: van.derive(() => getValue(props.edit_monitor_settings_dialog)?.table_group), schedule: van.derive(() => getValue(props.edit_monitor_settings_dialog)?.schedule), monitor_suite: van.derive(() => getValue(props.edit_monitor_settings_dialog)?.monitor_suite), cron_sample: van.derive(() => getValue(props.edit_monitor_settings_dialog)?.cron_sample), dialog: van.derive(() => getValue(props.edit_monitor_settings_dialog)?.dialog), }), - TableMonitoringTrend({ emit, + TableMonitoringTrend({ emit, freshness_events: van.derive(() => getValue(props.trends_dialog)?.freshness_events ?? []), volume_events: van.derive(() => getValue(props.trends_dialog)?.volume_events ?? []), schema_events: van.derive(() => getValue(props.trends_dialog)?.schema_events ?? []), @@ -534,14 +534,14 @@ const MonitorsDashboard = (/** @type Properties */ props) => { extended_history: van.derive(() => getValue(props.trends_dialog)?.extended_history), dialog: van.derive(() => getValue(props.trends_dialog)?.dialog), }), - EditTableMonitors({ emit, + EditTableMonitors({ emit, table_name: van.derive(() => getValue(props.edit_table_monitors_dialog)?.table_name), definitions: van.derive(() => getValue(props.edit_table_monitors_dialog)?.definitions ?? []), metric_test_type: van.derive(() => getValue(props.edit_table_monitors_dialog)?.metric_test_type), result: van.derive(() => getValue(props.edit_table_monitors_dialog)?.result), dialog: van.derive(() => getValue(props.edit_table_monitors_dialog)?.dialog), }), - SchemaChangesDialog({ emit, + SchemaChangesDialog({ emit, window_start: van.derive(() => getValue(props.schema_changes_dialog)?.window_start), window_end: van.derive(() => getValue(props.schema_changes_dialog)?.window_end), data_structure_logs: van.derive(() => getValue(props.schema_changes_dialog)?.data_structure_logs), @@ -568,9 +568,6 @@ const AnomalyTag = (anomalies, errorMessage = null, isTraining = false, isPendin const hasErrors = !!errorMessage; const content = van.derive(() => { - if (anomalies > 0) { - return span(anomalies); - } if (hasErrors) { return withTooltip( i({class: 'material-symbols-rounded'}, 'warning'), @@ -584,6 +581,9 @@ const AnomalyTag = (anomalies, errorMessage = null, isTraining = false, isPendin }, ); } + if (anomalies > 0) { + return span(anomalies); + } if (isTraining) { return withTooltip( i({class: 'material-symbols-rounded'}, 'more_horiz'), @@ -597,7 +597,7 @@ const AnomalyTag = (anomalies, errorMessage = null, isTraining = false, isPendin { class: `anomaly-tag-wrapper flex-row p-1 ${onClick ? 'clickable' : ''}`, onclick: onClick }, div( { - class: `anomaly-tag ${anomalies > 0 ? 'has-anomalies' : ''} ${hasErrors ? 'has-errors' : ''} ${isTraining ? 'is-training' : ''}`, + class: `anomaly-tag ${hasErrors ? 'has-errors' : anomalies > 0 ? 'has-anomalies' : ''} ${isTraining ? 'is-training' : ''}`, }, content, ), @@ -646,8 +646,8 @@ const ConditionalEmptyState = (projectSummary, userCanEdit, emit) => { }, }; } - - return EmptyState({ emit, + + return EmptyState({ emit, icon: 'apps_outage', ...args, }); From 4a611053c8cfb8fda8637fffd45651dcdad1ebb0 Mon Sep 17 00:00:00 2001 From: Ricardo Boni Date: Thu, 11 Jun 2026 10:30:01 -0400 Subject: [PATCH 47/78] feat(mcp): add export_tests and import_tests tools Thin wrappers over the TD export/import service. Export gated on view, import on edit with mode=preview default. Adds shared parse_enum helper to mcp/tools/common.py and migrates bulk_update_tests' action parsing. Co-Authored-By: Claude Fable 5 --- testgen/mcp/server.py | 4 + testgen/mcp/tools/common.py | 14 +- testgen/mcp/tools/test_definitions.py | 173 ++++++++++++++- tests/unit/mcp/test_tools_test_definitions.py | 204 +++++++++++++++++- 4 files changed, 388 insertions(+), 7 deletions(-) diff --git a/testgen/mcp/server.py b/testgen/mcp/server.py index cf64af9a..60cd6bf9 100644 --- a/testgen/mcp/server.py +++ b/testgen/mcp/server.py @@ -235,7 +235,9 @@ def build_mcp_server( create_test, create_test_note, delete_test_note, + export_tests, get_test, + import_tests, list_test_notes, list_test_types, list_tests, @@ -328,6 +330,8 @@ def safe_prompt(fn): safe_tool(update_test) safe_tool(validate_custom_test) safe_tool(bulk_update_tests) + safe_tool(export_tests) + safe_tool(import_tests) safe_tool(create_test_note) safe_tool(update_test_note) safe_tool(delete_test_note) diff --git a/testgen/mcp/tools/common.py b/testgen/mcp/tools/common.py index 2543bf3f..e038da6a 100644 --- a/testgen/mcp/tools/common.py +++ b/testgen/mcp/tools/common.py @@ -1,6 +1,7 @@ from dataclasses import dataclass, field from datetime import date, datetime from enum import StrEnum +from typing import NoReturn, TypeVar from uuid import UUID from sqlalchemy import select @@ -79,6 +80,17 @@ def parse_uuid(value: str, label: str = "ID") -> UUID: raise MCPUserError(f"Invalid {label}: `{value}` is not a valid UUID.") from err +_ParsedEnum = TypeVar("_ParsedEnum", bound=StrEnum) + + +def parse_enum(value: str, enum_cls: type[_ParsedEnum], label: str) -> _ParsedEnum: + try: + return enum_cls(value) + except ValueError as err: + valid = ", ".join(member.value for member in enum_cls) + raise MCPUserError(f"Invalid {label} `{value}`. Valid values: {valid}") from err + + def parse_result_status(value: str) -> TestResultStatus: try: return TestResultStatus(value) @@ -1385,7 +1397,7 @@ def effective_mode(connection: Connection, connection_mode: str | None) -> str | return str(inferred) if inferred is not None else None -def raise_validation_error(errors: list[str], header: str) -> None: +def raise_validation_error(errors: list[str], header: str) -> NoReturn: bullets = "\n".join(f"- {err}" for err in errors) raise MCPUserError(f"{header}\n\n{bullets}") diff --git a/testgen/mcp/tools/test_definitions.py b/testgen/mcp/tools/test_definitions.py index 67197f20..b0984c85 100644 --- a/testgen/mcp/tools/test_definitions.py +++ b/testgen/mcp/tools/test_definitions.py @@ -2,6 +2,7 @@ from enum import StrEnum from typing import NoReturn +from pydantic import ValidationError from sqlalchemy import update from testgen.common.custom_test_validation import validate_custom_query @@ -17,15 +18,32 @@ TestType, ) from testgen.common.models.test_result import TestResult +from testgen.common.test_definition_export_import_service import ( + ImportAction, + ImportConfig, + ImportMode, + ImportPayload, + ImportResponse, + ImportStrictViolation, + InvalidImportPayload, + OnAbsence, + OnMatch, + OnNew, + Origin, + export_definitions, + import_definitions, +) from testgen.mcp.exceptions import MCPUserError from testgen.mcp.permissions import get_authorized_mcp_user, get_project_permissions, mcp_permission from testgen.mcp.tools.common import ( DocGroup, format_page_footer, format_page_info, + parse_enum, parse_impact_dimension, parse_quality_dimension, parse_uuid, + raise_validation_error, resolve_test_definition, resolve_test_note, resolve_test_suite, @@ -698,11 +716,7 @@ def bulk_update_tests( table_name: Optional table-name filter. Case-sensitive. test_type: Optional test type name (e.g. ``Alpha Truncation``). """ - try: - bulk_action = BulkAction(action) - except ValueError as err: - valid = ", ".join(f"`{a.value}`" for a in BulkAction) - raise MCPUserError(f"`action` must be one of: {valid}.") from err + bulk_action = parse_enum(action, BulkAction, "action") suite = resolve_test_suite(test_suite_id) tt_code = resolve_test_type(test_type) if test_type else None @@ -742,3 +756,152 @@ def bulk_update_tests( doc.heading(1, f"{verb} {count} test(s) in suite `{suite.test_suite}`") doc.field("Filter", filter_str) return doc.render() + + +# --------------------------------------------------------------------------- +# Export / import +# +# Thin wrappers over ``testgen.common.test_definition_export_import_service`` — the +# same export document and import semantics as the REST API. +# --------------------------------------------------------------------------- + + +def _parse_import_payload(payload: str) -> ImportPayload: + try: + return ImportPayload.model_validate_json(payload) + except ValidationError as err: + errors = err.errors() + bullets = [f"`{'.'.join(str(part) for part in e['loc']) or 'payload'}`: {e['msg']}" for e in errors[:5]] + if len(errors) > 5: + bullets.append(f"…and {len(errors) - 5} more") + raise_validation_error(bullets, "`payload` is not a valid export document.") + + +def _append_import_result(doc: MdDoc, result: ImportResponse, preview: bool) -> None: + verb_suffix = " (projected)" if preview else "" + doc.field(f"Created{verb_suffix}", result.summary.created) + doc.field(f"Updated{verb_suffix}", result.summary.updated) + doc.field(f"Skipped{verb_suffix}", result.summary.skipped) + doc.field(f"Deleted{verb_suffix}", result.summary.deleted) + + rows_by_action: dict[ImportAction, list[list]] = {} + for item in result.items: + for td in item.tds: + rows_by_action.setdefault(item.action, []).append( + [item.reason.value, td.idx, td.target_id] + ) + + for action in ImportAction: + rows = rows_by_action.get(action) + if not rows: + continue + doc.heading(2, f"{action.value.capitalize()} ({len(rows)})") + doc.table(["Reason", "Index", "Test Definition ID"], rows, code=[2]) + + +@with_database_session +@mcp_permission("view") +def export_tests( + test_suite_id: str, + origin: str = "both", + table_name: str | None = None, + test_type: str | None = None, +) -> str: + """Export test definitions from a test suite as a portable JSON document. + + The returned JSON payload can be applied to another suite (or back to the same + suite) with ``import_tests`` — as-is for promotion or copying, or after editing + it for bulk changes such as threshold calibration. + + Args: + test_suite_id: UUID of the test suite to export from. + origin: Which tests to include: ``both`` (default), ``manual`` (only manually + created tests), or ``auto`` (only auto-generated tests). + table_name: Optional table-name filter. Case-sensitive. + test_type: Optional test type name filter, e.g. ``Alpha Truncation``. + """ + origin_enum = parse_enum(origin, Origin, "origin") + suite = resolve_test_suite(test_suite_id) + table_name = table_name or None + test_type_code = resolve_test_type(test_type) if test_type else None + + export = export_definitions(suite, origin_enum, table_name, test_type_code) + + filters = [] + if origin_enum is not Origin.both: + filters.append(f"{origin_enum.value} tests only") + if table_name: + filters.append(f"table `{table_name}`") + if test_type: + filters.append(f"type `{test_type}`") + + doc = MdDoc() + doc.heading(1, f"Test Definition Export from suite `{suite.test_suite}`") + doc.field("Project", export.source.project_code, code=True) + doc.field("Table Group", export.source.table_group) + doc.field("Schema", export.source.table_group_schema, code=True) + doc.field("Exported At", export.source.exported_at) + if filters: + doc.field("Filters", ", ".join(filters)) + doc.field("Tests Exported", len(export.definitions)) + doc.code_block(export.model_dump_json(indent=2, exclude_defaults=True), language="json") + return doc.render() + + +@with_database_session +@mcp_permission("edit") +def import_tests( + test_suite_id: str, + payload: str, + mode: str = "preview", + on_match: str = "overwrite_unlocked", + on_new: str = "create", + on_absence: str = "do_nothing", +) -> str: + """Apply an export document from ``export_tests`` to a test suite, defaulting to ``preview`` mode that reports the projected changes without persisting. + + Args: + test_suite_id: UUID of the target test suite. + payload: The export document as a JSON string (the fenced JSON returned by + ``export_tests``). + mode: ``preview`` (default) reports the projected changes without persisting; + ``apply`` persists, skipping entries that can't be applied; ``apply_strict`` + persists only if nothing would be skipped, otherwise applies nothing. + on_match: What to do when an incoming test matches an existing one: + ``overwrite_unlocked`` (default), ``overwrite_all``, or ``skip``. + on_new: What to do when an incoming test has no match: ``create`` (default), + ``create_and_lock``, or ``skip``. + on_absence: What to do with existing tests that are absent from the payload: + ``do_nothing`` (default), ``delete_all``, or ``delete_unlocked``. + """ + config = ImportConfig( + mode=parse_enum(mode, ImportMode, "mode"), + on_match=parse_enum(on_match, OnMatch, "on_match"), + on_new=parse_enum(on_new, OnNew, "on_new"), + on_absence=parse_enum(on_absence, OnAbsence, "on_absence"), + ) + suite = resolve_test_suite(test_suite_id) + parsed_payload = _parse_import_payload(payload) + + try: + result = import_definitions(suite, config, parsed_payload) + except InvalidImportPayload as err: + raise MCPUserError(str(err)) from err + except ImportStrictViolation as err: + breakdown = MdDoc() + _append_import_result(breakdown, err.result, preview=True) + raise MCPUserError( + f"Strict import failed: {err.result.summary.skipped} tests would be skipped. " + f"Nothing was applied.\n\n{breakdown.render()}\n\n" + "Retry with `mode='apply'` to import while skipping these, or adjust the payload or policies." + ) from err + + preview = config.mode is ImportMode.preview + doc = MdDoc() + if preview: + doc.heading(1, f"Test Definition Import Preview for suite `{suite.test_suite}`") + doc.text("Preview only — no changes were persisted. Re-run with `mode='apply'` to persist.") + else: + doc.heading(1, f"Test Definition Import into suite `{suite.test_suite}`") + _append_import_result(doc, result, preview=preview) + return doc.render() diff --git a/tests/unit/mcp/test_tools_test_definitions.py b/tests/unit/mcp/test_tools_test_definitions.py index 46488143..4f467ff5 100644 --- a/tests/unit/mcp/test_tools_test_definitions.py +++ b/tests/unit/mcp/test_tools_test_definitions.py @@ -1291,7 +1291,7 @@ def test_bulk_update_tests_invalid_action(mock_resolve_suite, mock_session, db_s from testgen.mcp.tools.test_definitions import bulk_update_tests - with pytest.raises(MCPUserError, match="`action`"): + with pytest.raises(MCPUserError, match="Invalid action `toggle`"): bulk_update_tests(test_suite_id=str(uuid4()), action="toggle") # Suite resolution happens before action validation in current code path? @@ -1313,3 +1313,205 @@ def test_bulk_update_tests_no_match(mock_resolve_suite, mock_session, db_session assert "No tests matched" in result assert "nonexistent" in result + + +# -- export_tests / import_tests ---------------------------------------------- + + +def _make_export_document(definitions=None): + from datetime import UTC, datetime + + from testgen.common.test_definition_export_import_service import ( + ExportDocument, + ExportSource, + TestDefinitionExport, + ) + + return ExportDocument( + version=1, + source=ExportSource( + project_code="demo", + test_suite="demo_suite", + table_group="tg", + table_group_schema="public", + exported_at=datetime(2026, 1, 1, tzinfo=UTC), + ), + definitions=definitions + if definitions is not None + else [TestDefinitionExport(test_type="Row_Ct", table_name="orders")], + ) + + +def _make_import_response(*, created=0, updated=0, skipped=0, deleted=0, items=None): + from testgen.common.test_definition_export_import_service import ImportResponse, ImportSummary + + return ImportResponse( + summary=ImportSummary(created=created, updated=updated, skipped=skipped, deleted=deleted), + items=items or [], + ) + + +@patch("testgen.mcp.tools.test_definitions.export_definitions") +@patch("testgen.mcp.tools.test_definitions.resolve_test_suite") +def test_export_tests_basic(mock_resolve_suite, mock_export, db_session_mock): + mock_resolve_suite.return_value = _make_suite() + mock_export.return_value = _make_export_document() + + from testgen.mcp.tools.test_definitions import export_tests + + result = export_tests(test_suite_id=str(uuid4())) + + assert "Test Definition Export from suite `demo_suite`" in result + assert "Tests Exported" in result + assert "```json" in result # payload rendered as a fenced JSON block + assert "Row_Ct" in result + # default origin=both → no Filters line + assert "Filters" not in result + + +@patch("testgen.mcp.tools.test_definitions.export_definitions") +@patch("testgen.mcp.tools.test_definitions.resolve_test_type") +@patch("testgen.mcp.tools.test_definitions.resolve_test_suite") +def test_export_tests_with_filters(mock_resolve_suite, mock_resolve_tt, mock_export, db_session_mock): + mock_resolve_suite.return_value = _make_suite() + mock_resolve_tt.return_value = "Row_Ct" + mock_export.return_value = _make_export_document() + + from testgen.mcp.tools.test_definitions import export_tests + + result = export_tests(test_suite_id=str(uuid4()), origin="manual", table_name="orders", test_type="Row Count") + + assert "Filters" in result + assert "manual tests only" in result + assert "table `orders`" in result + mock_resolve_tt.assert_called_once_with("Row Count") + # origin code, not label, is forwarded to the service + assert mock_export.call_args.args[1].value == "manual" + + +@patch("testgen.mcp.tools.test_definitions.resolve_test_suite") +def test_export_tests_invalid_origin(mock_resolve_suite, db_session_mock): + from testgen.mcp.tools.test_definitions import export_tests + + with pytest.raises(MCPUserError, match="Invalid origin `sideways`"): + export_tests(test_suite_id=str(uuid4()), origin="sideways") + + +@patch("testgen.mcp.tools.test_definitions.import_definitions") +@patch("testgen.mcp.tools.test_definitions.resolve_test_suite") +def test_import_tests_preview(mock_resolve_suite, mock_import, db_session_mock): + from testgen.common.test_definition_export_import_service import ( + ImportAction, + ImportItem, + ImportItemTD, + ImportReason, + ) + + mock_resolve_suite.return_value = _make_suite() + mock_import.return_value = _make_import_response( + created=1, + updated=1, + items=[ + ImportItem(action=ImportAction.create, reason=ImportReason.no_match, tds=[ImportItemTD(idx=0, target_id=uuid4())]), + ImportItem(action=ImportAction.update, reason=ImportReason.matched, tds=[ImportItemTD(idx=1, target_id=uuid4())]), + ], + ) + + from testgen.mcp.tools.test_definitions import import_tests + + result = import_tests(test_suite_id=str(uuid4()), payload='{"definitions": []}') + + assert "Import Preview" in result + assert "no changes were persisted" in result + assert "Created (projected)" in result + assert "Create (1)" in result + assert "Update (1)" in result + + +@patch("testgen.mcp.tools.test_definitions.import_definitions") +@patch("testgen.mcp.tools.test_definitions.resolve_test_suite") +def test_import_tests_apply(mock_resolve_suite, mock_import, db_session_mock): + mock_resolve_suite.return_value = _make_suite() + mock_import.return_value = _make_import_response(created=2) + + from testgen.mcp.tools.test_definitions import import_tests + + result = import_tests(test_suite_id=str(uuid4()), payload='{"definitions": []}', mode="apply") + + assert "Import into suite `demo_suite`" in result + assert "no changes were persisted" not in result + assert "Created" in result + # applied mode → no "(projected)" suffix + assert "(projected)" not in result + + +@patch("testgen.mcp.tools.test_definitions.resolve_test_suite") +def test_import_tests_invalid_mode(mock_resolve_suite, db_session_mock): + from testgen.mcp.tools.test_definitions import import_tests + + with pytest.raises(MCPUserError, match="Invalid mode `yolo`"): + import_tests(test_suite_id=str(uuid4()), payload='{"definitions": []}', mode="yolo") + + +@patch("testgen.mcp.tools.test_definitions.import_definitions") +@patch("testgen.mcp.tools.test_definitions.resolve_test_suite") +def test_import_tests_malformed_payload(mock_resolve_suite, mock_import, db_session_mock): + mock_resolve_suite.return_value = _make_suite() + + from testgen.mcp.tools.test_definitions import import_tests + + with pytest.raises(MCPUserError, match="not a valid export document"): + import_tests(test_suite_id=str(uuid4()), payload="{not json") + mock_import.assert_not_called() + + +@patch("testgen.mcp.tools.test_definitions.import_definitions") +@patch("testgen.mcp.tools.test_definitions.resolve_test_suite") +def test_import_tests_domain_error_surfaced(mock_resolve_suite, mock_import, db_session_mock): + from testgen.common.test_definition_export_import_service import ImportErrorCode, InvalidImportPayload + + mock_resolve_suite.return_value = _make_suite() + mock_import.side_effect = InvalidImportPayload( + ImportErrorCode.duplicate_natural_key, "Duplicate auto-gen key at index 1" + ) + + from testgen.mcp.tools.test_definitions import import_tests + + with pytest.raises(MCPUserError, match="Duplicate auto-gen key at index 1"): + import_tests(test_suite_id=str(uuid4()), payload='{"definitions": []}', mode="apply") + + +@patch("testgen.mcp.tools.test_definitions.import_definitions") +@patch("testgen.mcp.tools.test_definitions.resolve_test_suite") +def test_import_tests_strict_violation_embeds_breakdown(mock_resolve_suite, mock_import, db_session_mock): + from testgen.common.test_definition_export_import_service import ( + ImportAction, + ImportItem, + ImportItemTD, + ImportReason, + ImportStrictViolation, + ) + + mock_resolve_suite.return_value = _make_suite() + result = _make_import_response( + created=1, + skipped=1, + items=[ + ImportItem(action=ImportAction.create, reason=ImportReason.no_match, tds=[ImportItemTD(idx=0)]), + ImportItem(action=ImportAction.skip, reason=ImportReason.invalid_test_type, tds=[ImportItemTD(idx=1)]), + ], + ) + mock_import.side_effect = ImportStrictViolation(result) + + from testgen.mcp.tools.test_definitions import import_tests + + with pytest.raises(MCPUserError) as exc_info: + import_tests(test_suite_id=str(uuid4()), payload='{"definitions": []}', mode="apply_strict") + + msg = str(exc_info.value) + assert "Strict import failed" in msg + assert "1 tests would be skipped" in msg + assert "Nothing was applied" in msg + # the projected breakdown is embedded + assert "Skip (1)" in msg + assert "invalid_test_type" in msg From 3ce037323d7ba6fdd6841c9f9381a2082ae1654c Mon Sep 17 00:00:00 2001 From: Ricardo Boni Date: Thu, 18 Jun 2026 17:04:20 -0400 Subject: [PATCH 48/78] feat(api): add test run results endpoint MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add GET /api/v1/test-runs/{job_id}/results — paginated individual results for a test-execution run, using the standard {items, page, limit, total} envelope. Requires view permission on the run's project; not-found and not-accessible both return 404, and only test-execution runs resolve (other job kinds 404). Results are resource-shaped: test_type is the raw type code; result_status and disposition are lowercase snake_case StrEnums shared between request filters and responses. New api/enums.py maps the title-case DB values to the API surface. disposition exposes a no_decision value for results with no triage decision (NULL in the DB); omitting the filter returns active results (confirmed plus no_decision), excluding dismissed and muted. Co-Authored-By: Claude Opus 4.8 (1M context) --- testgen/api/enums.py | 52 ++++++ testgen/api/runs.py | 79 +++++++- testgen/api/schemas.py | 26 +++ testgen/common/models/test_result.py | 49 +++++ tests/unit/api/test_runs.py | 220 ++++++++++++++++++++++- tests/unit/mcp/test_model_test_result.py | 22 ++- 6 files changed, 442 insertions(+), 6 deletions(-) create mode 100644 testgen/api/enums.py diff --git a/testgen/api/enums.py b/testgen/api/enums.py new file mode 100644 index 00000000..7fc4eef0 --- /dev/null +++ b/testgen/api/enums.py @@ -0,0 +1,52 @@ +"""Lowercase presentation enums for API v1 request filters and responses. + +Shared home for StrEnums that map a DB-stored value to the lowercase snake_case +form exposed through the API. Several DB columns store title-case values +(``test_results.result_status``, ``*.disposition``); the mapping dicts here are the +single seam that normalizes them to the API surface — the DB is never changed. +""" + +from enum import StrEnum + +from testgen.common.enums import Disposition as DbDisposition +from testgen.common.models.test_result import TestResultStatus + + +class ResultStatus(StrEnum): + """Outcome of a single test result.""" + + passed = "passed" + failed = "failed" + warning = "warning" + error = "error" + log = "log" + + +class Disposition(StrEnum): + """Triage state of a test result. ``no_decision`` is the state of a result no one + has triaged yet. Omitting the filter returns active results (``confirmed`` and + ``no_decision``); pass an explicit value to filter to a single state.""" + + confirmed = "confirmed" + dismissed = "dismissed" + muted = "muted" + no_decision = "no_decision" + + +RESULT_STATUS_TO_DB: dict[ResultStatus, TestResultStatus] = { + ResultStatus.passed: TestResultStatus.Passed, + ResultStatus.failed: TestResultStatus.Failed, + ResultStatus.warning: TestResultStatus.Warning, + ResultStatus.error: TestResultStatus.Error, + ResultStatus.log: TestResultStatus.Log, +} +RESULT_STATUS_FROM_DB: dict[TestResultStatus, ResultStatus] = {v: k for k, v in RESULT_STATUS_TO_DB.items()} + +# ``no_decision`` has no stored DB value — it corresponds to a NULL ``disposition`` +# column, handled explicitly at the API boundary (not present in these dicts). +DISPOSITION_TO_DB: dict[Disposition, DbDisposition] = { + Disposition.confirmed: DbDisposition.CONFIRMED, + Disposition.dismissed: DbDisposition.DISMISSED, + Disposition.muted: DbDisposition.INACTIVE, +} +DISPOSITION_FROM_DB: dict[DbDisposition, Disposition] = {v: k for k, v in DISPOSITION_TO_DB.items()} diff --git a/testgen/api/runs.py b/testgen/api/runs.py index e16fbf8b..b32e5733 100644 --- a/testgen/api/runs.py +++ b/testgen/api/runs.py @@ -1,24 +1,35 @@ """API v1 — test run and profiling run retrieval.""" -from fastapi import APIRouter, Depends -from sqlalchemy import select +from fastapi import APIRouter, Depends, Query +from sqlalchemy import or_, select from testgen.api.deps import db_session, resolve_job +from testgen.api.enums import ( + DISPOSITION_FROM_DB, + DISPOSITION_TO_DB, + RESULT_STATUS_FROM_DB, + RESULT_STATUS_TO_DB, + Disposition, + ResultStatus, +) from testgen.api.schemas import ( ErrorResponse, IssueCounts, ProfilingRunResponse, ProfilingRunResult, ResultCounts, + TestResultItem, + TestResultListResponse, TestRunResponse, TestRunResult, ) +from testgen.common.enums import Disposition as DbDisposition from testgen.common.enums import JobKey from testgen.common.models import get_current_session from testgen.common.models.hygiene_issue import HygieneIssue from testgen.common.models.job_execution import JobExecution from testgen.common.models.profiling_run import ProfilingRun -from testgen.common.models.test_result import TestResult +from testgen.common.models.test_result import TestResult, TestRunResultRow from testgen.common.models.test_run import TestRun from testgen.common.models.test_suite import TestSuite @@ -63,6 +74,68 @@ def get_test_run(job: JobExecution = resolve_job("view", JobExecution.job_key == ) +def _to_item(row: TestRunResultRow) -> TestResultItem: + """Map a DB-valued result row to the API item, normalizing enum casing.""" + return TestResultItem( + test_definition_id=row.test_definition_id, + test_type=row.test_type, + schema_name=row.schema_name, + table_name=row.table_name, + column_names=row.column_names, + result_status=RESULT_STATUS_FROM_DB.get(row.status), + result_measure=row.result_measure, + threshold_value=row.threshold_value, + result_message=row.message, + test_time=row.test_time, + disposition=DISPOSITION_FROM_DB[DbDisposition(row.disposition)] if row.disposition else Disposition.no_decision, + ) + + +@router.get( + "/test-runs/{job_id}/results", + response_model=TestResultListResponse, +) +def list_test_run_results( + job: JobExecution = resolve_job("view", JobExecution.job_key == JobKey.run_tests), # noqa: B008 + status: ResultStatus | None = Query(default=None), # noqa: B008 + table_name: str | None = Query(default=None), + column_name: str | None = Query(default=None), + test_type: str | None = Query(default=None), + disposition: Disposition | None = Query(default=None), # noqa: B008 + page: int = Query(default=1, ge=1), + limit: int = Query(default=20, ge=1, le=100), +): + """List individual results for a test run. + + Omitting ``disposition`` returns active results — confirmed and no_decision + (excludes dismissed and muted); pass an explicit value to filter to one state. + """ + clauses = [] + if status: + clauses.append(TestResult.status == RESULT_STATUS_TO_DB[status]) + if table_name: + clauses.append(TestResult.table_name == table_name) + if column_name: + clauses.append(TestResult.column_names == column_name) + if test_type: + clauses.append(TestResult.test_type == test_type) + if disposition is None: + # Active: confirmed plus no_decision (NULL). Dismissed/muted excluded. + clauses.append(or_(TestResult.disposition.is_(None), TestResult.disposition == DbDisposition.CONFIRMED.value)) + elif disposition == Disposition.no_decision: + clauses.append(TestResult.disposition.is_(None)) + else: + clauses.append(TestResult.disposition == DISPOSITION_TO_DB[disposition].value) + + items, total = TestResult.list_for_run(job.id, *clauses, page=page, limit=limit) + return TestResultListResponse( + items=[_to_item(row) for row in items], + page=page, + limit=limit, + total=total, + ) + + @router.get( "/profiling-runs/{job_id}", response_model=ProfilingRunResponse, diff --git a/testgen/api/schemas.py b/testgen/api/schemas.py index 90deb0b7..4bfd0f40 100644 --- a/testgen/api/schemas.py +++ b/testgen/api/schemas.py @@ -5,6 +5,7 @@ from pydantic import BaseModel +from testgen.api.enums import Disposition, ResultStatus from testgen.common.enums import JobSource, JobStatus, PublicJobKey from testgen.common.test_definition_export_import_service import ImportConfig, ImportPayload, ImportResponse @@ -78,6 +79,31 @@ class TestRunResponse(BaseModel): result: TestRunResult | None = None +class TestResultItem(BaseModel): + """One individual test result within a test run.""" + + test_definition_id: UUID + test_type: str + schema_name: str + table_name: str | None = None + column_names: str | None = None + result_status: ResultStatus | None = None + result_measure: str | None = None + threshold_value: str | None = None + result_message: str | None = None + test_time: datetime | None = None + disposition: Disposition + + +class TestResultListResponse(BaseModel): + """Paginated list of individual test results.""" + + items: list[TestResultItem] + page: int + limit: int + total: int + + # --- Profiling Runs --- diff --git a/testgen/common/models/test_result.py b/testgen/common/models/test_result.py index 751e896f..4ffc3782 100644 --- a/testgen/common/models/test_result.py +++ b/testgen/common/models/test_result.py @@ -64,6 +64,23 @@ class TestResultSearchRow: result_message: str | None +@dataclass +class TestRunResultRow: + """One individual result within a single test run for the API results endpoint.""" + + test_definition_id: UUID + test_type: str + schema_name: str + table_name: str | None + column_names: str | None + status: TestResultStatus | None + result_measure: str | None + threshold_value: str | None + message: str | None + test_time: datetime | None + disposition: str | None + + @dataclass class TrendBucket: """One time-bucket of failure aggregates for ``get_failure_trend``.""" @@ -175,6 +192,38 @@ def select_results( query = query.order_by(cls.status, cls.table_name, cls.column_names).offset(offset).limit(limit) return get_current_session().scalars(query).all() + @classmethod + def list_for_run( + cls, + test_run_id: UUID, + *clauses, + page: int = 1, + limit: int = 20, + ) -> tuple[list[TestRunResultRow], int]: + """Paginated individual results for a single run, scoped by caller-supplied WHERE clauses. + + Monitor suites are always filtered out. + """ + query = ( + select( + cls.test_definition_id.label("test_definition_id"), + cls.test_type.label("test_type"), + cls.schema_name.label("schema_name"), + cls.table_name.label("table_name"), + cls.column_names.label("column_names"), + cls.status.label("status"), + cls.result_measure.label("result_measure"), + cls.threshold_value.label("threshold_value"), + cls.message.label("message"), + cls.test_time.label("test_time"), + cls.disposition.label("disposition"), + ) + .join(TestSuite, cls.test_suite_id == TestSuite.id) + .where(cls.test_run_id == test_run_id, TestSuite.is_monitor.isnot(True), *clauses) + .order_by(cls.status, cls.table_name, cls.column_names, cls.id) + ) + return cls._paginate(query, page=page, limit=limit, data_class=TestRunResultRow) + @classmethod def select_failures( cls, diff --git a/tests/unit/api/test_runs.py b/tests/unit/api/test_runs.py index 87c61ca1..b96ec05d 100644 --- a/tests/unit/api/test_runs.py +++ b/tests/unit/api/test_runs.py @@ -5,10 +5,15 @@ from uuid import uuid4 import pytest +from fastapi import FastAPI +from fastapi.testclient import TestClient +from sqlalchemy.dialects import postgresql -from testgen.api.runs import get_profiling_run, get_test_run +from testgen.api.deps import db_session, get_authorized_user +from testgen.api.enums import Disposition, ResultStatus +from testgen.api.runs import get_profiling_run, get_test_run, list_test_run_results, router from testgen.common.models.hygiene_issue import HygieneIssueCounts, IssueCounts, PotentialPiiCounts -from testgen.common.models.test_result import ResultStatusCounts +from testgen.common.models.test_result import ResultStatusCounts, TestResult, TestResultStatus, TestRunResultRow pytestmark = pytest.mark.unit @@ -18,6 +23,46 @@ TABLE_GROUP_ID = uuid4() +def _mock_result_row(**overrides): + defaults = { + "test_definition_id": uuid4(), + "test_type": "Unique", + "schema_name": "demo", + "table_name": "orders", + "column_names": "amount", + "status": TestResultStatus.Failed, + "result_measure": "3", + "threshold_value": "0", + "message": "Duplicate values: 3", + "test_time": datetime.now(UTC), + "disposition": None, + } + defaults.update(overrides) + return TestRunResultRow(**defaults) + + +def _result_clauses_sql(mock_list) -> str: + """Compile the WHERE clauses passed to list_for_run (args after test_run_id) to SQL.""" + clauses = mock_list.call_args.args[1:] + return " ".join( + str(c.compile(dialect=postgresql.dialect(), compile_kwargs={"literal_binds": True})) for c in clauses + ) + + +def _no_filters(**overrides): + kwargs = { + "status": None, + "table_name": None, + "column_name": None, + "test_type": None, + "disposition": None, + "page": 1, + "limit": 20, + } + kwargs.update(overrides) + return kwargs + + def _mock_job(**overrides): defaults = { "id": uuid4(), @@ -144,3 +189,174 @@ def test_get_profiling_run_pending_no_run(mock_pr_cls): assert result.status == "pending" assert result.table_group_id is None assert result.result is None + + +# --- list_test_run_results: envelope + field/enum mapping --- + + +@patch.object(TestResult, "list_for_run", return_value=([], 0)) +def test_list_results_envelope_and_mapping(mock_list): + job = _mock_job() + row = _mock_result_row(test_type="Unique", status=TestResultStatus.Failed, message="dup", disposition=None) + mock_list.return_value = ([row], 7) + + result = list_test_run_results(job, **_no_filters(page=2, limit=5)) + + assert result.total == 7 + assert result.page == 2 + assert result.limit == 5 + assert len(result.items) == 1 + item = result.items[0] + assert item.test_type == "Unique" # raw code, not a display name + assert item.result_status == ResultStatus.failed # title-case DB value -> lowercase API enum + assert item.result_message == "dup" # ORM attr `message` -> field `result_message` + assert item.disposition == Disposition.no_decision # NULL -> no_decision (not confirmed) + # test_run_id (job.id) is the scope; page/limit are forwarded as kwargs. + assert mock_list.call_args.args[0] == job.id + assert mock_list.call_args.kwargs == {"page": 2, "limit": 5} + + +@patch.object(TestResult, "list_for_run", return_value=([], 0)) +def test_list_results_empty_envelope(mock_list): + result = list_test_run_results(_mock_job(), **_no_filters()) + assert result.items == [] + assert result.total == 0 + assert result.page == 1 + + +@pytest.mark.parametrize( + "db_status,expected", + [ + (TestResultStatus.Passed, ResultStatus.passed), + (TestResultStatus.Failed, ResultStatus.failed), + (TestResultStatus.Warning, ResultStatus.warning), + (TestResultStatus.Error, ResultStatus.error), + (TestResultStatus.Log, ResultStatus.log), + (None, None), + ], +) +@patch.object(TestResult, "list_for_run") +def test_list_results_status_render(mock_list, db_status, expected): + mock_list.return_value = ([_mock_result_row(status=db_status)], 1) + item = list_test_run_results(_mock_job(), **_no_filters()).items[0] + assert item.result_status == expected + + +@pytest.mark.parametrize( + "db_disposition,expected", + [ + (None, Disposition.no_decision), + ("Confirmed", Disposition.confirmed), + ("Dismissed", Disposition.dismissed), + ("Inactive", Disposition.muted), + ], +) +@patch.object(TestResult, "list_for_run") +def test_list_results_disposition_render(mock_list, db_disposition, expected): + mock_list.return_value = ([_mock_result_row(disposition=db_disposition)], 1) + item = list_test_run_results(_mock_job(), **_no_filters()).items[0] + assert item.disposition == expected + + +# --- list_test_run_results: filter clause building --- + + +@patch.object(TestResult, "list_for_run", return_value=([], 0)) +def test_list_results_builds_filter_clauses(mock_list): + list_test_run_results( + _mock_job(), + **_no_filters(status=ResultStatus.failed, table_name="orders", column_name="amount", test_type="Unique"), + ) + sql = _result_clauses_sql(mock_list) + assert "'Failed'" in sql # status mapped to DB value + assert "'orders'" in sql + assert "'amount'" in sql # column_name -> column_names column + assert "'Unique'" in sql + + +@patch.object(TestResult, "list_for_run", return_value=([], 0)) +def test_disposition_omitted_returns_active(mock_list): + list_test_run_results(_mock_job(), **_no_filters(disposition=None)) + sql = _result_clauses_sql(mock_list) + # Active = confirmed + no_decision (NULL); dismissed/muted excluded. + assert "IS NULL" in sql + assert "'Confirmed'" in sql + assert "'Dismissed'" not in sql + assert "'Inactive'" not in sql + + +@patch.object(TestResult, "list_for_run", return_value=([], 0)) +def test_disposition_confirmed_excludes_nulls(mock_list): + list_test_run_results(_mock_job(), **_no_filters(disposition=Disposition.confirmed)) + sql = _result_clauses_sql(mock_list) + assert "'Confirmed'" in sql + assert "IS NULL" not in sql # explicit confirmed no longer sweeps in undecided rows + + +@patch.object(TestResult, "list_for_run", return_value=([], 0)) +def test_disposition_no_decision_is_null(mock_list): + list_test_run_results(_mock_job(), **_no_filters(disposition=Disposition.no_decision)) + sql = _result_clauses_sql(mock_list) + assert "IS NULL" in sql + assert "'Confirmed'" not in sql + + +@pytest.mark.parametrize( + "disposition,db_value", + [(Disposition.dismissed, "'Dismissed'"), (Disposition.muted, "'Inactive'")], +) +@patch.object(TestResult, "list_for_run", return_value=([], 0)) +def test_disposition_dismissed_muted_map_to_db(mock_list, disposition, db_value): + list_test_run_results(_mock_job(), **_no_filters(disposition=disposition)) + assert db_value in _result_clauses_sql(mock_list) + + +# --- list_test_run_results: HTTP-level query validation --- + + +def _client() -> TestClient: + app = FastAPI() + app.include_router(router, prefix="/api/v1") + app.dependency_overrides[db_session] = lambda: iter([None]) + app.dependency_overrides[get_authorized_user] = lambda: MagicMock(id=uuid4()) + return TestClient(app) + + +@patch("testgen.api.deps.has_project_permission", return_value=True) +@patch("testgen.api.deps.get_current_session") +@patch.object(TestResult, "list_for_run", return_value=([], 0)) +def test_http_rejects_unknown_status(_mock_list, mock_sess, _mock_perm): + mock_sess.return_value.scalars.return_value.first.return_value = _mock_job() + resp = _client().get(f"/api/v1/test-runs/{uuid4()}/results?status=BOGUS") + assert resp.status_code == 422 + assert resp.json()["detail"][0]["loc"] == ["query", "status"] + + +@patch("testgen.api.deps.has_project_permission", return_value=True) +@patch("testgen.api.deps.get_current_session") +@patch.object(TestResult, "list_for_run", return_value=([], 0)) +def test_http_rejects_unknown_disposition(_mock_list, mock_sess, _mock_perm): + mock_sess.return_value.scalars.return_value.first.return_value = _mock_job() + resp = _client().get(f"/api/v1/test-runs/{uuid4()}/results?disposition=foo") + assert resp.status_code == 422 + assert resp.json()["detail"][0]["loc"] == ["query", "disposition"] + + +@patch("testgen.api.deps.has_project_permission", return_value=True) +@patch("testgen.api.deps.get_current_session") +@patch.object(TestResult, "list_for_run", return_value=([], 0)) +def test_http_accepts_no_decision(mock_list, mock_sess, _mock_perm): + mock_sess.return_value.scalars.return_value.first.return_value = _mock_job() + resp = _client().get(f"/api/v1/test-runs/{uuid4()}/results?disposition=no_decision") + assert resp.status_code == 200 + mock_list.assert_called_once() + + +@pytest.mark.parametrize("query", ["page=0", "limit=0", "limit=101"]) +@patch("testgen.api.deps.has_project_permission", return_value=True) +@patch("testgen.api.deps.get_current_session") +@patch.object(TestResult, "list_for_run", return_value=([], 0)) +def test_http_rejects_out_of_range_pagination(_mock_list, mock_sess, _mock_perm, query): + mock_sess.return_value.scalars.return_value.first.return_value = _mock_job() + resp = _client().get(f"/api/v1/test-runs/{uuid4()}/results?{query}") + assert resp.status_code == 422 diff --git a/tests/unit/mcp/test_model_test_result.py b/tests/unit/mcp/test_model_test_result.py index ec037820..4b34bac8 100644 --- a/tests/unit/mcp/test_model_test_result.py +++ b/tests/unit/mcp/test_model_test_result.py @@ -4,7 +4,7 @@ import pytest from sqlalchemy.dialects import postgresql -from testgen.common.models.test_result import TestResult, TestResultStatus +from testgen.common.models.test_result import TestResult, TestResultStatus, TestRunResultRow @pytest.fixture @@ -128,6 +128,26 @@ def test_select_results_excludes_monitor_suites_with_project_codes(session_mock) assert "test_suites.project_code IN" in sql +def test_list_for_run_excludes_monitor_suites(): + with patch.object(TestResult, "_paginate", return_value=([], 0)) as mock_paginate: + TestResult.list_for_run(uuid4()) + + sql = _compiled_sql(mock_paginate.call_args.args[0]) + assert "test_suites.is_monitor IS NOT true" in sql + assert "JOIN test_suites" in sql + + +def test_list_for_run_applies_caller_clauses_and_pagination(): + with patch.object(TestResult, "_paginate", return_value=([], 0)) as mock_paginate: + TestResult.list_for_run(uuid4(), TestResult.table_name == "orders", page=3, limit=10) + + sql = _compiled_sql(mock_paginate.call_args.args[0]) + assert "test_results.table_name = " in sql + assert mock_paginate.call_args.kwargs["page"] == 3 + assert mock_paginate.call_args.kwargs["limit"] == 10 + assert mock_paginate.call_args.kwargs["data_class"] is TestRunResultRow + + def test_select_failures_excludes_monitor_suites(session_mock): session_mock.execute.return_value.all.return_value = [] From b73229958166914cb9783e75f06eaaebd6981bb3 Mon Sep 17 00:00:00 2001 From: Ricardo Boni Date: Fri, 19 Jun 2026 16:25:29 -0400 Subject: [PATCH 49/78] refactor(api): apply TG-1107 review feedback Degrade unmapped disposition values to no_decision instead of raising mid-serialization, matching result_status's tolerant mapping so one odd row never 500s the results page. Co-Authored-By: Claude Opus 4.8 (1M context) --- testgen/api/runs.py | 16 +++++++++++++++- tests/unit/api/test_runs.py | 2 ++ 2 files changed, 17 insertions(+), 1 deletion(-) diff --git a/testgen/api/runs.py b/testgen/api/runs.py index b32e5733..277e982b 100644 --- a/testgen/api/runs.py +++ b/testgen/api/runs.py @@ -74,6 +74,20 @@ def get_test_run(job: JobExecution = resolve_job("view", JobExecution.job_key == ) +def _disposition_from_db(value: str | None) -> Disposition: + """Map a stored ``disposition`` to the API enum, degrading unknown values to ``no_decision``. + + A NULL or unmapped value resolves to ``no_decision`` rather than raising, so a single odd + row never fails serialization of the whole results page. + """ + if not value: + return Disposition.no_decision + try: + return DISPOSITION_FROM_DB[DbDisposition(value)] + except (ValueError, KeyError): + return Disposition.no_decision + + def _to_item(row: TestRunResultRow) -> TestResultItem: """Map a DB-valued result row to the API item, normalizing enum casing.""" return TestResultItem( @@ -87,7 +101,7 @@ def _to_item(row: TestRunResultRow) -> TestResultItem: threshold_value=row.threshold_value, result_message=row.message, test_time=row.test_time, - disposition=DISPOSITION_FROM_DB[DbDisposition(row.disposition)] if row.disposition else Disposition.no_decision, + disposition=_disposition_from_db(row.disposition), ) diff --git a/tests/unit/api/test_runs.py b/tests/unit/api/test_runs.py index b96ec05d..1dc1a825 100644 --- a/tests/unit/api/test_runs.py +++ b/tests/unit/api/test_runs.py @@ -246,9 +246,11 @@ def test_list_results_status_render(mock_list, db_status, expected): "db_disposition,expected", [ (None, Disposition.no_decision), + ("", Disposition.no_decision), ("Confirmed", Disposition.confirmed), ("Dismissed", Disposition.dismissed), ("Inactive", Disposition.muted), + ("Bogus", Disposition.no_decision), ], ) @patch.object(TestResult, "list_for_run") From 34d8947a4dd302d97e522ba2ee597ba563e0befe Mon Sep 17 00:00:00 2001 From: Ricardo Boni Date: Mon, 22 Jun 2026 08:13:26 -0400 Subject: [PATCH 50/78] refactor(mcp): trim import_tests docstring to purpose line per review Co-Authored-By: Claude Fable 5 --- testgen/mcp/tools/test_definitions.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/testgen/mcp/tools/test_definitions.py b/testgen/mcp/tools/test_definitions.py index b0984c85..9c277c08 100644 --- a/testgen/mcp/tools/test_definitions.py +++ b/testgen/mcp/tools/test_definitions.py @@ -858,7 +858,7 @@ def import_tests( on_new: str = "create", on_absence: str = "do_nothing", ) -> str: - """Apply an export document from ``export_tests`` to a test suite, defaulting to ``preview`` mode that reports the projected changes without persisting. + """Apply an export document from ``export_tests`` to a test suite. Args: test_suite_id: UUID of the target test suite. From 2e42e435dbf5d8d7f2e0b43d74b9973b4f7d1f4c Mon Sep 17 00:00:00 2001 From: Astor Date: Tue, 9 Jun 2026 16:19:47 -0300 Subject: [PATCH 51/78] fix: use cross-platform date formatting for Excel export and schedule display Replaces Linux-only strftime directives (%-d, %-I) with f-string integers so date formatting works correctly on Windows, Linux, and macOS. Co-Authored-By: Claude Sonnet 4.6 --- testgen/ui/views/data_catalog.py | 2 +- testgen/ui/views/dialogs/manage_schedules.py | 2 +- testgen/ui/views/test_definitions.py | 2 +- 3 files changed, 3 insertions(+), 3 deletions(-) diff --git a/testgen/ui/views/data_catalog.py b/testgen/ui/views/data_catalog.py index 3b0b848f..2ad0df94 100644 --- a/testgen/ui/views/data_catalog.py +++ b/testgen/ui/views/data_catalog.py @@ -492,7 +492,7 @@ def get_excel_report_data(update_progress: PROGRESS_UPDATE_TYPE, table_group: Ta for key in ["min_date", "max_date", "add_date", "last_mod_date", "drop_date"]: data[key] = data[key].apply( - lambda val: val.strftime("%b %-d %Y, %-I:%M %p") if not pd.isna(val) and not isinstance(val, str) else val + lambda val: val.strftime(f"%b {val.day} %Y, {val.hour % 12 or 12}:%M %p") if not pd.isna(val) and not isinstance(val, str) else val ) for key in ["data_source", "source_system", "source_process", "business_domain", "stakeholder_group", "transform_level", "aggregation_level", "data_product"]: diff --git a/testgen/ui/views/dialogs/manage_schedules.py b/testgen/ui/views/dialogs/manage_schedules.py index 4d917f7c..d3faf1df 100644 --- a/testgen/ui/views/dialogs/manage_schedules.py +++ b/testgen/ui/views/dialogs/manage_schedules.py @@ -55,7 +55,7 @@ def build_data(self) -> dict: "readableExpr": cron_descriptor.get_description(job.cron_expr), "cronTz": job.cron_tz_str, "sample": [ - sample.strftime("%a %b %-d, %-I:%M %p") + sample.strftime(f"%a %b {sample.day}, {sample.hour % 12 or 12}:%M %p") for sample in job.get_sample_triggering_timestamps(CRON_SAMPLE_COUNT + 1) ], "active": job.active, diff --git a/testgen/ui/views/test_definitions.py b/testgen/ui/views/test_definitions.py index c2d5843c..8386f0c2 100644 --- a/testgen/ui/views/test_definitions.py +++ b/testgen/ui/views/test_definitions.py @@ -768,7 +768,7 @@ def get_excel_report_data( for key in ["profiling_as_of_date", "last_manual_update"]: data[key] = data[key].apply( - lambda val: datetime.strptime(val, "%Y-%m-%d %H:%M:%S").strftime("%b %-d %Y, %-I:%M %p") + lambda val: (lambda dt: dt.strftime(f"%b {dt.day} %Y, {dt.hour % 12 or 12}:%M %p"))(datetime.strptime(val, "%Y-%m-%d %H:%M:%S")) if (val and not pd.isna(val) and val != "NaT") else None ) From 7da561a46efdd55e7f58d667d4772d717f79e6ff Mon Sep 17 00:00:00 2001 From: Astor Date: Wed, 17 Jun 2026 17:07:28 -0300 Subject: [PATCH 52/78] fix: extend cross-platform date formatting to all remaining sites Adds format_friendly_datetime() to date_service and routes all strftime calls that used Linux-only %-d/%-I through it, covering cron preview, UI timestamps, email templates, and PDF reports. Co-Authored-By: Claude Sonnet 4.6 --- testgen/common/cron_service.py | 5 ++++- testgen/common/date_service.py | 9 ++++++++- testgen/common/notifications/notifications.py | 3 ++- testgen/ui/views/data_catalog.py | 3 ++- testgen/ui/views/dialogs/manage_schedules.py | 3 ++- testgen/ui/views/test_definitions.py | 2 +- 6 files changed, 19 insertions(+), 6 deletions(-) diff --git a/testgen/common/cron_service.py b/testgen/common/cron_service.py index afe75485..c6abc5c0 100644 --- a/testgen/common/cron_service.py +++ b/testgen/common/cron_service.py @@ -5,6 +5,9 @@ import cron_converter import cron_descriptor +from testgen.common import date_service + + class CronSample(TypedDict, total=False): id: str | None @@ -26,7 +29,7 @@ def get_cron_sample( cron_schedule = cron_obj.schedule(reference_time or datetime.now(zoneinfo.ZoneInfo(cron_tz))) readable_cron_schedule = cron_descriptor.get_description(cron_expr) if formatted: - samples = [cron_schedule.next().strftime("%a %b %-d, %-I:%M %p") for _ in range(sample_count)] + samples = [date_service.format_friendly_datetime(cron_schedule.next(), "%a %b %-d, %-I:%M %p") for _ in range(sample_count)] else: samples = [int(cron_schedule.next().timestamp()) for _ in range(sample_count)] except zoneinfo.ZoneInfoNotFoundError: diff --git a/testgen/common/date_service.py b/testgen/common/date_service.py index eefaf131..54c04faf 100644 --- a/testgen/common/date_service.py +++ b/testgen/common/date_service.py @@ -97,12 +97,19 @@ def accommodate_dataframe_to_timezone(df, streamlit_session, time_columns=None): df[time_column] = df[time_column].dt.strftime("%Y-%m-%d %H:%M:%S") +def format_friendly_datetime(dt: datetime, fmt: str) -> str: + """Cross-platform strftime: replaces Linux-only %-d and %-I before calling strftime.""" + return dt.strftime(fmt.replace("%-d", str(dt.day)).replace("%-I", str(dt.hour % 12 or 12))) + + def get_timezoned_timestamp(streamlit_session, value, dateformat="%b %-d, %-I:%M %p"): ret = None if value and "browser_timezone" in streamlit_session: data = {"value": [value]} df = pd.DataFrame(data) timezone = streamlit_session["browser_timezone"] - df["value"] = df["value"].dt.tz_localize("UTC").dt.tz_convert(timezone).dt.strftime(dateformat) + df["value"] = df["value"].dt.tz_localize("UTC").dt.tz_convert(timezone).apply( + lambda dt: format_friendly_datetime(dt, dateformat) if pd.notna(dt) else "" + ) ret = df.iloc[0, 0] return ret diff --git a/testgen/common/notifications/notifications.py b/testgen/common/notifications/notifications.py index 23a9f896..52226950 100644 --- a/testgen/common/notifications/notifications.py +++ b/testgen/common/notifications/notifications.py @@ -1,6 +1,7 @@ import math from datetime import datetime +from testgen.common import date_service from testgen.common.notifications.base import BaseEmailTemplate from testgen.utils import friendly_score @@ -14,7 +15,7 @@ def format_number_helper(self, number: int) -> str: return "" if number is None else f"{number:,}" def format_dt_helper(self, dt: datetime) -> str: - return "" if dt is None else dt.strftime("%b %d, %-I:%M %p UTC") + return "" if dt is None else date_service.format_friendly_datetime(dt, "%b %d, %-I:%M %p UTC") def format_duration_helper(self, start_time: datetime, end_time: datetime) -> str: total_seconds = abs(end_time - start_time).total_seconds() diff --git a/testgen/ui/views/data_catalog.py b/testgen/ui/views/data_catalog.py index 2ad0df94..ae8dcdd2 100644 --- a/testgen/ui/views/data_catalog.py +++ b/testgen/ui/views/data_catalog.py @@ -11,6 +11,7 @@ from sqlalchemy.sql.expression import func as sa_func from streamlit.delta_generator import DeltaGenerator +from testgen.common import date_service from testgen.common.data_catalog_service import ( apply_column_metadata, apply_table_metadata, @@ -492,7 +493,7 @@ def get_excel_report_data(update_progress: PROGRESS_UPDATE_TYPE, table_group: Ta for key in ["min_date", "max_date", "add_date", "last_mod_date", "drop_date"]: data[key] = data[key].apply( - lambda val: val.strftime(f"%b {val.day} %Y, {val.hour % 12 or 12}:%M %p") if not pd.isna(val) and not isinstance(val, str) else val + lambda val: date_service.format_friendly_datetime(val, "%b %-d %Y, %-I:%M %p") if not pd.isna(val) and not isinstance(val, str) else val ) for key in ["data_source", "source_system", "source_process", "business_domain", "stakeholder_group", "transform_level", "aggregation_level", "data_product"]: diff --git a/testgen/ui/views/dialogs/manage_schedules.py b/testgen/ui/views/dialogs/manage_schedules.py index d3faf1df..e22d3dd8 100644 --- a/testgen/ui/views/dialogs/manage_schedules.py +++ b/testgen/ui/views/dialogs/manage_schedules.py @@ -6,6 +6,7 @@ from sqlalchemy import select from sqlalchemy.exc import IntegrityError +from testgen.common import date_service from testgen.common.models import Session, with_database_session from testgen.common.models.scheduler import JobSchedule from testgen.ui.session import session @@ -55,7 +56,7 @@ def build_data(self) -> dict: "readableExpr": cron_descriptor.get_description(job.cron_expr), "cronTz": job.cron_tz_str, "sample": [ - sample.strftime(f"%a %b {sample.day}, {sample.hour % 12 or 12}:%M %p") + date_service.format_friendly_datetime(sample, "%a %b %-d, %-I:%M %p") for sample in job.get_sample_triggering_timestamps(CRON_SAMPLE_COUNT + 1) ], "active": job.active, diff --git a/testgen/ui/views/test_definitions.py b/testgen/ui/views/test_definitions.py index 8386f0c2..b4cac73d 100644 --- a/testgen/ui/views/test_definitions.py +++ b/testgen/ui/views/test_definitions.py @@ -768,7 +768,7 @@ def get_excel_report_data( for key in ["profiling_as_of_date", "last_manual_update"]: data[key] = data[key].apply( - lambda val: (lambda dt: dt.strftime(f"%b {dt.day} %Y, {dt.hour % 12 or 12}:%M %p"))(datetime.strptime(val, "%Y-%m-%d %H:%M:%S")) + lambda val: date_service.format_friendly_datetime(datetime.strptime(val, "%Y-%m-%d %H:%M:%S"), "%b %-d %Y, %-I:%M %p") if (val and not pd.isna(val) and val != "NaT") else None ) From 0bd1764baee0ada92c1f1e749f26e6b58cccfeae Mon Sep 17 00:00:00 2001 From: Astor Date: Mon, 22 Jun 2026 12:04:25 -0300 Subject: [PATCH 53/78] fix: remove extra blank line in cron_service imports to satisfy ruff I001 Co-Authored-By: Claude Sonnet 4.6 --- testgen/common/cron_service.py | 1 - 1 file changed, 1 deletion(-) diff --git a/testgen/common/cron_service.py b/testgen/common/cron_service.py index c6abc5c0..fd9a2984 100644 --- a/testgen/common/cron_service.py +++ b/testgen/common/cron_service.py @@ -8,7 +8,6 @@ from testgen.common import date_service - class CronSample(TypedDict, total=False): id: str | None error: str | None From 1fe6947a01227ff288351f98877f643ba97b438d Mon Sep 17 00:00:00 2001 From: Ricardo Boni Date: Tue, 23 Jun 2026 15:36:09 -0400 Subject: [PATCH 54/78] refactor(mcp): apply TG-1094 review feedback Co-Authored-By: Claude Opus 4.8 (1M context) --- testgen/mcp/tools/monitors.py | 26 +++++++++++++++----------- 1 file changed, 15 insertions(+), 11 deletions(-) diff --git a/testgen/mcp/tools/monitors.py b/testgen/mcp/tools/monitors.py index 9aeef2f6..f1493585 100644 --- a/testgen/mcp/tools/monitors.py +++ b/testgen/mcp/tools/monitors.py @@ -337,10 +337,11 @@ def _format_row_count_change(row: MonitorTableSummary) -> str | None: @with_database_session @mcp_permission("edit") def enable_monitors(table_group_id: str, cron_expression: str, cron_tz: str = "UTC") -> str: - """Turn on monitoring for a table group: create its monitors and put them on a schedule. + """Turn on monitoring for a table group: create its monitors and schedule them to run. Sets up the initial Volume and Schema monitors with default settings; adjust them afterward - with ``update_monitor_settings``. Fails if monitoring is already on for the group. + with ``update_monitor_settings``. Freshness monitors are added automatically once the table + group has profiling data. Fails if monitoring is already on for the group. Args: table_group_id: UUID of the table group, e.g. from ``list_table_groups``. @@ -404,12 +405,12 @@ def update_monitor_settings( Args: table_group_id: UUID of the table group, e.g. from ``list_table_groups``. - sensitivity: Anomaly-detection sensitivity. One of ``low`` / ``medium`` / ``high``. - lookback_runs: Monitor runs aggregated for dashboard summaries (1-200). - min_training_lookback: Minimum runs before predictions activate (20-1000). - exclude_weekends: Whether to exclude weekends from prediction baselines. - holiday_codes: Holiday calendars to exclude from baselines — ISO country codes (e.g. ``US``, ``GB``) or financial-market codes (e.g. ``NYSE``, ``ECB``); see https://holidays.readthedocs.io/en/latest/#available-countries. Pass an empty list to clear. - regenerate_freshness: Whether to auto-regenerate Freshness monitors when the schema shifts. + sensitivity: How readily monitors flag a deviation. ``high`` flags smaller deviations (more alerts), ``low`` only large ones (fewer alerts), ``medium`` is balanced. + lookback_runs: Monitor runs aggregated for dashboard summaries. Display only — does not affect detection (1-200). + min_training_lookback: Minimum monitor runs required to train the prediction model (20-1000). + exclude_weekends: Whether to exclude weekends from the model's training data. + holiday_codes: Holiday calendars to exclude from the model's training data — ISO country codes (e.g. ``US``, ``GB``) or financial-market codes (e.g. ``NYSE``, ``ECB``); see https://holidays.readthedocs.io/en/latest/#available-countries. Pass an empty list to clear. + regenerate_freshness: Whether to automatically reconfigure Freshness monitors with new fingerprints after each profiling run. cron_expression: New cron expression for the schedule, e.g. ``0 6 * * *``. cron_tz: New IANA timezone for the schedule, e.g. ``America/New_York``. active: ``True`` to resume the schedule, ``False`` to pause it. @@ -543,14 +544,16 @@ def _last_monitor_run(schedule_id: UUID) -> datetime | None: def _render_monitor_settings(doc: MdDoc, monitor_suite: TestSuite, schedule: JobSchedule) -> None: + doc.field("Lookback runs", monitor_suite.monitor_lookback) + doc.field("Regenerate freshness", monitor_suite.monitor_regenerate_freshness) + + doc.heading(2, "Prediction Model") sensitivity = monitor_suite.predict_sensitivity.value if monitor_suite.predict_sensitivity else None doc.field("Sensitivity", sensitivity) - doc.field("Lookback runs", monitor_suite.monitor_lookback) doc.field("Min training lookback", monitor_suite.predict_min_lookback) doc.field("Exclude weekends", monitor_suite.predict_exclude_weekends) holidays = monitor_suite.holiday_codes_list doc.field("Holiday codes", ", ".join(holidays) if holidays else None) - doc.field("Regenerate freshness", monitor_suite.monitor_regenerate_freshness) doc.heading(2, "Schedule") doc.field("Cron expression", schedule.cron_expr, code=True) @@ -568,4 +571,5 @@ def _render_monitor_settings(doc: MdDoc, monitor_suite: TestSuite, schedule: Job # rendered " UTC" suffix is accurate. doc.field("Next run", next_runs[0].astimezone(UTC)) if (last_run := _last_monitor_run(schedule.id)) is not None: - doc.field("Last run", last_run) + # job_executions timestamps are tz-aware UTC; convert so the rendered " UTC" suffix is accurate. + doc.field("Last run", last_run.astimezone(UTC)) From 3a5b1f2c3f3758c4625e53f6263d1e97ab932259 Mon Sep 17 00:00:00 2001 From: Aarthy Adityan Date: Tue, 23 Jun 2026 23:42:14 -0400 Subject: [PATCH 55/78] fix(monitors): event-space SARIMAX for freshness-gated volume/metric prediction MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Freshness-gated Volume_Trend/Metric_Trend monitors still flagged legitimate refreshes: get_sarimax_forecast resampled the (irregular) freshness-filtered series onto a calendar grid and linearly interpolated the gaps, collapsing the refresh-jump variance and biasing the forecast low. Fit the gated series in event-space (one step per refresh, no interpolation) instead — preserving the SE floor and volume>=0 floor. Walk-forward on the real series: 6.1% -> 0% false positives, true over-loads still caught. Display: gated monitors render a flat baseline up to the predicted next-update window then a coupled step to the forecast, instead of a rising cone; the "show more history" toggle moved into the dialog header (new headerActions slot). TG-1131 Co-Authored-By: Claude Opus 4.8 (1M context) --- .../commands/test_thresholds_prediction.py | 8 + testgen/common/time_series_service.py | 57 ++++-- .../js/pages/table_monitoring_trends.js | 29 +-- testgen/ui/static/js/components/dialog.js | 10 +- testgen/ui/views/monitors_dashboard.py | 185 ++++++++++++++---- .../commands/test_thresholds_prediction.py | 30 +++ tests/unit/common/test_time_series_service.py | 34 ++++ tests/unit/ui/test_monitors_dashboard.py | 155 +++++++++++++++ 8 files changed, 426 insertions(+), 82 deletions(-) create mode 100644 tests/unit/ui/test_monitors_dashboard.py diff --git a/testgen/commands/test_thresholds_prediction.py b/testgen/commands/test_thresholds_prediction.py index d357693b..9089551e 100644 --- a/testgen/commands/test_thresholds_prediction.py +++ b/testgen/commands/test_thresholds_prediction.py @@ -299,9 +299,12 @@ def compute_sarimax_threshold( exclude_weekends: bool = False, holiday_codes: list[str] | None = None, schedule_tz: str | None = None, + event_space: bool = False, ) -> tuple[float | None, float | None, str | None]: """Compute SARIMAX-based thresholds for the next forecast point. + `event_space=True` fits in event-space (one step per observation, no interpolation). + Returns (lower, upper, forecast_json) or (None, None, None) if insufficient data. """ try: @@ -311,6 +314,7 @@ def compute_sarimax_threshold( exclude_weekends=exclude_weekends, holiday_codes=holiday_codes, tz=schedule_tz, + event_space=event_space, ) num_points = len(history) @@ -359,6 +363,9 @@ def compute_volume_or_metric_threshold( Freshness_Trend detected a fingerprint change. """ filtered_history = history.loc[history["test_run_id"].astype(str).isin(freshness_updates)] + # Event-space fit: the filtered series has one point per refresh and is irregularly spaced. + # Calendar resampling + interpolation would smooth the refresh jumps into small uniform + # increments and collapse the jump variance, biasing the forecast low and the band too tight. lower, upper, prediction = compute_sarimax_threshold( filtered_history, sensitivity=sensitivity, @@ -366,6 +373,7 @@ def compute_volume_or_metric_threshold( exclude_weekends=exclude_weekends, holiday_codes=holiday_codes, schedule_tz=schedule_tz, + event_space=True, ) if prediction is not None: # Pull the baseline value from the most-recent filtered row. diff --git a/testgen/common/time_series_service.py b/testgen/common/time_series_service.py index 7a68fb35..ddf329c0 100644 --- a/testgen/common/time_series_service.py +++ b/testgen/common/time_series_service.py @@ -23,6 +23,7 @@ def get_sarimax_forecast( exclude_weekends: bool = False, holiday_codes: list[str] | None = None, tz: str | None = None, + event_space: bool = False, ) -> pd.DataFrame: """ # Parameters @@ -33,6 +34,14 @@ def get_sarimax_forecast( :param exclude_weekends: Whether weekends should be considered exogenous when training the model and forecasting. :param holiday_codes: List of country or financial market codes defining holidays to be considered exogenous when training the model and forecasting. :param tz: IANA timezone (e.g. "America/New_York") for day-of-week/holiday checks. Naive timestamps are treated as UTC and converted to this timezone before determining weekday/holiday status. + :param event_space: When False (default), resample the series onto a regular calendar grid and + linearly interpolate gaps before fitting — appropriate for series sampled at a + steady cadence, and where weekend/holiday exogenous flags apply. When True, fit + in event-space: one model step per observation, no interpolation, no calendar + exog. Required for irregularly-spaced series (e.g. the freshness-update points of + a refresh-driven table) where interpolation would smooth multi-period jumps into + small uniform increments and collapse the very jump variance the forecast must + capture. # Return value Returns a Pandas dataframe with forecast DatetimeIndex, "mean" column, and "se" (standard error) column. @@ -40,23 +49,33 @@ def get_sarimax_forecast( if len(history) < MIN_TRAIN_VALUES: raise NotEnoughData("Not enough data points in history.") - # statsmodels requires DatetimeIndex with a regular frequency - # Resample the data to get a regular time series - datetimes = history.index.to_series() - frequency = infer_frequency(datetimes) - resampled_history = history.resample(frequency).mean().interpolate(method="linear") - - if len(resampled_history) < MIN_TRAIN_VALUES: - raise NotEnoughData("Not enough data points after resampling.") + if event_space: + # Event-space: one step per observation. Map the values onto a synthetic regular grid + # stepped by the median observed interval, so forecast timestamps remain realistic while + # the model sees each observation as a single step (no interpolated points between them). + median_step = history.index.to_series().diff().median() + if pd.isna(median_step) or median_step <= pd.Timedelta(0): + median_step = pd.Timedelta(days=1) + step = median_step + synthetic_index = pd.date_range(end=history.index[-1], periods=len(history), freq=step) + train_history = pd.DataFrame(history.iloc[:, 0].values, index=synthetic_index, columns=[history.columns[0]]) + else: + # statsmodels requires DatetimeIndex with a regular frequency + # Resample the data to get a regular time series + datetimes = history.index.to_series() + frequency = infer_frequency(datetimes) + train_history = history.resample(frequency).mean().interpolate(method="linear") + if len(train_history) < MIN_TRAIN_VALUES: + raise NotEnoughData("Not enough data points after resampling.") + step = pd.to_timedelta(frequency) # Generate DatetimeIndex with future dates - forecast_start = resampled_history.index[-1] + pd.to_timedelta(frequency) - forecast_index = pd.date_range(start=forecast_start, periods=num_forecast, freq=frequency) + forecast_index = pd.date_range(start=train_history.index[-1] + step, periods=num_forecast, freq=step) - # Detect holidays in entire date range + # Detect holidays in entire date range (calendar-aware path only) holiday_dates = None - if holiday_codes: - all_dates_index = resampled_history.index.append(forecast_index) + if not event_space and holiday_codes: + all_dates_index = train_history.index.append(forecast_index) holiday_dates = get_holiday_dates(holiday_codes, all_dates_index) def get_exog_flags(index: pd.DatetimeIndex) -> pd.DataFrame: @@ -71,11 +90,13 @@ def get_exog_flags(index: pd.DatetimeIndex) -> pd.DataFrame: exog.loc[pd.Index(check_index.date).isin(holiday_dates), "is_excluded"] = 1 return exog - exog_train = get_exog_flags(resampled_history.index) + # Calendar exogenous flags only apply when fitting against real calendar time. + exog_train = None if event_space else get_exog_flags(train_history.index) + exog_forecast = None if event_space else get_exog_flags(forecast_index) # When seasonal_order is not specified, this is effectively the ARIMAX model model = SARIMAX( - resampled_history.iloc[:, 0], + train_history.iloc[:, 0], exog=exog_train, # This is a good starting point according to Gemini - tune if needed order=(1, 1, 1), @@ -85,12 +106,6 @@ def get_exog_flags(index: pd.DatetimeIndex) -> pd.DataFrame: ) fitted_model = model.fit(disp=False) - forecast_index = pd.date_range( - start=resampled_history.index[-1] + pd.to_timedelta(frequency), - periods=num_forecast, - freq=frequency - ) - exog_forecast = get_exog_flags(forecast_index) forecast = fitted_model.get_forecast(steps=num_forecast, exog=exog_forecast) results = pd.DataFrame(index=forecast_index) diff --git a/testgen/ui/components/frontend/js/pages/table_monitoring_trends.js b/testgen/ui/components/frontend/js/pages/table_monitoring_trends.js index 8e0c86cd..94180e2d 100644 --- a/testgen/ui/components/frontend/js/pages/table_monitoring_trends.js +++ b/testgen/ui/components/frontend/js/pages/table_monitoring_trends.js @@ -123,18 +123,6 @@ const TableMonitoringTrend = (props) => { }, div( { class: '', style: 'width: 100%;' }, - () => { - const extendedHistory = getValue(props.extended_history) ?? false; - return div( - { class: 'extended-history-toggle' }, - Button({ - label: extendedHistory ? 'Show default view' : 'Show more history', - icon: extendedHistory ? 'history_toggle_off' : 'history', - width: 'auto', - onclick: () => emit('ToggleExtendedHistory', { payload: {} }), - }), - ); - }, () => { if (!getValue(props.dialog)?.open) return div(); return ChartsSection(props, { schemaChartSelection, getDataStructureLogs }); @@ -229,12 +217,22 @@ const TableMonitoringTrend = (props) => { ); const dialogTitle = van.derive(() => getValue(props.dialog)?.title ?? ''); + const historyToggle = () => { + const extendedHistory = getValue(props.extended_history) ?? false; + return Button({ + label: extendedHistory ? 'Show default view' : 'Show more history', + icon: extendedHistory ? 'history_toggle_off' : 'history', + width: 'auto', + onclick: () => emit('ToggleExtendedHistory', { payload: {} }), + }); + }; return Dialog( { title: dialogTitle, open: dialogOpen, onClose: () => { dialogOpen.val = false; emit('CloseTrendsDialog', {}); }, width: '75rem', + headerActions: historyToggle, }, content, ); @@ -825,13 +823,6 @@ stylesheet.replace(` position: relative; } - .extended-history-toggle { - position: absolute; - top: -70px; - right: 48px; - z-index: 1; - } - .table-monitoring-trend-wrapper:not(.has-sidebar) > .tg-dualpane-divider { display: none; } diff --git a/testgen/ui/static/js/components/dialog.js b/testgen/ui/static/js/components/dialog.js index 465c3faf..0f285fd3 100644 --- a/testgen/ui/static/js/components/dialog.js +++ b/testgen/ui/static/js/components/dialog.js @@ -5,6 +5,7 @@ * @property {import('../van.min.js').State} open - Reactive open state * @property {Function} onClose - Called when the dialog is closed (backdrop click or X button) * @property {string} [width] - CSS width value, default '30rem' + * @property {(Element | Function)} [headerActions] - Optional node rendered right-aligned in header row */ import van from '../van.min.js'; import { getValue, loadStylesheet } from '../utils.js'; @@ -28,7 +29,7 @@ const { button, div, i, span } = van.tags; * @param {DialogProps} props * @param {...(Element | string)} children - Content rendered in the dialog body */ -const Dialog = ({ title, open, onClose, width = '30rem' }, ...children) => { +const Dialog = ({ title, open, onClose, width = '30rem', headerActions }, ...children) => { loadStylesheet('dialog', stylesheet); const overlay = div( @@ -50,6 +51,7 @@ const Dialog = ({ title, open, onClose, width = '30rem' }, ...children) => { div( { class: 'tg-dialog-header' }, span({ 'data-testid': 'dialog-title', class: 'tg-dialog-title' }, title), + headerActions ? div({ class: 'tg-dialog-header-actions' }, headerActions) : null, ), div({ 'data-testid': 'dialog-content', class: 'tg-dialog-content' }, ...children), button( @@ -117,6 +119,12 @@ stylesheet.replace(` flex-shrink: 0; } +.tg-dialog-header-actions { + margin-left: auto; + font-size: initial; + font-weight: initial; +} + .tg-dialog-content { padding: 0.75rem 1.5rem 1.5rem; overflow-y: auto; diff --git a/testgen/ui/views/monitors_dashboard.py b/testgen/ui/views/monitors_dashboard.py index 8e44d4ba..9bc4e687 100644 --- a/testgen/ui/views/monitors_dashboard.py +++ b/testgen/ui/views/monitors_dashboard.py @@ -516,6 +516,96 @@ def _resolve_holiday_dates(test_suite: TestSuite) -> set[date] | None: return resolve_holiday_dates(test_suite.holiday_codes_list, idx) +def _freshness_next_update_window( + freshness_definition: TestDefinition | None, + events: dict, + test_suite: TestSuite, + monitor_schedule, +) -> dict | None: + """Predicted next freshness-update window as {"start", "end"} epoch-ms, or None. + + The schedule-derived business-time interval from the last detected update out to the + lower/upper staleness tolerance. Drives the Freshness_Trend display window and couples the + freshness-gated Volume/Metric forecast to the expected next refresh. + """ + if ( + freshness_definition is None + or freshness_definition.history_calculation != "PREDICT" + or (freshness_definition.prediction and not freshness_definition.prediction.get("schedule_stage")) + or freshness_definition.upper_tolerance is None + ): + return None + + last_update_events = [ + e for e in events["freshness_events"] + if e["changed"] and not e["is_training"] and not e["is_pending"] + ] + if not last_update_events: + return None + + last_detection_time = max(e["time"] for e in last_update_events) + holiday_dates = _resolve_holiday_dates(test_suite) + tz = monitor_schedule.cron_tz or "UTC" if monitor_schedule else None + sched = get_schedule_params(freshness_definition.prediction) + + window_end = add_business_minutes( + pd.Timestamp(last_detection_time), + float(freshness_definition.upper_tolerance), + test_suite.predict_exclude_weekends, + holiday_dates, tz, + excluded_days=sched.excluded_days, + ) + window_start = None + if lower_minutes := (float(freshness_definition.lower_tolerance) if freshness_definition.lower_tolerance else None): + window_start = add_business_minutes( + pd.Timestamp(last_detection_time), + lower_minutes, + test_suite.predict_exclude_weekends, + holiday_dates, tz, + excluded_days=sched.excluded_days, + ) + + return { + "start": int(window_start.timestamp() * 1000) if window_start else None, + "end": int(window_end.timestamp() * 1000), + } + + +def _build_gated_forecast_prediction( + definition: TestDefinition, + freshness_window: dict | None, + last_run_time: datetime | None, +) -> dict | None: + """Coupled forecast payload for a freshness-gated Volume/Metric monitor. + + Holds a flat baseline line from the latest run up to the predicted next-update window, then + steps to the forecast's next-refresh value with the band opening to its tolerance. The anchor + is never earlier than the latest run, so the forecast extends forward rather than back over + history. Returns None when there is no usable forward window — the caller then falls back to a + flat band — i.e. when freshness has no predicted window, the window has already elapsed, or the + gated prediction carries no baseline. + """ + baseline = definition.prediction.get("baseline_value") if definition.prediction else None + window_end = freshness_window.get("end") if freshness_window else None + now_ms = int(pd.Timestamp(last_run_time).timestamp() * 1000) if last_run_time is not None else None + if window_end is None or now_ms is None or window_end <= now_ms or baseline is None: + return None + + forecast_means = (definition.prediction.get("mean") if definition.prediction else None) or {} + next_refresh_mean = forecast_means[min(forecast_means, key=lambda k: int(k))] if forecast_means else baseline + flat_anchor = max(freshness_window.get("start") or now_ms, now_ms) + # lower/upper_tolerance are VARCHAR columns — coerce to float so the band dicts are numerically + # typed throughout (baseline is already a float). + lower_tol = float(definition.lower_tolerance) if definition.lower_tolerance is not None else None + upper_tol = float(definition.upper_tolerance) if definition.upper_tolerance is not None else None + return { + "method": "predict", + "mean": {flat_anchor: baseline, window_end: next_refresh_mean}, + "lower_tolerance": {flat_anchor: baseline, window_end: lower_tol}, + "upper_tolerance": {flat_anchor: baseline, window_end: upper_tol}, + } + + @with_database_session def build_table_trends_data( table_group: TableGroupMinimal, payload: dict, dialog: dict | None = None, @@ -571,9 +661,20 @@ def on_close_trends(_payload=None): metric_definition_id = metric_group["test_definition_id"] last_run_time_per_test_key[f"metric:{metric_definition_id}"] = max(e["time"] for e in metric_group["events"]) + # Predicted next freshness-update window for the table — shared by the Freshness_Trend + # display window and the freshness-gated Volume/Metric forecast (which expects no change + # until a refresh lands in this window). + freshness_definition = next((d for d in definitions if d.test_type == "Freshness_Trend"), None) + freshness_window = _freshness_next_update_window(freshness_definition, events, test_suite, monitor_schedule) + for definition in definitions: test_key = f"metric:{definition.id}" if definition.test_type == "Metric_Trend" else definition.test_type.lower() - if definition.history_calculation == "PREDICT" and definition.prediction and (base_mean_predictions := definition.prediction.get("mean")): + if ( + definition.history_calculation == "PREDICT" + and definition.prediction + and not definition.prediction.get("freshness_gated") + and (base_mean_predictions := definition.prediction.get("mean")) + ): predicted_times = sorted([datetime.fromtimestamp(int(timestamp) / 1000.0, UTC) for timestamp in base_mean_predictions.keys()]) # Limit predictions to 1/3 of the lookback, with minimum 3 points predicted_times = [str(int(t.timestamp() * 1000)) for idx, t in enumerate(predicted_times) if idx < 3 or idx < monitor_lookback / 3] @@ -592,6 +693,43 @@ def on_close_trends(_payload=None): "lower_tolerance": lower_tolerance_predictions, "upper_tolerance": upper_tolerance_predictions, } + elif ( + definition.history_calculation == "PREDICT" + and definition.prediction + and definition.prediction.get("freshness_gated") + and (definition.lower_tolerance is not None or definition.upper_tolerance is not None) + ): + # A freshness-gated monitor holds at its baseline between refreshes (the stale-period + # check is value == baseline), so it must never render the rising forecast cone. + gated_prediction = _build_gated_forecast_prediction( + definition, freshness_window, last_run_time_per_test_key.get(test_key), + ) + if gated_prediction is not None: + predictions[test_key] = gated_prediction + else: + # No freshness window available — fall back to a flat band at the next-refresh + # tolerance sampled across upcoming scheduled runs. + cron_sample = get_cron_sample( + monitor_schedule.cron_expr, + monitor_schedule.cron_tz, + sample_count=ceil(min(max(3, monitor_lookback / 3), 10)), + reference_time=last_run_time_per_test_key.get(test_key), + ) + mean_predictions: dict = {} + lower_tolerance_predictions: dict = {} + upper_tolerance_predictions: dict = {} + sample_next_runs = [timestamp * 1000 for timestamp in (cron_sample.get("samples") or [])] + for timestamp in sample_next_runs: + mean_predictions[timestamp] = None + lower_tolerance_predictions[timestamp] = definition.lower_tolerance + upper_tolerance_predictions[timestamp] = definition.upper_tolerance + + predictions[test_key] = { + "method": "static", + "mean": mean_predictions, + "lower_tolerance": lower_tolerance_predictions, + "upper_tolerance": upper_tolerance_predictions, + } elif definition.history_calculation is None and (definition.lower_tolerance is not None or definition.upper_tolerance is not None): cron_sample = get_cron_sample( monitor_schedule.cron_expr, @@ -614,46 +752,11 @@ def on_close_trends(_payload=None): "lower_tolerance": lower_tolerance_predictions, "upper_tolerance": upper_tolerance_predictions, } - elif ( - definition.test_type == "Freshness_Trend" - and definition.history_calculation == "PREDICT" - and (not definition.prediction or definition.prediction.get("schedule_stage")) - and definition.upper_tolerance is not None - ): - last_update_events = [ - e for e in events["freshness_events"] - if e["changed"] and not e["is_training"] and not e["is_pending"] - ] - if last_update_events: - last_detection_time = max(e["time"] for e in last_update_events) - holiday_dates = _resolve_holiday_dates(test_suite) - tz = monitor_schedule.cron_tz or "UTC" if monitor_schedule else None - sched = get_schedule_params(definition.prediction) - - window_end = add_business_minutes( - pd.Timestamp(last_detection_time), - float(definition.upper_tolerance), - test_suite.predict_exclude_weekends, - holiday_dates, tz, - excluded_days=sched.excluded_days, - ) - window_start = None - if lower_minutes := float(definition.lower_tolerance) if definition.lower_tolerance else None: - window_start = add_business_minutes( - pd.Timestamp(last_detection_time), - lower_minutes, - test_suite.predict_exclude_weekends, - holiday_dates, tz, - excluded_days=sched.excluded_days, - ) - - predictions["freshness_trend"] = { - "method": "freshness_window", - "window": { - "start": int(window_start.timestamp() * 1000) if window_start else None, - "end": int(window_end.timestamp() * 1000), - }, - } + elif definition.test_type == "Freshness_Trend" and freshness_window is not None: + predictions["freshness_trend"] = { + "method": "freshness_window", + "window": freshness_window, + } data = { **make_json_safe(events), diff --git a/tests/unit/commands/test_thresholds_prediction.py b/tests/unit/commands/test_thresholds_prediction.py index 37891a8c..d9f1f2ae 100644 --- a/tests/unit/commands/test_thresholds_prediction.py +++ b/tests/unit/commands/test_thresholds_prediction.py @@ -369,6 +369,36 @@ def test_freshness_gating_fits_on_filtered_series(mock_forecast): assert len(fitted_history) == len(freshness_updates) +@patch(MOCK_TARGET) +def test_freshness_gating_filtered_fit_uses_event_space(mock_forecast): + """The freshness-filtered fit must run in event-space (event_space=True): the filtered + series has one point per refresh and is irregularly spaced, so calendar resampling would + interpolate the refresh jumps into uniform increments and bias the forecast low.""" + mock_forecast.return_value = _make_forecast([220.0], [1.0]) + timestamps = [f"2026-01-{day:02d}" for day in range(1, 21)] + run_ids = [f"run_{i:02d}" for i in range(len(timestamps))] + history = _history_with_run_ids(timestamps, run_ids, value=220.0) + + compute_volume_or_metric_threshold(history, run_ids[:8], PredictSensitivity.medium) + + assert mock_forecast.call_args.kwargs["event_space"] is True + + +@patch(MOCK_TARGET) +def test_freshness_gating_raw_fallback_uses_calendar(mock_forecast): + """The raw-history fallback (when the filtered fit fails) keeps calendar regularization.""" + mock_forecast.side_effect = [NotEnoughData("not enough"), _make_forecast([220.0], [1.0])] + timestamps = [f"2026-01-{day:02d}" for day in range(1, 21)] + run_ids = [f"run_{i:02d}" for i in range(len(timestamps))] + history = _history_with_run_ids(timestamps, run_ids, value=220.0) + + compute_volume_or_metric_threshold(history, run_ids[:5], PredictSensitivity.medium) + + assert mock_forecast.call_count == 2 # filtered (event-space) failed, raw retried + assert mock_forecast.call_args_list[0].kwargs["event_space"] is True + assert mock_forecast.call_args_list[1].kwargs["event_space"] is False + + @patch(MOCK_TARGET) def test_freshness_gating_baseline_from_filtered_when_events_extend_past_history(mock_forecast): """When freshness_updates includes runs beyond the (retention-trimmed) history window, diff --git a/tests/unit/common/test_time_series_service.py b/tests/unit/common/test_time_series_service.py index c8a5d66a..b34bea76 100644 --- a/tests/unit/common/test_time_series_service.py +++ b/tests/unit/common/test_time_series_service.py @@ -617,3 +617,37 @@ def test_without_exclusions_timezone_has_no_effect(self): forecast_with_tz = get_sarimax_forecast(history, num_forecast=3, exclude_weekends=False, tz="America/New_York") pd.testing.assert_frame_equal(forecast_no_tz, forecast_with_tz) + + +class Test_GetSarimaxForecast_EventSpace: + """Event-space (event_space=True) vs default calendar-regularized fitting for irregular series.""" + + @staticmethod + def _irregular_refresh_history() -> pd.DataFrame: + """A refresh-driven series: irregular gaps (1-3 days), jumps that scale with the gap. + This is the shape that calendar interpolation smooths into uniform daily increments.""" + times = pd.to_datetime([ + "2026-01-01", "2026-01-02", "2026-01-03", "2026-01-05", "2026-01-06", + "2026-01-09", "2026-01-10", "2026-01-12", "2026-01-15", "2026-01-16", + "2026-01-18", "2026-01-21", + ]) + values = [100.0, 110, 121, 150, 162, 200, 212, 240, 290, 303, 330, 375] + return pd.DataFrame({"result_signal": values}, index=times) + + def test_event_space_centers_on_full_refresh_jump(self): + history = self._irregular_refresh_history() + + regularized = get_sarimax_forecast(history, num_forecast=1, event_space=False) + event_space = get_sarimax_forecast(history, num_forecast=1, event_space=True) + + last = history["result_signal"].iloc[-1] + # Calendar interpolation forecasts a single per-day increment; event-space forecasts a + # full refresh-to-refresh jump, so its next-point prediction sits meaningfully higher. + assert (event_space["mean"].iloc[0] - last) > (regularized["mean"].iloc[0] - last) + + def test_event_space_returns_requested_forecast_length(self): + history = self._irregular_refresh_history() + forecast = get_sarimax_forecast(history, num_forecast=5, event_space=True) + assert len(forecast) == 5 + assert forecast["mean"].notna().all() + assert forecast["se"].notna().all() diff --git a/tests/unit/ui/test_monitors_dashboard.py b/tests/unit/ui/test_monitors_dashboard.py new file mode 100644 index 00000000..84e2fd01 --- /dev/null +++ b/tests/unit/ui/test_monitors_dashboard.py @@ -0,0 +1,155 @@ +from datetime import datetime +from unittest.mock import MagicMock, patch + +import pandas as pd +import pytest + +from testgen.ui.views.monitors_dashboard import ( + _build_gated_forecast_prediction, + _freshness_next_update_window, +) + +pytestmark = pytest.mark.unit + +MODULE = "testgen.ui.views.monitors_dashboard" + + +def _freshness_def(history_calculation="PREDICT", prediction=None, upper="1000", lower="500"): + d = MagicMock() + d.history_calculation = history_calculation + d.prediction = {"schedule_stage": "active"} if prediction is None else prediction + d.upper_tolerance = upper + d.lower_tolerance = lower + return d + + +def _events(changed=True, is_training=False, is_pending=False, time=datetime(2026, 6, 22, 16, 0)): + return {"freshness_events": [{"changed": changed, "is_training": is_training, "is_pending": is_pending, "time": time}]} + + +def _suite(): + s = MagicMock() + s.predict_exclude_weekends = False + s.holiday_codes_list = None + return s + + +def _schedule(): + s = MagicMock() + s.cron_tz = "UTC" + return s + + +# --- _freshness_next_update_window: guard branches return None --- + + +def test_window_none_when_no_freshness_definition(): + assert _freshness_next_update_window(None, _events(), _suite(), _schedule()) is None + + +def test_window_none_when_not_predict(): + d = _freshness_def(history_calculation=None) + assert _freshness_next_update_window(d, _events(), _suite(), _schedule()) is None + + +def test_window_none_when_prediction_lacks_schedule_stage(): + # Non-empty prediction without schedule_stage (e.g. still learning) → no window + d = _freshness_def(prediction={"frequency": "daily"}) + assert _freshness_next_update_window(d, _events(), _suite(), _schedule()) is None + + +def test_window_none_when_no_upper_tolerance(): + d = _freshness_def(upper=None) + assert _freshness_next_update_window(d, _events(), _suite(), _schedule()) is None + + +def test_window_none_when_no_qualifying_update_events(): + # Only training / pending / unchanged events → nothing to anchor the window to + events = {"freshness_events": [ + {"changed": True, "is_training": True, "is_pending": False, "time": datetime(2026, 6, 22)}, + {"changed": False, "is_training": False, "is_pending": False, "time": datetime(2026, 6, 22)}, + ]} + assert _freshness_next_update_window(_freshness_def(), events, _suite(), _schedule()) is None + + +@patch(f"{MODULE}.get_schedule_params", return_value=MagicMock(excluded_days=[])) +@patch(f"{MODULE}.add_business_minutes") +def test_window_returns_start_end_when_valid(mock_abm, _mock_sched): + # add_business_minutes is called for the end (upper_tolerance) then the start (lower_tolerance) + mock_abm.side_effect = [pd.Timestamp("2026-06-23 19:00"), pd.Timestamp("2026-06-23 04:00")] + window = _freshness_next_update_window(_freshness_def(), _events(), _suite(), _schedule()) + assert window == { + "start": int(pd.Timestamp("2026-06-23 04:00").timestamp() * 1000), + "end": int(pd.Timestamp("2026-06-23 19:00").timestamp() * 1000), + } + + +@patch(f"{MODULE}.get_schedule_params", return_value=MagicMock(excluded_days=[])) +@patch(f"{MODULE}.add_business_minutes") +def test_window_start_is_none_when_no_lower_tolerance(mock_abm, _mock_sched): + mock_abm.return_value = pd.Timestamp("2026-06-23 19:00") + window = _freshness_next_update_window(_freshness_def(lower=None), _events(), _suite(), _schedule()) + assert window["start"] is None + assert window["end"] == int(pd.Timestamp("2026-06-23 19:00").timestamp() * 1000) + + +# --- _build_gated_forecast_prediction --- + +LAST_RUN = datetime(2026, 6, 23, 16, 0) +NOW_MS = int(pd.Timestamp(LAST_RUN).timestamp() * 1000) +HOUR = 3600 * 1000 + + +def _gated_def(baseline=1000.0, mean=None, lower="950", upper="1400"): + d = MagicMock() + d.prediction = {"freshness_gated": True, "baseline_value": baseline} + if mean is not None: + d.prediction["mean"] = mean + d.lower_tolerance = lower + d.upper_tolerance = upper + return d + + +def test_gated_prediction_none_without_window(): + assert _build_gated_forecast_prediction(_gated_def(), None, LAST_RUN) is None + + +def test_gated_prediction_none_when_window_elapsed(): + window = {"start": NOW_MS - 12 * HOUR, "end": NOW_MS - HOUR} # window_end already in the past + assert _build_gated_forecast_prediction(_gated_def(), window, LAST_RUN) is None + + +def test_gated_prediction_none_without_baseline(): + window = {"start": NOW_MS - HOUR, "end": NOW_MS + 3 * HOUR} + assert _build_gated_forecast_prediction(_gated_def(baseline=None), window, LAST_RUN) is None + + +def test_gated_prediction_none_without_last_run(): + window = {"start": NOW_MS - HOUR, "end": NOW_MS + 3 * HOUR} + assert _build_gated_forecast_prediction(_gated_def(), window, None) is None + + +def test_gated_prediction_anchors_at_now_when_window_started(): + # window_start is before the latest run → anchor clamps to now (forecast never draws backward) + window = {"start": NOW_MS - 12 * HOUR, "end": NOW_MS + 3 * HOUR} + mean = {str(NOW_MS + 3 * HOUR): 1200.0, str(NOW_MS + 27 * HOUR): 1300.0} + result = _build_gated_forecast_prediction(_gated_def(mean=mean), window, LAST_RUN) + assert result["method"] == "predict" + assert result["mean"] == {NOW_MS: 1000.0, NOW_MS + 3 * HOUR: 1200.0} + # tolerances coerced from VARCHAR to float + assert result["lower_tolerance"] == {NOW_MS: 1000.0, NOW_MS + 3 * HOUR: 950.0} + assert result["upper_tolerance"] == {NOW_MS: 1000.0, NOW_MS + 3 * HOUR: 1400.0} + + +def test_gated_prediction_anchors_at_window_start_when_future(): + # window opens after the latest run → flat segment runs out to window_start + window = {"start": NOW_MS + HOUR, "end": NOW_MS + 3 * HOUR} + mean = {str(NOW_MS + 3 * HOUR): 1200.0} + result = _build_gated_forecast_prediction(_gated_def(mean=mean), window, LAST_RUN) + assert set(result["mean"].keys()) == {NOW_MS + HOUR, NOW_MS + 3 * HOUR} + + +def test_gated_prediction_next_mean_falls_back_to_baseline_without_forecast(): + window = {"start": NOW_MS - HOUR, "end": NOW_MS + 3 * HOUR} + result = _build_gated_forecast_prediction(_gated_def(mean=None), window, LAST_RUN) + assert result["mean"][NOW_MS + 3 * HOUR] == 1000.0 # baseline used as the step value From f60c55f51b7de5b96ff665fbc3eddec072a2eb2c Mon Sep 17 00:00:00 2001 From: Astor Date: Wed, 24 Jun 2026 15:46:54 -0300 Subject: [PATCH 56/78] fix(TG-1113): show Managed Identity option on Synapse connection form SynapseMSSQLForm was aliased to RedshiftForm, which only renders host/username/password with no Managed Identity radio. Alias it to AzureMSSQLForm instead so the connect_with_identity field and the Connect with Managed Identity radio are exposed for Synapse connections, matching MCP parity (MR !529). Co-Authored-By: Claude Sonnet 4.6 --- .gitignore | 2 ++ testgen/ui/static/js/components/connection_form.js | 2 +- 2 files changed, 3 insertions(+), 1 deletion(-) diff --git a/.gitignore b/.gitignore index 62b713e7..f1b1b4ab 100644 --- a/.gitignore +++ b/.gitignore @@ -7,6 +7,8 @@ __pycache__/ work_area/ .envrc environment.list +.claude/settings.local.json +.claude/repo.yml # C extensions *.so diff --git a/testgen/ui/static/js/components/connection_form.js b/testgen/ui/static/js/components/connection_form.js index 1fadaacf..c0d3b1d7 100644 --- a/testgen/ui/static/js/components/connection_form.js +++ b/testgen/ui/static/js/components/connection_form.js @@ -784,7 +784,7 @@ const AzureMSSQLForm = ( ); }; -const SynapseMSSQLForm = RedshiftForm; +const SynapseMSSQLForm = AzureMSSQLForm; const MSSQLForm = RedshiftForm; From 5ef153763fd49660f52d5946b9e19637ae295ee8 Mon Sep 17 00:00:00 2001 From: Luis Date: Tue, 16 Jun 2026 18:28:08 -0400 Subject: [PATCH 57/78] :feat(test-definitions): faceted add-test picker dialog Add a faceted picker to the Add-Test dialog so users browse and search test types by taxonomy (algorithm, statistical technique) and column type before advancing to the existing parameter form. Persists the taxonomy fields on test types and surfaces them, with column-type metadata, to the dialog. Co-Authored-By: Claude Opus 4.8 (1M context) --- testgen/common/models/test_definition.py | 98 ++++ .../030_initialize_new_schema_structure.sql | 2 + .../test_types_Aggregate_Balance.yaml | 1 + .../test_types_Aggregate_Balance_Percent.yaml | 1 + .../test_types_Aggregate_Balance_Range.yaml | 1 + .../test_types_Aggregate_Minimum.yaml | 1 + .../test_types_Alpha_Trunc.yaml | 1 + .../test_types_Avg_Shift.yaml | 2 + .../dbsetup_test_types/test_types_CUSTOM.yaml | 1 + .../test_types_Combo_Match.yaml | 1 + .../test_types_Condition_Flag.yaml | 1 + .../test_types_Constant.yaml | 1 + .../test_types_Daily_Record_Ct.yaml | 1 + .../test_types_Dec_Trunc.yaml | 1 + .../test_types_Distinct_Date_Ct.yaml | 3 +- .../test_types_Distinct_Value_Ct.yaml | 1 + .../test_types_Distribution_Shift.yaml | 2 + .../test_types_Dupe_Rows.yaml | 1 + .../test_types_Email_Format.yaml | 1 + .../test_types_Freshness_Trend.yaml | 3 +- .../test_types_Future_Date.yaml | 3 +- .../test_types_Future_Date_1Y.yaml | 3 +- .../test_types_Incr_Avg_Shift.yaml | 2 + .../test_types_LOV_All.yaml | 1 + .../test_types_LOV_Match.yaml | 1 + .../test_types_Metric_Trend.yaml | 3 +- .../test_types_Min_Date.yaml | 1 + .../test_types_Min_Val.yaml | 1 + .../test_types_Missing_Pct.yaml | 2 + .../test_types_Monthly_Rec_Ct.yaml | 1 + .../test_types_Outlier_Pct_Above.yaml | 4 +- .../test_types_Outlier_Pct_Below.yaml | 4 +- .../test_types_Pattern_Match.yaml | 1 + .../test_types_Recency.yaml | 4 +- .../test_types_Required.yaml | 1 + .../dbsetup_test_types/test_types_Row_Ct.yaml | 1 + .../test_types_Row_Ct_Pct.yaml | 1 + .../test_types_Schema_Drift.yaml | 1 + .../test_types_Street_Addr_Pattern.yaml | 1 + .../test_types_Table_Freshness.yaml | 3 +- .../test_types_Timeframe_Combo_Gain.yaml | 1 + .../test_types_Timeframe_Combo_Match.yaml | 1 + .../test_types_US_State.yaml | 1 + .../dbsetup_test_types/test_types_Unique.yaml | 1 + .../test_types_Unique_Pct.yaml | 2 + .../test_types_Valid_Characters.yaml | 1 + .../test_types_Valid_Month.yaml | 1 + .../test_types_Valid_US_Zip.yaml | 1 + .../test_types_Valid_US_Zip3.yaml | 1 + .../test_types_Variability_Decrease.yaml | 2 + .../test_types_Variability_Increase.yaml | 2 + .../test_types_Volume_Trend.yaml | 2 + .../test_types_Weekly_Rec_Ct.yaml | 1 + .../dbupgrade/0196_incremental_upgrade.sql | 5 + .../frontend/js/pages/test_definitions.js | 553 +++++++++++++++--- testgen/ui/static/css/shared.css | 27 + .../js/components/test_picker_taxonomy.js | 108 ++++ testgen/ui/views/test_definitions.py | 23 +- .../common/models/test_test_type_columns.py | 49 ++ tests/unit/test_test_criteria.py | 104 ++++ tests/unit/test_test_type_taxonomy.py | 45 ++ 61 files changed, 998 insertions(+), 96 deletions(-) create mode 100644 testgen/template/dbupgrade/0196_incremental_upgrade.sql create mode 100644 testgen/ui/static/js/components/test_picker_taxonomy.js create mode 100644 tests/unit/common/models/test_test_type_columns.py create mode 100644 tests/unit/test_test_criteria.py create mode 100644 tests/unit/test_test_type_taxonomy.py diff --git a/testgen/common/models/test_definition.py b/testgen/common/models/test_definition.py index ec9ddde0..13fe7705 100644 --- a/testgen/common/models/test_definition.py +++ b/testgen/common/models/test_definition.py @@ -9,6 +9,7 @@ from sqlalchemy import ( Boolean, Column, + Enum, ForeignKey, String, Text, @@ -40,6 +41,87 @@ class Severity(StrEnum): WARNING = "Warning" +class TestAlgorithm(StrEnum): + """SQL-derived algorithm family for a test type (faceted picker axis).""" + + BOUNDARY_CHECK = "Boundary check" + COUNTING = "Counting" + PATTERN_REGEX = "Pattern / regex" + SET_LOOKUP = "Set / lookup" + STATISTICAL_DRIFT = "Statistical drift" + AGGREGATE_RECONCILIATION = "Aggregate reconciliation" + FRESHNESS_TIME = "Freshness / time" + SCHEMA_METADATA = "Schema / metadata" + CUSTOM_SQL = "Custom SQL" + + +class StatisticalTechnique(StrEnum): + """Named statistical technique a test type uses to evaluate its measure.""" + + COHENS_D = "Cohen's D" + COHENS_H = "Cohen's H" + OUTLIER_DETECTION = "Outlier Detection" + SD_SHIFT = "SD Shift" + JENSEN_SHANNON_DIVERGENCE = "Jensen-Shannon Divergence" + PREDICTIVE_MODEL = "Predictive Model" + + +class TestCriteria(StrEnum): + """What a test type needs to be set up (faceted picker axis). + + Derived, not stored: ``derive_test_criteria`` is the single source of truth, shared by the + UI lookup and MCP so the value never drifts between surfaces. + """ + + DEFINED_RULE = "Defined Rule" + DEFINED_THRESHOLD = "Defined Threshold" + DEFINED_VALUE = "Defined Value" + LIST_OF_VALUES = "List of Values" + REFERENCE_DATASET = "Reference Dataset" + CUSTOM_CRITERIA = "Custom Criteria" + + +# Predefined validity/integrity rules — the user just enables them; the rule itself is fixed +# (dedup/uniqueness and standard format validators). Not separable from structural attributes +# alone: e.g. Valid_Month is structurally identical to the Defined-Value test Pattern_Match +# (both column-scoped, Pattern / regex, with baseline_value + threshold_value params). +_DEFINED_RULE_TESTS = frozenset({ + "Dupe_Rows", "Unique", "Email_Format", "Street_Addr_Pattern", + "Valid_Characters", "Valid_Month", "Valid_US_Zip", "Valid_US_Zip3", +}) +# The user asserts the expected value/pattern the column should hold (a constant, a regex +# baseline, or simply that a value is present). Enumerated for the same reason as above. +_DEFINED_VALUE_TESTS = frozenset({ + "Constant", "Pattern_Match", "Required", +}) + + +def derive_test_criteria( + test_type: str, + test_scope: str | None, + algorithm: str | None, +) -> TestCriteria: + """Classify a test type by the kind of setup it requires. + + Single source of truth for the Criteria facet — call from both the UI lookup and MCP rather + than reproducing the rules. Scope and algorithm resolve the cleanly-typed buckets (referential + is checked before Set / lookup so Combo_Match stays Reference Dataset). Defined Rule, Defined + Value, and Defined Threshold can't be told apart from structural attributes, so the first two + are enumerated and Defined Threshold is the fallthrough. + """ + if test_scope == "custom": + return TestCriteria.CUSTOM_CRITERIA + if test_scope == "referential": + return TestCriteria.REFERENCE_DATASET + if algorithm == TestAlgorithm.SET_LOOKUP: + return TestCriteria.LIST_OF_VALUES + if test_type in _DEFINED_RULE_TESTS: + return TestCriteria.DEFINED_RULE + if test_type in _DEFINED_VALUE_TESTS: + return TestCriteria.DEFINED_VALUE + return TestCriteria.DEFINED_THRESHOLD + + class InvalidTestDefinitionFields(ValueError): """Aggregated field-level validation errors. ``errors``: ``dict[field_name, reason]``.""" @@ -175,6 +257,15 @@ def process_bind_param(self, value: str | None, _dialect) -> str | None: return value or None +def _enum_by_value(enum_cls: type[StrEnum]) -> Enum: + """Map a StrEnum to its VARCHAR column by member value (not name) so reads return enum members. + + These columns store the display value (e.g. ``"Boundary check"``), which differs from the enum + member name, so the default name-based mapping would not round-trip. + """ + return Enum(enum_cls, native_enum=False, values_callable=lambda cls: [member.value for member in cls]) + + class TestType(ParamFieldsMixin, Entity): __tablename__ = "test_types" @@ -204,12 +295,19 @@ class TestType(ParamFieldsMixin, Entity): dq_dimension: str = Column(String) impact_dimension: str = Column(String) health_dimension: str = Column(String) + algorithm: TestAlgorithm | None = Column(_enum_by_value(TestAlgorithm)) + statistical_technique: StatisticalTechnique | None = Column(_enum_by_value(StatisticalTechnique)) threshold_description: str = Column(String) usage_notes: str = Column(String) active: str = Column(String) # Unmapped columns: generation_template, result_visualization, result_visualization_params + @property + def criteria(self) -> TestCriteria: + """Setup-kind facet, derived via the shared classifier (see ``derive_test_criteria``).""" + return derive_test_criteria(self.test_type, self.test_scope, self.algorithm) + _summary_columns = ( *[key for key in TestTypeSummary.__annotations__.keys() if key not in ("default_test_description", "default_impact_dimension")], test_description.label("default_test_description"), diff --git a/testgen/template/dbsetup/030_initialize_new_schema_structure.sql b/testgen/template/dbsetup/030_initialize_new_schema_structure.sql index a71e6718..814f57b9 100644 --- a/testgen/template/dbsetup/030_initialize_new_schema_structure.sql +++ b/testgen/template/dbsetup/030_initialize_new_schema_structure.sql @@ -576,6 +576,8 @@ CREATE TABLE test_types ( dq_dimension VARCHAR(50), impact_dimension VARCHAR(20), health_dimension VARCHAR(50), + algorithm VARCHAR(64), + statistical_technique VARCHAR(64), threshold_description VARCHAR(200), result_visualization VARCHAR(50) DEFAULT 'line_chart', result_visualization_params TEXT DEFAULT NULL, diff --git a/testgen/template/dbsetup_test_types/test_types_Aggregate_Balance.yaml b/testgen/template/dbsetup_test_types/test_types_Aggregate_Balance.yaml index f38b89f4..15cacaf1 100644 --- a/testgen/template/dbsetup_test_types/test_types_Aggregate_Balance.yaml +++ b/testgen/template/dbsetup_test_types/test_types_Aggregate_Balance.yaml @@ -37,6 +37,7 @@ test_types: usage_notes: |- This test compares sums or counts of a column rolled up to one or more category combinations across two different tables. Both tables must be accessible at the same time. It's ideal for confirming that two datasets exactly match -- that the sum of a measure or count of a value hasn't changed or shifted between categories. Use this test to compare a raw and processed version of the same dataset, or to confirm that an aggregated table exactly matches the detail table that it's built from. An error here means that one or more value combinations fail to match. New categories or combinations will cause failure. active: Y + algorithm: Aggregate reconciliation cat_test_conditions: [] target_data_lookups: - id: '1400' diff --git a/testgen/template/dbsetup_test_types/test_types_Aggregate_Balance_Percent.yaml b/testgen/template/dbsetup_test_types/test_types_Aggregate_Balance_Percent.yaml index 1415731d..47882a0f 100644 --- a/testgen/template/dbsetup_test_types/test_types_Aggregate_Balance_Percent.yaml +++ b/testgen/template/dbsetup_test_types/test_types_Aggregate_Balance_Percent.yaml @@ -37,6 +37,7 @@ test_types: usage_notes: |- This test compares sums or counts of a column rolled up to one or more category combinations across two different tables. Both tables must be accessible at the same time. Use it to confirm that two datasets closely match within the tolerance you set -- that the sum of a measure or count of a value remains sufficiently consistent between categories. You could use this test compare sales per product within one month to another, when you want to be alerted if the difference for any product falls outside of the range defined as 5% below to 10% above the prior month. An error here means that one or more value combinations fail to match within the set tolerances. New categories or combinations will cause failure. active: Y + algorithm: Aggregate reconciliation cat_test_conditions: [] target_data_lookups: - id: '1404' diff --git a/testgen/template/dbsetup_test_types/test_types_Aggregate_Balance_Range.yaml b/testgen/template/dbsetup_test_types/test_types_Aggregate_Balance_Range.yaml index 84f20602..c578c084 100644 --- a/testgen/template/dbsetup_test_types/test_types_Aggregate_Balance_Range.yaml +++ b/testgen/template/dbsetup_test_types/test_types_Aggregate_Balance_Range.yaml @@ -37,6 +37,7 @@ test_types: usage_notes: |- This test compares sums or counts of a column rolled up to one or more category combinations across two different tables. Both tables must be accessible at the same time. Use it to confirm that two datasets closely match within the tolerances you define as specific values above or below the aggregate measure for the same categories in the reference dataset -- that the sum of a measure or count of a value remains sufficiently consistent between categories. For instance, you can use this test to compare sales per product within one month to another, when you want to be alerted if the difference for any product falls outside of the range defined as 10000 dollars above or below the prior week. An error here means that one or more value combinations fail to match within the set tolerances. New categories or combinations will cause failure. active: Y + algorithm: Aggregate reconciliation cat_test_conditions: [] target_data_lookups: - id: '1405' diff --git a/testgen/template/dbsetup_test_types/test_types_Aggregate_Minimum.yaml b/testgen/template/dbsetup_test_types/test_types_Aggregate_Minimum.yaml index 425a72e2..4fa29cee 100644 --- a/testgen/template/dbsetup_test_types/test_types_Aggregate_Minimum.yaml +++ b/testgen/template/dbsetup_test_types/test_types_Aggregate_Minimum.yaml @@ -37,6 +37,7 @@ test_types: usage_notes: |- This test compares sums or counts of a column rolled up to one or more category combinations, but requires a match or increase in the aggregate value, rather than an exact match, across two different tables. Both tables must be accessible at the same time. Use this to confirm that aggregate values have not dropped for any set of categories, even if some values may rise. This test is useful to compare an older and newer version of a cumulative dataset. An error here means that one or more values per category set fail to match or exceed the prior dataset. New categories or combinations are allowed (but can be restricted independently with a Combo_Match test). Both tables must be present to run this test. active: Y + algorithm: Aggregate reconciliation cat_test_conditions: [] target_data_lookups: - id: '1401' diff --git a/testgen/template/dbsetup_test_types/test_types_Alpha_Trunc.yaml b/testgen/template/dbsetup_test_types/test_types_Alpha_Trunc.yaml index 88577f60..8cd8f223 100644 --- a/testgen/template/dbsetup_test_types/test_types_Alpha_Trunc.yaml +++ b/testgen/template/dbsetup_test_types/test_types_Alpha_Trunc.yaml @@ -36,6 +36,7 @@ test_types: usage_notes: |- Alpha Truncation tests that the longest text value in a column hasn't become shorter than the defined threshold, initially 95% of the longest value at baseline. This could indicate a problem in a cumulative dataset, where prior values should still exist unchanged. A failure here would suggest that some process changed data that you would still expect to be present and matching its value when the column was profiled. This test would not be appropriate for an incremental or windowed dataset. active: Y + algorithm: Boundary check cat_test_conditions: - id: '7001' test_type: Alpha_Trunc diff --git a/testgen/template/dbsetup_test_types/test_types_Avg_Shift.yaml b/testgen/template/dbsetup_test_types/test_types_Avg_Shift.yaml index 08f801a7..30f8c7c7 100644 --- a/testgen/template/dbsetup_test_types/test_types_Avg_Shift.yaml +++ b/testgen/template/dbsetup_test_types/test_types_Avg_Shift.yaml @@ -37,6 +37,8 @@ test_types: usage_notes: |- Average Shift tests that the average of a numeric column has not significantly changed since baseline, when profiling was done. A significant shift may indicate errors in processing, differences in source data, or valid changes that may nevertheless impact assumptions in downstream data products. The test uses Cohen's D, a statistical technique to identify significant shifts in a value. Cohen's D measures the difference between the two averages, reporting results on a standardized scale, which can be interpreted via a rule-of-thumb from small to huge. Depending on your data, some difference may be expected, so it's reasonable to adjust the threshold value that triggers test failure. This test works well for measures, or even for identifiers if you expect them to increment consistently. You may want to periodically adjust the expected threshold, or even the expected average value if you expect shifting over time. Consider this test along with Variability Increase. If variability rises too, process or measurement flaws could be at work. If variability remains consistent, the issue is more likely to be with the source data itself. active: Y + algorithm: Statistical drift + statistical_technique: Cohen's D cat_test_conditions: - id: '7002' test_type: Avg_Shift diff --git a/testgen/template/dbsetup_test_types/test_types_CUSTOM.yaml b/testgen/template/dbsetup_test_types/test_types_CUSTOM.yaml index 6005806c..099e4b05 100644 --- a/testgen/template/dbsetup_test_types/test_types_CUSTOM.yaml +++ b/testgen/template/dbsetup_test_types/test_types_CUSTOM.yaml @@ -38,6 +38,7 @@ test_types: usage_notes: |- This business-rule test is highly flexible, covering any error state that can be expressed by a SQL query against one or more tables in the database. In operation, the user-defined query is embedded within a parent query returning the count of error rows identified. Any row returned by the query is interpreted as a single error condition in the test. Note that this query is run independently of other tests, and that performance will be slower, depending in large part on the efficiency of the query you write. Interpretation is based on the user-defined meaning of the test. Your query might be written to return errors in individual rows identified by joining tables. Or it might return an error based on a multi-column aggregate condition returning a single row if an error is found. This query is run separately when you click `Review Source Data` from Test Results, so be sure to include enough data in your results to follow-up. Interpretation is based on the user-defined meaning of the test. active: Y + algorithm: Custom SQL cat_test_conditions: [] target_data_lookups: [] test_templates: diff --git a/testgen/template/dbsetup_test_types/test_types_Combo_Match.yaml b/testgen/template/dbsetup_test_types/test_types_Combo_Match.yaml index f4016136..fef1ba37 100644 --- a/testgen/template/dbsetup_test_types/test_types_Combo_Match.yaml +++ b/testgen/template/dbsetup_test_types/test_types_Combo_Match.yaml @@ -37,6 +37,7 @@ test_types: usage_notes: |- This test verifies that values, or combinations of values, that are present in the main table are also found in a reference table. This is a useful test for referential integrity between fact and dimension tables. You can also use it to confirm the validity of a code or category, or of combinations of values that should only be found together within each record, such as product/size/color. An error here means that one or more category combinations in the main table are not found in the reference table. Both tables must be present to run this test. active: Y + algorithm: Set / lookup cat_test_conditions: [] target_data_lookups: - id: '1402' diff --git a/testgen/template/dbsetup_test_types/test_types_Condition_Flag.yaml b/testgen/template/dbsetup_test_types/test_types_Condition_Flag.yaml index 9c63f169..c2b08622 100644 --- a/testgen/template/dbsetup_test_types/test_types_Condition_Flag.yaml +++ b/testgen/template/dbsetup_test_types/test_types_Condition_Flag.yaml @@ -37,6 +37,7 @@ test_types: usage_notes: |- Custom Condition is a business-rule test for a user-defined error condition based on the value of one or more columns. The condition is applied to each record within the table, and the count of records failing the condition is added up. If that count exceeds a threshold of errors, the test as a whole is failed. This test is ideal for error conditions that TestGen cannot automatically infer, and any condition that involves the values of more than one column in the same record. Performance of this test is fast, since it is performed together with other aggregate tests. Interpretation is based on the user-defined meaning of the test. active: Y + algorithm: Custom SQL cat_test_conditions: - id: '7003' test_type: Condition_Flag diff --git a/testgen/template/dbsetup_test_types/test_types_Constant.yaml b/testgen/template/dbsetup_test_types/test_types_Constant.yaml index fff389cd..0bd3351e 100644 --- a/testgen/template/dbsetup_test_types/test_types_Constant.yaml +++ b/testgen/template/dbsetup_test_types/test_types_Constant.yaml @@ -36,6 +36,7 @@ test_types: usage_notes: |- Constant Match tests that a single value determined to be a constant in baseline profiling is still the only value for the column that appears in subsequent versions of the dataset. Sometimes new data or business knowledge may reveal that the value is not a constant at all, even though only one value was present at profiling. In this case, you will want to disable this test. Alternatively, you can use the Value Match test to provide a limited number of valid values for the column. active: Y + algorithm: Boundary check cat_test_conditions: - id: '7004' test_type: Constant diff --git a/testgen/template/dbsetup_test_types/test_types_Daily_Record_Ct.yaml b/testgen/template/dbsetup_test_types/test_types_Daily_Record_Ct.yaml index df372cd6..2926c8c4 100644 --- a/testgen/template/dbsetup_test_types/test_types_Daily_Record_Ct.yaml +++ b/testgen/template/dbsetup_test_types/test_types_Daily_Record_Ct.yaml @@ -40,6 +40,7 @@ test_types: \ of days identified without data. You can adjust the threshold to accept a number\ \ of days that you know legitimately have no records. " active: Y + algorithm: Counting cat_test_conditions: - id: '7005' test_type: Daily_Record_Ct diff --git a/testgen/template/dbsetup_test_types/test_types_Dec_Trunc.yaml b/testgen/template/dbsetup_test_types/test_types_Dec_Trunc.yaml index 770d1175..bd6d5168 100644 --- a/testgen/template/dbsetup_test_types/test_types_Dec_Trunc.yaml +++ b/testgen/template/dbsetup_test_types/test_types_Dec_Trunc.yaml @@ -37,6 +37,7 @@ test_types: usage_notes: |- Decimal Truncation tests that the fractional (decimal) part of a numeric column has not been truncated since Baseline. This works by summing all the fractional values after the decimal point and confirming that the total is at least equal to the fractional total at baseline. This could indicate a problem in a cumulative dataset, where prior values should still exist unchanged. A failure here would suggest that some process changed data that you would still expect to be present and matching its value when the column was profiled. This test would not be appropriate for an incremental or windowed dataset. active: Y + algorithm: Boundary check cat_test_conditions: - id: '7006' test_type: Dec_Trunc diff --git a/testgen/template/dbsetup_test_types/test_types_Distinct_Date_Ct.yaml b/testgen/template/dbsetup_test_types/test_types_Distinct_Date_Ct.yaml index 23967398..b77887a2 100644 --- a/testgen/template/dbsetup_test_types/test_types_Distinct_Date_Ct.yaml +++ b/testgen/template/dbsetup_test_types/test_types_Distinct_Date_Ct.yaml @@ -29,7 +29,7 @@ test_types: test_scope: column dq_dimension: Timeliness impact_dimension: Reliability - health_dimension: Recency + health_dimension: Freshness threshold_description: |- Minimum distinct date count expected result_visualization: line_chart @@ -37,6 +37,7 @@ test_types: usage_notes: |- Date Count tests that the count of distinct dates present in the column has not dropped since baseline. The test is relevant for cumulative datasets, where old records are retained. A failure here would indicate missing records, which could be caused by a processing error or changed upstream data sources. active: Y + algorithm: Counting cat_test_conditions: - id: '7007' test_type: Distinct_Date_Ct diff --git a/testgen/template/dbsetup_test_types/test_types_Distinct_Value_Ct.yaml b/testgen/template/dbsetup_test_types/test_types_Distinct_Value_Ct.yaml index 47e609ec..630cb901 100644 --- a/testgen/template/dbsetup_test_types/test_types_Distinct_Value_Ct.yaml +++ b/testgen/template/dbsetup_test_types/test_types_Distinct_Value_Ct.yaml @@ -36,6 +36,7 @@ test_types: usage_notes: |- Value Count tests that the count of unique values present in the column has not dropped since baseline. The test is relevant for cumulative datasets, where old records are retained, or for any dataset where you would expect a set number of distinct values should be present. A failure here would indicate missing records or a change in categories or value assignment. active: Y + algorithm: Counting cat_test_conditions: - id: '7008' test_type: Distinct_Value_Ct diff --git a/testgen/template/dbsetup_test_types/test_types_Distribution_Shift.yaml b/testgen/template/dbsetup_test_types/test_types_Distribution_Shift.yaml index 666bf095..9f5efdce 100644 --- a/testgen/template/dbsetup_test_types/test_types_Distribution_Shift.yaml +++ b/testgen/template/dbsetup_test_types/test_types_Distribution_Shift.yaml @@ -38,6 +38,8 @@ test_types: usage_notes: |- This test measures the similarity of two sets of counts per categories, by using their proportional counts as probability distributions. Using Jensen-Shannon divergence, a measure of relative entropy or difference between two distributions, the test assigns a score ranging from 0, meaning that the distributions are identical, to 1, meaning that the distributions are completely unrelated. This test can be used to compare datasets that may not match exactly, but should have similar distributions. For example, it is a useful sanity check for data from different sources that you would expect to have a consistent spread, such as shipment of building materials per state and construction projects by state. Scores can be compared over time even if the distributions are not identical -- a dataset can be expected to maintain a comparable divergence score with a reference dataset over time. Both tables must be present to run this test. active: Y + algorithm: Statistical drift + statistical_technique: Jensen-Shannon Divergence cat_test_conditions: [] target_data_lookups: - id: '1403' diff --git a/testgen/template/dbsetup_test_types/test_types_Dupe_Rows.yaml b/testgen/template/dbsetup_test_types/test_types_Dupe_Rows.yaml index 1ef27125..5eeb9b81 100644 --- a/testgen/template/dbsetup_test_types/test_types_Dupe_Rows.yaml +++ b/testgen/template/dbsetup_test_types/test_types_Dupe_Rows.yaml @@ -38,6 +38,7 @@ test_types: usage_notes: |- This test verifies that combinations of values are not repeated within the table. By default when auto-generated, the test considers all columns to protect against duplication of entire rows. If you know the minimum columns that should constitute a unique record, such as a set of ID's, you should use those to make the test as sensitive as possible. Alternatively, if you know of columns you can always exclude, such as file_date or refresh_snapshot_id, remove them to tighten the test somewhat. active: Y + algorithm: Counting cat_test_conditions: [] target_data_lookups: - id: '1409' diff --git a/testgen/template/dbsetup_test_types/test_types_Email_Format.yaml b/testgen/template/dbsetup_test_types/test_types_Email_Format.yaml index 6d6573b4..3a625be1 100644 --- a/testgen/template/dbsetup_test_types/test_types_Email_Format.yaml +++ b/testgen/template/dbsetup_test_types/test_types_Email_Format.yaml @@ -36,6 +36,7 @@ test_types: result_visualization_params: null usage_notes: null active: Y + algorithm: Pattern / regex cat_test_conditions: - id: '7009' test_type: Email_Format diff --git a/testgen/template/dbsetup_test_types/test_types_Freshness_Trend.yaml b/testgen/template/dbsetup_test_types/test_types_Freshness_Trend.yaml index c8d1e3a2..15250c0a 100644 --- a/testgen/template/dbsetup_test_types/test_types_Freshness_Trend.yaml +++ b/testgen/template/dbsetup_test_types/test_types_Freshness_Trend.yaml @@ -30,7 +30,7 @@ test_types: test_scope: table dq_dimension: Recency impact_dimension: Reliability - health_dimension: Recency + health_dimension: Freshness threshold_description: |- Expected time window result_visualization: binary_chart @@ -38,6 +38,7 @@ test_types: usage_notes: |- This test compares the current table fingerprint, calculated signature of column contents, to confirm that the table has been updated within the expectd time window. The table fingerprint is derived from a set of values and aggregates from columns most likely to change. This test allows you to track the schedule and frequency of updates and refreshes to the table. active: Y + algorithm: Freshness / time cat_test_conditions: [] target_data_lookups: [] test_templates: diff --git a/testgen/template/dbsetup_test_types/test_types_Future_Date.yaml b/testgen/template/dbsetup_test_types/test_types_Future_Date.yaml index c2327843..78217c1b 100644 --- a/testgen/template/dbsetup_test_types/test_types_Future_Date.yaml +++ b/testgen/template/dbsetup_test_types/test_types_Future_Date.yaml @@ -28,13 +28,14 @@ test_types: test_scope: column dq_dimension: Timeliness impact_dimension: Conformance - health_dimension: Recency + health_dimension: Freshness threshold_description: |- Expected count of future dates result_visualization: line_chart result_visualization_params: null usage_notes: null active: Y + algorithm: Boundary check cat_test_conditions: - id: '7010' test_type: Future_Date diff --git a/testgen/template/dbsetup_test_types/test_types_Future_Date_1Y.yaml b/testgen/template/dbsetup_test_types/test_types_Future_Date_1Y.yaml index c1ec7d6d..5e4254f9 100644 --- a/testgen/template/dbsetup_test_types/test_types_Future_Date_1Y.yaml +++ b/testgen/template/dbsetup_test_types/test_types_Future_Date_1Y.yaml @@ -28,7 +28,7 @@ test_types: test_scope: column dq_dimension: Timeliness impact_dimension: Conformance - health_dimension: Recency + health_dimension: Freshness threshold_description: |- Expected count of future dates beyond one year result_visualization: line_chart @@ -36,6 +36,7 @@ test_types: usage_notes: |- Future Year looks for date values in the column that extend beyond one year after the test date. This would be appropriate for transactional dates where you would expect to find dates in the near future, but not beyond one year ahead. Errors could indicate invalid entries or possibly dummy dates representing blank values. active: Y + algorithm: Boundary check cat_test_conditions: - id: '7011' test_type: Future_Date_1Y diff --git a/testgen/template/dbsetup_test_types/test_types_Incr_Avg_Shift.yaml b/testgen/template/dbsetup_test_types/test_types_Incr_Avg_Shift.yaml index 00b6b9f6..69b35b8a 100644 --- a/testgen/template/dbsetup_test_types/test_types_Incr_Avg_Shift.yaml +++ b/testgen/template/dbsetup_test_types/test_types_Incr_Avg_Shift.yaml @@ -37,6 +37,8 @@ test_types: usage_notes: |- This is a more sensitive test than Average Shift, because it calculates an incremental difference in the average of new values compared to the average of values at baseline. This is appropriate for a cumulative dataset only, because it calculates the average of new entries based on the assumption that the count and average of records present at baseline are still present at the time of the test. This test compares the mean of new values with the standard deviation of the baseline average to calculate a Z-score. If the new mean falls outside the Z-score threshold, a shift is detected. Potential Z-score thresholds may range from 0 to 3, depending on the sensitivity you prefer. A failed test could indicate a quality issue or a legitimate shift in new data that should be noted and assessed by business users. Consider this test along with Variability Increase. If variability rises too, process, methodology or measurement flaws could be at issue. If variability remains consistent, the problem is more likely to be with the source data itself. active: Y + algorithm: Statistical drift + statistical_technique: Cohen's D cat_test_conditions: - id: '7012' test_type: Incr_Avg_Shift diff --git a/testgen/template/dbsetup_test_types/test_types_LOV_All.yaml b/testgen/template/dbsetup_test_types/test_types_LOV_All.yaml index cdd4bfda..e6b7145d 100644 --- a/testgen/template/dbsetup_test_types/test_types_LOV_All.yaml +++ b/testgen/template/dbsetup_test_types/test_types_LOV_All.yaml @@ -34,6 +34,7 @@ test_types: usage_notes: |- This is a more restrictive form of Value Match, testing that all values in the dataset match the list provided, and also that all values present in the list appear at least once in the dataset. This would be appropriate for tables where all category values in the column are represented at least once. active: Y + algorithm: Set / lookup cat_test_conditions: - id: '7013' test_type: LOV_All diff --git a/testgen/template/dbsetup_test_types/test_types_LOV_Match.yaml b/testgen/template/dbsetup_test_types/test_types_LOV_Match.yaml index 38b8040c..463503f2 100644 --- a/testgen/template/dbsetup_test_types/test_types_LOV_Match.yaml +++ b/testgen/template/dbsetup_test_types/test_types_LOV_Match.yaml @@ -140,6 +140,7 @@ test_types: usage_notes: |- This tests that all values in the column match the hard-coded list provided. This is relevant when the list of allowable values is small and not expected to change often. Even if new values might occasionally be added, this test is useful for downstream data products to provide warning that assumptions and logic may need to change. active: Y + algorithm: Set / lookup cat_test_conditions: - id: '7014' test_type: LOV_Match diff --git a/testgen/template/dbsetup_test_types/test_types_Metric_Trend.yaml b/testgen/template/dbsetup_test_types/test_types_Metric_Trend.yaml index 7675e9d1..2e0f56f1 100644 --- a/testgen/template/dbsetup_test_types/test_types_Metric_Trend.yaml +++ b/testgen/template/dbsetup_test_types/test_types_Metric_Trend.yaml @@ -25,7 +25,7 @@ test_types: test_scope: table dq_dimension: Validity impact_dimension: Regularity - health_dimension: null + health_dimension: Data Drift threshold_description: |- Expected aggregate metric range. result_visualization: line_chart @@ -33,6 +33,7 @@ test_types: usage_notes: |- This test compares the aggregate metric of all or a subset of records in a table against a derived tolerance range. active: Y + algorithm: Custom SQL cat_test_conditions: - id: '2516' test_type: Metric_Trend diff --git a/testgen/template/dbsetup_test_types/test_types_Min_Date.yaml b/testgen/template/dbsetup_test_types/test_types_Min_Date.yaml index 1ade805e..80dfeffa 100644 --- a/testgen/template/dbsetup_test_types/test_types_Min_Date.yaml +++ b/testgen/template/dbsetup_test_types/test_types_Min_Date.yaml @@ -36,6 +36,7 @@ test_types: usage_notes: |- This test is appropriate for a cumulative dataset only, because it assumes all prior values are still present. It's appropriate where new records are added with more recent dates, but old dates dates do not change. active: Y + algorithm: Boundary check cat_test_conditions: - id: '7015' test_type: Min_Date diff --git a/testgen/template/dbsetup_test_types/test_types_Min_Val.yaml b/testgen/template/dbsetup_test_types/test_types_Min_Val.yaml index 90f107f4..49a3e5d4 100644 --- a/testgen/template/dbsetup_test_types/test_types_Min_Val.yaml +++ b/testgen/template/dbsetup_test_types/test_types_Min_Val.yaml @@ -36,6 +36,7 @@ test_types: usage_notes: |- This test is appropriate for a cumulative dataset only, assuming all prior values are still present. It is also appropriate for any measure that has an absolute, definable minimum value, or a heuristic that makes senes for valid data. active: Y + algorithm: Boundary check cat_test_conditions: - id: '7016' test_type: Min_Val diff --git a/testgen/template/dbsetup_test_types/test_types_Missing_Pct.yaml b/testgen/template/dbsetup_test_types/test_types_Missing_Pct.yaml index f4dd0a0a..8669d6c4 100644 --- a/testgen/template/dbsetup_test_types/test_types_Missing_Pct.yaml +++ b/testgen/template/dbsetup_test_types/test_types_Missing_Pct.yaml @@ -37,6 +37,8 @@ test_types: usage_notes: |- This test uses Cohen's H, a statistical test to identify a significant difference between two ratios. Results are reported on a standardized scale, which can be interpreted via a rule-of-thumb from small to huge. An uptick in missing data may indicate a collection issue at the source. A larger change may indicate a processing failure. A drop in missing data may also be significant, if it affects assumptions built into analytic products downstream. You can refine the expected threshold value as you view legitimate results of the measure over time. active: Y + algorithm: Statistical drift + statistical_technique: Cohen's H cat_test_conditions: - id: '7017' test_type: Missing_Pct diff --git a/testgen/template/dbsetup_test_types/test_types_Monthly_Rec_Ct.yaml b/testgen/template/dbsetup_test_types/test_types_Monthly_Rec_Ct.yaml index 35580b34..29df97dc 100644 --- a/testgen/template/dbsetup_test_types/test_types_Monthly_Rec_Ct.yaml +++ b/testgen/template/dbsetup_test_types/test_types_Monthly_Rec_Ct.yaml @@ -37,6 +37,7 @@ test_types: usage_notes: |- Monthly Records tests that at least one record is present for every calendar month within the minimum and maximum date range for the column. The test is relevant for transactional data, where you would expect at least one transaction to be recorded each month. A failure here would suggest missing records for the number of months identified without data. You can adjust the threshold to accept a number of month that you know legitimately have no records. active: Y + algorithm: Counting cat_test_conditions: - id: '7018' test_type: Monthly_Rec_Ct diff --git a/testgen/template/dbsetup_test_types/test_types_Outlier_Pct_Above.yaml b/testgen/template/dbsetup_test_types/test_types_Outlier_Pct_Above.yaml index 5cf1493f..2112c1b7 100644 --- a/testgen/template/dbsetup_test_types/test_types_Outlier_Pct_Above.yaml +++ b/testgen/template/dbsetup_test_types/test_types_Outlier_Pct_Above.yaml @@ -28,7 +28,7 @@ test_types: test_scope: column dq_dimension: Accuracy impact_dimension: Regularity - health_dimension: Data Drift + health_dimension: Statistical Drift threshold_description: |- Expected maximum pct records over upper 2 SD limit result_visualization: line_chart @@ -41,6 +41,8 @@ test_types: \ you expect to see. This test uses the baseline mean rather than the mean for\ \ the latest dataset to capture systemic shift as well as individual outliers. " active: Y + algorithm: Statistical drift + statistical_technique: Outlier Detection cat_test_conditions: - id: '7019' test_type: Outlier_Pct_Above diff --git a/testgen/template/dbsetup_test_types/test_types_Outlier_Pct_Below.yaml b/testgen/template/dbsetup_test_types/test_types_Outlier_Pct_Below.yaml index fa88ab8c..fc4d7262 100644 --- a/testgen/template/dbsetup_test_types/test_types_Outlier_Pct_Below.yaml +++ b/testgen/template/dbsetup_test_types/test_types_Outlier_Pct_Below.yaml @@ -28,7 +28,7 @@ test_types: test_scope: column dq_dimension: Accuracy impact_dimension: Regularity - health_dimension: Data Drift + health_dimension: Statistical Drift threshold_description: |- Expected maximum pct records over lower 2 SD limit result_visualization: line_chart @@ -41,6 +41,8 @@ test_types: \ you expect to see. This test uses the baseline mean rather than the mean for\ \ the latest dataset to capture systemic shift as well as individual outliers. " active: Y + algorithm: Statistical drift + statistical_technique: Outlier Detection cat_test_conditions: - id: '7020' test_type: Outlier_Pct_Below diff --git a/testgen/template/dbsetup_test_types/test_types_Pattern_Match.yaml b/testgen/template/dbsetup_test_types/test_types_Pattern_Match.yaml index 6998da47..36f81737 100644 --- a/testgen/template/dbsetup_test_types/test_types_Pattern_Match.yaml +++ b/testgen/template/dbsetup_test_types/test_types_Pattern_Match.yaml @@ -36,6 +36,7 @@ test_types: usage_notes: |- This test is appropriate for character fields that are expected to appear in a consistent format. It uses pattern matching syntax as appropriate for your database: REGEX matching if available, otherwise LIKE expressions. The expected threshold is the number of records that fail to match the defined pattern. active: Y + algorithm: Pattern / regex cat_test_conditions: - id: '7021' test_type: Pattern_Match diff --git a/testgen/template/dbsetup_test_types/test_types_Recency.yaml b/testgen/template/dbsetup_test_types/test_types_Recency.yaml index 34945e9b..b1d656c7 100644 --- a/testgen/template/dbsetup_test_types/test_types_Recency.yaml +++ b/testgen/template/dbsetup_test_types/test_types_Recency.yaml @@ -29,7 +29,8 @@ test_types: test_scope: column dq_dimension: Timeliness impact_dimension: Reliability - health_dimension: Recency + health_dimension: Freshness + statistical_technique: Predictive Model threshold_description: |- Expected maximum count of days preceding test date result_visualization: line_chart @@ -37,6 +38,7 @@ test_types: usage_notes: |- This test evaluates recency based on the latest referenced dates in the column. The test is appropriate for transactional dates and timestamps. The test can be especially valuable because timely data deliveries themselves may not assure that the most recent data is present. You can adjust the expected threshold to the maximum number of days that you expect the data to age before the dataset is refreshed. active: Y + algorithm: Freshness / time cat_test_conditions: - id: '7022' test_type: Recency diff --git a/testgen/template/dbsetup_test_types/test_types_Required.yaml b/testgen/template/dbsetup_test_types/test_types_Required.yaml index cb294860..6f785df2 100644 --- a/testgen/template/dbsetup_test_types/test_types_Required.yaml +++ b/testgen/template/dbsetup_test_types/test_types_Required.yaml @@ -35,6 +35,7 @@ test_types: result_visualization_params: null usage_notes: null active: Y + algorithm: Counting cat_test_conditions: - id: '7023' test_type: Required diff --git a/testgen/template/dbsetup_test_types/test_types_Row_Ct.yaml b/testgen/template/dbsetup_test_types/test_types_Row_Ct.yaml index 06c3d62e..fee0fb9c 100644 --- a/testgen/template/dbsetup_test_types/test_types_Row_Ct.yaml +++ b/testgen/template/dbsetup_test_types/test_types_Row_Ct.yaml @@ -34,6 +34,7 @@ test_types: usage_notes: |- Because this tests the row count against a constant minimum threshold, it's appropriate for any dataset, as long as the number of rows doesn't radically change from refresh to refresh. But it's not responsive to change over time. You may want to adjust the threshold periodically if you are dealing with a cumulative dataset. active: Y + algorithm: Counting cat_test_conditions: - id: '7024' test_type: Row_Ct diff --git a/testgen/template/dbsetup_test_types/test_types_Row_Ct_Pct.yaml b/testgen/template/dbsetup_test_types/test_types_Row_Ct_Pct.yaml index 47cbf379..85588851 100644 --- a/testgen/template/dbsetup_test_types/test_types_Row_Ct_Pct.yaml +++ b/testgen/template/dbsetup_test_types/test_types_Row_Ct_Pct.yaml @@ -35,6 +35,7 @@ test_types: usage_notes: |- This test is better than Row Count for an incremental or windowed dataset where you would expect the row count to range within a percentage of baseline. active: Y + algorithm: Counting cat_test_conditions: - id: '7025' test_type: Row_Ct_Pct diff --git a/testgen/template/dbsetup_test_types/test_types_Schema_Drift.yaml b/testgen/template/dbsetup_test_types/test_types_Schema_Drift.yaml index 4992ba2c..d492f211 100644 --- a/testgen/template/dbsetup_test_types/test_types_Schema_Drift.yaml +++ b/testgen/template/dbsetup_test_types/test_types_Schema_Drift.yaml @@ -32,6 +32,7 @@ test_types: usage_notes: |- This test compares the current table column types with previous data, to check whether the table schema has changed. This test allows you to track any changes to the table structure. active: Y + algorithm: Schema / metadata cat_test_conditions: [] target_data_lookups: [] test_templates: diff --git a/testgen/template/dbsetup_test_types/test_types_Street_Addr_Pattern.yaml b/testgen/template/dbsetup_test_types/test_types_Street_Addr_Pattern.yaml index 36b83009..ff5e2025 100644 --- a/testgen/template/dbsetup_test_types/test_types_Street_Addr_Pattern.yaml +++ b/testgen/template/dbsetup_test_types/test_types_Street_Addr_Pattern.yaml @@ -37,6 +37,7 @@ test_types: usage_notes: |- The street address pattern used in this test should match the vast majority of USA addresses. You can adjust the threshold percent of matches based on the results you are getting -- you may well want to tighten it to make the test more sensitive to invalid entries. active: Y + algorithm: Pattern / regex cat_test_conditions: - id: '7026' test_type: Street_Addr_Pattern diff --git a/testgen/template/dbsetup_test_types/test_types_Table_Freshness.yaml b/testgen/template/dbsetup_test_types/test_types_Table_Freshness.yaml index d81e834e..11658f2a 100644 --- a/testgen/template/dbsetup_test_types/test_types_Table_Freshness.yaml +++ b/testgen/template/dbsetup_test_types/test_types_Table_Freshness.yaml @@ -29,7 +29,7 @@ test_types: test_scope: table dq_dimension: Recency impact_dimension: Reliability - health_dimension: Recency + health_dimension: Freshness threshold_description: |- Most recent prior table fingerprint result_visualization: binary_chart @@ -37,6 +37,7 @@ test_types: usage_notes: |- This test compares the current table fingerprint, calculated signature of column contents, to confirm that the table has been updated. The table fingerprint is derived from a set of values and aggregates from columns most likely to change. This test allows you to track the schedule and frequency of updates and refreshes to the table. active: Y + algorithm: Freshness / time cat_test_conditions: [] target_data_lookups: [] test_templates: diff --git a/testgen/template/dbsetup_test_types/test_types_Timeframe_Combo_Gain.yaml b/testgen/template/dbsetup_test_types/test_types_Timeframe_Combo_Gain.yaml index 34329e26..55273486 100644 --- a/testgen/template/dbsetup_test_types/test_types_Timeframe_Combo_Gain.yaml +++ b/testgen/template/dbsetup_test_types/test_types_Timeframe_Combo_Gain.yaml @@ -38,6 +38,7 @@ test_types: usage_notes: |- This test checks a single transactional table to verify that categorical values or combinations that are present in the most recent time window you define include at least all those found in the prior time window of the same duration. Missing values in the latest time window will trigger the test to fail. New values are permitted. Use this test to confirm that codes or categories are not lost across successive time periods in a transactional table. active: Y + algorithm: Aggregate reconciliation cat_test_conditions: [] target_data_lookups: - id: '1406' diff --git a/testgen/template/dbsetup_test_types/test_types_Timeframe_Combo_Match.yaml b/testgen/template/dbsetup_test_types/test_types_Timeframe_Combo_Match.yaml index 6b10231d..c378cfd3 100644 --- a/testgen/template/dbsetup_test_types/test_types_Timeframe_Combo_Match.yaml +++ b/testgen/template/dbsetup_test_types/test_types_Timeframe_Combo_Match.yaml @@ -36,6 +36,7 @@ test_types: usage_notes: |- This test checks a single transactional table (such as a fact table) to verify that categorical values or combinations that are present in the most recent time window you define match those found in the prior time window of the same duration. New or missing values in the latest time window will trigger the test to fail. Use this test to confirm the consistency in the occurrence of codes or categories across successive time periods in a transactional table. active: Y + algorithm: Aggregate reconciliation cat_test_conditions: [] target_data_lookups: - id: '1407' diff --git a/testgen/template/dbsetup_test_types/test_types_US_State.yaml b/testgen/template/dbsetup_test_types/test_types_US_State.yaml index 397611df..d936d0f8 100644 --- a/testgen/template/dbsetup_test_types/test_types_US_State.yaml +++ b/testgen/template/dbsetup_test_types/test_types_US_State.yaml @@ -37,6 +37,7 @@ test_types: usage_notes: |- This test validates entries against a fixed list of two-character US state codes and related Armed Forces codes. active: Y + algorithm: Set / lookup cat_test_conditions: - id: '7027' test_type: US_State diff --git a/testgen/template/dbsetup_test_types/test_types_Unique.yaml b/testgen/template/dbsetup_test_types/test_types_Unique.yaml index e1b5b661..6dee77c4 100644 --- a/testgen/template/dbsetup_test_types/test_types_Unique.yaml +++ b/testgen/template/dbsetup_test_types/test_types_Unique.yaml @@ -37,6 +37,7 @@ test_types: usage_notes: |- This test is ideal when the database itself does not enforce a primary key constraint on the table. It serves as an independent check on uniqueness. If's also useful when there are a small number of exceptions to uniqueness, which can be reflected in the expected threshold count of duplicates. active: Y + algorithm: Counting cat_test_conditions: - id: '7028' test_type: Unique diff --git a/testgen/template/dbsetup_test_types/test_types_Unique_Pct.yaml b/testgen/template/dbsetup_test_types/test_types_Unique_Pct.yaml index 6cb5908a..d3ca67bc 100644 --- a/testgen/template/dbsetup_test_types/test_types_Unique_Pct.yaml +++ b/testgen/template/dbsetup_test_types/test_types_Unique_Pct.yaml @@ -37,6 +37,8 @@ test_types: usage_notes: |- You can think of this as a test of similarity that measures whether the percentage of unique values is consistent with the percentage at baseline. A significant change might indicate duplication or a telling shift in cardinality between entities. The test uses Cohen's H, a statistical test to identify a significant difference between two ratios. Results are reported on a standardized scale, which can be interpreted via a rule-of-thumb from small to huge. You can refine the expected threshold value as you view legitimate results of the measure over time. active: Y + algorithm: Statistical drift + statistical_technique: Cohen's H cat_test_conditions: - id: '7029' test_type: Unique_Pct diff --git a/testgen/template/dbsetup_test_types/test_types_Valid_Characters.yaml b/testgen/template/dbsetup_test_types/test_types_Valid_Characters.yaml index cd73a08c..c3437dcf 100644 --- a/testgen/template/dbsetup_test_types/test_types_Valid_Characters.yaml +++ b/testgen/template/dbsetup_test_types/test_types_Valid_Characters.yaml @@ -37,6 +37,7 @@ test_types: usage_notes: |- This test looks for the presence of non-printing ASCII characters that are considered non-standard in basic text processing. It also identifies leading spaces and values enclosed in quotes. Values that fail this test may be artifacts of data conversion, or just more difficult to process or analyze downstream. active: N + algorithm: Pattern / regex cat_test_conditions: - id: '7036' test_type: Valid_Characters diff --git a/testgen/template/dbsetup_test_types/test_types_Valid_Month.yaml b/testgen/template/dbsetup_test_types/test_types_Valid_Month.yaml index fab14ff1..f1e46ae8 100644 --- a/testgen/template/dbsetup_test_types/test_types_Valid_Month.yaml +++ b/testgen/template/dbsetup_test_types/test_types_Valid_Month.yaml @@ -36,6 +36,7 @@ test_types: result_visualization_params: null usage_notes: null active: N + algorithm: Pattern / regex cat_test_conditions: - id: '7033' test_type: Valid_Month diff --git a/testgen/template/dbsetup_test_types/test_types_Valid_US_Zip.yaml b/testgen/template/dbsetup_test_types/test_types_Valid_US_Zip.yaml index e380caef..2aad44e1 100644 --- a/testgen/template/dbsetup_test_types/test_types_Valid_US_Zip.yaml +++ b/testgen/template/dbsetup_test_types/test_types_Valid_US_Zip.yaml @@ -35,6 +35,7 @@ test_types: result_visualization_params: null usage_notes: null active: Y + algorithm: Pattern / regex cat_test_conditions: - id: '7034' test_type: Valid_US_Zip diff --git a/testgen/template/dbsetup_test_types/test_types_Valid_US_Zip3.yaml b/testgen/template/dbsetup_test_types/test_types_Valid_US_Zip3.yaml index 45218af9..af148a48 100644 --- a/testgen/template/dbsetup_test_types/test_types_Valid_US_Zip3.yaml +++ b/testgen/template/dbsetup_test_types/test_types_Valid_US_Zip3.yaml @@ -36,6 +36,7 @@ test_types: usage_notes: |- This test looks for the presence of values that fail to match the three-digit numeric code expected for US Zip Code regional prefixes. These prefixes are often used to roll up Zip Code data to a regional level, and may be critical to anonymize detailed data and protect PID. Depending on your needs and regulatory requirements, longer zip codes could place PID at risk. active: Y + algorithm: Pattern / regex cat_test_conditions: - id: '7035' test_type: Valid_US_Zip3 diff --git a/testgen/template/dbsetup_test_types/test_types_Variability_Decrease.yaml b/testgen/template/dbsetup_test_types/test_types_Variability_Decrease.yaml index bb671fd8..4ed8791f 100644 --- a/testgen/template/dbsetup_test_types/test_types_Variability_Decrease.yaml +++ b/testgen/template/dbsetup_test_types/test_types_Variability_Decrease.yaml @@ -41,6 +41,8 @@ test_types: \ process, better precision in measurement, the elimination of outliers, or a\ \ more homogeneous cohort. " active: Y + algorithm: Statistical drift + statistical_technique: SD Shift cat_test_conditions: - id: '7032' test_type: Variability_Decrease diff --git a/testgen/template/dbsetup_test_types/test_types_Variability_Increase.yaml b/testgen/template/dbsetup_test_types/test_types_Variability_Increase.yaml index 54e11245..a0072b37 100644 --- a/testgen/template/dbsetup_test_types/test_types_Variability_Increase.yaml +++ b/testgen/template/dbsetup_test_types/test_types_Variability_Increase.yaml @@ -45,6 +45,8 @@ test_types: \ that should be noted and assessed by business users. If the average does not\ \ shift, this may point to a data quality or data collection problem. " active: Y + algorithm: Statistical drift + statistical_technique: SD Shift cat_test_conditions: - id: '7031' test_type: Variability_Increase diff --git a/testgen/template/dbsetup_test_types/test_types_Volume_Trend.yaml b/testgen/template/dbsetup_test_types/test_types_Volume_Trend.yaml index 67c9ff29..0565a1e3 100644 --- a/testgen/template/dbsetup_test_types/test_types_Volume_Trend.yaml +++ b/testgen/template/dbsetup_test_types/test_types_Volume_Trend.yaml @@ -27,6 +27,7 @@ test_types: dq_dimension: Completeness impact_dimension: Reliability health_dimension: Volume + statistical_technique: Predictive Model threshold_description: |- Expected row count range. result_visualization: line_chart @@ -34,6 +35,7 @@ test_types: usage_notes: |- This test compares the row count of all or a subset of records in a table against a derived tolerance range. active: Y + algorithm: Custom SQL cat_test_conditions: - id: '2515' test_type: Volume_Trend diff --git a/testgen/template/dbsetup_test_types/test_types_Weekly_Rec_Ct.yaml b/testgen/template/dbsetup_test_types/test_types_Weekly_Rec_Ct.yaml index bf0e91df..74be242d 100644 --- a/testgen/template/dbsetup_test_types/test_types_Weekly_Rec_Ct.yaml +++ b/testgen/template/dbsetup_test_types/test_types_Weekly_Rec_Ct.yaml @@ -37,6 +37,7 @@ test_types: usage_notes: |- Weekly Records tests that at least one record is present for every calendar week within the minimum and maximum date range for the column. The test is relevant for transactional data, where you would expect at least one transaction to be recorded each week. A failure here would suggest missing records for the number of weeks identified without data. You can adjust the threshold to accept a number of weeks that you know legitimately have no records. active: Y + algorithm: Counting cat_test_conditions: - id: '7030' test_type: Weekly_Rec_Ct diff --git a/testgen/template/dbupgrade/0196_incremental_upgrade.sql b/testgen/template/dbupgrade/0196_incremental_upgrade.sql new file mode 100644 index 00000000..7b712aec --- /dev/null +++ b/testgen/template/dbupgrade/0196_incremental_upgrade.sql @@ -0,0 +1,5 @@ +SET SEARCH_PATH TO {SCHEMA_NAME}; + +ALTER TABLE test_types + ADD COLUMN IF NOT EXISTS algorithm VARCHAR(64), + ADD COLUMN IF NOT EXISTS statistical_technique VARCHAR(64); diff --git a/testgen/ui/components/frontend/js/pages/test_definitions.js b/testgen/ui/components/frontend/js/pages/test_definitions.js index 7e83fbcc..2132b56f 100644 --- a/testgen/ui/components/frontend/js/pages/test_definitions.js +++ b/testgen/ui/components/frontend/js/pages/test_definitions.js @@ -17,8 +17,12 @@ import { TestDefinitionNotes } from './test_definition_notes.js'; import { withTooltip } from '/app/static/js/components/tooltip.js'; import { Icon } from '/app/static/js/components/icon.js'; import { ProfilingResultsDialog } from '../shared/profiling_results_dialog.js'; +import { AXES, FACET_AXES, GROUP_BY_AXES, EMPTY, appliesToSelectedColumn } from '/app/static/js/components/test_picker_taxonomy.js'; +import { enterPage, exitPage, getPageSignal } from '/app/static/js/page_lifecycle.js'; -const { button: btn, div, i: icon, span, strong, input, label } = van.tags; +const { button: btn, div, i: icon, span, strong } = van.tags; + +const PAGE_KEY = 'testDefinitions'; const TABLE_COLUMNS = [ { name: 'table_name', label: 'Table', width: 180, sortable: true, overflow: 'hidden' }, @@ -40,8 +44,6 @@ const SEVERITY_OPTIONS = [ { label: 'Fail', value: 'Fail' }, ]; -const SCOPE_LABELS = { referential: 'Referential', table: 'Table', column: 'Column', custom: 'Custom' }; - // Blank test definition field defaults for add mode const BLANK_PARAM_FIELDS = { custom_query: null, @@ -826,7 +828,7 @@ const DetailPanel = (row) => { Attribute({ label: 'Schema Name', value: row.schema_name }), Attribute({ label: 'Table Name', value: row.table_name }), Attribute({ label: 'Test Focus', value: row.column_name }), - Attribute({ label: 'Test Type', value: row.test_type }), + Attribute({ label: 'Test Type', value: row.test_name_short }), Attribute({ label: 'Test Active', value: row.test_active_display }), Attribute({ label: 'Validation Status', value: row.test_definition_status }), Attribute({ label: 'Lock Refresh', value: row.lock_refresh_display }), @@ -847,32 +849,62 @@ const DetailPanel = (row) => { ); }; -// Add dialog — mounted once, state persists across Python reruns -const AddDialogComponent = ({ open, info, validateResult: validateResultProp, onClose }, emit) => { +const TestPickerChip = (text, color) => { + const { span } = van.tags; + return span( + { + class: 'tg-test-chip', + style: `border:1px solid color-mix(in srgb, ${color} 25%, transparent);color:${color};`, + }, + text, + ); +}; + +// A keyboard-shortcut hint for the picker footer: a key chip followed by its action. +const KeyHint = (key, action) => { + const { span } = van.tags; + return span( + { class: 'tg-key-hint' }, + span({ class: 'tg-kbd' }, key), + action, + ); +}; + +// Faceted add-test picker dialog (step 1) reusing the param form (step 2) +const AddDialogComponent = ({ open, info, validateResult, onClose }, emit) => { + const { div, span, input, label, h4 } = van.tags; + const testTypes = van.derive(() => getValue(info)?.test_types ?? []); + const tableColumns = van.derive(() => getValue(info)?.table_columns ?? []); + const testSuite = van.derive(() => getValue(info)?.test_suite ?? {}); const tableGroupSchema = van.derive(() => getValue(info)?.table_group_schema ?? ''); const tableGroupsId = van.derive(() => getValue(info)?.table_groups_id ?? ''); - const testSuite = van.derive(() => getValue(info)?.test_suite ?? {}); - const tableColumns = van.derive(() => getValue(info)?.table_columns ?? []); const qualifiesTableRefsWithSchema = van.derive(() => getValue(info)?.qualifies_table_refs_with_schema ?? true); - const validateResult = van.derive(() => getValue(validateResultProp) ?? null); + const prefillColumn = van.derive(() => getValue(info)?.prefill_column ?? null); + + // selectedColumn carries the general_type needed by the type-aware filter; resolve it from + // the tableColumns entry so a prefill and a manual pick produce the identical shape. + const columnFromEntry = (c) => ({ + table_name: c.table_name, + column_name: c.column_name, + general_type: c.general_type ?? null, + }); - const scopeFilter = { - referential: van.state(true), - table: van.state(true), - column: van.state(true), - custom: van.state(true), - }; + // ---- Step + selection state ---- + const step = van.state(1); + const selectedTestType = van.state(null); + const formValues = van.state({}); - const filteredTestTypeOptions = van.derive(() => - testTypes.val - .filter(tt => tt.test_scope !== 'tablegroup' && (scopeFilter[tt.test_scope]?.val ?? true)) - .map(tt => ({ label: tt.select_name ?? tt.test_name_short, value: tt.test_type })) - ); + // ---- Picker state ---- + const searchQuery = van.state(''); + const groupBy = van.state('impact'); + const focusIndex = van.state(-1); // -1 = no row highlighted until keyboard nav + const selectedColumn = van.state(null); // { table_name, column_name, general_type } | null - const selectedTestType = van.state(null); - const formValues = van.state(null); + const facetSel = {}; + Object.keys(AXES).forEach((ax) => { facetSel[ax] = van.state([]); }); + // Build blank form values for the selected test type. const buildFormValues = (testType) => { if (!testType) return null; const tt = testTypes.rawVal.find(t => t.test_type === testType); @@ -900,80 +932,378 @@ const AddDialogComponent = ({ open, info, validateResult: validateResultProp, on }; }; - const selectTestType = (testType) => { - selectedTestType.val = testType; - formValues.val = buildFormValues(testType); + const toggleFacet = (ax, value) => { + const cur = facetSel[ax].val; + facetSel[ax].val = cur.includes(value) ? cur.filter((v) => v !== value) : [...cur, value]; + focusIndex.val = -1; + }; + const clearAll = () => { + Object.values(facetSel).forEach((s) => { s.val = []; }); + searchQuery.val = ''; + focusIndex.val = -1; }; - // Reset form state when dialog opens (transitions from closed→open) + // Reset to a fresh step-1 picker on each closed->open transition. The dialog is mounted + // once and its state otherwise persists across reopens, which would show the stale step-2 + // form. Also re-seeds the locked column per open. const wasOpen = van.state(false); van.derive(() => { const isOpen = open.val; if (isOpen && !wasOpen.val) { - selectTestType(null); + step.val = 1; + selectedTestType.val = null; + formValues.val = {}; + searchQuery.val = ''; + focusIndex.val = -1; + groupBy.val = 'impact'; + Object.values(facetSel).forEach((s) => { s.val = []; }); + const prefill = prefillColumn.rawVal; + const match = prefill + ? tableColumns.rawVal.find((c) => c.table_name === prefill.table_name && c.column_name === prefill.column_name) + : null; + selectedColumn.val = match ? columnFromEntry(match) : null; wasOpen.val = true; } else if (!isOpen) { wasOpen.val = false; } }); - return Dialog( - { title: 'Add Test', open, onClose, width: '52rem' }, - div( - { class: 'flex-column fx-gap-4 td-form-dialog' }, + const matchesSearch = (t, q) => { + if (!q) return true; + const hay = `${t.test_name_short} ${t.test_name_long} ${t.test_description}`.toLowerCase(); + return q.toLowerCase().split(/\s+/).filter(Boolean).every((tok) => hay.includes(tok)); + }; - // Test type picker — always visible - div( - { class: 'flex-column fx-gap-3' }, + // Search + column relevance only -- drives facet counts (PRD F9). A selected column filters + // the list to tests applicable to its type; clearing the column shows all tests. + const baseVisible = van.derive(() => { + const q = searchQuery.val; + const col = selectedColumn.val; + return testTypes.val.filter((t) => + matchesSearch(t, q) && (!col || appliesToSelectedColumn(t, col.general_type))); + }); + + const passesFacets = (t) => Object.entries(facetSel).every(([ax, s]) => { + const sel = s.val; + if (!sel.length) return true; + const v = AXES[ax].value(t); + return v != null && sel.includes(v); + }); + + const filtered = van.derive(() => baseVisible.val.filter(passesFacets)); + + // Group the result list by the chosen axis; null -> EMPTY bucket. + const grouped = van.derive(() => { + const axis = AXES[groupBy.val]; + const buckets = new Map(); + filtered.val.forEach((t) => { + const key = axis.value(t) || EMPTY; + if (!buckets.has(key)) buckets.set(key, []); + buckets.get(key).push(t); + }); + const order = axis.order; // sparse axes have a canonical order + const entries = [...buckets.entries()]; + entries.sort((a, b) => { + if (order) { + const ia = order.indexOf(a[0]); const ib = order.indexOf(b[0]); + return (ia < 0 ? Infinity : ia) - (ib < 0 ? Infinity : ib); // EMPTY bucket sorts last + } + return b[1].length - a[1].length; + }); + return entries; + }); + + const flatList = van.derive(() => grouped.val.flatMap(([, tests]) => tests)); + + // Per-axis value -> count, over baseVisible (not other facet selections). + const counts = (ax) => { + const m = new Map(); + baseVisible.val.forEach((t) => { + const v = AXES[ax].value(t); + if (v != null) m.set(v, (m.get(v) || 0) + 1); + }); + return m; + }; + + const selectTest = (t) => { + if (!t) return; + selectedTestType.val = t; + const fv = buildFormValues(t.test_type); + const col = selectedColumn.rawVal; + if (col) { + const scope = fv.test_scope ?? 'column'; + if (scope !== 'tablegroup') { + fv.table_name = col.table_name; + } + if (scope === 'column' || scope === 'referential' || scope === 'custom') { + fv.column_name = col.column_name; + } + } + formValues.val = fv; + step.val = 2; + }; + + // ---- Keyboard: Cmd/Ctrl+K or "/" focus search, Up/Down move, Enter add, Esc close ---- + let searchEl = null; + const onKeyDown = (e) => { + if (step.val !== 1) return; + if ((e.key === 'k' && (e.metaKey || e.ctrlKey)) || (e.key === '/' && document.activeElement !== searchEl)) { + e.preventDefault(); searchEl?.focus(); return; + } + if (e.key === 'Escape') { onClose(); return; } + const list = flatList.val; + if (e.key === 'ArrowDown') { e.preventDefault(); focusIndex.val = Math.min(focusIndex.val + 1, Math.max(list.length - 1, 0)); } + else if (e.key === 'ArrowUp') { e.preventDefault(); focusIndex.val = Math.max(focusIndex.val - 1, 0); } + // Enter adds the focused row, or the only match when nothing is explicitly focused. + else if (e.key === 'Enter') { e.preventDefault(); selectTest(list[focusIndex.val] ?? (list.length === 1 ? list[0] : null)); } + }; + + // Listen at the document level while step 1 is open so arrow keys work without first + // clicking the picker. Attach/detach reactively; guard prevents duplicate registration. + // The page's AbortSignal removes the listener on teardown — without it, a teardown + // while the dialog is open (e.g. browser Back) would orphan onKeyDown on the document, + // since nothing toggles open/step on unmount to run the detach branch below. + let keydownAttached = false; + van.derive(() => { + const active = open.val && step.val === 1; + if (active && !keydownAttached) { + document.addEventListener('keydown', onKeyDown, { signal: getPageSignal(PAGE_KEY) ?? undefined }); + keydownAttached = true; + } else if (!active && keydownAttached) { + document.removeEventListener('keydown', onKeyDown); + keydownAttached = false; + } + }); + + // ---- Renderers ---- + const FacetGroup = (ax) => { + const axis = AXES[ax]; + // Whole group is reactive: its title is hidden when no option has a count (PRD facet review). + return () => { + const m = counts(ax); + let keys = [...m.keys()]; + if (axis.order) keys = axis.order.filter((k) => m.has(k)); + else keys.sort((a, b) => m.get(b) - m.get(a)); + if (!keys.length) return ''; + return div( + { class: 'tg-facet-group', 'data-testid': 'facet-group', 'data-axis': ax }, + h4({ class: 'tg-facet-title' }, axis.label), div( - { class: 'flex-row fx-gap-4 fx-align-flex-center fx-flex-wrap' }, - span({ class: 'text-caption' }, 'Show Types:'), - ...Object.entries(SCOPE_LABELS).map(([scope, scopeLabel]) => - Checkbox({ - label: scopeLabel, - checked: scopeFilter[scope], - onChange: (v) => { scopeFilter[scope].val = v; }, - }) - ), + { class: 'flex-column fx-gap-1' }, + ...keys.map((value) => { + const checked = van.derive(() => facetSel[ax].val.includes(value)); + // Display-only Checkbox (pointer-events disabled via .tg-facet-checkbox); the + // row's onclick is the single toggle source, so the box can't double-fire. + return div( + { + class: 'tg-facet-option', + 'data-testid': 'facet-option', + 'data-axis': ax, + 'data-value': value, + onclick: () => toggleFacet(ax, value), + }, + span({ class: 'tg-facet-checkbox' }, Checkbox({ label: '', checked })), + span({ class: 'tg-facet-dot', style: `background:${axis.color(value)}` }), + span({ class: 'tg-facet-label' }, value), + span({ class: 'tg-facet-count' }, m.get(value)), + ); + }), ), - () => Select({ - label: 'Test Type', - value: selectedTestType.val, - options: filteredTestTypeOptions.val, - allowNull: true, - filterable: true, - onChange: (value) => { selectTestType(value); }, - }), - ), + ); + }; + }; - // Form (shown after test type selected) — imperative update - // because VanJS binding replacement doesn't work inside Dialog portals - () => { - open.val; + const ResultRow = (t, isFocused) => { + const row = div( + { + class: () => `tg-test-row${isFocused.val ? ' tg-test-row--focused' : ''}`, + 'data-testid': 'test-picker-row', + 'data-test-type': t.test_type, + onclick: () => selectTest(t), + }, + div( + { class: 'flex-row fx-justify-space-between' }, + div({ class: 'tg-test-row-title' }, t.test_name_short), + ), + div({ class: 'tg-test-desc' }, t.test_description), + div( + { class: 'flex-row fx-flex-wrap fx-gap-1' }, + t.impact_dimension ? TestPickerChip(t.impact_dimension, AXES.impact.color(t.impact_dimension)) : null, + t.algorithm ? TestPickerChip(t.algorithm, AXES.algorithm.color(t.algorithm)) : null, + t.statistical_technique ? TestPickerChip(t.statistical_technique, AXES.technique.color(t.statistical_technique)) : null, + t.health_dimension ? TestPickerChip(t.health_dimension, AXES.health.color(t.health_dimension)) : null, + t.criteria ? TestPickerChip(t.criteria, AXES.criteria.color(t.criteria)) : null, + ), + ); + // Keep the keyboard-focused row visible within the scrolling list. + van.derive(() => { if (isFocused.val) row.scrollIntoView({ block: 'nearest' }); }); + return row; + }; - selectedTestType.val; - const fv = formValues.val; - const vr = validateResult.val; + // One filter row: column relevance + search on the left, group-by pushed to the right. + const ContextBar = (searchField) => div( + { class: 'tg-picker-context flex-row fx-gap-3 fx-align-flex-end' }, + () => { + // Column relevance filter. Option VALUES are the array index (string) so that + // table/column names containing spaces or dots can never break parsing. + const cols = tableColumns.val; + const options = cols.map((c, i) => ({ label: `${c.table_name}.${c.column_name}`, value: String(i) })); + const sel = selectedColumn.val; + const currentIdx = sel ? cols.findIndex((c) => c.table_name === sel.table_name && c.column_name === sel.column_name) : -1; + return Select({ + label: 'Show tests relevant to a column', + allowNull: true, + // Cap growth so a long table.column value ellipsizes instead of crowding search. + style: 'max-width: 320px;', + value: currentIdx >= 0 ? String(currentIdx) : null, + options, + onChange: (val) => { + if (val == null || val === '') { selectedColumn.val = null; return; } + selectedColumn.val = columnFromEntry(tableColumns.rawVal[Number(val)]); + }, + }); + }, + div({ class: 'tg-picker-search fx-flex', 'data-testid': 'test-picker-search' }, searchField), + Select({ + label: 'Group by', + value: groupBy.val, + options: GROUP_BY_AXES.map((ax) => ({ label: AXES[ax].label, value: ax })), + onChange: (val) => { groupBy.val = val; focusIndex.val = -1; }, + }), + ); - if (!fv) return ''; + const ActiveFilters = () => div( + { class: 'flex-row fx-flex-wrap fx-gap-1', 'data-testid': 'test-picker-active-filters' }, + () => { + const chips = []; + Object.entries(facetSel).forEach(([ax, s]) => s.val.forEach((value) => { + chips.push(span( + { class: 'tg-active-chip flex-row', onclick: () => toggleFacet(ax, value) }, + span(`${AXES[ax].label}: ${value}`), + Icon({ size: 14 }, 'close'), + )); + })); + if (!chips.length) return ''; + chips.push(span( + { 'data-testid': 'test-picker-clear' }, + Button({ type: 'basic', label: 'Clear all', onclick: clearAll }), + )); + return div({ class: 'flex-row fx-flex-wrap fx-gap-1' }, ...chips); + }, + ); - return TestDefFormContent({ - formValues: fv, - tableColumns: tableColumns.rawVal, - testSuite: testSuite.rawVal, - qualifiesTableRefsWithSchema: qualifiesTableRefsWithSchema.rawVal, - validateResult: vr, - mode: 'add', - onFormChange: (changes) => { - formValues.val = { ...formValues.rawVal, ...changes }; - }, - onValidate: () => emit('ValidateTest', { payload: formValues.rawVal }), - onSave: () => emit('AddTestSaved', { payload: formValues.rawVal }), - onCancel: onClose, - }); + const ResultsPane = () => div( + { class: 'tg-picker-results' }, + // Fixed header — count stays put while the list below scrolls. + div( + { class: 'tg-picker-results-header' }, + span( + { 'data-testid': 'test-picker-count', class: 'tg-picker-count' }, + () => `${filtered.val.length} matching test types${!!selectedColumn.val ? ' for selected column' : ''}` + ), + ), + div( + { class: 'tg-picker-list' }, + // One-result hint (PRD F13): when exactly one test matches, prompt Enter-to-add. + () => filtered.val.length === 1 + ? div({ class: 'tg-one-result-hint mb-2', 'data-testid': 'test-picker-one-hint' }, + `${filtered.val[0].test_name_short} — press Enter to add`) + : '', + () => { + const groups = grouped.val; + if (!groups.length) { + return div( + { class: 'tg-picker-empty', 'data-testid': 'test-picker-empty' }, + span('No tests match this filter combination.'), + Button({ type: 'stroked', label: 'Clear filters and search', onclick: clearAll }), + ); + } + let runningIndex = -1; + return div( + { class: 'flex-column fx-gap-2' }, + ...groups.map(([groupName, tests]) => div( + { class: 'tg-result-group' }, + div( + { class: 'tg-group-band' }, + span({ class: 'tg-group-band-name' }, groupName), + span({ class: 'tg-group-band-count' }, `${tests.length} test types`), + ), + ...tests.map((t) => { + runningIndex += 1; + const myIndex = runningIndex; + const isFocused = van.derive(() => focusIndex.val === myIndex); + return ResultRow(t, isFocused); + }), + )), + ); }, ), ); + + const PickerView = () => { + // Reusable Input component for search. The keyboard handler focuses the inner , + // resolved by id below (Input renders its label with this id). + const searchField = Input({ + id: 'tg-picker-search', + icon: 'search', + clearable: true, + style: 'width: 100%;', + placeholder: 'Search name or description', + value: searchQuery, + onChange: (value) => { searchQuery.val = value; focusIndex.val = -1; }, + }); + // Load-bearing on every step-1 (re)render, not just first open: refocuses search when + // returning from the step-2 form. Do not hoist out of PickerView. + requestAnimationFrame(() => { + searchEl = document.getElementById('tg-picker-search')?.querySelector('input') ?? null; + searchEl?.focus(); + }); + + return div( + { class: 'tg-test-picker', 'data-testid': 'test-picker' }, + // Fixed header: column relevance + search + group-by in one row, then active filters. + div( + { class: 'tg-picker-header' }, + ContextBar(searchField), + ActiveFilters(), + ), + div( + { class: 'tg-picker-body fx-gap-3' }, + div( + { class: 'tg-picker-rail' }, + ...FACET_AXES.map(FacetGroup), + ), + ResultsPane(), + ), + div( + { class: 'tg-picker-footer', 'data-testid': 'test-picker-footer' }, + KeyHint('↑ ↓', 'navigate'), + KeyHint('Enter', 'select'), + KeyHint('/', 'focus search'), + ), + ); + }; + + const FormView = () => TestDefFormContent({ + formValues: formValues.val, + tableColumns: tableColumns.rawVal, + testSuite: testSuite.rawVal, + qualifiesTableRefsWithSchema: qualifiesTableRefsWithSchema.rawVal, + validateResult: getValue(validateResult), + mode: 'add', + onFormChange: (changes) => { formValues.val = { ...formValues.rawVal, ...changes }; }, + onValidate: () => emit('ValidateTest', { payload: formValues.rawVal }), + onSave: () => emit('AddTestSaved', { payload: formValues.rawVal }), + onBack: () => { step.val = 1; selectedTestType.val = null; }, + onCancel: onClose, + }); + + const dialogTitle = van.derive(() => `Add Test - ${step.val === 1 ? 'Pick a Test Type' : 'Configure Test'}`); + + return Dialog( + { title: dialogTitle, open, onClose, width: '68rem' }, + () => step.val === 1 ? PickerView() : FormView(), + ); }; // Edit dialog — mounted once, state persists across Python reruns @@ -1041,7 +1371,7 @@ const EditDialogComponent = ({ open, info, validateResult: validateResultProp, o }; // Shared form content for add/edit dialogs -const TestDefFormContent = ({ formValues, tableColumns, testSuite, validateResult, mode, qualifiesTableRefsWithSchema, onFormChange, onValidate, onSave, onCancel }) => { +const TestDefFormContent = ({ formValues, tableColumns, testSuite, validateResult, mode, qualifiesTableRefsWithSchema, onFormChange, onValidate, onSave, onCancel, onBack }) => { const testScope = formValues.test_scope ?? 'column'; const runType = formValues.run_type ?? 'CAT'; const testType = formValues.test_type ?? ''; @@ -1278,15 +1608,28 @@ const TestDefFormContent = ({ formValues, tableColumns, testSuite, validateResul // Buttons div( { class: 'flex-row fx-justify-space-between fx-gap-2' }, - isValidatable - ? Button({ - type: 'stroked', - color: 'basic', - label: 'Validate', - width: 'auto', - onclick: onValidate, - }) - : span(''), + div( + { class: 'flex-row fx-gap-2' }, + onBack + ? Button({ + type: 'stroked', + color: 'basic', + icon: 'arrow_back', + label: 'Back', + width: 'auto', + onclick: onBack, + }) + : null, + isValidatable + ? Button({ + type: 'stroked', + color: 'basic', + label: 'Validate', + width: 'auto', + onclick: onValidate, + }) + : null, + ), div( { class: 'flex-row fx-gap-2' }, Button({ @@ -1555,6 +1898,42 @@ stylesheet.replace(` padding-top: 12px; margin-top: 4px; } + +.tg-test-picker { display: flex; flex-direction: column; gap: 12px; height: 70vh; min-height: 0; } +.tg-picker-header { flex: none; display: flex; flex-direction: column; gap: 12px; } +.tg-picker-body { flex: 1; min-height: 0; display: flex; flex-direction: row; align-items: stretch; } +.tg-picker-rail { flex: 0 0 260px; min-height: 0; overflow-y: auto; overflow-x: hidden; padding-right: 8px; border-right: 1px solid var(--border-color, #e0e0e0); } +.tg-picker-results { flex: 1; min-height: 0; display: flex; flex-direction: column; } +.tg-picker-results-header { flex: none; padding-bottom: 8px; margin-bottom: 8px; border-bottom: 1px solid var(--border-color, #e0e0e0); } +.tg-picker-count { font-size: 13px; color: var(--secondary-text-color, #666); } +.tg-picker-list { flex: 1; overflow-y: auto; } +.tg-facet-group { margin-bottom: 14px; } +.tg-facet-title { margin: 0 0 6px; font-size: 12px; text-transform: uppercase; letter-spacing: .06em; color: var(--secondary-text-color, #666); } +.tg-facet-option { display: flex; align-items: center; gap: 8px; cursor: pointer; font-size: 13px; color: var(--primary-text-color); } +.tg-facet-dot { width: 8px; height: 8px; border-radius: 50%; flex-shrink: 0; } +.tg-facet-label { color: var(--primary-text-color); } +.tg-facet-count { margin-left: auto; color: var(--secondary-text-color, #888); font-variant-numeric: tabular-nums; } +.tg-facet-checkbox { display: inline-flex; pointer-events: none; } +/* --empty stays distinct from the --table-hover-color row hover so the band reads as a divider. */ +.tg-group-band { position: sticky; top: 0; z-index: 1; display: flex; justify-content: space-between; align-items: center; gap: 8px; padding: 6px 10px; margin-bottom: 4px; border-radius: 6px; background: var(--empty, #e0e0e0); } +.tg-group-band-name { font-weight: 600; } +.tg-group-band-count { font-size: 12px; color: var(--secondary-text-color, #888); font-variant-numeric: tabular-nums; } +.tg-result-group { display: flex; flex-direction: column; gap: 6px; } +.tg-test-row { position: relative; padding: 10px; border: 2px solid transparent; border-radius: 8px; cursor: pointer; } +/* Divider sits in the gap below the row, not on its border, so it never overlaps the focus outline. */ +.tg-result-group .tg-test-row:not(:last-child)::after { content: ''; position: absolute; left: 8px; right: 8px; bottom: -4px; border-bottom: 1px dashed var(--border-color, #dddfe2); } +.tg-test-row:hover { background: var(--table-hover-color, #f3f4f6); } +.tg-test-row--focused { border-color: var(--primary-color, #1976d2); background: var(--table-hover-color, #f3f4f6); } +.tg-test-row-title { font-weight: 600; } +.tg-test-desc { font-size: 13px; color: var(--secondary-text-color, #555); margin: 2px 0 6px; } +.tg-test-chip { font-size: 11px; padding: 1px 7px; border-radius: 999px; white-space: nowrap; } +.tg-active-chip { font-size: 12px; line-height: 1; gap: 3px; padding: 2px 6px 2px 8px; border-radius: 999px; background: var(--table-hover-color, #eee); cursor: pointer; } +.tg-active-chip .material-symbols-rounded { cursor: pointer; } +.tg-one-result-hint { padding: 8px 10px; border-radius: 6px; background: var(--table-hover-color, #eef2ff); font-size: 13px; } +.tg-picker-empty { display: flex; flex-direction: column; gap: 12px; align-items: flex-start; padding: 40px 8px; } +.tg-picker-footer { flex: none; display: flex; gap: 16px; align-items: center; padding-top: 10px; border-top: 1px solid var(--border-color, #e0e0e0); font-size: 12px; color: var(--secondary-text-color, #666); } +.tg-key-hint { display: inline-flex; align-items: center; gap: 6px; } +.tg-kbd { display: inline-flex; align-items: center; min-width: 18px; justify-content: center; padding: 1px 6px; border: 1px solid var(--border-color, #d0d0d0); border-radius: 4px; background: var(--empty, #e0e0e0); font-family: monospace; font-size: 11px; color: var(--primary-text-color); } `); export { TestDefinitions, EditDialogComponent }; @@ -1570,6 +1949,7 @@ export default (component) => { } parentElement.state = componentState; componentState.emit = createEmitter(setTriggerValue); + componentState.signal = enterPage(PAGE_KEY); van.add(parentElement, TestDefinitions(componentState)); } else { for (const [key, value] of Object.entries(data)) { @@ -1579,5 +1959,8 @@ export default (component) => { } } - return () => { parentElement.state = null; }; + return () => { + exitPage(PAGE_KEY); + parentElement.state = null; + }; }; diff --git a/testgen/ui/static/css/shared.css b/testgen/ui/static/css/shared.css index 4d5b1f0b..4a0ff545 100644 --- a/testgen/ui/static/css/shared.css +++ b/testgen/ui/static/css/shared.css @@ -26,6 +26,20 @@ body { --empty-dark: #BDBDBD; --empty-teal: #E7F1F0; + /* Faceted test-picker taxonomy palette (test_picker_taxonomy.js). Deep, saturated hues tuned + for light backgrounds; overridden with lighter, higher-contrast variants in dark mode below. */ + --facet-crimson: #9F1239; + --facet-blue: #1E40AF; + --facet-amber: #B45309; + --facet-purple: #6B21A8; + --facet-rust: #7C2D12; + --facet-teal: #0E5F4F; + --facet-bronze: #92400E; + --facet-indigo: #312E81; + --facet-orange: #C75D2C; + --facet-emerald: #1B7C68; + --facet-grey: #6A6F7E; + --primary-text-color: #000000de; --secondary-text-color: #0000008a; --disabled-text-color: #00000042; @@ -97,6 +111,19 @@ body { --empty-dark: #757575; --empty-teal: #242E2D; + /* Lighter, higher-contrast facet hues for dark backgrounds (chip text + facet dots). */ + --facet-crimson: #FB7185; + --facet-blue: #60A5FA; + --facet-amber: #FBBF24; + --facet-purple: #C084FC; + --facet-rust: #F4845F; + --facet-teal: #2DD4BF; + --facet-bronze: #D9A05B; + --facet-indigo: #A5B4FC; + --facet-orange: #FB923C; + --facet-emerald: #34D399; + --facet-grey: #9CA3AF; + --primary-text-color: rgba(255, 255, 255); --secondary-text-color: rgba(255, 255, 255, .7); --disabled-text-color: rgba(255, 255, 255, .5); diff --git a/testgen/ui/static/js/components/test_picker_taxonomy.js b/testgen/ui/static/js/components/test_picker_taxonomy.js new file mode 100644 index 00000000..64673b65 --- /dev/null +++ b/testgen/ui/static/js/components/test_picker_taxonomy.js @@ -0,0 +1,108 @@ +// Faceted Add-Test picker taxonomy: colors, axis definitions, and the column-aware pre-filter. +// All seven facets read a column on test_types (no derivation here): `impact_dimension`, +// `dq_dimension`, `health_dimension`, `algorithm`, `statistical_technique`, `test_scope`, and +// `criteria`. `criteria` is derived once in Python (derive_test_criteria) and shared by UI + MCP, +// so this file holds presentation only. Colors reference the --facet-* CSS variables defined in +// shared.css, which carry light values plus higher-contrast dark-mode overrides. + +const IMPACT_COLOR = { + 'Conformance': 'var(--facet-crimson)', 'Reliability': 'var(--facet-blue)', 'Regularity': 'var(--facet-amber)', 'Usability': 'var(--facet-purple)', +}; + +const QD_COLOR = { + 'Validity': 'var(--facet-crimson)', 'Completeness': 'var(--facet-blue)', 'Consistency': 'var(--facet-rust)', 'Accuracy': 'var(--facet-teal)', + 'Timeliness': 'var(--facet-amber)', 'Uniqueness': 'var(--facet-purple)', 'Recency': 'var(--facet-bronze)', +}; + +const HD_COLOR = { + 'Schema Drift': 'var(--facet-indigo)', 'Data Drift': 'var(--facet-orange)', 'Volume': 'var(--facet-teal)', 'Freshness': 'var(--facet-amber)', +}; + +const ALGO_COLOR = { + 'Boundary check': 'var(--facet-rust)', + 'Counting': 'var(--facet-blue)', + 'Pattern / regex': 'var(--facet-crimson)', + 'Set / lookup': 'var(--facet-purple)', + 'Statistical drift': 'var(--facet-orange)', + 'Aggregate reconciliation': 'var(--facet-bronze)', + 'Freshness / time': 'var(--facet-amber)', + 'Schema / metadata': 'var(--facet-indigo)', + 'Custom SQL': 'var(--facet-teal)', +}; + +const TECH_COLOR = { + "Cohen's D": 'var(--facet-blue)', + "Cohen's H": 'var(--facet-purple)', + 'Outlier Detection': 'var(--facet-orange)', + 'SD Shift': 'var(--facet-rust)', + 'Jensen-Shannon Divergence': 'var(--facet-teal)', +}; + +// Canonical order for the sparse Statistical Technique facet. +const TECHNIQUE_ORDER = ["Cohen's D", "Cohen's H", 'Outlier Detection', 'SD Shift', 'Jensen-Shannon Divergence']; + +const CRITERIA_COLOR = { + 'Defined Rule': 'var(--facet-teal)', 'Defined Threshold': 'var(--facet-blue)', 'Defined Value': 'var(--facet-emerald)', + 'List of Values': 'var(--facet-purple)', 'Reference Dataset': 'var(--facet-bronze)', 'Custom Criteria': 'var(--facet-crimson)', +}; +// Canonical order for the Criteria facet (least to most setup effort). +const CRITERIA_ORDER = [ + 'Defined Rule', 'Defined Threshold', 'Defined Value', + 'List of Values', 'Reference Dataset', 'Custom Criteria', +]; + +const SCOPE_LABEL = { + column: 'Column', table: 'Table', referential: 'Referential', custom: 'Custom', tablegroup: 'Table Group', +}; +const SCOPE_ORDER = ['Column', 'Table', 'Referential', 'Custom', 'Table Group']; + +const FALLBACK = 'var(--facet-grey)'; +const EMPTY = '—'; + +// ── Column-aware pre-filter (PRD §9) ────────────────────────────────────────── +// Not a facet: the separate "hide tests that don't apply to this column" toggle. +// `generalType` is data_column_chars.general_type: 'N' Numeric, 'D' Datetime, 'T' Time, +// 'A' Alpha, 'B' Boolean, 'X' Other. Lists enumerate type affinity by test code; gating on +// scope first excludes table/referential-scoped codes regardless. +const NUMERIC_ONLY = ['Min_Val', 'Avg_Shift', 'Incr_Avg_Shift', 'Variability_Increase', 'Variability_Decrease', + 'Outlier_Pct_Above', 'Outlier_Pct_Below', 'Dec_Trunc', 'Distribution_Shift', 'Aggregate_Minimum']; +const DATE_ONLY = ['Min_Date', 'Future_Date', 'Future_Date_1Y', 'Recency', 'Distinct_Date_Ct', + 'Daily_Record_Ct', 'Weekly_Rec_Ct', 'Monthly_Rec_Ct', 'Valid_Month', 'Freshness_Trend', 'Table_Freshness']; + +function appliesToSelectedColumn(t, generalType) { + if (['table', 'tablegroup', 'referential'].includes(t.test_scope)) return false; + // Unknown column type (null/undefined/empty): we can't assert a test is inapplicable, so + // don't hide it. Without this, an absent general_type fails every type check below and + // silently hides all numeric/date tests. Known non-numeric/non-date codes ('A','B','X') + // are truthy and still filter correctly. + if (!generalType) return true; + if (NUMERIC_ONLY.includes(t.test_type) && generalType !== 'N') return false; + const isDate = generalType === 'D' || generalType === 'T'; + if (DATE_ONLY.includes(t.test_type) && !isDate) return false; + return true; +} + +// ── Axis registry ───────────────────────────────────────────────────────────── +// `value(t)` returns the facet value for a test (or null when absent). `color(key)` returns a +// hex color. `sparse` axes filter out null-valued tests when active. `order`, when present, +// fixes the facet's display order. + +const AXES = { + impact: { label: 'Impact Dimension', value: (t) => t.impact_dimension || null, color: (k) => IMPACT_COLOR[k] || FALLBACK }, + dq: { label: 'Quality Dimension', value: (t) => t.dq_dimension || null, color: (k) => QD_COLOR[k] || FALLBACK }, + health: { label: 'Health Dimension', value: (t) => t.health_dimension || null, color: (k) => HD_COLOR[k] || FALLBACK }, + algorithm: { label: 'Algorithm', value: (t) => t.algorithm || null, color: (k) => ALGO_COLOR[k] || FALLBACK }, + technique: { label: 'Statistical Technique', value: (t) => t.statistical_technique || null, color: (k) => TECH_COLOR[k] || FALLBACK, sparse: true, order: TECHNIQUE_ORDER }, + scope: { label: 'Test Scope', value: (t) => SCOPE_LABEL[t.test_scope] || 'Other', color: () => FALLBACK, order: SCOPE_ORDER }, + criteria: { label: 'Criteria', value: (t) => t.criteria || null, color: (k) => CRITERIA_COLOR[k] || FALLBACK, order: CRITERIA_ORDER }, +}; + +// Facet order in the rail (PRD facet review). All facets are always visible. +const FACET_AXES = ['impact', 'dq', 'health', 'algorithm', 'technique', 'scope', 'criteria']; +// Group-by options offered in the result-list dropdown. +const GROUP_BY_AXES = ['impact', 'dq', 'health', 'algorithm', 'technique', 'scope', 'criteria']; + +export { + AXES, FACET_AXES, GROUP_BY_AXES, TECHNIQUE_ORDER, CRITERIA_ORDER, SCOPE_LABEL, EMPTY, + appliesToSelectedColumn, +}; diff --git a/testgen/ui/views/test_definitions.py b/testgen/ui/views/test_definitions.py index b4cac73d..ef28ff6f 100644 --- a/testgen/ui/views/test_definitions.py +++ b/testgen/ui/views/test_definitions.py @@ -20,6 +20,7 @@ TestDefinitionNote, TestDefinitionSummary, TestType, + derive_test_criteria, ) from testgen.common.models.test_suite import TestSuite from testgen.common.pii_masking import get_pii_columns, mask_profiling_pii @@ -195,6 +196,15 @@ def render( add_dialog = None if st.session_state.get(TD_ADD_DIALOG_KEY): + # Pre-fill the picker's column filter from the page's table+column filter when a + # specific column is filtered. The frontend resolves general_type from table_columns, + # so only the column identity is passed here. The prefilled column stays editable -- + # the user can change or clear it inside the picker. + prefill_column = ( + {"table_name": table_name, "column_name": column_name} + if table_name and column_name + else None + ) add_dialog = { "open": True, "test_types": test_types, @@ -203,6 +213,7 @@ def render( "table_group_schema": table_group.table_group_schema, "test_suite": test_suite_info, "qualifies_table_refs_with_schema": qualifies_table_refs_with_schema, + "prefill_column": prefill_column, } edit_dialog = None @@ -805,6 +816,7 @@ def run_test_type_lookup_query(test_type: str | None = None) -> pd.DataFrame: tt.measure_uom, COALESCE(tt.measure_uom_description, '') as measure_uom_description, tt.default_parm_columns, tt.default_severity, tt.run_type, tt.test_scope, tt.dq_dimension, tt.impact_dimension, tt.threshold_description, + tt.health_dimension, tt.algorithm, tt.statistical_technique, tt.column_name_prompt, tt.column_name_help, tt.default_parm_prompts, tt.default_parm_help, tt.usage_notes, CASE tt.test_scope @@ -836,7 +848,14 @@ def run_test_type_lookup_query(test_type: str | None = None) -> pd.DataFrame: END, tt.test_name_short; """ - return fetch_df_from_db(query, {"test_type": test_type}) + df = fetch_df_from_db(query, {"test_type": test_type}) + if not df.empty: + # Criteria facet is derived (not stored) via the shared classifier so UI and MCP agree. + df["criteria"] = df.apply( + lambda row: str(derive_test_criteria(row["test_type"], row["test_scope"], row["algorithm"])), + axis=1, + ) + return df @st.cache_data(show_spinner=False) @@ -994,7 +1013,7 @@ def get_test_definitions_collision( def get_columns(table_groups_id: str) -> list[dict]: results = fetch_all_from_db( """ - SELECT table_name, column_name + SELECT table_name, column_name, general_type FROM data_column_chars WHERE table_groups_id = :table_groups_id AND drop_date IS NULL diff --git a/tests/unit/common/models/test_test_type_columns.py b/tests/unit/common/models/test_test_type_columns.py new file mode 100644 index 00000000..f7b97095 --- /dev/null +++ b/tests/unit/common/models/test_test_type_columns.py @@ -0,0 +1,49 @@ +import pytest +from sqlalchemy.dialects.postgresql import dialect as pg_dialect + +from testgen.common.models.test_definition import StatisticalTechnique, TestAlgorithm, TestType + +pytestmark = pytest.mark.unit + +DIALECT = pg_dialect() + + +@pytest.mark.parametrize( + ("column_name", "enum_cls", "member"), + [ + ("algorithm", TestAlgorithm, TestAlgorithm.BOUNDARY_CHECK), + ("statistical_technique", StatisticalTechnique, StatisticalTechnique.COHENS_D), + ], +) +def test_enum_column_reads_db_value_as_enum_member(column_name, enum_cls, member): + column_type = TestType.__table__.c[column_name].type + process = column_type.result_processor(DIALECT, None) + + result = process(member.value) + + assert result is member + assert isinstance(result, enum_cls) + + +@pytest.mark.parametrize( + ("column_name", "member"), + [ + ("algorithm", TestAlgorithm.STATISTICAL_DRIFT), + ("statistical_technique", StatisticalTechnique.JENSEN_SHANNON_DIVERGENCE), + ], +) +def test_enum_column_writes_enum_as_db_value(column_name, member): + column_type = TestType.__table__.c[column_name].type + process = column_type.bind_processor(DIALECT) + + assert process(member) == member.value + + +@pytest.mark.parametrize("column_name", ["algorithm", "statistical_technique"]) +def test_enum_column_round_trips_none(column_name): + column_type = TestType.__table__.c[column_name].type + read = column_type.result_processor(DIALECT, None) + write = column_type.bind_processor(DIALECT) + + assert read(None) is None + assert write(None) is None diff --git a/tests/unit/test_test_criteria.py b/tests/unit/test_test_criteria.py new file mode 100644 index 00000000..582c7d94 --- /dev/null +++ b/tests/unit/test_test_criteria.py @@ -0,0 +1,104 @@ +from pathlib import Path + +import pytest +from yaml import safe_load + +import testgen +from testgen.common.models.test_definition import TestCriteria, derive_test_criteria + +YAML_DIR = Path(testgen.__file__).parent / "template" / "dbsetup_test_types" + +# Authoritative test-type -> Criteria mapping. The derivation is pure Python (no column), +# shared by the UI lookup and MCP. This dict pins the expected facet value for every test type; +# redline here and the derivation together. +EXPECTED_CRITERIA = { + # Custom Criteria: user authors the test in SQL (custom scope) + "CUSTOM": TestCriteria.CUSTOM_CRITERIA, + "Condition_Flag": TestCriteria.CUSTOM_CRITERIA, + # Reference Dataset: compared against another dataset (referential scope) + "Aggregate_Balance": TestCriteria.REFERENCE_DATASET, + "Aggregate_Balance_Percent": TestCriteria.REFERENCE_DATASET, + "Aggregate_Balance_Range": TestCriteria.REFERENCE_DATASET, + "Aggregate_Minimum": TestCriteria.REFERENCE_DATASET, + "Combo_Match": TestCriteria.REFERENCE_DATASET, + "Distribution_Shift": TestCriteria.REFERENCE_DATASET, + "Timeframe_Combo_Gain": TestCriteria.REFERENCE_DATASET, + "Timeframe_Combo_Match": TestCriteria.REFERENCE_DATASET, + # List of Values: tested against a set/lookup of allowed values (Set / lookup algorithm) + "LOV_All": TestCriteria.LIST_OF_VALUES, + "LOV_Match": TestCriteria.LIST_OF_VALUES, + "US_State": TestCriteria.LIST_OF_VALUES, + # Defined Rule: predefined validity/integrity rule, no user-supplied value (enumerated) + "Dupe_Rows": TestCriteria.DEFINED_RULE, + "Unique": TestCriteria.DEFINED_RULE, + "Email_Format": TestCriteria.DEFINED_RULE, + "Street_Addr_Pattern": TestCriteria.DEFINED_RULE, + "Valid_Characters": TestCriteria.DEFINED_RULE, + "Valid_Month": TestCriteria.DEFINED_RULE, + "Valid_US_Zip": TestCriteria.DEFINED_RULE, + "Valid_US_Zip3": TestCriteria.DEFINED_RULE, + # Defined Value: user asserts the expected value/pattern (enumerated) + "Constant": TestCriteria.DEFINED_VALUE, + "Pattern_Match": TestCriteria.DEFINED_VALUE, + "Required": TestCriteria.DEFINED_VALUE, + # Defined Threshold: user tunes a numeric threshold/tolerance (fallthrough) + "Alpha_Trunc": TestCriteria.DEFINED_THRESHOLD, + "Avg_Shift": TestCriteria.DEFINED_THRESHOLD, + "Daily_Record_Ct": TestCriteria.DEFINED_THRESHOLD, + "Dec_Trunc": TestCriteria.DEFINED_THRESHOLD, + "Distinct_Date_Ct": TestCriteria.DEFINED_THRESHOLD, + "Distinct_Value_Ct": TestCriteria.DEFINED_THRESHOLD, + "Freshness_Trend": TestCriteria.DEFINED_THRESHOLD, + "Future_Date": TestCriteria.DEFINED_THRESHOLD, + "Future_Date_1Y": TestCriteria.DEFINED_THRESHOLD, + "Incr_Avg_Shift": TestCriteria.DEFINED_THRESHOLD, + "Metric_Trend": TestCriteria.DEFINED_THRESHOLD, + "Min_Date": TestCriteria.DEFINED_THRESHOLD, + "Min_Val": TestCriteria.DEFINED_THRESHOLD, + "Missing_Pct": TestCriteria.DEFINED_THRESHOLD, + "Monthly_Rec_Ct": TestCriteria.DEFINED_THRESHOLD, + "Outlier_Pct_Above": TestCriteria.DEFINED_THRESHOLD, + "Outlier_Pct_Below": TestCriteria.DEFINED_THRESHOLD, + "Recency": TestCriteria.DEFINED_THRESHOLD, + "Row_Ct": TestCriteria.DEFINED_THRESHOLD, + "Row_Ct_Pct": TestCriteria.DEFINED_THRESHOLD, + "Table_Freshness": TestCriteria.DEFINED_THRESHOLD, + "Unique_Pct": TestCriteria.DEFINED_THRESHOLD, + "Variability_Increase": TestCriteria.DEFINED_THRESHOLD, + "Variability_Decrease": TestCriteria.DEFINED_THRESHOLD, + "Volume_Trend": TestCriteria.DEFINED_THRESHOLD, + "Weekly_Rec_Ct": TestCriteria.DEFINED_THRESHOLD, +} + +# Test types that never surface in the Add-Test picker, so they carry no Criteria facet at all. +NON_PUBLIC_TEST_TYPES = frozenset({"Schema_Drift"}) + + +def _test_type_docs(): + for path in sorted([*YAML_DIR.glob("*.yaml"), *YAML_DIR.glob("*.yml")]): + yield safe_load(path.read_text())["test_types"] + + +@pytest.mark.unit +def test_every_test_type_has_expected_criteria(): + seen = set() + for tt in _test_type_docs(): + test_type = tt["test_type"] + if test_type in NON_PUBLIC_TEST_TYPES: + continue + seen.add(test_type) + criteria = derive_test_criteria(test_type, tt.get("test_scope"), tt.get("algorithm")) + assert criteria == EXPECTED_CRITERIA[test_type], ( + f"{test_type}: derived {criteria!r}, expected {EXPECTED_CRITERIA[test_type]!r}" + ) + missing = set(EXPECTED_CRITERIA) - seen + assert not missing, f"expected-criteria mapping references unknown test types: {sorted(missing)}" + + +@pytest.mark.unit +def test_derive_returns_a_valid_criteria_for_every_test_type(): + for tt in _test_type_docs(): + if tt["test_type"] in NON_PUBLIC_TEST_TYPES: + continue + criteria = derive_test_criteria(tt["test_type"], tt.get("test_scope"), tt.get("algorithm")) + assert isinstance(criteria, TestCriteria) diff --git a/tests/unit/test_test_type_taxonomy.py b/tests/unit/test_test_type_taxonomy.py new file mode 100644 index 00000000..123eed95 --- /dev/null +++ b/tests/unit/test_test_type_taxonomy.py @@ -0,0 +1,45 @@ +from pathlib import Path + +import pytest +from yaml import safe_load + +import testgen +from testgen.common.models.test_definition import StatisticalTechnique, TestAlgorithm + +YAML_DIR = Path(testgen.__file__).parent / "template" / "dbsetup_test_types" + +ALGORITHMS = {a.value for a in TestAlgorithm} +TECHNIQUES = {t.value for t in StatisticalTechnique} +HEALTH_DIMENSIONS = {"Schema Drift", "Data Drift", "Statistical Drift", "Volume", "Freshness"} + + +def _test_type_docs(): + for path in sorted([*YAML_DIR.glob("*.yaml"), *YAML_DIR.glob("*.yml")]): + data = safe_load(path.read_text())["test_types"] + yield path.name, data + + +@pytest.mark.unit +def test_every_test_type_declares_a_valid_algorithm(): + for filename, tt in _test_type_docs(): + algorithm = tt.get("algorithm") + assert algorithm in ALGORITHMS, f"{filename}: invalid/missing algorithm {algorithm!r}" + + +@pytest.mark.unit +def test_statistical_technique_is_valid_when_present(): + for filename, tt in _test_type_docs(): + algorithm = tt.get("algorithm") + technique = tt.get("statistical_technique") + if technique is not None: + assert technique in TECHNIQUES, f"{filename}: invalid statistical_technique {technique!r}" + if algorithm == TestAlgorithm.STATISTICAL_DRIFT.value: + assert technique in TECHNIQUES, f"{filename}: drift test needs a valid technique, got {technique!r}" + + +@pytest.mark.unit +def test_health_dimension_uses_freshness_not_recency(): + for filename, tt in _test_type_docs(): + health = tt.get("health_dimension") + assert health != "Recency", f"{filename}: health_dimension must be 'Freshness', not 'Recency'" + assert health is None or health in HEALTH_DIMENSIONS, f"{filename}: unexpected health_dimension {health!r}" From 5a06c3a803a7bd581df23dc7c94f6cf534bb0cd0 Mon Sep 17 00:00:00 2001 From: Aarthy Adityan Date: Fri, 26 Jun 2026 02:50:59 -0400 Subject: [PATCH 58/78] feat(test-definitions): add external URL and custom metadata for click-through Add two optional fields to test definitions so a test can link out to the code, pipeline step, or system that produces the data under test: - external_url: a user-supplied link, rendered as a clickable link in the test-results view and a trailing icon-only column in the results list. - custom_metadata: an arbitrary JSON object of key-value pairs, edited as validated JSON text in the authoring form and displayed as a key-value list. Both fields are optional, apply to every test type, travel with the test in the export/import JSON, and are surfaced through the test-definition and test-result MCP read tools (create/update accept custom_metadata as an object or a JSON string). Co-Authored-By: Claude Opus 4.8 --- testgen/common/models/test_definition.py | 8 ++ .../test_definition_export_import_service.py | 4 + testgen/mcp/tools/test_definitions.py | 30 +++++++ testgen/mcp/tools/test_results.py | 15 +++- .../030_initialize_new_schema_structure.sql | 2 + .../dbupgrade/0197_incremental_upgrade.sql | 5 ++ .../js/pages/test_definition_summary.js | 57 +++++++++++- .../frontend/js/pages/test_definitions.js | 34 ++++++++ .../frontend/js/pages/test_results.js | 24 ++++++ testgen/ui/queries/test_result_queries.py | 1 + testgen/ui/static/js/form_validators.js | 23 +++++ testgen/ui/views/test_results.py | 2 + tests/unit/api/test_td_export_import.py | 86 +++++++++++++++++++ .../common/models/test_test_definition.py | 25 ++++++ tests/unit/mcp/test_tools_test_definitions.py | 30 +++++++ tests/unit/mcp/test_tools_test_results.py | 16 +++- 16 files changed, 356 insertions(+), 6 deletions(-) create mode 100644 testgen/template/dbupgrade/0197_incremental_upgrade.sql diff --git a/testgen/common/models/test_definition.py b/testgen/common/models/test_definition.py index 13fe7705..e5c9c4d6 100644 --- a/testgen/common/models/test_definition.py +++ b/testgen/common/models/test_definition.py @@ -224,6 +224,8 @@ class TestDefinitionSummary(TestTypeSummary): prediction: dict[str, dict[str, float]] | None flagged: bool impact_dimension: str | None + external_url: str + custom_metadata: dict | None @property def display_name(self) -> str: @@ -397,6 +399,8 @@ class TestDefinition(Entity): flagged: bool = Column(Boolean, default=False, nullable=False) external_id: UUID | None = Column(postgresql.UUID(as_uuid=True)) impact_dimension: str | None = Column(String, nullable=True) + external_url: str = Column(NullIfEmptyString) + custom_metadata: dict | None = Column(postgresql.JSONB) _default_order_by = ( asc(func.lower(schema_name)), @@ -555,6 +559,7 @@ def select_page( # Fields editable on every test type regardless of param_columns. EDITABLE_BASE_FIELDS: ClassVar[frozenset[str]] = frozenset({ "test_active", "severity", "lock_refresh", "flagged", "test_description", + "external_url", "custom_metadata", }) def editable_fields(self, test_type: TestType) -> set[str]: @@ -601,6 +606,9 @@ def validate(self, test_type: TestType) -> None: f"test type `{test_type.test_type}` does not accept a custom query" ) + if self.custom_metadata is not None and not isinstance(self.custom_metadata, dict): + errors["custom_metadata"] = "must be a JSON object of key-value pairs" + for required in _required_fields_for(test_type): if _is_blank(getattr(self, required, None)): errors[required] = f"required for test type `{test_type.test_type}`" diff --git a/testgen/common/test_definition_export_import_service.py b/testgen/common/test_definition_export_import_service.py index 1e831d6b..7e69bc9e 100644 --- a/testgen/common/test_definition_export_import_service.py +++ b/testgen/common/test_definition_export_import_service.py @@ -129,6 +129,10 @@ class TestDefinitionExport(BaseModel): history_calculation_upper: str | None = None history_lookback: int = 0 + # URL / metadata + external_url: str | None = None + custom_metadata: dict[str, Any] | None = None + @field_validator("skip_errors", "window_days", "history_lookback", mode="before") @classmethod def _coerce_none_to_zero(cls, v: int | None) -> int: diff --git a/testgen/mcp/tools/test_definitions.py b/testgen/mcp/tools/test_definitions.py index 9c277c08..56d183f4 100644 --- a/testgen/mcp/tools/test_definitions.py +++ b/testgen/mcp/tools/test_definitions.py @@ -1,3 +1,4 @@ +import json from datetime import UTC, datetime from enum import StrEnum from typing import NoReturn @@ -242,6 +243,14 @@ def _append_td_summary(doc: MdDoc, td: TestDefinitionSummary) -> None: elif td.last_auto_gen_date: doc.field("Last Updated", f"{td.last_auto_gen_date} (auto-generated)") + # URL and metadata + if td.external_url: + doc.field("External URL", td.external_url) + if td.custom_metadata: + doc.heading(2, "Custom Metadata") + for key, value in td.custom_metadata.items(): + doc.field(key, value) + # Parameters (editable fields from test type metadata) _append_parameters_section(doc, td) @@ -506,6 +515,25 @@ def _raise_validation_errors(err: InvalidTestDefinitionFields, header: str) -> N raise MCPUserError(f"{header}\n\n{bullets}") from err +def _coerce_custom_metadata(fields: dict) -> None: + """Accept ``custom_metadata`` as a JSON string for convenience, parsing it to an object in place. + + A blank string becomes ``None`` (clears the field). Non-object JSON is left for model + validation to reject with a consistent message. + """ + value = fields.get("custom_metadata") + if not isinstance(value, str): + return + stripped = value.strip() + if not stripped: + fields["custom_metadata"] = None + return + try: + fields["custom_metadata"] = json.loads(stripped) + except json.JSONDecodeError as e: + raise MCPUserError("`custom_metadata` must be a JSON object of key-value pairs.") from e + + @with_database_session @mcp_permission("edit") def create_test( @@ -547,6 +575,7 @@ def create_test( ) fields = fields or {} + _coerce_custom_metadata(fields) accepted = td.editable_fields(tt) rejected = sorted(set(fields) - accepted) if rejected: @@ -591,6 +620,7 @@ def update_test(test_definition_id: str, fields: dict) -> str: if not fields: raise MCPUserError("No fields supplied to update.") + _coerce_custom_metadata(fields) accepted = td.editable_fields(tt) rejected = sorted(set(fields) - accepted) if rejected: diff --git a/testgen/mcp/tools/test_results.py b/testgen/mcp/tools/test_results.py index f38c574b..dbdb719c 100644 --- a/testgen/mcp/tools/test_results.py +++ b/testgen/mcp/tools/test_results.py @@ -6,7 +6,7 @@ from testgen.common.enums import JobStatus from testgen.common.models import get_current_session, with_database_session from testgen.common.models.job_execution import JobExecution -from testgen.common.models.test_definition import TestType +from testgen.common.models.test_definition import TestDefinition, TestType from testgen.common.models.test_result import BucketInterval, TestResult, TestResultStatus from testgen.common.models.test_run import TestRun, TestRunSummary from testgen.common.models.test_suite import TestSuite @@ -130,6 +130,12 @@ def list_test_results( type_names = {tt.test_type: tt.test_name_short for tt in TestType.select_where(TestType.active == "Y")} + td_ids = {r.test_definition_id for r in results if r.test_definition_id} + source_by_td = { + td.id: (td.external_url, td.custom_metadata) + for td in (TestDefinition.select_where(TestDefinition.id.in_(td_ids)) if td_ids else []) + } + doc = MdDoc() doc.heading(1, f"Test Results for run `{run_id_label}`") if resolved_via_suite: @@ -153,6 +159,13 @@ def list_test_results( doc.field("Threshold", r.threshold_value) if r.message: doc.field("Message", r.message) + external_url, custom_metadata = source_by_td.get(r.test_definition_id, (None, None)) + if external_url: + doc.field("External URL", external_url) + if custom_metadata: + doc.heading(3, "Custom Metadata") + for key, value in custom_metadata.items(): + doc.field(key, value) return doc.render() diff --git a/testgen/template/dbsetup/030_initialize_new_schema_structure.sql b/testgen/template/dbsetup/030_initialize_new_schema_structure.sql index 814f57b9..a4c84fc4 100644 --- a/testgen/template/dbsetup/030_initialize_new_schema_structure.sql +++ b/testgen/template/dbsetup/030_initialize_new_schema_structure.sql @@ -241,6 +241,8 @@ CREATE TABLE test_definitions ( flagged BOOLEAN DEFAULT FALSE NOT NULL, external_id UUID, impact_dimension VARCHAR(20), + external_url VARCHAR, + custom_metadata JSONB, CONSTRAINT test_definitions_test_suites_test_suite_id_fk FOREIGN KEY (test_suite_id) REFERENCES test_suites ); diff --git a/testgen/template/dbupgrade/0197_incremental_upgrade.sql b/testgen/template/dbupgrade/0197_incremental_upgrade.sql new file mode 100644 index 00000000..b8a37d41 --- /dev/null +++ b/testgen/template/dbupgrade/0197_incremental_upgrade.sql @@ -0,0 +1,5 @@ +SET SEARCH_PATH TO {SCHEMA_NAME}; + +ALTER TABLE test_definitions + ADD COLUMN IF NOT EXISTS external_url VARCHAR, + ADD COLUMN IF NOT EXISTS custom_metadata JSONB; diff --git a/testgen/ui/components/frontend/js/pages/test_definition_summary.js b/testgen/ui/components/frontend/js/pages/test_definition_summary.js index c3cb23e5..b075dd84 100644 --- a/testgen/ui/components/frontend/js/pages/test_definition_summary.js +++ b/testgen/ui/components/frontend/js/pages/test_definition_summary.js @@ -17,8 +17,10 @@ * @property {string} export_to_observability * @property {string?} last_manual_update * @property {string?} usage_notes + * @property {string?} external_url + * @property {object?} custom_metadata * @property {Array} attributes - * + * * @typedef Properties * @type {object} * @property {TestDefinition} test_definition @@ -27,9 +29,15 @@ import van from '/app/static/js/van.min.js'; import { createEmitter, getValue, isEqual, loadStylesheet } from '/app/static/js/utils.js'; import { Alert } from '/app/static/js/components/alert.js'; import { Attribute } from '/app/static/js/components/attribute.js'; +import { Link } from '/app/static/js/components/link.js'; const { div, strong } = van.tags; +const isHttpUrl = (value) => typeof value === 'string' && /^https?:\/\//i.test(value.trim()); + +const metadataDisplayValue = (value) => + (value !== null && typeof value === 'object') ? JSON.stringify(value) : value; + /** * @param {Properties} props * @returns @@ -112,6 +120,38 @@ const TestDefinitionSummary = (props) => { ), ), ), + testDefinition.external_url + ? Attribute({ + label: 'External URL', + value: isHttpUrl(testDefinition.external_url) + ? Link({ + href: testDefinition.external_url, + label: testDefinition.external_url, + open_new: true, + underline: true, + right_icon: 'open_in_new', + right_icon_size: 16, + }) + : testDefinition.external_url, + class: 'mt-4 external-url-attribute', + }) + : '', + testDefinition.custom_metadata && Object.keys(testDefinition.custom_metadata).length + ? div( + { class: 'flex-column fx-gap-3 mt-4' }, + strong({}, 'Custom Metadata'), + div( + { class: 'flex-row fx-flex-wrap fx-gap-4 test-definition-attributes' }, + Object.entries(testDefinition.custom_metadata).map(([key, value]) => + Attribute({ + label: key, + value: metadataDisplayValue(value), + class: 'fx-flex', + }) + ), + ), + ) + : '', testDefinition.usage_notes ? Alert( { type: 'info', class: 'mt-4' }, @@ -132,6 +172,21 @@ stylesheet.replace(` .test-definition-attributes > div .attribute-value { font-size: 16px; } +.external-url-attribute .attribute-value { + overflow-wrap: anywhere; +} +.external-url-attribute .tg-link { + width: 100%; + max-width: 100%; +} +.external-url-attribute .tg-link--wrapper { + flex-wrap: wrap; +} +.external-url-attribute .tg-link--text { + overflow-wrap: anywhere; + word-break: break-all; + min-width: 0; +} `); export { TestDefinitionSummary }; diff --git a/testgen/ui/components/frontend/js/pages/test_definitions.js b/testgen/ui/components/frontend/js/pages/test_definitions.js index 2132b56f..79de48dd 100644 --- a/testgen/ui/components/frontend/js/pages/test_definitions.js +++ b/testgen/ui/components/frontend/js/pages/test_definitions.js @@ -19,6 +19,7 @@ import { Icon } from '/app/static/js/components/icon.js'; import { ProfilingResultsDialog } from '../shared/profiling_results_dialog.js'; import { AXES, FACET_AXES, GROUP_BY_AXES, EMPTY, appliesToSelectedColumn } from '/app/static/js/components/test_picker_taxonomy.js'; import { enterPage, exitPage, getPageSignal } from '/app/static/js/page_lifecycle.js'; +import { jsonObject } from '/app/static/js/form_validators.js'; const { button: btn, div, i: icon, span, strong } = van.tags; @@ -1384,6 +1385,13 @@ const TestDefFormContent = ({ formValues, tableColumns, testSuite, validateResul onFormChange({ [key]: value }); }; + // Custom Metadata is stored as a JSON object (JSONB). Keep the raw text the user edits in a + // local state, and only commit a parsed object to the form when the JSON is valid. + const metadataText = van.state( + formValues.custom_metadata ? JSON.stringify(formValues.custom_metadata, null, 2) : '' + ); + const metadataValid = van.state(true); + const inheritedSeverity = testSuite.severity ?? formValues.default_severity ?? 'Warning'; const severityOptions = [ { label: `Inherited (${inheritedSeverity})`, value: null }, @@ -1459,6 +1467,31 @@ const TestDefFormContent = ({ formValues, tableColumns, testSuite, validateResul onChange: (value) => updateField('test_description', value || null), }), + // External link + custom metadata + Input({ + name: 'external_url', + label: 'External URL', + help: 'Optional link to the code, pipeline step, or system that produces this data.', + value: () => fv.val.external_url ?? '', + onChange: (value) => updateField('external_url', value || null), + }), + Textarea({ + name: 'custom_metadata', + label: 'Custom Metadata', + help: 'Optional JSON object of key-value pairs, e.g. {"pipeline": "daily_load", "task": "transform_orders"}.', + value: metadataText, + placeholder: '{\n "pipeline": "daily_load",\n "task": "transform_orders"\n}', + height: 120, + validators: [jsonObject], + onChange: (value, state) => { + metadataText.val = value; + metadataValid.val = state.valid; + if (state.valid) { + updateField('custom_metadata', value && value.trim() ? JSON.parse(value) : null); + } + }, + }), + // Checkboxes div( { class: 'flex-row fx-gap-4' }, @@ -1644,6 +1677,7 @@ const TestDefFormContent = ({ formValues, tableColumns, testSuite, validateResul color: 'primary', label: mode === 'edit' ? 'Save' : 'Add', width: 'auto', + disabled: () => !metadataValid.val, onclick: onSave, }), ), diff --git a/testgen/ui/components/frontend/js/pages/test_results.js b/testgen/ui/components/frontend/js/pages/test_results.js index 8fdc3bab..7dda87b4 100644 --- a/testgen/ui/components/frontend/js/pages/test_results.js +++ b/testgen/ui/components/frontend/js/pages/test_results.js @@ -20,6 +20,7 @@ * @property {string?} table_groups_id * @property {string?} severity * @property {string} test_type + * @property {string?} external_url * * @typedef Properties * @type {object} @@ -49,6 +50,7 @@ import { Button } from '/app/static/js/components/button.js'; import { Checkbox } from '/app/static/js/components/checkbox.js'; import { DropdownButton } from '/app/static/js/components/dropdown_button.js'; import { Icon } from '/app/static/js/components/icon.js'; +import { Link } from '/app/static/js/components/link.js'; import { SummaryBar } from '/app/static/js/components/summary_bar.js'; import { Dialog } from '/app/static/js/components/dialog.js'; import { Toggle } from '/app/static/js/components/toggle.js'; @@ -106,6 +108,7 @@ const DATA_COLUMNS = [ { name: 'flagged_display', label: 'Flagged', width: 80, align: 'center' }, { name: 'notes_count', label: 'Notes', width: 70, align: 'center' }, { name: 'result_message', label: 'Details', width: 200, overflow: 'hidden' }, + { name: 'external_link', label: 'Link', width: 60, align: 'center' }, ]; const HISTORY_COLUMNS = [ @@ -149,6 +152,26 @@ const formatNumber = (v) => { return n.toLocaleString(undefined, { maximumFractionDigits: 5 }); }; +const isHttpUrl = (value) => typeof value === 'string' && /^https?:\/\//i.test(value.trim()); + +const buildExternalLinkCell = (url) => { + if (!isHttpUrl(url)) { + return ''; + } + // Stop the click from also selecting the row; the anchor still opens in a new tab. + return span( + { class: 'flex-row fx-justify-center', onclick: (event) => event.stopPropagation() }, + Link({ + href: url, + label: '', + open_new: true, + left_icon: 'open_in_new', + left_icon_size: 18, + tooltip: 'Open external link', + }), + ); +}; + const buildTableRow = (item) => ({ id: item.test_result_id, table_name: item.table_name ?? '', @@ -171,6 +194,7 @@ const buildTableRow = (item) => ({ span(item.notes_count), ) : '', result_message: item.result_message ?? '', + external_link: buildExternalLinkCell(item.external_url), }); const ExportMenu = (statusFilter, tableFilter, columnFilter, testTypeFilter, actionFilter, flaggedFilter, hasSelection, getSelectedIds, emit) => { diff --git a/testgen/ui/queries/test_result_queries.py b/testgen/ui/queries/test_result_queries.py index 4e0b0c58..92f38673 100644 --- a/testgen/ui/queries/test_result_queries.py +++ b/testgen/ui/queries/test_result_queries.py @@ -115,6 +115,7 @@ def get_test_results( r.test_definition_id::VARCHAR, r.auto_gen, td.flagged, + td.external_url, (SELECT COUNT(*) FROM test_definition_notes tdn WHERE tdn.test_definition_id = td.id) as notes_count, -- These are used in the PDF report diff --git a/testgen/ui/static/js/form_validators.js b/testgen/ui/static/js/form_validators.js index 58c085bb..799292aa 100644 --- a/testgen/ui/static/js/form_validators.js +++ b/testgen/ui/static/js/form_validators.js @@ -140,6 +140,28 @@ function notIn(values, options) { }; } +/** + * Validate that the value is empty (optional) or a well-formed JSON object of key-value pairs. + * + * @param {any} value + * @returns {string} + */ +function jsonObject(value) { + if (!value || !String(value).trim()) { + return null; + } + let parsed; + try { + parsed = JSON.parse(value); + } catch { + return 'Must be valid JSON.'; + } + if (typeof parsed !== 'object' || parsed === null || Array.isArray(parsed)) { + return 'Must be a JSON object of key-value pairs.'; + } + return null; +} + export { maxLength, minLength, @@ -149,4 +171,5 @@ export { requiredIf, sizeLimit, notIn, + jsonObject, }; diff --git a/testgen/ui/views/test_results.py b/testgen/ui/views/test_results.py index b62473aa..1814b08c 100644 --- a/testgen/ui/views/test_results.py +++ b/testgen/ui/views/test_results.py @@ -720,6 +720,8 @@ def readable_boolean(v: bool) -> str: "locked": readable_boolean(test_definition.lock_refresh), "active": readable_boolean(test_definition.test_active), "usage_notes": test_definition.usage_notes, + "external_url": test_definition.external_url, + "custom_metadata": test_definition.custom_metadata, "last_manual_update": ( test_definition.last_manual_update.isoformat() if test_definition.last_manual_update else None ), diff --git a/tests/unit/api/test_td_export_import.py b/tests/unit/api/test_td_export_import.py index 5a875136..e56d262f 100644 --- a/tests/unit/api/test_td_export_import.py +++ b/tests/unit/api/test_td_export_import.py @@ -6,6 +6,7 @@ import pytest from fastapi import HTTPException +from pydantic import ValidationError from testgen.api.schemas import ImportRequest from testgen.common.test_definition_export_import_service import ( @@ -1081,3 +1082,88 @@ def test_defaults_match_expected_values(self): assert td.window_days == 0 assert td.history_lookback == 0 + def test_click_through_fields_default_none(self): + td = TestDefinitionExport(test_type="Alpha") + assert td.external_url is None + assert td.custom_metadata is None + assert "external_url" not in td.model_fields_set + assert "custom_metadata" not in td.model_fields_set + + def test_custom_metadata_accepts_object(self): + td = TestDefinitionExport(test_type="Alpha", custom_metadata={"pipeline": "daily_load", "task": "t1"}) + assert td.custom_metadata == {"pipeline": "daily_load", "task": "t1"} + + @pytest.mark.parametrize("bad_value", ["not-json", ["a", "b"], 42]) + def test_custom_metadata_rejects_non_object(self, bad_value): + with pytest.raises(ValidationError): + TestDefinitionExport(test_type="Alpha", custom_metadata=bad_value) + + +# --- Click-through fields: export + import --- + + +class Test_click_through_fields: + + @patch(f"{SERVICE_MODULE}.settings") + @patch(f"{SERVICE_MODULE}.TableGroup") + @patch(f"{SERVICE_MODULE}.get_current_session") + def test_export_includes_external_url_and_custom_metadata(self, mock_session_fn, mock_tg_cls, mock_settings): + mock_settings.VERSION = "5.12.0" + session = MagicMock() + mock_session_fn.return_value = session + mock_tg_cls.get.return_value = _make_table_group() + + ts = _make_test_suite() + + td_obj = MagicMock(spec=[]) + for field in TestDefinitionExport.model_fields: + setattr(td_obj, field, None) + td_obj.test_type = "Alpha" + td_obj.test_active = True + td_obj.lock_refresh = False + td_obj.skip_errors = 0 + td_obj.window_days = 0 + td_obj.history_lookback = 0 + td_obj.external_url = "https://example.com/notebook" + td_obj.custom_metadata = {"pipeline": "daily_load", "task": "transform_orders"} + + session.scalars.return_value.all.return_value = [td_obj] + + from testgen.common.test_definition_export_import_service import export_definitions + + result = export_definitions(ts, Origin.both, None, None) + + exported = result.definitions[0] + assert exported.external_url == "https://example.com/notebook" + assert exported.custom_metadata == {"pipeline": "daily_load", "task": "transform_orders"} + + @patch(f"{SERVICE_MODULE}.DataTable") + @patch(f"{SERVICE_MODULE}.TableGroup") + @patch(f"{SERVICE_MODULE}.get_current_session") + def test_create_sets_click_through_fields(self, mock_session_fn, mock_tg_cls, mock_dt_cls): + session = MagicMock() + mock_session_fn.return_value = session + mock_tg_cls.get.return_value = _make_table_group() + mock_dt_cls.select_table_names.return_value = ["t1"] + session.execute.return_value.all.return_value = [] + + ts = _make_test_suite() + td = _make_import_td( + test_type="Alpha", + table_name="t1", + external_url="https://example.com/task", + custom_metadata={"node": "n1"}, + ) + config = _make_config(mode=ImportMode.apply, on_new=OnNew.create) + + from testgen.common.test_definition_export_import_service import import_definitions + + with patch(f"{SERVICE_MODULE}._load_valid_test_types", return_value={"Alpha"}): + with patch(f"{SERVICE_MODULE}.TestDefinition") as mock_td_cls: + created_td = MagicMock(id=uuid4()) + mock_td_cls.return_value = created_td + import_definitions(ts, config, _make_payload(td)) + + assert created_td.external_url == "https://example.com/task" + assert created_td.custom_metadata == {"node": "n1"} + diff --git a/tests/unit/common/models/test_test_definition.py b/tests/unit/common/models/test_test_definition.py index f733d1b6..660f2205 100644 --- a/tests/unit/common/models/test_test_definition.py +++ b/tests/unit/common/models/test_test_definition.py @@ -99,6 +99,13 @@ def test_editable_fields_includes_param_columns(): assert {"threshold_value", "baseline_value"} <= accepted +def test_editable_fields_includes_click_through_fields(): + tt = make_test_type(param_columns=set(), default_parm_columns=None) + td = make_td() + accepted = td.editable_fields(tt) + assert {"external_url", "custom_metadata"} <= accepted + + def test_editable_fields_includes_impact_dimension_only_for_custom_or_referential_scope(): """impact_dimension is overridable only for user-defined-semantic scopes.""" td = make_td() @@ -220,6 +227,22 @@ def test_validate_severity_empty_string_treated_as_unset(): td.validate(tt) # empty severity is OK — falls back to test type default +def test_validate_custom_metadata_accepts_object_and_none(): + tt = make_test_type() + make_td(column_name="email", threshold_value="10", custom_metadata={"pipeline": "p1"}).validate(tt) + make_td(column_name="email", threshold_value="10", custom_metadata=None).validate(tt) + make_td(column_name="email", threshold_value="10", custom_metadata={}).validate(tt) + + +@pytest.mark.parametrize("bad_value", ["a string", ["a", "b"], 42]) +def test_validate_custom_metadata_rejects_non_object(bad_value): + tt = make_test_type() + td = make_td(column_name="email", threshold_value="10", custom_metadata=bad_value) + with pytest.raises(InvalidTestDefinitionFields) as exc_info: + td.validate(tt) + assert "custom_metadata" in exc_info.value.errors + + def test_validate_aggregates_errors(): tt = make_test_type(scope="column") td = make_td(severity="critical", custom_query="SELECT 1") # no column_name @@ -294,6 +317,8 @@ def _make_summary_row(table_name: str = "my_table") -> dict: "prediction": None, "flagged": False, "impact_dimension": None, + "external_url": None, + "custom_metadata": None, "test_name_short": "Custom", "default_test_description": "A test", "measure_uom": "", diff --git a/tests/unit/mcp/test_tools_test_definitions.py b/tests/unit/mcp/test_tools_test_definitions.py index 4f467ff5..ebd814a6 100644 --- a/tests/unit/mcp/test_tools_test_definitions.py +++ b/tests/unit/mcp/test_tools_test_definitions.py @@ -1099,6 +1099,36 @@ def test_update_test_multi_field(mock_resolve_td, mock_tt_model, db_session_mock td.save.assert_called_once() +@patch("testgen.mcp.tools.test_definitions.TestType") +@patch("testgen.mcp.tools.test_definitions.resolve_test_definition") +def test_update_test_custom_metadata_json_string_coerced(mock_resolve_td, mock_tt_model, db_session_mock): + td = _make_td_orm() + td.editable_fields.return_value = td.editable_fields.return_value | {"custom_metadata"} + mock_resolve_td.return_value = td + mock_tt_model.get.return_value = _make_test_type() + + from testgen.mcp.tools.test_definitions import update_test + + update_test(str(td.id), fields={"custom_metadata": '{"pipeline": "p1", "task": "t1"}'}) + assert td.custom_metadata == {"pipeline": "p1", "task": "t1"} + td.save.assert_called_once() + + +@patch("testgen.mcp.tools.test_definitions.TestType") +@patch("testgen.mcp.tools.test_definitions.resolve_test_definition") +def test_update_test_custom_metadata_invalid_json_rejected(mock_resolve_td, mock_tt_model, db_session_mock): + td = _make_td_orm() + td.editable_fields.return_value = td.editable_fields.return_value | {"custom_metadata"} + mock_resolve_td.return_value = td + mock_tt_model.get.return_value = _make_test_type() + + from testgen.mcp.tools.test_definitions import update_test + + with pytest.raises(MCPUserError, match="custom_metadata"): + update_test(str(td.id), fields={"custom_metadata": "{not valid json"}) + td.save.assert_not_called() + + # -- validate_custom_test ----------------------------------------------------- diff --git a/tests/unit/mcp/test_tools_test_results.py b/tests/unit/mcp/test_tools_test_results.py index 47e04f6c..64991e20 100644 --- a/tests/unit/mcp/test_tools_test_results.py +++ b/tests/unit/mcp/test_tools_test_results.py @@ -21,8 +21,9 @@ def _mock_test_run(test_run_id=None): @patch("testgen.mcp.tools.test_results.TestSuite") @patch("testgen.mcp.tools.test_results.TestRun") @patch("testgen.mcp.tools.test_results.TestType") +@patch("testgen.mcp.tools.test_results.TestDefinition") @patch("testgen.mcp.tools.test_results.TestResult") -def test_list_test_results_basic(mock_result, mock_tt_cls, mock_test_run_cls, mock_suite_cls, db_session_mock): +def test_list_test_results_basic(mock_result, mock_td_cls, mock_tt_cls, mock_test_run_cls, mock_suite_cls, db_session_mock): job_id = str(uuid4()) mock_test_run_cls.get.return_value = _mock_test_run() mock_suite_cls.get_regular.return_value = _mock_test_suite() @@ -37,6 +38,7 @@ def test_list_test_results_basic(mock_result, mock_tt_cls, mock_test_run_cls, mo r1.threshold_value = "10.0" r1.message = "Truncation detected" mock_result.select_results.return_value = [r1] + mock_td_cls.select_where.return_value = [] tt = MagicMock() tt.test_type = "Alpha_Trunc" @@ -57,8 +59,9 @@ def test_list_test_results_basic(mock_result, mock_tt_cls, mock_test_run_cls, mo @patch("testgen.mcp.tools.test_results.TestSuite") @patch("testgen.mcp.tools.test_results.TestRun") @patch("testgen.mcp.tools.test_results.TestType") +@patch("testgen.mcp.tools.test_results.TestDefinition") @patch("testgen.mcp.tools.test_results.TestResult") -def test_list_test_results_emits_test_result_id(mock_result, mock_tt_cls, mock_test_run_cls, mock_suite_cls, db_session_mock): +def test_list_test_results_emits_test_result_id(mock_result, mock_td_cls, mock_tt_cls, mock_test_run_cls, mock_suite_cls, db_session_mock): mock_test_run_cls.get.return_value = _mock_test_run() mock_suite_cls.get_regular.return_value = _mock_test_suite() @@ -74,6 +77,7 @@ def test_list_test_results_emits_test_result_id(mock_result, mock_tt_cls, mock_t r1.threshold_value = "10.0" r1.message = "Truncation detected" mock_result.select_results.return_value = [r1] + mock_td_cls.select_where.return_value = [] tt = MagicMock() tt.test_type = "Alpha_Trunc" @@ -91,8 +95,9 @@ def test_list_test_results_emits_test_result_id(mock_result, mock_tt_cls, mock_t @patch("testgen.mcp.tools.test_results.TestSuite") @patch("testgen.mcp.tools.test_results.TestRun") @patch("testgen.mcp.tools.test_results.TestType") +@patch("testgen.mcp.tools.test_results.TestDefinition") @patch("testgen.mcp.tools.test_results.TestResult") -def test_list_test_results_table_level_title(mock_result, mock_tt_cls, mock_test_run_cls, mock_suite_cls, db_session_mock): +def test_list_test_results_table_level_title(mock_result, mock_td_cls, mock_tt_cls, mock_test_run_cls, mock_suite_cls, db_session_mock): mock_test_run_cls.get.return_value = _mock_test_run() mock_suite_cls.get_regular.return_value = _mock_test_suite() @@ -106,6 +111,7 @@ def test_list_test_results_table_level_title(mock_result, mock_tt_cls, mock_test r1.threshold_value = "500" r1.message = None mock_result.select_results.return_value = [r1] + mock_td_cls.select_where.return_value = [] tt = MagicMock() tt.test_type = "Row_Ct" @@ -324,9 +330,10 @@ def test_list_test_results_by_suite_id_no_completed_runs(mock_suite_cls, db_sess @patch("testgen.mcp.tools.test_results.TestSuite") @patch("testgen.mcp.tools.test_results.TestRun") @patch("testgen.mcp.tools.test_results.TestType") +@patch("testgen.mcp.tools.test_results.TestDefinition") @patch("testgen.mcp.tools.test_results.TestResult") def test_list_test_results_by_suite_id_resolves_latest_run( - mock_result, mock_tt_cls, mock_test_run_cls, mock_suite_cls, db_session_mock + mock_result, mock_td_cls, mock_tt_cls, mock_test_run_cls, mock_suite_cls, db_session_mock ): last_run_id = uuid4() mock_suite_cls.get_regular.return_value = _mock_test_suite(last_complete_test_run_id=last_run_id) @@ -345,6 +352,7 @@ def test_list_test_results_by_suite_id_resolves_latest_run( r1.threshold_value = "1" r1.message = None mock_result.select_results.return_value = [r1] + mock_td_cls.select_where.return_value = [] tt = MagicMock() tt.test_type = "Alpha_Trunc" From bdaa16a5cfac9091c88e84024ab75a8d9115e2dd Mon Sep 17 00:00:00 2001 From: Aarthy Adityan Date: Fri, 26 Jun 2026 02:52:14 -0400 Subject: [PATCH 59/78] fix(test-definitions): allow column_name on referential tests in validation Referential test types (Aggregate_Balance, Combo_Match, ...) use column_name for the aggregate expression or categorical column list under test, but TestDefinition.validate() and editable_fields() only permitted column_name on column- and custom-scoped tests. This rejected any MCP create/update of a referential test that carries a column_name. Treat referential like column and custom for column_name; table scope still excludes it. Co-Authored-By: Claude Opus 4.8 --- testgen/common/models/test_definition.py | 14 ++++++++------ .../unit/common/models/test_test_definition.py | 18 +++++++++++++----- 2 files changed, 21 insertions(+), 11 deletions(-) diff --git a/testgen/common/models/test_definition.py b/testgen/common/models/test_definition.py index e5c9c4d6..23c4e447 100644 --- a/testgen/common/models/test_definition.py +++ b/testgen/common/models/test_definition.py @@ -565,9 +565,10 @@ def select_page( def editable_fields(self, test_type: TestType) -> set[str]: """Fields a caller may set or change on this test definition under the given test type.""" fields = self.EDITABLE_BASE_FIELDS | test_type.param_columns - # column_name is meaningful for column-scoped tests (the column under test) and - # custom-scoped tests (a "Test Focus" label). Other scopes don't use it. - if test_type.test_scope in ("column", "custom"): + # column_name is meaningful for column-scoped tests (the column under test), + # custom-scoped tests (a "Test Focus" label), and referential tests (the aggregate + # expression or categorical column list under test). Table-scoped tests don't use it. + if test_type.test_scope in ("column", "custom", "referential"): fields = fields | {"column_name"} # impact_dimension is overridable only for user-defined-semantic scopes # (custom-scope = user-authored SQL; referential-scope = comparison-based tests). @@ -593,9 +594,10 @@ def validate(self, test_type: TestType) -> None: f"(got `{self.severity}`)" ) - # column_name applies to column-scoped tests (the column under test) and - # custom-scoped tests (a "Test Focus" label). Other scopes don't use it. - if test_type.test_scope not in ("column", "custom") and not _is_blank(self.column_name): + # column_name applies to column-scoped tests (the column under test), + # custom-scoped tests (a "Test Focus" label), and referential tests (the aggregate + # expression or categorical column list under test). Table-scoped tests don't use it. + if test_type.test_scope not in ("column", "custom", "referential") and not _is_blank(self.column_name): errors["column_name"] = ( f"test type `{test_type.test_type}` has scope `{test_type.test_scope}`; " f"column_name does not apply to this scope" diff --git a/tests/unit/common/models/test_test_definition.py b/tests/unit/common/models/test_test_definition.py index 660f2205..d3466751 100644 --- a/tests/unit/common/models/test_test_definition.py +++ b/tests/unit/common/models/test_test_definition.py @@ -123,8 +123,9 @@ def test_editable_fields_includes_impact_dimension_only_for_custom_or_referentia assert "impact_dimension" not in td.editable_fields(table_tt) -def test_editable_fields_includes_column_name_only_for_column_or_custom_scope(): - """column_name is meaningful for column-scope (column under test) and custom-scope (label).""" +def test_editable_fields_includes_column_name_except_for_table_scope(): + """column_name is meaningful for column (column under test), custom (label), and referential + (aggregate expression / categorical column list) scopes — but not table scope.""" td = make_td() column_tt = make_test_type(scope="column", param_columns={"threshold_value"}) @@ -133,12 +134,12 @@ def test_editable_fields_includes_column_name_only_for_column_or_custom_scope(): custom_tt = make_test_type(scope="custom", param_columns={"custom_query"}) assert "column_name" in td.editable_fields(custom_tt) + referential_tt = make_test_type(scope="referential", param_columns={"match_column_names"}) + assert "column_name" in td.editable_fields(referential_tt) + table_tt = make_test_type(scope="table", param_columns=set()) assert "column_name" not in td.editable_fields(table_tt) - referential_tt = make_test_type(scope="referential", param_columns={"match_column_names"}) - assert "column_name" not in td.editable_fields(referential_tt) - def test_editable_fields_does_not_leak_identity_or_internal_columns(): tt = make_test_type(param_columns={"threshold_value"}) @@ -177,6 +178,13 @@ def test_validate_wrong_scope_column_name_rejected(): assert "column_name" in exc_info.value.errors +def test_validate_referential_scope_accepts_column_name(): + # Referential tests use column_name as the aggregate expression / categorical column list. + tt = make_test_type(code="Aggregate_Balance", scope="referential", param_columns={"match_column_names"}) + td = make_td(column_name="SUM(total_amount)", match_column_names="SUM(total_amount)") + td.validate(tt) # no raise + + def test_validate_custom_scope_accepts_column_name_as_label(): # CUSTOM uses column_name as a "Test Focus" label — must be accepted. tt = make_test_type( From 68ab291c9e0b91c9dfe191f71c15c1d38a94f1d7 Mon Sep 17 00:00:00 2001 From: Aarthy Adityan Date: Fri, 26 Jun 2026 03:13:50 -0400 Subject: [PATCH 60/78] refactor(test-definitions): address review feedback on click-through fields - Extract the http(s) link guard to a shared `isHttpUrl` in utils.js; drop the duplicate copies in the summary and results-list pages and trim the href. - Type `external_url` as `str | None` on the model and summary dataclass to match the nullable column. - Bound `custom_metadata` (<=50 keys, <=10KB serialized) via a shared `validate_custom_metadata`, enforced in the model's validate(), the import schema (the import path's only guard), and a maxLength on the authoring textarea. - Escape user-supplied metadata keys in MCP output; note why the list_test_results re-fetch needs no is_monitor filter. Co-Authored-By: Claude Opus 4.8 --- testgen/common/models/test_definition.py | 32 ++++++++++++++++--- .../test_definition_export_import_service.py | 13 +++++++- testgen/mcp/tools/test_definitions.py | 2 +- testgen/mcp/tools/test_results.py | 2 +- .../js/pages/test_definition_summary.js | 8 ++--- .../frontend/js/pages/test_definitions.js | 4 +-- .../frontend/js/pages/test_results.js | 6 ++-- testgen/ui/static/js/utils.js | 8 ++++- tests/unit/api/test_td_export_import.py | 8 +++++ .../common/models/test_test_definition.py | 20 ++++++++++++ 10 files changed, 84 insertions(+), 19 deletions(-) diff --git a/testgen/common/models/test_definition.py b/testgen/common/models/test_definition.py index 23c4e447..878bdc8c 100644 --- a/testgen/common/models/test_definition.py +++ b/testgen/common/models/test_definition.py @@ -1,3 +1,4 @@ +import json from collections.abc import Iterable from dataclasses import dataclass from datetime import UTC, datetime @@ -135,6 +136,28 @@ def _is_blank(value: object) -> bool: return value is None or value == "" +CUSTOM_METADATA_MAX_KEYS = 50 +CUSTOM_METADATA_MAX_BYTES = 10_240 + + +def validate_custom_metadata(value: object) -> str | None: + """Return an error message if ``value`` is not a valid ``custom_metadata`` payload, else ``None``. + + ``custom_metadata`` must be a JSON object (key-value pairs), bounded in key count and serialized + size. Shared by every write path — ``TestDefinition.validate`` (UI/CLI/MCP) and the + ``TestDefinitionExport`` import schema — so the rule has a single definition. + """ + if value is None: + return None + if not isinstance(value, dict): + return "must be a JSON object of key-value pairs" + if len(value) > CUSTOM_METADATA_MAX_KEYS: + return f"must have at most {CUSTOM_METADATA_MAX_KEYS} keys" + if len(json.dumps(value)) > CUSTOM_METADATA_MAX_BYTES: + return f"must be at most {CUSTOM_METADATA_MAX_BYTES} bytes when serialized as JSON" + return None + + class ParamFieldsMixin: """Parsed access to default_parm_columns/prompts/help metadata. @@ -224,7 +247,7 @@ class TestDefinitionSummary(TestTypeSummary): prediction: dict[str, dict[str, float]] | None flagged: bool impact_dimension: str | None - external_url: str + external_url: str | None custom_metadata: dict | None @property @@ -399,7 +422,7 @@ class TestDefinition(Entity): flagged: bool = Column(Boolean, default=False, nullable=False) external_id: UUID | None = Column(postgresql.UUID(as_uuid=True)) impact_dimension: str | None = Column(String, nullable=True) - external_url: str = Column(NullIfEmptyString) + external_url: str | None = Column(NullIfEmptyString) custom_metadata: dict | None = Column(postgresql.JSONB) _default_order_by = ( @@ -608,8 +631,9 @@ def validate(self, test_type: TestType) -> None: f"test type `{test_type.test_type}` does not accept a custom query" ) - if self.custom_metadata is not None and not isinstance(self.custom_metadata, dict): - errors["custom_metadata"] = "must be a JSON object of key-value pairs" + metadata_error = validate_custom_metadata(self.custom_metadata) + if metadata_error: + errors["custom_metadata"] = metadata_error for required in _required_fields_for(test_type): if _is_blank(getattr(self, required, None)): diff --git a/testgen/common/test_definition_export_import_service.py b/testgen/common/test_definition_export_import_service.py index 7e69bc9e..fb9a4e57 100644 --- a/testgen/common/test_definition_export_import_service.py +++ b/testgen/common/test_definition_export_import_service.py @@ -14,7 +14,7 @@ from testgen.common.models import get_current_session from testgen.common.models.data_table import DataTable from testgen.common.models.table_group import TableGroup -from testgen.common.models.test_definition import TestDefinition, TestType +from testgen.common.models.test_definition import TestDefinition, TestType, validate_custom_metadata from testgen.common.models.test_suite import TestSuite EXPORT_FORMAT_VERSION = 1 @@ -138,6 +138,17 @@ class TestDefinitionExport(BaseModel): def _coerce_none_to_zero(cls, v: int | None) -> int: return v if v is not None else 0 + @field_validator("custom_metadata") + @classmethod + def _check_custom_metadata(cls, v: dict[str, Any] | None) -> dict[str, Any] | None: + # The import path does not run TestDefinition.validate(), so this is the only enforcement + # of the custom_metadata shape/size bound on import. Shares validate_custom_metadata with + # the model so the rule has a single definition. + error = validate_custom_metadata(v) + if error: + raise ValueError(f"custom_metadata {error}") + return v + class ExportSource(BaseModel): project_code: str diff --git a/testgen/mcp/tools/test_definitions.py b/testgen/mcp/tools/test_definitions.py index 56d183f4..c33c1edd 100644 --- a/testgen/mcp/tools/test_definitions.py +++ b/testgen/mcp/tools/test_definitions.py @@ -249,7 +249,7 @@ def _append_td_summary(doc: MdDoc, td: TestDefinitionSummary) -> None: if td.custom_metadata: doc.heading(2, "Custom Metadata") for key, value in td.custom_metadata.items(): - doc.field(key, value) + doc.field(MdDoc.escape(key), value) # Parameters (editable fields from test type metadata) _append_parameters_section(doc, td) diff --git a/testgen/mcp/tools/test_results.py b/testgen/mcp/tools/test_results.py index dbdb719c..bb571954 100644 --- a/testgen/mcp/tools/test_results.py +++ b/testgen/mcp/tools/test_results.py @@ -165,7 +165,7 @@ def list_test_results( if custom_metadata: doc.heading(3, "Custom Metadata") for key, value in custom_metadata.items(): - doc.field(key, value) + doc.field(MdDoc.escape(key), value) return doc.render() diff --git a/testgen/ui/components/frontend/js/pages/test_definition_summary.js b/testgen/ui/components/frontend/js/pages/test_definition_summary.js index b075dd84..e8933b98 100644 --- a/testgen/ui/components/frontend/js/pages/test_definition_summary.js +++ b/testgen/ui/components/frontend/js/pages/test_definition_summary.js @@ -26,15 +26,13 @@ * @property {TestDefinition} test_definition */ import van from '/app/static/js/van.min.js'; -import { createEmitter, getValue, isEqual, loadStylesheet } from '/app/static/js/utils.js'; +import { createEmitter, getValue, isEqual, isHttpUrl, loadStylesheet } from '/app/static/js/utils.js'; import { Alert } from '/app/static/js/components/alert.js'; import { Attribute } from '/app/static/js/components/attribute.js'; import { Link } from '/app/static/js/components/link.js'; const { div, strong } = van.tags; -const isHttpUrl = (value) => typeof value === 'string' && /^https?:\/\//i.test(value.trim()); - const metadataDisplayValue = (value) => (value !== null && typeof value === 'object') ? JSON.stringify(value) : value; @@ -125,8 +123,8 @@ const TestDefinitionSummary = (props) => { label: 'External URL', value: isHttpUrl(testDefinition.external_url) ? Link({ - href: testDefinition.external_url, - label: testDefinition.external_url, + href: testDefinition.external_url.trim(), + label: testDefinition.external_url.trim(), open_new: true, underline: true, right_icon: 'open_in_new', diff --git a/testgen/ui/components/frontend/js/pages/test_definitions.js b/testgen/ui/components/frontend/js/pages/test_definitions.js index 79de48dd..d8def8b9 100644 --- a/testgen/ui/components/frontend/js/pages/test_definitions.js +++ b/testgen/ui/components/frontend/js/pages/test_definitions.js @@ -19,7 +19,7 @@ import { Icon } from '/app/static/js/components/icon.js'; import { ProfilingResultsDialog } from '../shared/profiling_results_dialog.js'; import { AXES, FACET_AXES, GROUP_BY_AXES, EMPTY, appliesToSelectedColumn } from '/app/static/js/components/test_picker_taxonomy.js'; import { enterPage, exitPage, getPageSignal } from '/app/static/js/page_lifecycle.js'; -import { jsonObject } from '/app/static/js/form_validators.js'; +import { jsonObject, maxLength } from '/app/static/js/form_validators.js'; const { button: btn, div, i: icon, span, strong } = van.tags; @@ -1482,7 +1482,7 @@ const TestDefFormContent = ({ formValues, tableColumns, testSuite, validateResul value: metadataText, placeholder: '{\n "pipeline": "daily_load",\n "task": "transform_orders"\n}', height: 120, - validators: [jsonObject], + validators: [jsonObject, maxLength(10240)], onChange: (value, state) => { metadataText.val = value; metadataValid.val = state.valid; diff --git a/testgen/ui/components/frontend/js/pages/test_results.js b/testgen/ui/components/frontend/js/pages/test_results.js index 7dda87b4..5c986d0e 100644 --- a/testgen/ui/components/frontend/js/pages/test_results.js +++ b/testgen/ui/components/frontend/js/pages/test_results.js @@ -42,7 +42,7 @@ * @property {object} filter_options */ import van from '/app/static/js/van.min.js'; -import { createEmitter, getValue, isEqual, loadStylesheet, parseDate } from '/app/static/js/utils.js'; +import { createEmitter, getValue, isEqual, isHttpUrl, loadStylesheet, parseDate } from '/app/static/js/utils.js'; import { Table } from '/app/static/js/components/table.js'; import { Select } from '/app/static/js/components/select.js'; import { Tabs, Tab } from '/app/static/js/components/tabs.js'; @@ -152,8 +152,6 @@ const formatNumber = (v) => { return n.toLocaleString(undefined, { maximumFractionDigits: 5 }); }; -const isHttpUrl = (value) => typeof value === 'string' && /^https?:\/\//i.test(value.trim()); - const buildExternalLinkCell = (url) => { if (!isHttpUrl(url)) { return ''; @@ -162,7 +160,7 @@ const buildExternalLinkCell = (url) => { return span( { class: 'flex-row fx-justify-center', onclick: (event) => event.stopPropagation() }, Link({ - href: url, + href: url.trim(), label: '', open_new: true, left_icon: 'open_in_new', diff --git a/testgen/ui/static/js/utils.js b/testgen/ui/static/js/utils.js index 71eb7403..bcd45d3b 100644 --- a/testgen/ui/static/js/utils.js +++ b/testgen/ui/static/js/utils.js @@ -189,6 +189,12 @@ function isDataURL(/** @type string */ url) { return url.startsWith('data:'); } +// Guards a user-supplied value before it is rendered as a link href, blocking +// javascript:/data:/vbscript: URIs. Only http(s) URLs may become clickable links. +function isHttpUrl(/** @type any */ value) { + return typeof value === 'string' && /^https?:\/\//i.test(value.trim()); +} + function checkIsRequired(validators) { let isRequired = validators.some(v => v.name === 'required'); if (!isRequired) { @@ -228,4 +234,4 @@ function createEmitter(setTriggerValue) { }; } -export { afterMount, createEmitter, debounce, fillViewportHeight, getRandomId, getValue, getParents, isEqual, isState, loadStylesheet, friendlyPercent, slugify, isDataURL, checkIsRequired, onFrameResized, parseDate }; +export { afterMount, createEmitter, debounce, fillViewportHeight, getRandomId, getValue, getParents, isEqual, isState, loadStylesheet, friendlyPercent, slugify, isDataURL, isHttpUrl, checkIsRequired, onFrameResized, parseDate }; diff --git a/tests/unit/api/test_td_export_import.py b/tests/unit/api/test_td_export_import.py index e56d262f..ba448956 100644 --- a/tests/unit/api/test_td_export_import.py +++ b/tests/unit/api/test_td_export_import.py @@ -1098,6 +1098,14 @@ def test_custom_metadata_rejects_non_object(self, bad_value): with pytest.raises(ValidationError): TestDefinitionExport(test_type="Alpha", custom_metadata=bad_value) + def test_custom_metadata_rejects_too_many_keys(self): + with pytest.raises(ValidationError): + TestDefinitionExport(test_type="Alpha", custom_metadata={f"k{i}": "v" for i in range(51)}) + + def test_custom_metadata_rejects_oversized(self): + with pytest.raises(ValidationError): + TestDefinitionExport(test_type="Alpha", custom_metadata={"blob": "x" * 10_241}) + # --- Click-through fields: export + import --- diff --git a/tests/unit/common/models/test_test_definition.py b/tests/unit/common/models/test_test_definition.py index d3466751..8a689b00 100644 --- a/tests/unit/common/models/test_test_definition.py +++ b/tests/unit/common/models/test_test_definition.py @@ -7,6 +7,8 @@ import pytest from testgen.common.models.test_definition import ( + CUSTOM_METADATA_MAX_BYTES, + CUSTOM_METADATA_MAX_KEYS, InvalidTestDefinitionFields, Severity, TestDefinition, @@ -251,6 +253,24 @@ def test_validate_custom_metadata_rejects_non_object(bad_value): assert "custom_metadata" in exc_info.value.errors +def test_validate_custom_metadata_rejects_too_many_keys(): + tt = make_test_type() + too_many = {f"k{i}": "v" for i in range(CUSTOM_METADATA_MAX_KEYS + 1)} + td = make_td(column_name="email", threshold_value="10", custom_metadata=too_many) + with pytest.raises(InvalidTestDefinitionFields) as exc_info: + td.validate(tt) + assert "custom_metadata" in exc_info.value.errors + + +def test_validate_custom_metadata_rejects_oversized(): + tt = make_test_type() + oversized = {"blob": "x" * (CUSTOM_METADATA_MAX_BYTES + 1)} + td = make_td(column_name="email", threshold_value="10", custom_metadata=oversized) + with pytest.raises(InvalidTestDefinitionFields) as exc_info: + td.validate(tt) + assert "custom_metadata" in exc_info.value.errors + + def test_validate_aggregates_errors(): tt = make_test_type(scope="column") td = make_td(severity="critical", custom_query="SELECT 1") # no column_name From 13d6e747dc550a61ad793a90cca0a6f1a50fdec7 Mon Sep 17 00:00:00 2001 From: Aarthy Adityan Date: Fri, 26 Jun 2026 14:57:20 -0400 Subject: [PATCH 61/78] refactor(test-definitions): split add/edit dialog into Parameters and Settings tabs The dialog had grown long and led with optional configuration ahead of the core test parameters. Restructure it: - Keep the test type name, description, and usage notes as a header above the content. - Add "Parameters" and "Settings" tabs. Parameters holds schema/table/column and the test-type parameters; Settings holds the description override, the external URL / custom metadata, the active/lock flags, and the severity/observability/impact-dimension overrides. - Drop the now-redundant test type name/description that repeated inside the form (via a new opt-in hideHeader on TestDefinitionForm; the monitors editor keeps its header). - Add an optional external activeTab state to the shared Tabs component and hoist it to the dialog so the selected tab survives the form's re-render on each field change. Co-Authored-By: Claude Opus 4.8 --- .../frontend/js/pages/test_definitions.js | 365 +++++++++--------- testgen/ui/static/js/components/tabs.js | 6 +- .../js/components/test_definition_form.js | 17 +- 3 files changed, 202 insertions(+), 186 deletions(-) diff --git a/testgen/ui/components/frontend/js/pages/test_definitions.js b/testgen/ui/components/frontend/js/pages/test_definitions.js index d8def8b9..e1bfdbfb 100644 --- a/testgen/ui/components/frontend/js/pages/test_definitions.js +++ b/testgen/ui/components/frontend/js/pages/test_definitions.js @@ -11,6 +11,7 @@ import { Attribute } from '/app/static/js/components/attribute.js'; import { TestDefinitionForm } from '/app/static/js/components/test_definition_form.js'; import { RunTestsDialog } from '/app/static/js/components/run_tests_dialog.js'; import { Textarea } from '/app/static/js/components/textarea.js'; +import { Tabs, Tab } from '/app/static/js/components/tabs.js'; import { Checkbox } from '/app/static/js/components/checkbox.js'; import { DropdownButton } from '/app/static/js/components/dropdown_button.js'; import { TestDefinitionNotes } from './test_definition_notes.js'; @@ -895,6 +896,8 @@ const AddDialogComponent = ({ open, info, validateResult, onClose }, emit) => { const step = van.state(1); const selectedTestType = van.state(null); const formValues = van.state({}); + // Hoisted so the active form tab survives the form re-render on each field change. + const formActiveTab = van.state(0); // ---- Picker state ---- const searchQuery = van.state(''); @@ -1041,6 +1044,7 @@ const AddDialogComponent = ({ open, info, validateResult, onClose }, emit) => { } } formValues.val = fv; + formActiveTab.val = 0; step.val = 2; }; @@ -1292,6 +1296,7 @@ const AddDialogComponent = ({ open, info, validateResult, onClose }, emit) => { qualifiesTableRefsWithSchema: qualifiesTableRefsWithSchema.rawVal, validateResult: getValue(validateResult), mode: 'add', + activeTab: formActiveTab, onFormChange: (changes) => { formValues.val = { ...formValues.rawVal, ...changes }; }, onValidate: () => emit('ValidateTest', { payload: formValues.rawVal }), onSave: () => emit('AddTestSaved', { payload: formValues.rawVal }), @@ -1316,6 +1321,8 @@ const EditDialogComponent = ({ open, info, validateResult: validateResultProp, o const validateResult = van.derive(() => getValue(validateResultProp) ?? null); const formValues = van.state(null); + // Hoisted so the active form tab survives the form re-render on each field change. + const formActiveTab = van.state(0); const initFormFromInfo = () => { const di = dialogInfo.rawVal; @@ -1328,6 +1335,7 @@ const EditDialogComponent = ({ open, info, validateResult: validateResultProp, o column_name_prompt: ttRow.column_name_prompt ?? null, column_name_help: ttRow.column_name_help ?? null, }; + formActiveTab.val = 0; }; // Reset form when dialog opens (closed→open), clear when it closes @@ -1359,6 +1367,7 @@ const EditDialogComponent = ({ open, info, validateResult: validateResultProp, o qualifiesTableRefsWithSchema: qualifiesTableRefsWithSchema.rawVal, validateResult: vr, mode: 'edit', + activeTab: formActiveTab, onFormChange: (changes) => { formValues.val = { ...formValues.rawVal, ...changes }; }, @@ -1372,7 +1381,7 @@ const EditDialogComponent = ({ open, info, validateResult: validateResultProp, o }; // Shared form content for add/edit dialogs -const TestDefFormContent = ({ formValues, tableColumns, testSuite, validateResult, mode, qualifiesTableRefsWithSchema, onFormChange, onValidate, onSave, onCancel, onBack }) => { +const TestDefFormContent = ({ formValues, tableColumns, testSuite, validateResult, mode, qualifiesTableRefsWithSchema, activeTab, onFormChange, onValidate, onSave, onCancel, onBack }) => { const testScope = formValues.test_scope ?? 'column'; const runType = formValues.run_type ?? 'CAT'; const testType = formValues.test_type ?? ''; @@ -1433,8 +1442,8 @@ const TestDefFormContent = ({ formValues, tableColumns, testSuite, validateResul return div( { class: 'flex-column fx-gap-3' }, - // Test type header (add mode) or read-only test type (edit mode) - mode === 'add' && formValues.test_name_short + // Header — test type identity, shown above the tabs in both add and edit modes + formValues.test_name_short ? div( { class: 'mb-1' }, div({ class: 'text-large' }, formValues.test_name_short), @@ -1444,195 +1453,197 @@ const TestDefFormContent = ({ formValues, tableColumns, testSuite, validateResul ) : null, - mode === 'edit' - ? Input({ - name: 'test_type_display', - label: 'Test Type', - value: formValues.test_name_short ?? formValues.test_type ?? '', - disabled: true, - }) - : null, - formValues.usage_notes ? Alert({ type: 'info' }, strong({ class: 'mb-1' }, 'Usage Notes'), div({}, formValues.usage_notes)) : null, - // Description override - Textarea({ - name: 'test_description', - label: 'Test Description Override', - value: () => fv.val.test_description ?? '', - placeholder: `Inherited (${formValues.default_test_description ?? ''})`, - height: 72, - onChange: (value) => updateField('test_description', value || null), - }), + Tabs( + { activeTab }, + Tab( + { label: 'Parameters' }, + + // Schema (read-only) + qualifiesTableRefsWithSchema + ? Input({ + name: 'schema_name', + label: 'Schema', + value: formValues.schema_name ?? '', + disabled: true, + }) + : null, - // External link + custom metadata - Input({ - name: 'external_url', - label: 'External URL', - help: 'Optional link to the code, pipeline step, or system that produces this data.', - value: () => fv.val.external_url ?? '', - onChange: (value) => updateField('external_url', value || null), - }), - Textarea({ - name: 'custom_metadata', - label: 'Custom Metadata', - help: 'Optional JSON object of key-value pairs, e.g. {"pipeline": "daily_load", "task": "transform_orders"}.', - value: metadataText, - placeholder: '{\n "pipeline": "daily_load",\n "task": "transform_orders"\n}', - height: 120, - validators: [jsonObject, maxLength(10240)], - onChange: (value, state) => { - metadataText.val = value; - metadataValid.val = state.valid; - if (state.valid) { - updateField('custom_metadata', value && value.trim() ? JSON.parse(value) : null); - } - }, - }), + // Table name + testScope !== 'tablegroup' + ? testScope === 'custom' + ? Input({ + name: 'table_name', + label: 'Table', + value: () => fv.val.table_name ?? '', + onChange: (value) => updateField('table_name', value || null), + }) + : () => Select({ + label: 'Table', + value: fv.val.table_name ?? null, + options: tableNameOptions, + allowNull: true, + filterable: true, + disabled: mode === 'edit', + onChange: (value) => { + updateField('table_name', value); + updateField('column_name', null); + }, + }) + : null, - // Checkboxes - div( - { class: 'flex-row fx-gap-4' }, - Checkbox({ - label: 'Test Active', - checked: () => fv.val.test_active ?? true, - onChange: (v) => updateField('test_active', v), - }), - Checkbox({ - label: 'Lock Refresh', - checked: () => fv.val.lock_refresh ?? false, - onChange: (v) => updateField('lock_refresh', v), - }), - ), + // Column name (scope-dependent) + testScope === 'column' + ? () => Select({ + label: 'Column', + value: fv.val.column_name ?? null, + options: columnNameOptions.val, + allowNull: true, + filterable: true, + onChange: (value) => updateField('column_name', value), + }) + : testScope === 'referential' || testScope === 'custom' + ? Input({ + name: 'column_name', + label: columnLabel, + help: columnHelp, + value: () => fv.val.column_name ?? '', + onChange: (value) => updateField('column_name', value || null), + }) + : null, + + // Validation status (edit mode only) + mode === 'edit' && formValues.test_definition_status + ? Input({ + name: 'test_definition_status', + label: 'Validation Status', + value: formValues.test_definition_status || 'OK', + disabled: true, + }) + : null, - // Severity + Observability + Impact Dimension selects - div( - { class: 'flex-row fx-gap-3 fx-flex-wrap' }, - div( - { style: 'flex: calc(50% - 8px) 0 0;' }, - () => Select({ - label: 'Urgency Override', - value: fv.val.severity ?? null, - options: severityOptions, - allowNull: false, - onChange: (value) => updateField('severity', value), - }), - ), - div( - { style: 'flex: calc(50% - 8px) 0 0;' }, - () => Select({ - label: 'Send to Observability - Override', - value: fv.val.export_to_observability ?? null, - options: obsOptions, - allowNull: false, - onChange: (value) => updateField('export_to_observability', value), - }), + // Dynamic parameter fields + div( + { class: 'td-form-params-section' }, + TestDefinitionForm({ + definition: formValues, + qualifiesTableRefsWithSchema, + hideHeader: true, + onChange: (changes) => { + if (Object.keys(changes).length === 0) return; + const updated = { ...fv.rawVal, ...changes }; + fv.val = updated; + onFormChange(changes); + }, + }), + ), + + // Skip errors (QUERY run type only) + runType === 'QUERY' + ? Input({ + name: 'skip_errors', + label: 'Threshold Error Count', + type: 'number', + value: () => fv.val.skip_errors ?? 0, + step: 1, + onChange: (value) => updateField('skip_errors', value ?? 0), + }) + : null, ), - showImpactDimensionOverride ? div( - { style: 'flex: calc(50% - 8px) 0 0;' }, - () => Select({ - label: 'Impact Dimension Override', - value: fv.val.impact_dimension ?? null, - options: impactDimensionOptions, - allowNull: false, - helpText: 'Override the default impact classification for this test. Affects how the test result is categorized in score breakdowns.', - onChange: (value) => updateField('impact_dimension', value), + Tab( + { label: 'Settings' }, + + // Description override + Textarea({ + name: 'test_description', + label: 'Test Description Override', + value: () => fv.val.test_description ?? '', + placeholder: `Inherited (${formValues.default_test_description ?? ''})`, + height: 72, + onChange: (value) => updateField('test_description', value || null), }), - ) : null, - ), - // Schema (read-only) - qualifiesTableRefsWithSchema - ? Input({ - name: 'schema_name', - label: 'Schema', - value: formValues.schema_name ?? '', - disabled: true, - }) - : null, - - // Table name - testScope !== 'tablegroup' - ? testScope === 'custom' - ? Input({ - name: 'table_name', - label: 'Table', - value: () => fv.val.table_name ?? '', - onChange: (value) => updateField('table_name', value || null), - }) - : () => Select({ - label: 'Table', - value: fv.val.table_name ?? null, - options: tableNameOptions, - allowNull: true, - filterable: true, - disabled: mode === 'edit', - onChange: (value) => { - updateField('table_name', value); - updateField('column_name', null); + // External link + custom metadata + Input({ + name: 'external_url', + label: 'External URL', + help: 'Optional link to the code, pipeline step, or system that produces this data.', + value: () => fv.val.external_url ?? '', + onChange: (value) => updateField('external_url', value || null), + }), + Textarea({ + name: 'custom_metadata', + label: 'Custom Metadata', + help: 'Optional JSON object of key-value pairs, e.g. {"pipeline": "daily_load", "task": "transform_orders"}.', + value: metadataText, + placeholder: '{\n "pipeline": "daily_load",\n "task": "transform_orders"\n}', + height: 120, + validators: [jsonObject, maxLength(10240)], + onChange: (value, state) => { + metadataText.val = value; + metadataValid.val = state.valid; + if (state.valid) { + updateField('custom_metadata', value && value.trim() ? JSON.parse(value) : null); + } }, - }) - : null, + }), - // Column name (scope-dependent) - testScope === 'column' - ? () => Select({ - label: 'Column', - value: fv.val.column_name ?? null, - options: columnNameOptions.val, - allowNull: true, - filterable: true, - onChange: (value) => updateField('column_name', value), - }) - : testScope === 'referential' || testScope === 'custom' - ? Input({ - name: 'column_name', - label: columnLabel, - help: columnHelp, - value: () => fv.val.column_name ?? '', - onChange: (value) => updateField('column_name', value || null), - }) - : null, - - // Validation status (edit mode only) - mode === 'edit' && formValues.test_definition_status - ? Input({ - name: 'test_definition_status', - label: 'Validation Status', - value: formValues.test_definition_status || 'OK', - disabled: true, - }) - : null, + // Checkboxes + div( + { class: 'flex-row fx-gap-4' }, + Checkbox({ + label: 'Test Active', + checked: () => fv.val.test_active ?? true, + onChange: (v) => updateField('test_active', v), + }), + Checkbox({ + label: 'Lock Refresh', + checked: () => fv.val.lock_refresh ?? false, + onChange: (v) => updateField('lock_refresh', v), + }), + ), - // Dynamic parameter fields - div( - { class: 'td-form-params-section' }, - TestDefinitionForm({ - definition: formValues, - qualifiesTableRefsWithSchema, - onChange: (changes) => { - if (Object.keys(changes).length === 0) return; - const updated = { ...fv.rawVal, ...changes }; - fv.val = updated; - onFormChange(changes); - }, - }), + // Severity + Observability + Impact Dimension selects + div( + { class: 'flex-row fx-gap-3 fx-flex-wrap' }, + div( + { style: 'flex: calc(50% - 8px) 0 0;' }, + () => Select({ + label: 'Urgency Override', + value: fv.val.severity ?? null, + options: severityOptions, + allowNull: false, + onChange: (value) => updateField('severity', value), + }), + ), + div( + { style: 'flex: calc(50% - 8px) 0 0;' }, + () => Select({ + label: 'Send to Observability - Override', + value: fv.val.export_to_observability ?? null, + options: obsOptions, + allowNull: false, + onChange: (value) => updateField('export_to_observability', value), + }), + ), + showImpactDimensionOverride ? div( + { style: 'flex: calc(50% - 8px) 0 0;' }, + () => Select({ + label: 'Impact Dimension Override', + value: fv.val.impact_dimension ?? null, + options: impactDimensionOptions, + allowNull: false, + helpText: 'Override the default impact classification for this test. Affects how the test result is categorized in score breakdowns.', + onChange: (value) => updateField('impact_dimension', value), + }), + ) : null, + ), + ), ), - // Skip errors (QUERY run type only) - runType === 'QUERY' - ? Input({ - name: 'skip_errors', - label: 'Threshold Error Count', - type: 'number', - value: () => fv.val.skip_errors ?? 0, - step: 1, - onChange: (value) => updateField('skip_errors', value ?? 0), - }) - : null, - // Validate feedback validateResult ? Alert({ type: validateResult.success ? 'success' : 'error' }, validateResult.message) diff --git a/testgen/ui/static/js/components/tabs.js b/testgen/ui/static/js/components/tabs.js index d6cbc51c..ae23f7b1 100644 --- a/testgen/ui/static/js/components/tabs.js +++ b/testgen/ui/static/js/components/tabs.js @@ -20,6 +20,8 @@ const Tab = ({ label }, ...children) => ({ /** * @typedef {Object} TabsProps * @property {string?} class + * @property {van.State?} activeTab Optional external state for the active tab index. Pass one + * to control or persist the selection across re-renders; omit to use internal state. * * @param {TabsProps} props * @param {...Tab} tabs @@ -27,9 +29,9 @@ const Tab = ({ label }, ...children) => ({ const Tabs = (props, ...tabs) => { loadStylesheet('tabs', stylesheet); - const { ...restProps } = props; + const { activeTab: activeTabProp, ...restProps } = props; - const activeTab = van.state(0); + const activeTab = activeTabProp ?? van.state(0); let labelsContainerEl; const highlightEl = span({ class: "tg-tabs--highlight" }); diff --git a/testgen/ui/static/js/components/test_definition_form.js b/testgen/ui/static/js/components/test_definition_form.js index e7fa14d6..ec3265b3 100644 --- a/testgen/ui/static/js/components/test_definition_form.js +++ b/testgen/ui/static/js/components/test_definition_form.js @@ -61,6 +61,7 @@ * @property {TestDefinition} definition * @property {string?} class * @property {boolean} qualifiesTableRefsWithSchema + * @property {boolean?} hideHeader Hide the test-type name/description header (when the host already shows it). * @property {(changes: object, valid: boolean) => void} onChange */ @@ -143,13 +144,15 @@ const TestDefinitionForm = (/** @type Properties */ props) => { return div( { class: props.class }, - div( - { class: 'mb-2' }, - div({ class: 'text-large' }, definition.test_name_short), - definition.test_description || definition.default_test_description - ? span({ class: 'text-caption mt-2' }, definition.test_description ?? definition.default_test_description) - : null, - ), + getValue(props.hideHeader) + ? null + : div( + { class: 'mb-2' }, + div({ class: 'text-large' }, definition.test_name_short), + definition.test_description || definition.default_test_description + ? span({ class: 'text-caption mt-2' }, definition.test_description ?? definition.default_test_description) + : null, + ), () => div( { class: 'flex-row fx-flex-wrap fx-gap-3' }, dynamicParamColumns.map(config => { From a66c85fc6126640e494b466140d80f22f012bb7c Mon Sep 17 00:00:00 2001 From: Aarthy Adityan Date: Fri, 26 Jun 2026 15:50:59 -0400 Subject: [PATCH 62/78] fix(test-definitions): restore field spacing in add/edit dialog tabs The tab panels render into the Tabs component's content container, which has no flex gap, so the per-field spacing the dialog had as a flex-column was lost. Wrap each tab's fields in a flex-column fx-gap-3 container to restore it. Co-Authored-By: Claude Opus 4.8 --- .../ui/components/frontend/js/pages/test_definitions.js | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/testgen/ui/components/frontend/js/pages/test_definitions.js b/testgen/ui/components/frontend/js/pages/test_definitions.js index e1bfdbfb..af8d602d 100644 --- a/testgen/ui/components/frontend/js/pages/test_definitions.js +++ b/testgen/ui/components/frontend/js/pages/test_definitions.js @@ -1461,6 +1461,8 @@ const TestDefFormContent = ({ formValues, tableColumns, testSuite, validateResul { activeTab }, Tab( { label: 'Parameters' }, + div( + { class: 'flex-column fx-gap-3' }, // Schema (read-only) qualifiesTableRefsWithSchema @@ -1552,9 +1554,12 @@ const TestDefFormContent = ({ formValues, tableColumns, testSuite, validateResul onChange: (value) => updateField('skip_errors', value ?? 0), }) : null, + ), ), Tab( { label: 'Settings' }, + div( + { class: 'flex-column fx-gap-3' }, // Description override Textarea({ @@ -1641,6 +1646,7 @@ const TestDefFormContent = ({ formValues, tableColumns, testSuite, validateResul }), ) : null, ), + ), ), ), @@ -1941,7 +1947,7 @@ stylesheet.replace(` .td-form-params-section { border-top: 1px solid var(--border-color); padding-top: 12px; - margin-top: 4px; + margin-top: 16px; } .tg-test-picker { display: flex; flex-direction: column; gap: 12px; height: 70vh; min-height: 0; } From 48fc2c366f01685e99e6f7103de038e183659bcd Mon Sep 17 00:00:00 2001 From: Aarthy Adityan Date: Fri, 26 Jun 2026 16:16:56 -0400 Subject: [PATCH 63/78] fix(test-definitions): import Tooltip in Link and drop external-URL underline Address review feedback on the click-through link: - Link referenced Tooltip without importing it, so any Link given a `tooltip` (the test-results external-link cell) threw at render. Import it from tooltip.js. - In the result test-definition summary, the external URL used `underline` plus a full-width link, so the hover underline spanned the panel and read as a divider. Drop the underline (the open_in_new icon is enough affordance) and size the link to content (max-width instead of width) so long URLs still wrap. Co-Authored-By: Claude Opus 4.8 --- .../ui/components/frontend/js/pages/test_definition_summary.js | 2 -- testgen/ui/static/js/components/link.js | 1 + 2 files changed, 1 insertion(+), 2 deletions(-) diff --git a/testgen/ui/components/frontend/js/pages/test_definition_summary.js b/testgen/ui/components/frontend/js/pages/test_definition_summary.js index e8933b98..7f25cba7 100644 --- a/testgen/ui/components/frontend/js/pages/test_definition_summary.js +++ b/testgen/ui/components/frontend/js/pages/test_definition_summary.js @@ -126,7 +126,6 @@ const TestDefinitionSummary = (props) => { href: testDefinition.external_url.trim(), label: testDefinition.external_url.trim(), open_new: true, - underline: true, right_icon: 'open_in_new', right_icon_size: 16, }) @@ -174,7 +173,6 @@ stylesheet.replace(` overflow-wrap: anywhere; } .external-url-attribute .tg-link { - width: 100%; max-width: 100%; } .external-url-attribute .tg-link--wrapper { diff --git a/testgen/ui/static/js/components/link.js b/testgen/ui/static/js/components/link.js index c78e821c..e709f601 100644 --- a/testgen/ui/static/js/components/link.js +++ b/testgen/ui/static/js/components/link.js @@ -20,6 +20,7 @@ * @property {((event: any) => void)?} onClick */ import { getValue, loadStylesheet } from '../utils.js'; +import { Tooltip } from './tooltip.js'; import van from '../van.min.js'; const { a, div, i, span } = van.tags; From 57cf36e962a6b5a06716e08e2ee3f4947562f5f6 Mon Sep 17 00:00:00 2001 From: Ricardo Boni Date: Fri, 19 Jun 2026 16:09:55 -0400 Subject: [PATCH 64/78] feat(oauth): personal access token schema and model --- testgen/api/oauth/models.py | 45 ---------- testgen/api/oauth/routes.py | 2 +- testgen/api/oauth/server.py | 23 +---- testgen/common/auth.py | 8 +- testgen/common/models/oauth.py | 90 +++++++++++++++++++ testgen/settings.py | 11 +++ .../030_initialize_new_schema_structure.sql | 7 +- .../dbupgrade/0198_incremental_upgrade.sql | 15 ++++ testgen/ui/app.py | 1 + testgen/ui/bootstrap.py | 4 +- testgen/ui/components/widgets/sidebar.py | 2 + testgen/ui/navigation/page.py | 1 + testgen/ui/static/js/sidebar.js | 83 +++++++++++++++-- tests/unit/api/oauth/test_server.py | 61 +------------ .../models/test_oauth.py} | 32 ++++++- 15 files changed, 241 insertions(+), 144 deletions(-) delete mode 100644 testgen/api/oauth/models.py create mode 100644 testgen/common/models/oauth.py create mode 100644 testgen/template/dbupgrade/0198_incremental_upgrade.sql rename tests/unit/{api/oauth/test_models.py => common/models/test_oauth.py} (53%) diff --git a/testgen/api/oauth/models.py b/testgen/api/oauth/models.py deleted file mode 100644 index 49b2c961..00000000 --- a/testgen/api/oauth/models.py +++ /dev/null @@ -1,45 +0,0 @@ -import time - -from authlib.integrations.sqla_oauth2 import ( - OAuth2AuthorizationCodeMixin, - OAuth2ClientMixin, - OAuth2TokenMixin, -) -from sqlalchemy import Column, ForeignKey, String -from sqlalchemy.dialects import postgresql - -from testgen import settings -from testgen.common.models import Base - - -class OAuth2Client(Base, OAuth2ClientMixin): - __tablename__ = "oauth2_clients" - - id = Column(postgresql.UUID(as_uuid=True), primary_key=True, server_default="gen_random_uuid()") - user_id = Column(postgresql.UUID(as_uuid=True), ForeignKey("auth_users.id", ondelete="SET NULL"), nullable=True) - - # Override to widen — JWTs can exceed 255 chars - # (the mixin defines client_id as VARCHAR(48) which is fine) - - -class OAuth2AuthorizationCode(Base, OAuth2AuthorizationCodeMixin): - __tablename__ = "oauth2_authorization_codes" - - id = Column(postgresql.UUID(as_uuid=True), primary_key=True, server_default="gen_random_uuid()") - user_id = Column(postgresql.UUID(as_uuid=True), ForeignKey("auth_users.id", ondelete="CASCADE"), nullable=False) - - -class OAuth2Token(Base, OAuth2TokenMixin): - __tablename__ = "oauth2_tokens" - - id = Column(postgresql.UUID(as_uuid=True), primary_key=True, server_default="gen_random_uuid()") - user_id = Column(postgresql.UUID(as_uuid=True), ForeignKey("auth_users.id", ondelete="CASCADE"), nullable=True) - - # Override to allow longer JWTs as access tokens - access_token = Column(String(2048), unique=True, nullable=False) - - def is_refresh_token_active(self) -> bool: - if self.refresh_token_revoked_at: - return False - expires_at = self.issued_at + settings.REFRESH_TOKEN_EXPIRES_IN - return expires_at >= time.time() diff --git a/testgen/api/oauth/routes.py b/testgen/api/oauth/routes.py index d2ea828b..da228383 100644 --- a/testgen/api/oauth/routes.py +++ b/testgen/api/oauth/routes.py @@ -21,10 +21,10 @@ from testgen import settings from testgen.api.deps import db_session from testgen.api.oauth.login import render_login_page -from testgen.api.oauth.models import OAuth2Client from testgen.api.oauth.server import TestGenAuthorizationServer from testgen.common.auth import create_jwt_token, decode_jwt_token, verify_password from testgen.common.models import get_current_session +from testgen.common.models.oauth import OAuth2Client from testgen.common.models.user import User LOG = logging.getLogger("testgen") diff --git a/testgen/api/oauth/server.py b/testgen/api/oauth/server.py index 32ca120a..0d7859cb 100644 --- a/testgen/api/oauth/server.py +++ b/testgen/api/oauth/server.py @@ -2,7 +2,6 @@ Grant types: - Authorization Code + PKCE (for MCP clients) -- Client Credentials (for automation scripts) - Refresh Token (for token renewal) All DB operations use get_current_session() for thread-local session access. @@ -13,15 +12,14 @@ from typing import ClassVar from authlib.oauth2.rfc6749 import AuthorizationServer, JsonRequest, OAuth2Request, grants -from authlib.oauth2.rfc6749.errors import InvalidGrantError from authlib.oauth2.rfc7009 import RevocationEndpoint from authlib.oauth2.rfc7636 import CodeChallenge from sqlalchemy import select from testgen import settings -from testgen.api.oauth.models import OAuth2AuthorizationCode, OAuth2Client, OAuth2Token from testgen.common.auth import create_jwt_token from testgen.common.models import get_current_session +from testgen.common.models.oauth import OAuth2AuthorizationCode, OAuth2Client, OAuth2Token from testgen.common.models.user import User @@ -85,24 +83,6 @@ def revoke_old_credential(self, credential): credential.access_token_revoked_at = int(time.time()) -class ClientCredentialsGrant(grants.ClientCredentialsGrant): - """Client credentials grant that resolves the client's owner as the token user. - - Ensures every token has a real User identity — no "ghost" usernames. - """ - - def validate_token_request(self): - super().validate_token_request() - client = self.request.client - if not client.user_id: - raise InvalidGrantError(description="Client has no registered owner.") - session = get_current_session() - owner = session.scalars(select(User).where(User.id == client.user_id)).first() - if owner is None: - raise InvalidGrantError(description="Client owner no longer exists.") - self.request.user = owner - - class TestGenRevocationEndpoint(RevocationEndpoint): def query_token(self, token_string, token_type_hint): session = get_current_session() @@ -183,7 +163,6 @@ def create_authorization_server() -> TestGenAuthorizationServer: """Create and configure the authorization server with all grant types.""" server = TestGenAuthorizationServer() server.register_grant(AuthorizationCodeGrant, [CodeChallenge(required=True)]) - server.register_grant(ClientCredentialsGrant) server.register_grant(RefreshTokenGrant) server.register_endpoint(TestGenRevocationEndpoint) server.register_token_generator("default", _generate_bearer_token) diff --git a/testgen/common/auth.py b/testgen/common/auth.py index 141b032e..1df03684 100644 --- a/testgen/common/auth.py +++ b/testgen/common/auth.py @@ -4,8 +4,11 @@ import bcrypt import jwt +from sqlalchemy import func, select from testgen import settings +from testgen.common.models.oauth import OAuth2Token +from testgen.common.models.user import User LOG = logging.getLogger("testgen") @@ -45,11 +48,6 @@ def authorize_token(token_str: str, username: str, session): Shared implementation for API and MCP authorization. """ - from sqlalchemy import func, select - - from testgen.api.oauth.models import OAuth2Token - from testgen.common.models.user import User - user = session.scalars(select(User).where(func.lower(User.username) == func.lower(username))).first() if user is None: raise AuthError("User not found") diff --git a/testgen/common/models/oauth.py b/testgen/common/models/oauth.py new file mode 100644 index 00000000..f8f6856a --- /dev/null +++ b/testgen/common/models/oauth.py @@ -0,0 +1,90 @@ +import time +from enum import StrEnum + +from authlib.integrations.sqla_oauth2 import ( + OAuth2AuthorizationCodeMixin, + OAuth2ClientMixin, + OAuth2TokenMixin, +) +from sqlalchemy import Column, ForeignKey, String, case, func +from sqlalchemy.dialects import postgresql +from sqlalchemy.ext.hybrid import hybrid_property + +from testgen import settings +from testgen.common.models import Base + + +class OAuth2ClientType(StrEnum): + """How an OAuth2 client was provisioned. + + PAT clients are created by the authenticated personal-access-token path and owned by + a user; EXTERNAL clients are dynamically registered (MCP apps, automation scripts). + """ + + PAT = "pat" + EXTERNAL = "external" + + +class PersonalAccessTokenStatus(StrEnum): + """Display status of a personal access token. Revoked takes precedence over expired.""" + + ACTIVE = "active" + REVOKED = "revoked" + EXPIRED = "expired" + + +class OAuth2Client(Base, OAuth2ClientMixin): + __tablename__ = "oauth2_clients" + + id = Column(postgresql.UUID(as_uuid=True), primary_key=True, server_default="gen_random_uuid()") + client_type = Column(String(20), nullable=False, default=OAuth2ClientType.EXTERNAL) + + # Override to widen — JWTs can exceed 255 chars + # (the mixin defines client_id as VARCHAR(48) which is fine) + + +class OAuth2AuthorizationCode(Base, OAuth2AuthorizationCodeMixin): + __tablename__ = "oauth2_authorization_codes" + + id = Column(postgresql.UUID(as_uuid=True), primary_key=True, server_default="gen_random_uuid()") + user_id = Column(postgresql.UUID(as_uuid=True), ForeignKey("auth_users.id", ondelete="CASCADE"), nullable=False) + + +class OAuth2Token(Base, OAuth2TokenMixin): + __tablename__ = "oauth2_tokens" + + id = Column(postgresql.UUID(as_uuid=True), primary_key=True, server_default="gen_random_uuid()") + user_id = Column(postgresql.UUID(as_uuid=True), ForeignKey("auth_users.id", ondelete="CASCADE"), nullable=True) + + # Override to allow longer JWTs as access tokens + access_token = Column(String(2048), unique=True, nullable=False) + + # Set only for personal access tokens; NULL for tokens from the auth-code / MCP flows + name = Column(String(255), nullable=True) + + def is_refresh_token_active(self) -> bool: + if self.refresh_token_revoked_at: + return False + expires_at = self.issued_at + settings.REFRESH_TOKEN_EXPIRES_IN + return expires_at >= time.time() + + @hybrid_property + def status(self) -> PersonalAccessTokenStatus: + """Personal access token status, derived from revocation and expiry. + + Instance side reads the app clock; the SQL expression reads the DB clock — + both are wall-clock "now", evaluated independently. + """ + if self.access_token_revoked_at: + return PersonalAccessTokenStatus.REVOKED + if self.issued_at + self.expires_in < time.time(): + return PersonalAccessTokenStatus.EXPIRED + return PersonalAccessTokenStatus.ACTIVE + + @status.expression # type: ignore[no-redef] + def status(cls): + return case( + (cls.access_token_revoked_at > 0, PersonalAccessTokenStatus.REVOKED.value), + (cls.issued_at + cls.expires_in < func.extract("epoch", func.now()), PersonalAccessTokenStatus.EXPIRED.value), + else_=PersonalAccessTokenStatus.ACTIVE.value, + ) diff --git a/testgen/settings.py b/testgen/settings.py index ad3ed162..cf760cba 100644 --- a/testgen/settings.py +++ b/testgen/settings.py @@ -533,6 +533,17 @@ def _ssl_files_present() -> bool: Lifetime of OAuth access and refresh tokens. """ +PAT_DEFAULT_LIFETIME_SECONDS: int = 31_536_000 # 365 days +""" +Default maximum lifetime a personal access token can be created with. +""" + +PAT_MAX_LIFETIME_SECONDS: int = 63_072_000 # 2 years (730 days) +""" +Absolute ceiling on the admin-configurable personal access token maximum lifetime. +The configured maximum cannot exceed this. +""" + JWT_HASHING_KEY_B64: str = getenv("TG_JWT_HASHING_KEY") """ Random key used to sign/verify the authentication token diff --git a/testgen/template/dbsetup/030_initialize_new_schema_structure.sql b/testgen/template/dbsetup/030_initialize_new_schema_structure.sql index a4c84fc4..75f37b15 100644 --- a/testgen/template/dbsetup/030_initialize_new_schema_structure.sql +++ b/testgen/template/dbsetup/030_initialize_new_schema_structure.sql @@ -742,12 +742,12 @@ CREATE INDEX ix_pm_role ON project_memberships(role); CREATE TABLE oauth2_clients ( id UUID PRIMARY KEY DEFAULT gen_random_uuid(), - user_id UUID REFERENCES auth_users(id) ON DELETE SET NULL, client_id VARCHAR(48) NOT NULL UNIQUE, client_secret VARCHAR(120), client_id_issued_at INTEGER NOT NULL DEFAULT 0, client_secret_expires_at INTEGER NOT NULL DEFAULT 0, - client_metadata TEXT NOT NULL DEFAULT '{}' + client_metadata TEXT NOT NULL DEFAULT '{}', + client_type VARCHAR(20) NOT NULL DEFAULT 'external' ); CREATE INDEX idx_oauth2_clients_client_id ON oauth2_clients(client_id); @@ -778,7 +778,8 @@ CREATE TABLE oauth2_tokens ( issued_at INTEGER NOT NULL DEFAULT EXTRACT(EPOCH FROM NOW())::INTEGER, access_token_revoked_at INTEGER NOT NULL DEFAULT 0, refresh_token_revoked_at INTEGER NOT NULL DEFAULT 0, - expires_in INTEGER NOT NULL DEFAULT 0 + expires_in INTEGER NOT NULL DEFAULT 0, + name VARCHAR(255) ); CREATE INDEX idx_oauth2_tokens_refresh_token ON oauth2_tokens(refresh_token); diff --git a/testgen/template/dbupgrade/0198_incremental_upgrade.sql b/testgen/template/dbupgrade/0198_incremental_upgrade.sql new file mode 100644 index 00000000..1677eb3e --- /dev/null +++ b/testgen/template/dbupgrade/0198_incremental_upgrade.sql @@ -0,0 +1,15 @@ +SET SEARCH_PATH TO {SCHEMA_NAME}; + +-- Personal access tokens are stored as oauth2_tokens rows; the name labels each +-- one in the user's token list. NULL marks a non-PAT token (auth-code / MCP flows). +ALTER TABLE oauth2_tokens ADD COLUMN name VARCHAR(255); + +-- Classify clients so personal-access-token clients are distinguishable from +-- dynamically-registered external clients (MCP apps, automation). Existing rows +-- all predate PATs and are externally registered, so they default to 'external'. +ALTER TABLE oauth2_clients ADD COLUMN client_type VARCHAR(20) NOT NULL DEFAULT 'external'; + +-- Drop the client owner: it was only ever read by the client-credentials grant +-- (removed). User identity now rides the token (oauth2_tokens.user_id + the JWT +-- username), so a client needs no owner. The FK constraint drops with the column. +ALTER TABLE oauth2_clients DROP COLUMN user_id; diff --git a/testgen/ui/app.py b/testgen/ui/app.py index 5bbd67eb..dc26e6cf 100644 --- a/testgen/ui/app.py +++ b/testgen/ui/app.py @@ -84,6 +84,7 @@ def render(log_level: int = logging.INFO): support_email=settings.SUPPORT_EMAIL, global_context=is_global_context, is_global_admin=session.auth.user_has_permission("global_admin") and bool(application.global_admin_paths), + account_path=application.account_path, ) application.router.run() diff --git a/testgen/ui/bootstrap.py b/testgen/ui/bootstrap.py index bea240d5..d5e5db54 100644 --- a/testgen/ui/bootstrap.py +++ b/testgen/ui/bootstrap.py @@ -49,13 +49,14 @@ class Application(singleton.Singleton): - def __init__(self, auth_class: Authentication, logo: plugins.Logo, router: Router, menu: Menu, logger: logging.Logger, global_admin_paths: frozenset[str]) -> None: + def __init__(self, auth_class: Authentication, logo: plugins.Logo, router: Router, menu: Menu, logger: logging.Logger, global_admin_paths: frozenset[str], account_path: str | None) -> None: self.auth_class = auth_class self.logo = logo self.router = router self.menu = menu self.logger = logger self.global_admin_paths = global_admin_paths + self.account_path = account_path def run(log_level: int = logging.INFO) -> Application: @@ -91,4 +92,5 @@ def run(log_level: int = logging.INFO) -> Application: ), logger=LOG, global_admin_paths=frozenset(page.path for page in pages if page.permission == "global_admin"), + account_path=next((page.path for page in pages if page.is_account_page), None), ) diff --git a/testgen/ui/components/widgets/sidebar.py b/testgen/ui/components/widgets/sidebar.py index 6806d22c..2b68a3b6 100644 --- a/testgen/ui/components/widgets/sidebar.py +++ b/testgen/ui/components/widgets/sidebar.py @@ -25,6 +25,7 @@ def sidebar( support_email: str | None = None, global_context: bool = False, is_global_admin: bool = False, + account_path: str | None = None, ) -> None: """ Testgen custom component to display a styled menu over streamlit's @@ -53,6 +54,7 @@ def sidebar( "support_email": support_email, "global_context": global_context, "is_global_admin": is_global_admin, + "account_path": account_path, }, on_Navigate_change=_on_navigate, ) diff --git a/testgen/ui/navigation/page.py b/testgen/ui/navigation/page.py index 11e93f06..9f4de9a8 100644 --- a/testgen/ui/navigation/page.py +++ b/testgen/ui/navigation/page.py @@ -21,6 +21,7 @@ class Page(abc.ABC): menu_item: MenuItem | None = None permission: Permission | None = "view" can_activate: typing.ClassVar[list[CanActivateGuard] | None] = None + is_account_page: bool = False def __init__(self, router: testgen.ui.navigation.router.Router) -> None: self.router = router diff --git a/testgen/ui/static/js/sidebar.js b/testgen/ui/static/js/sidebar.js index 9ebd4b2f..89006e5d 100644 --- a/testgen/ui/static/js/sidebar.js +++ b/testgen/ui/static/js/sidebar.js @@ -36,6 +36,7 @@ * @property {string} support_email * @property {boolean} global_context * @property {boolean} is_global_admin + * @property {(string|null)} account_path */ const van = window.top.van; const { a, button, div, i, img, label, option, select, span } = van.tags; @@ -93,11 +94,7 @@ const Sidebar = (/** @type {Properties} */ props) => { }, ), div( - div( - { class: 'menu--user' }, - span({class: 'menu--username', title: props.username}, props.username), - span({class: 'menu--role'}, () => props.role.val?.replace('_', ' ')), - ), + () => UserBlock(props, currentProject), div( { class: 'menu--buttons' }, button( @@ -167,6 +164,36 @@ const ProjectSelect = (/** @type Project[] */ projects, /** @type string */ curr ); }; +const UserBlock = (/** @type Properties */ props, currentProject) => { + const roleAndProject = () => { + const role = props.role.val?.replace('_', ' '); + const projectName = currentProject.val?.name; + return [role, projectName].filter(Boolean).join(' · '); + }; + + const userInfo = div( + { class: 'menu--user--info' }, + span({ class: 'menu--username', title: props.username }, props.username), + span({ class: 'menu--role' }, roleAndProject), + ); + + const accountPath = props.account_path?.val; + if (!accountPath) { + return div({ class: 'menu--user' }, userInfo); + } + + return a( + { + class: () => isCurrentPage(accountPath, props.current_page?.val) ? 'menu--user--link active' : 'menu--user--link', + href: `/${accountPath}?${PROJECT_CODE_QUERY_PARAM}=${currentProject.val?.code ?? ''}`, + onclick: (event) => navigate(event, accountPath, { [PROJECT_CODE_QUERY_PARAM]: currentProject.val?.code }), + }, + i({ class: 'menu--user--icon material-symbols-rounded' }, 'manage_accounts'), + userInfo, + i({ class: 'menu--user--chevron material-symbols-rounded' }, 'chevron_right'), + ); +}; + const MenuSection = ( /** @type {MenuItem} */ item, /** @type {string} */ currentPage, @@ -336,6 +363,49 @@ stylesheet.replace(` padding: 16px; } +.menu .menu--user--link { + display: flex; + flex-direction: row; + align-items: center; + gap: 12px; + margin: 4px 8px; + padding: 8px 12px; + border-radius: 8px; + color: var(--primary-text-color); + text-decoration: none; + cursor: pointer; +} + +.menu .menu--user--link:hover { + background: var(--sidebar-item-hover-color); +} + +.menu .menu--user--link.active { + color: var(--primary-color); + background: var(--sidebar-active-item-color); +} + +.menu .menu--user--info { + display: flex; + flex-direction: column; + flex: 1; + min-width: 0; +} + +.menu .menu--user--icon { + flex: none; + font-size: 24px; + line-height: 24px; +} + +.menu .menu--user--chevron { + flex: none; + margin-left: auto; + font-size: 18px; + line-height: 18px; + color: var(--secondary-text-color); +} + .menu .menu--username { overflow-x: hidden; text-overflow: ellipsis; @@ -346,6 +416,9 @@ stylesheet.replace(` text-transform: uppercase; font-size: 12px; color: var(--secondary-text-color); + overflow-x: hidden; + text-overflow: ellipsis; + text-wrap: nowrap; } .menu .content > .menu--section > .menu--section--label { diff --git a/tests/unit/api/oauth/test_server.py b/tests/unit/api/oauth/test_server.py index c9810e15..83e5bf20 100644 --- a/tests/unit/api/oauth/test_server.py +++ b/tests/unit/api/oauth/test_server.py @@ -5,12 +5,10 @@ from uuid import uuid4 import pytest -from authlib.oauth2.rfc6749 import grants from testgen import settings from testgen.api.oauth.server import ( AuthorizationCodeGrant, - ClientCredentialsGrant, RefreshTokenGrant, TestGenAuthorizationServer, TestGenRevocationEndpoint, @@ -335,63 +333,6 @@ def test_generate_bearer_token_no_scope_omits_field(mock_jwt): assert "scope" not in token -# --- ClientCredentialsGrant --- - - -@patch("testgen.api.oauth.server.get_current_session") -def test_client_credentials_resolves_owner(mock_get_session): - mock_session = MagicMock() - mock_get_session.return_value = mock_session - - mock_owner = MagicMock() - mock_owner.username = "owner_user" - mock_session.scalars.return_value.first.return_value = mock_owner - - grant = ClientCredentialsGrant.__new__(ClientCredentialsGrant) - grant.request = MagicMock() - grant.request.client.user_id = uuid4() - grant.request.client.check_grant_type.return_value = True - - with patch.object(grants.ClientCredentialsGrant, "validate_token_request"): - grant.validate_token_request() - - assert grant.request.user is mock_owner - - -@patch("testgen.api.oauth.server.get_current_session") -def test_client_credentials_rejects_client_without_owner(mock_get_session): - from authlib.oauth2.rfc6749.errors import InvalidGrantError - - grant = ClientCredentialsGrant.__new__(ClientCredentialsGrant) - grant.request = MagicMock() - grant.request.client.user_id = None - - with ( - patch.object(grants.ClientCredentialsGrant, "validate_token_request"), - pytest.raises(InvalidGrantError), - ): - grant.validate_token_request() - - -@patch("testgen.api.oauth.server.get_current_session") -def test_client_credentials_rejects_deleted_owner(mock_get_session): - from authlib.oauth2.rfc6749.errors import InvalidGrantError - - mock_session = MagicMock() - mock_get_session.return_value = mock_session - mock_session.scalars.return_value.first.return_value = None - - grant = ClientCredentialsGrant.__new__(ClientCredentialsGrant) - grant.request = MagicMock() - grant.request.client.user_id = uuid4() - - with ( - patch.object(grants.ClientCredentialsGrant, "validate_token_request"), - pytest.raises(InvalidGrantError), - ): - grant.validate_token_request() - - # --- RevocationEndpoint --- @@ -460,6 +401,6 @@ def test_create_authorization_server_returns_configured_server(): assert isinstance(server, TestGenAuthorizationServer) token_grant_classes = [entry[0] for entry in server._token_grants] assert RefreshTokenGrant in token_grant_classes - assert ClientCredentialsGrant in token_grant_classes + assert AuthorizationCodeGrant in token_grant_classes assert "revocation" in server._endpoints diff --git a/tests/unit/api/oauth/test_models.py b/tests/unit/common/models/test_oauth.py similarity index 53% rename from tests/unit/api/oauth/test_models.py rename to tests/unit/common/models/test_oauth.py index d546feb6..bf9a5f7f 100644 --- a/tests/unit/api/oauth/test_models.py +++ b/tests/unit/common/models/test_oauth.py @@ -1,8 +1,8 @@ -"""Tests for testgen.api.oauth.models — OAuth2 ORM model business logic.""" +"""Tests for testgen.common.models.oauth — OAuth2 ORM model business logic.""" import time -from testgen.api.oauth.models import OAuth2Token +from testgen.common.models.oauth import OAuth2Token, PersonalAccessTokenStatus def _make_token(**overrides): @@ -37,3 +37,31 @@ def test_is_refresh_token_active_ignores_access_revocation(): def test_is_refresh_token_active_returns_false_when_expired(): token = _make_token(issued_at=int(time.time()) - (31 * 86400)) assert token.is_refresh_token_active() is False + + +# --- status (hybrid_property) --- + + +def test_status_active_when_unrevoked_and_unexpired(): + token = _make_token(issued_at=int(time.time()) - 10, expires_in=3600) + assert token.status == PersonalAccessTokenStatus.ACTIVE + + +def test_status_expired_when_past_expiry(): + token = _make_token(issued_at=int(time.time()) - 3600, expires_in=60) + assert token.status == PersonalAccessTokenStatus.EXPIRED + + +def test_status_revoked_when_access_revoked(): + token = _make_token(access_token_revoked_at=int(time.time())) + assert token.status == PersonalAccessTokenStatus.REVOKED + + +def test_status_revoked_takes_precedence_over_expired(): + # Past expiry AND revoked: revoked wins. + token = _make_token( + issued_at=int(time.time()) - 3600, + expires_in=60, + access_token_revoked_at=int(time.time()), + ) + assert token.status == PersonalAccessTokenStatus.REVOKED From fbd021fc9939924ed4b63a80bf384ac5fb900de0 Mon Sep 17 00:00:00 2001 From: Diogo Basto Date: Fri, 26 Jun 2026 16:20:39 +0100 Subject: [PATCH 65/78] feat(mcp): add CRUD write tools for projects and test suites (TG-1070) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Three new MCP write tools in core and one in the enterprise project-management plugin, so an LLM can manage projects and test suites without falling back to the UI: - update_project (core, administer permission): updates the full per-project settings surface exposed in the UI — display name, weighted DQ scoring toggle, DataOps Observability URL + key, data retention enable/days, and retention schedule cron + timezone. Side effects mirror the UI: weights flip triggers a background score recalc, retention toggles upsert/delete the scheduled cleanup job. observability_api_key is consumed but never echoed back (mcp-patterns secrets rule). - create_test_suite (core, edit permission): under a table group. Accepts the full UI-editable surface in one call — name, description, default severity, observability export toggle, DataOps Observability component fields, and scoring exclusion — so the LLM doesn't need a follow-up update_test_suite to finish configuring a new suite. Always non-monitor. - update_test_suite (core, edit permission): same surface, partial updates. Resolves through the existing monitor-filtered resolve_test_suite so a monitor-suite ID surfaces the unified "not found or not accessible" wording. Empty-string args on NullIfEmptyString columns normalize to None on the way in so the diff never reports a phantom change. create_project (enterprise plugin, global_admin) registers through the TG-1124 plugin MCP-tools hook in testgen_project_management. Mirrors the UI's "Add Project" handler: persists the row and schedules the data-retention cleanup job at the project's defaults so MCP-created projects behave the same as UI-created ones. Project creation is enterprise-only because it belongs to the project-management feature. global_admin is a User flag, not a per-project role, so the core mcp_permission decorator was extended to special-case it. Every other call site stays unchanged. Co-Authored-By: Claude Opus 4.7 --- testgen/mcp/permissions.py | 22 +- testgen/mcp/server.py | 5 + testgen/mcp/tools/common.py | 11 +- testgen/mcp/tools/projects.py | 262 +++++++++++++ testgen/mcp/tools/test_suites.py | 280 +++++++++++++ tests/unit/mcp/test_permissions.py | 48 +++ tests/unit/mcp/test_tools_common.py | 28 ++ tests/unit/mcp/test_tools_projects.py | 430 ++++++++++++++++++++ tests/unit/mcp/test_tools_test_suites.py | 479 +++++++++++++++++++++++ 9 files changed, 1559 insertions(+), 6 deletions(-) create mode 100644 testgen/mcp/tools/projects.py create mode 100644 testgen/mcp/tools/test_suites.py create mode 100644 tests/unit/mcp/test_tools_projects.py create mode 100644 tests/unit/mcp/test_tools_test_suites.py diff --git a/testgen/mcp/permissions.py b/testgen/mcp/permissions.py index 77a597ba..d65483a1 100644 --- a/testgen/mcp/permissions.py +++ b/testgen/mcp/permissions.py @@ -138,17 +138,29 @@ def mcp_permission(permission: str) -> Callable: Raises ``MCPPermissionDenied`` if the user has no projects with the required permission. Other ``MCPPermissionDenied`` exceptions from tool code propagate through — the ``safe_tool`` error boundary handles conversion to text. + + ``global_admin`` is a user-level flag, not a per-project role: the check + consults ``User.is_global_admin`` and the resulting ``ProjectPermissions`` + has empty ``memberships``. Tools gated on ``global_admin`` operate above + project scope and should not call ``get_project_permissions()``. """ def decorator(fn: Callable) -> Callable: @functools.wraps(fn) def wrapper(*args, **kwargs): user = get_authorized_mcp_user() - perms = _compute_project_permissions(user, permission) - if not perms.allowed_codes: - raise MCPPermissionDenied( - "Your role does not include the necessary permission for this operation on any project." - ) + if permission == "global_admin": + if not user.is_global_admin: + raise MCPPermissionDenied( + "Your role does not include the necessary permission for this operation." + ) + perms = ProjectPermissions(memberships={}, permission=permission, username=user.username) + else: + perms = _compute_project_permissions(user, permission) + if not perms.allowed_codes: + raise MCPPermissionDenied( + "Your role does not include the necessary permission for this operation on any project." + ) tok = _mcp_project_permissions.set(perms) try: return fn(*args, **kwargs) diff --git a/testgen/mcp/server.py b/testgen/mcp/server.py index 45ed6102..cf8ac6ff 100644 --- a/testgen/mcp/server.py +++ b/testgen/mcp/server.py @@ -204,6 +204,7 @@ def build_mcp_server( list_profiling_summaries, search_columns, ) + from testgen.mcp.tools.projects import update_project from testgen.mcp.tools.quality_scores import ( create_scorecard, delete_scorecard, @@ -263,6 +264,7 @@ def build_mcp_server( update_test_result, ) from testgen.mcp.tools.test_runs import get_test_run, list_test_runs + from testgen.mcp.tools.test_suites import create_test_suite, update_test_suite if server_url is None: server_url = f"{api_base_url}/mcp" @@ -375,6 +377,9 @@ def safe_prompt(fn): safe_tool(create_table_group) safe_tool(update_table_group) safe_tool(preview_table_group) + safe_tool(update_project) + safe_tool(create_test_suite) + safe_tool(update_test_suite) safe_tool(update_catalog_metadata) # Resources diff --git a/testgen/mcp/tools/common.py b/testgen/mcp/tools/common.py index e038da6a..581ef7e6 100644 --- a/testgen/mcp/tools/common.py +++ b/testgen/mcp/tools/common.py @@ -38,7 +38,7 @@ from testgen.common.models.scheduler import SCHEDULABLE_JOB_KEYS, JobSchedule from testgen.common.models.scores import ScoreCategory, ScoreDefinition from testgen.common.models.table_group import TableGroup -from testgen.common.models.test_definition import TestDefinition, TestDefinitionNote, TestType +from testgen.common.models.test_definition import Severity, TestDefinition, TestDefinitionNote, TestType from testgen.common.models.test_result import TestResult, TestResultStatus from testgen.common.models.test_suite import TestSuite from testgen.mcp.exceptions import MCPResourceNotAccessible, MCPUserError @@ -132,6 +132,15 @@ def parse_quality_dimension(value: str) -> QualityDimension: raise MCPUserError(f"Invalid quality_dimension `{value}`. Valid values: {valid}") from err +def parse_severity(value: str) -> Severity: + """Validate a test-suite default severity. Accepts ``Fail`` or ``Warning``.""" + try: + return Severity(value) + except ValueError as err: + valid = ", ".join(s.value for s in Severity) + raise MCPUserError(f"Invalid severity `{value}`. Valid values: {valid}") from err + + class ScoreGroupBy(StrEnum): """User-facing values accepted for the ``group_by`` argument on quality-score rollups.""" diff --git a/testgen/mcp/tools/projects.py b/testgen/mcp/tools/projects.py new file mode 100644 index 00000000..2e60e3eb --- /dev/null +++ b/testgen/mcp/tools/projects.py @@ -0,0 +1,262 @@ +"""MCP write tools for projects. + +Project creation and deletion live in the enterprise project-management plugin +(gated on ``global_admin``); editing a project's own settings is a per-project +``administer`` operation that ships in core. This module holds the latter. +""" + +from __future__ import annotations + +from typing import Any + +from testgen.common.enums import JobKey, JobSource +from testgen.common.models import with_database_session +from testgen.common.models.job_execution import JobExecution +from testgen.common.models.project import Project +from testgen.common.models.scheduler import ( + DEFAULT_DATA_CLEANUP_CRON, + DEFAULT_RETENTION_CRON_TZ, + JobSchedule, +) +from testgen.mcp.exceptions import MCPResourceNotAccessible, MCPUserError +from testgen.mcp.permissions import get_project_permissions, mcp_permission +from testgen.mcp.tools.common import DocGroup, raise_validation_error +from testgen.mcp.tools.markdown import MdDoc + +_DOC_GROUP = DocGroup.MANAGE + +DEFAULT_RETENTION_DAYS = 180 + + +def _empty_to_none(value: str | None) -> str | None: + return value if value else None + + +@with_database_session +@mcp_permission("administer") +def update_project( + project_code: str, + *, + project_name: str | None = None, + use_dq_score_weights: bool | None = None, + observability_api_url: str | None = None, + observability_api_key: str | None = None, + data_retention_enabled: bool | None = None, + data_retention_days: int | None = None, + retention_cron_expr: str | None = None, + retention_cron_tz: str | None = None, +) -> str: + """Update a project's settings. + + Args: + project_code: The project code, e.g. from `list_projects`. + project_name: New display name. + use_dq_score_weights: Whether quality scoring uses the configured weights. + Changing this re-runs project-wide score recalculation in the background. + observability_api_url: DataOps Observability API URL. Pass an empty string to clear. + observability_api_key: DataOps Observability API key. Pass an empty string to clear. + Never echoed back in tool output. + data_retention_enabled: Whether old profiling and test history is automatically deleted. + data_retention_days: How many days of history to keep (only when retention is enabled). + Defaults to 180 when retention is being enabled without specifying days. + retention_cron_expr: Cron expression for the retention cleanup job + (only when retention is enabled). Defaults to daily at 01:00. + retention_cron_tz: Timezone for the retention cleanup cron + (only when retention is enabled). Defaults to UTC. + """ + supplied = { + "project_name": project_name, + "use_dq_score_weights": use_dq_score_weights, + "observability_api_url": observability_api_url, + "observability_api_key": observability_api_key, + "data_retention_enabled": data_retention_enabled, + "data_retention_days": data_retention_days, + "retention_cron_expr": retention_cron_expr, + "retention_cron_tz": retention_cron_tz, + } + if all(value is None for value in supplied.values()): + raise MCPUserError("No fields supplied to update.") + + perms = get_project_permissions() + perms.verify_access(project_code, not_found=MCPResourceNotAccessible("Project", project_code)) + + project = Project.get(project_code) + if project is None: + raise MCPResourceNotAccessible("Project", project_code) + + schedule = JobSchedule.get( + JobSchedule.project_code == project_code, + JobSchedule.key == JobKey.run_data_cleanup, + ) + + # Compute effective retention state from supplied args + current state. Errors below. + effective_enabled = ( + data_retention_enabled if data_retention_enabled is not None else project.data_retention_enabled + ) + effective_days = ( + data_retention_days if data_retention_days is not None else project.data_retention_days + ) + # When retention is being newly enabled without explicit days, fall back to the system default. + if effective_enabled and effective_days is None: + effective_days = DEFAULT_RETENTION_DAYS + + # Validate field-by-field; surface all errors at once. + errors: list[str] = [] + cleaned_name: str | None = None + if project_name is not None: + cleaned_name = project_name.strip() + if not cleaned_name: + errors.append("project_name: must not be empty.") + + if data_retention_days is not None and data_retention_days < 1: + errors.append("data_retention_days: must be a positive integer.") + + # Setting retention schedule args only makes sense when retention will be enabled. + if not effective_enabled: + if data_retention_days is not None: + errors.append("data_retention_days: cannot be set when data_retention_enabled is False.") + if retention_cron_expr is not None: + errors.append("retention_cron_expr: cannot be set when data_retention_enabled is False.") + if retention_cron_tz is not None: + errors.append("retention_cron_tz: cannot be set when data_retention_enabled is False.") + + if errors: + raise_validation_error(errors, "Update rejected. No changes saved.") + + weights_were = project.use_dq_score_weights + + # Snapshot the editable surface (including schedule cron + tz, which live on JobSchedule). + before = _snapshot(project, schedule) + + # Apply project changes. + if cleaned_name is not None: + project.project_name = cleaned_name + if use_dq_score_weights is not None: + project.use_dq_score_weights = use_dq_score_weights + if observability_api_url is not None: + project.observability_api_url = _empty_to_none(observability_api_url) + if observability_api_key is not None: + project.observability_api_key = _empty_to_none(observability_api_key) + project.data_retention_enabled = effective_enabled + project.data_retention_days = effective_days if effective_enabled else None + + # Compute the effective cron values now so the snapshot reflects what the schedule will be. + schedule_supplied = any( + v is not None + for v in (data_retention_enabled, data_retention_days, retention_cron_expr, retention_cron_tz) + ) + effective_cron_expr: str | None + effective_cron_tz: str | None + if effective_enabled: + effective_cron_expr = ( + retention_cron_expr or (schedule.cron_expr if schedule else None) or DEFAULT_DATA_CLEANUP_CRON + ) + effective_cron_tz = ( + retention_cron_tz or (schedule.cron_tz if schedule else None) or DEFAULT_RETENTION_CRON_TZ + ) + else: + effective_cron_expr = None + effective_cron_tz = None + + after = _snapshot( + project, + _ScheduleSnapshot(cron_expr=effective_cron_expr, cron_tz=effective_cron_tz) + if effective_enabled + else None, + ) + changed = {attr for attr in before if before[attr] != after[attr]} + + doc = MdDoc() + doc.heading(1, f"Project `{project_code}` updated") + + if not changed: + doc.text("No fields changed — supplied values matched the current state.") + return doc.render() + + project.save() + + # Schedule side effects: only touch JobSchedule when a retention-related arg was supplied. + if schedule_supplied: + if effective_enabled: + JobSchedule.upsert_for_retention( + project_code=project_code, + retention_days=effective_days, # type: ignore[arg-type] + cron_expr=effective_cron_expr, # type: ignore[arg-type] + cron_tz=effective_cron_tz, # type: ignore[arg-type] + ) + else: + JobSchedule.delete_for_retention(project_code) + + # Weights side effect: submit a background recalculation job, same as the UI. + if use_dq_score_weights is not None and use_dq_score_weights != weights_were: + JobExecution.submit( + job_key=JobKey.recalculate_project_scores, + kwargs={"project_code": project_code}, + source=JobSource.mcp, + project_code=project_code, + ) + + rows: list[list[object]] = [] + for attr in _DIFF_ORDER: + if attr not in changed: + continue + rows.append([_DIFF_LABELS[attr], _render_field_value(attr, before[attr]), _render_field_value(attr, after[attr])]) + doc.table(["Field", "Before", "After"], rows, code=[0]) + return doc.render() + + +class _ScheduleSnapshot: + """Stand-in for a JobSchedule row, used while computing the post-update snapshot + before any DB write has happened.""" + + def __init__(self, cron_expr: str | None, cron_tz: str | None) -> None: + self.cron_expr = cron_expr + self.cron_tz = cron_tz + + +def _snapshot(project: Project, schedule: Any) -> dict[str, Any]: + return { + "project_name": project.project_name, + "use_dq_score_weights": project.use_dq_score_weights, + "observability_api_url": project.observability_api_url, + "observability_api_key": project.observability_api_key, + "data_retention_enabled": project.data_retention_enabled, + "data_retention_days": project.data_retention_days, + "retention_cron_expr": schedule.cron_expr if schedule is not None else None, + "retention_cron_tz": schedule.cron_tz if schedule is not None else None, + } + + +_DIFF_ORDER: tuple[str, ...] = ( + "project_name", + "use_dq_score_weights", + "observability_api_url", + "observability_api_key", + "data_retention_enabled", + "data_retention_days", + "retention_cron_expr", + "retention_cron_tz", +) + +_DIFF_LABELS: dict[str, str] = { + "project_name": "Name", + "use_dq_score_weights": "Weighted quality scoring", + "observability_api_url": "DataOps Observability API URL", + "observability_api_key": "DataOps Observability API key", + "data_retention_enabled": "Data retention", + "data_retention_days": "Retention days", + "retention_cron_expr": "Retention cron expression", + "retention_cron_tz": "Retention timezone", +} + + +def _render_field_value(attr: str, value: Any) -> str | None: + # Redact the API key — mcp-patterns "Secrets in inputs" rule. The diff still tells the LLM + # the field changed; it just never echoes the value (cleartext OR masked). + if attr == "observability_api_key" and value: + return "[secret]" + if isinstance(value, bool): + return "Yes" if value else "No" + if value is None or value == "": + return None + return str(value) diff --git a/testgen/mcp/tools/test_suites.py b/testgen/mcp/tools/test_suites.py new file mode 100644 index 00000000..a5773de3 --- /dev/null +++ b/testgen/mcp/tools/test_suites.py @@ -0,0 +1,280 @@ +"""MCP write tools for test suites — create and update regular (non-monitor) suites. + +Monitor suites (``is_monitor=True``) are managed exclusively through the dedicated +monitors surface. ``resolve_test_suite`` already filters them out, so an +``update_test_suite`` call against a monitor-suite ID surfaces the unified +``not found or not accessible`` error. +""" + +from __future__ import annotations + +from typing import Any + +from testgen.common.models import get_current_session, with_database_session +from testgen.common.models.test_definition import Severity +from testgen.common.models.test_suite import TestSuite +from testgen.mcp.exceptions import MCPUserError +from testgen.mcp.permissions import mcp_permission +from testgen.mcp.tools.common import ( + DocGroup, + parse_severity, + raise_validation_error, + resolve_table_group, + resolve_test_suite, +) +from testgen.mcp.tools.markdown import MdDoc + +_DOC_GROUP = DocGroup.MANAGE + +# Free-text + component fields are stored as ``NullIfEmptyString`` columns: writing ``""`` +# to one becomes ``NULL`` in the DB. Normalize on the way in so the diff-renderer doesn't +# flag a phantom change (``before=None`` → ``after=""``) on a value that the DB will read +# back as identical to before. +_OPTIONAL_TEXT_ATTRS = ( + "test_suite_description", + "component_key", + "component_type", + "component_name", +) + + +def _empty_to_none(value: str | None) -> str | None: + return value if value else None + + +@with_database_session +@mcp_permission("edit") +def create_test_suite( + table_group_id: str, + test_suite_name: str, + *, + description: str | None = None, + severity_default: str | None = None, + export_to_observability: bool = False, + component_key: str | None = None, + component_type: str | None = "dataset", + component_name: str | None = None, + dq_score_exclude: bool = False, +) -> str: + """Create a test suite under a table group. + + The new suite inherits the table group's project and connection. + + Args: + table_group_id: UUID of the table group, e.g. from `list_table_groups`. + test_suite_name: Display name for the suite. + description: Optional free-text description. + severity_default: Optional default severity applied to tests that do + not set their own. Accepts `Fail` or `Warning`. + export_to_observability: Whether to export test results to the + configured DataOps Observability backend. Defaults to False. + component_key: Component identifier in DataOps Observability. + component_type: Component type in DataOps Observability + (e.g. `dataset`). Defaults to `dataset`. + component_name: Component display name in DataOps Observability. + dq_score_exclude: Whether to exclude this suite's results from data + quality scoring. Defaults to False. + """ + name = test_suite_name.strip() if test_suite_name else "" + errors: list[str] = [] + if not name: + errors.append("test_suite_name: must not be empty.") + parsed_severity: Severity | None = _parse_severity_field(severity_default, errors) + if errors: + raise_validation_error(errors, "Test suite creation rejected. No changes saved.") + + table_group = resolve_table_group(table_group_id) + + suite = TestSuite( + project_code=table_group.project_code, + connection_id=table_group.connection_id, + table_groups_id=table_group.id, + test_suite=name, + test_suite_description=_empty_to_none(description), + severity=parsed_severity.value if parsed_severity is not None else None, + export_to_observability=export_to_observability, + component_key=_empty_to_none(component_key), + component_type=_empty_to_none(component_type), + component_name=_empty_to_none(component_name), + dq_score_exclude=dq_score_exclude, + is_monitor=False, + ) + session = get_current_session() + session.add(suite) + # Flush so the autogenerated ``id`` is available for the rendered response. + session.flush() + + return _render_created_suite(suite, table_group_name=table_group.table_groups_name) + + +@with_database_session +@mcp_permission("edit") +def update_test_suite( + test_suite_id: str, + *, + test_suite_name: str | None = None, + description: str | None = None, + severity_default: str | None = None, + export_to_observability: bool | None = None, + component_key: str | None = None, + component_type: str | None = None, + component_name: str | None = None, + dq_score_exclude: bool | None = None, +) -> str: + """Update fields on a test suite. Atomic — nothing saved unless every supplied value is valid. + + Args: + test_suite_id: UUID of the test suite, e.g. from `list_test_suites`. + test_suite_name: New display name. + description: New free-text description. Pass an empty string to clear. + severity_default: New default severity for tests that do not set their + own. Accepts `Fail` or `Warning`. + export_to_observability: Whether test results are exported to the + configured DataOps Observability backend. + component_key: Component identifier in DataOps Observability. + Pass an empty string to clear. + component_type: Component type in DataOps Observability + (e.g. `dataset`). Pass an empty string to clear. + component_name: Component display name in DataOps Observability. + Pass an empty string to clear. + dq_score_exclude: Whether this suite's results are excluded from data + quality scoring. + """ + supplied = { + "test_suite_name": test_suite_name, + "description": description, + "severity_default": severity_default, + "export_to_observability": export_to_observability, + "component_key": component_key, + "component_type": component_type, + "component_name": component_name, + "dq_score_exclude": dq_score_exclude, + } + if all(value is None for value in supplied.values()): + raise MCPUserError("No fields supplied to update.") + + suite = resolve_test_suite(test_suite_id) + + # Validate every supplied field first; surface all errors at once. + errors: list[str] = [] + updates: dict[str, Any] = {} + + if test_suite_name is not None: + cleaned = test_suite_name.strip() + if not cleaned: + errors.append("test_suite_name: must not be empty.") + else: + updates["test_suite"] = cleaned + + if description is not None: + updates["test_suite_description"] = _empty_to_none(description) + + if severity_default is not None: + parsed = _parse_severity_field(severity_default, errors) + if parsed is not None: + updates["severity"] = parsed.value + + if export_to_observability is not None: + updates["export_to_observability"] = export_to_observability + if component_key is not None: + updates["component_key"] = _empty_to_none(component_key) + if component_type is not None: + updates["component_type"] = _empty_to_none(component_type) + if component_name is not None: + updates["component_name"] = _empty_to_none(component_name) + if dq_score_exclude is not None: + updates["dq_score_exclude"] = dq_score_exclude + + if errors: + raise_validation_error(errors, "Update rejected. No changes saved.") + + # Diff scoped to the attributes actually touched by this call. + before = {attr: getattr(suite, attr, None) for attr in updates} + for attr, value in updates.items(): + setattr(suite, attr, value) + after = {attr: getattr(suite, attr, None) for attr in updates} + changed = {attr for attr in before if before[attr] != after[attr]} + + doc = MdDoc() + doc.heading(1, f"Test Suite `{suite.test_suite}` updated") + doc.field("ID", str(suite.id), code=True) + + if not changed: + doc.text("No fields changed — supplied values matched the current state.") + return doc.render() + + suite.save() + + rows: list[list[object]] = [] + for attr in _DIFF_ORDER: + if attr not in changed: + continue + rows.append([_DIFF_LABELS[attr], _render_field_value(before[attr]), _render_field_value(after[attr])]) + doc.table(["Field", "Before", "After"], rows, code=[0]) + return doc.render() + + +def _parse_severity_field(value: str | None, errors: list[str]) -> Severity | None: + """Validate a ``severity_default`` arg; append a field-scoped bullet on failure.""" + if value is None: + return None + try: + return parse_severity(value) + except MCPUserError: + valid = ", ".join(s.value for s in Severity) + errors.append(f"severity_default: must be one of {valid} (got `{value}`).") + return None + + +def _render_created_suite(suite: TestSuite, *, table_group_name: str) -> str: + doc = MdDoc() + doc.heading(1, f"Test Suite `{suite.test_suite}` created") + doc.field("ID", str(suite.id), code=True) + doc.field("Project", suite.project_code, code=True) + doc.field("Table Group", f"{table_group_name} (`{suite.table_groups_id}`)") + if suite.test_suite_description: + doc.field("Description", suite.test_suite_description) + if suite.severity: + doc.field("Default severity", suite.severity) + doc.field("Export to Observability", suite.export_to_observability) + if suite.component_key: + doc.field("Component key", suite.component_key) + if suite.component_type: + doc.field("Component type", suite.component_type) + if suite.component_name: + doc.field("Component name", suite.component_name) + doc.field("Exclude from quality scoring", suite.dq_score_exclude) + return doc.render() + + +# Order drives the diff-table rendering. Mirrors the update_test_suite arg order: +# observability + components grouped together, dq_score_exclude separate. +_DIFF_ORDER: tuple[str, ...] = ( + "test_suite", + "test_suite_description", + "severity", + "export_to_observability", + "component_key", + "component_type", + "component_name", + "dq_score_exclude", +) + +_DIFF_LABELS: dict[str, str] = { + "test_suite": "Name", + "test_suite_description": "Description", + "severity": "Default severity", + "export_to_observability": "Export to Observability", + "component_key": "Component key", + "component_type": "Component type", + "component_name": "Component name", + "dq_score_exclude": "Exclude from quality scoring", +} + + +def _render_field_value(value: Any) -> str | None: + if isinstance(value, bool): + return "Yes" if value else "No" + if value is None or value == "": + return None + return str(value) diff --git a/tests/unit/mcp/test_permissions.py b/tests/unit/mcp/test_permissions.py index bad81379..f0815d4a 100644 --- a/tests/unit/mcp/test_permissions.py +++ b/tests/unit/mcp/test_permissions.py @@ -332,3 +332,51 @@ def my_tool(x: int, y: str = "default") -> str: assert my_tool.__name__ == "my_tool" assert my_tool.__doc__ == "Tool docstring." + + +# --- mcp_permission("global_admin") --- + + +def test_mcp_permission_global_admin_allows_when_flag_set(mcp_user): + """global_admin gates on User.is_global_admin, not on project memberships.""" + set_mcp_username("test") + mcp_user.is_global_admin = True + + captured = {} + + @mcp_permission("global_admin") + def tool_fn(): + captured["perms"] = get_project_permissions() + return "ok" + + assert tool_fn() == "ok" + assert captured["perms"].permission == "global_admin" + assert captured["perms"].memberships == {} + + +def test_mcp_permission_global_admin_denies_when_flag_false(mcp_user): + set_mcp_username("test") + mcp_user.is_global_admin = False + + @mcp_permission("global_admin") + def tool_fn(): + raise AssertionError("Should not be called") + + with pytest.raises(MCPPermissionDenied, match="necessary permission"): + tool_fn() + + +def test_mcp_permission_global_admin_allows_with_zero_memberships(mcp_user): + """A global admin with no project memberships still passes the gate — the standard + 'no allowed_codes' bailout would otherwise refuse a fresh super-user.""" + set_mcp_username("test") + mcp_user.is_global_admin = True + + # global_admin path must NOT call _compute_project_permissions — the tool still runs + # even though the conftest's membership mock would yield zero allowed_codes for an + # unknown permission. + @mcp_permission("global_admin") + def tool_fn(): + return "ok" + + assert tool_fn() == "ok" diff --git a/tests/unit/mcp/test_tools_common.py b/tests/unit/mcp/test_tools_common.py index a1d81f9a..eb9ef43d 100644 --- a/tests/unit/mcp/test_tools_common.py +++ b/tests/unit/mcp/test_tools_common.py @@ -5,6 +5,7 @@ from testgen.common.enums import Disposition, ImpactDimension, IssueLikelihood, PiiRisk, QualityDimension from testgen.common.models.scores import ScoreCategory +from testgen.common.models.test_definition import Severity from testgen.common.models.test_result import TestResultStatus from testgen.mcp.exceptions import MCPResourceNotAccessible, MCPUserError from testgen.mcp.tools.common import ( @@ -28,6 +29,7 @@ parse_score_filter_field, parse_score_group_by, parse_score_type, + parse_severity, parse_uuid, resolve_hygiene_issue, resolve_issue_type, @@ -88,6 +90,32 @@ def test_parse_result_status_invalid_lists_valid_values(): assert status.value in str(exc_info.value) +# --- parse_severity --- + + +def test_parse_severity_valid(): + assert parse_severity("Fail") == Severity.FAIL + assert parse_severity("Warning") == Severity.WARNING + + +def test_parse_severity_invalid_names_value(): + with pytest.raises(MCPUserError, match="Invalid severity `Critical`"): + parse_severity("Critical") + + +def test_parse_severity_invalid_lists_valid_values(): + with pytest.raises(MCPUserError, match="Valid values:") as exc_info: + parse_severity("nope") + for severity in Severity: + assert severity.value in str(exc_info.value) + + +def test_parse_severity_case_sensitive(): + """Severity stores 'Fail' / 'Warning' verbatim; lowercase rejected by design.""" + with pytest.raises(MCPUserError, match="Invalid severity `fail`"): + parse_severity("fail") + + # --- validate_page --- diff --git a/tests/unit/mcp/test_tools_projects.py b/tests/unit/mcp/test_tools_projects.py new file mode 100644 index 00000000..616d7890 --- /dev/null +++ b/tests/unit/mcp/test_tools_projects.py @@ -0,0 +1,430 @@ +"""Tests for the MCP project write tools — currently just update_project. + +Project create/delete live in the enterprise plugin (gated on ``global_admin``); +this module owns the per-project ``administer`` slice that ships in core. +""" + +from unittest.mock import MagicMock, patch + +import pytest + +from testgen.mcp.exceptions import MCPPermissionDenied, MCPResourceNotAccessible, MCPUserError +from testgen.mcp.permissions import ProjectPermissions + +pytestmark = pytest.mark.unit + +MODULE = "testgen.mcp.tools.projects" + + +def _patch_perms(allowed=("demo",), memberships=None, permission="administer", role="role_a"): + memberships = memberships or dict.fromkeys(allowed, role) + return patch( + "testgen.mcp.permissions._compute_project_permissions", + return_value=ProjectPermissions(memberships=memberships, permission=permission, username="test_user"), + ) + + +def _mock_project(**overrides) -> MagicMock: + project = MagicMock() + project.project_code = overrides.get("project_code", "demo") + project.project_name = overrides.get("project_name", "Demo Project") + project.use_dq_score_weights = overrides.get("use_dq_score_weights", True) + project.observability_api_url = overrides.get("observability_api_url", None) + project.observability_api_key = overrides.get("observability_api_key", None) + project.data_retention_enabled = overrides.get("data_retention_enabled", True) + project.data_retention_days = overrides.get("data_retention_days", 180) + return project + + +def _mock_schedule(cron_expr: str = "0 1 * * *", cron_tz: str = "UTC") -> MagicMock: + schedule = MagicMock() + schedule.cron_expr = cron_expr + schedule.cron_tz = cron_tz + return schedule + + +# --------------------------------------------------------------------------- +# update_project — guards +# --------------------------------------------------------------------------- + + +def test_update_project_no_fields_supplied(db_session_mock): + from testgen.mcp.tools.projects import update_project + + with _patch_perms(), pytest.raises(MCPUserError) as exc: + update_project(project_code="demo") + assert str(exc.value) == "No fields supplied to update." + + +def test_update_project_inaccessible_uses_unified_wording(db_session_mock): + from testgen.mcp.tools.projects import update_project + + with _patch_perms(allowed=("demo",)), pytest.raises(MCPResourceNotAccessible) as exc: + update_project(project_code="secret", project_name="Anything") + assert "Project `secret` not found or not accessible" in str(exc.value) + + +@patch(f"{MODULE}.JobSchedule") +@patch(f"{MODULE}.Project") +def test_update_project_not_found_uses_unified_wording(mock_project_cls, mock_schedule_cls, db_session_mock): + """User has the project code in scope, but Project.get returns None — still the unified error.""" + mock_project_cls.get.return_value = None + mock_schedule_cls.get.return_value = None + + from testgen.mcp.tools.projects import update_project + + with _patch_perms(), pytest.raises(MCPResourceNotAccessible) as exc: + update_project(project_code="demo", project_name="Anything") + assert "Project `demo` not found or not accessible" in str(exc.value) + + +def test_update_project_requires_administer(db_session_mock): + """role_d has edit but NOT administer (per conftest matrix).""" + from testgen.mcp.tools.projects import update_project + + with _patch_perms(memberships={"demo": "role_d"}), pytest.raises(MCPPermissionDenied): + update_project(project_code="demo", project_name="anything") + + +# --------------------------------------------------------------------------- +# update_project — rename +# --------------------------------------------------------------------------- + + +@patch(f"{MODULE}.JobSchedule") +@patch(f"{MODULE}.Project") +def test_update_project_renames(mock_project_cls, mock_schedule_cls, db_session_mock): + project = _mock_project(project_name="Demo Project") + mock_project_cls.get.return_value = project + mock_schedule_cls.get.return_value = _mock_schedule() + + from testgen.mcp.tools.projects import update_project + + with _patch_perms(): + out = update_project(project_code="demo", project_name="Demo Renamed") + + project.save.assert_called_once() + assert "Project `demo` updated" in out + assert "| Field | Before | After |" in out + assert "Demo Project" in out + assert "Demo Renamed" in out + + +@patch(f"{MODULE}.JobSchedule") +@patch(f"{MODULE}.Project") +def test_update_project_no_op_when_value_unchanged(mock_project_cls, mock_schedule_cls, db_session_mock): + project = _mock_project(project_name="Demo Project") + mock_project_cls.get.return_value = project + mock_schedule_cls.get.return_value = _mock_schedule() + + from testgen.mcp.tools.projects import update_project + + with _patch_perms(): + out = update_project(project_code="demo", project_name="Demo Project") + + project.save.assert_not_called() + assert "No fields changed" in out + + +@patch(f"{MODULE}.JobSchedule") +@patch(f"{MODULE}.Project") +def test_update_project_empty_name_rejected(mock_project_cls, mock_schedule_cls, db_session_mock): + project = _mock_project(project_name="Demo") + mock_project_cls.get.return_value = project + mock_schedule_cls.get.return_value = _mock_schedule() + + from testgen.mcp.tools.projects import update_project + + with _patch_perms(), pytest.raises(MCPUserError) as exc: + update_project(project_code="demo", project_name=" ") + assert "project_name: must not be empty" in str(exc.value) + project.save.assert_not_called() + + +@patch(f"{MODULE}.JobSchedule") +@patch(f"{MODULE}.Project") +def test_update_project_strips_whitespace(mock_project_cls, mock_schedule_cls, db_session_mock): + project = _mock_project(project_name="Demo Project") + mock_project_cls.get.return_value = project + mock_schedule_cls.get.return_value = _mock_schedule() + + from testgen.mcp.tools.projects import update_project + + with _patch_perms(): + update_project(project_code="demo", project_name=" Demo Renamed ") + + assert project.project_name == "Demo Renamed" + + +# --------------------------------------------------------------------------- +# update_project — weights toggle (recalc-scores side effect) +# --------------------------------------------------------------------------- + + +@patch(f"{MODULE}.JobExecution") +@patch(f"{MODULE}.JobSchedule") +@patch(f"{MODULE}.Project") +def test_update_project_weights_toggle_submits_recalc( + mock_project_cls, mock_schedule_cls, mock_job_exec_cls, db_session_mock, +): + project = _mock_project(use_dq_score_weights=True) + mock_project_cls.get.return_value = project + mock_schedule_cls.get.return_value = _mock_schedule() + + from testgen.mcp.tools.projects import update_project + + with _patch_perms(): + update_project(project_code="demo", use_dq_score_weights=False) + + project.save.assert_called_once() + # Background score recalc is submitted, matching the UI behaviour. + mock_job_exec_cls.submit.assert_called_once() + _, kwargs = mock_job_exec_cls.submit.call_args + assert kwargs["project_code"] == "demo" + + +@patch(f"{MODULE}.JobExecution") +@patch(f"{MODULE}.JobSchedule") +@patch(f"{MODULE}.Project") +def test_update_project_weights_unchanged_no_recalc( + mock_project_cls, mock_schedule_cls, mock_job_exec_cls, db_session_mock, +): + project = _mock_project(use_dq_score_weights=True) + mock_project_cls.get.return_value = project + mock_schedule_cls.get.return_value = _mock_schedule() + + from testgen.mcp.tools.projects import update_project + + with _patch_perms(): + # Supplied but unchanged — no side effect. + update_project(project_code="demo", use_dq_score_weights=True) + + mock_job_exec_cls.submit.assert_not_called() + + +# --------------------------------------------------------------------------- +# update_project — observability fields +# --------------------------------------------------------------------------- + + +@patch(f"{MODULE}.JobSchedule") +@patch(f"{MODULE}.Project") +def test_update_project_observability_url_diff_visible(mock_project_cls, mock_schedule_cls, db_session_mock): + project = _mock_project(observability_api_url=None) + mock_project_cls.get.return_value = project + mock_schedule_cls.get.return_value = _mock_schedule() + + from testgen.mcp.tools.projects import update_project + + with _patch_perms(): + out = update_project(project_code="demo", observability_api_url="https://obs.example/api") + + project.save.assert_called_once() + assert "DataOps Observability API URL" in out + assert "https://obs.example/api" in out + + +@patch(f"{MODULE}.JobSchedule") +@patch(f"{MODULE}.Project") +def test_update_project_observability_key_redacted_in_diff(mock_project_cls, mock_schedule_cls, db_session_mock): + """Per mcp-patterns 'Secrets in inputs': the key is consumed but never echoed back.""" + project = _mock_project(observability_api_key=None) + mock_project_cls.get.return_value = project + mock_schedule_cls.get.return_value = _mock_schedule() + + from testgen.mcp.tools.projects import update_project + + with _patch_perms(): + out = update_project(project_code="demo", observability_api_key="super-secret-key-value") + + project.save.assert_called_once() + assert "DataOps Observability API key" in out + assert "[secret]" in out + assert "super-secret-key-value" not in out + + +@patch(f"{MODULE}.JobSchedule") +@patch(f"{MODULE}.Project") +def test_update_project_observability_url_empty_string_clears( + mock_project_cls, mock_schedule_cls, db_session_mock, +): + """Empty string clears the field (NullIfEmptyString column).""" + project = _mock_project(observability_api_url="https://existing.example") + mock_project_cls.get.return_value = project + mock_schedule_cls.get.return_value = _mock_schedule() + + from testgen.mcp.tools.projects import update_project + + with _patch_perms(): + update_project(project_code="demo", observability_api_url="") + + # After normalization "" → None on the way in. + assert project.observability_api_url is None + + +# --------------------------------------------------------------------------- +# update_project — data retention +# --------------------------------------------------------------------------- + + +@patch(f"{MODULE}.JobSchedule") +@patch(f"{MODULE}.Project") +def test_update_project_disable_retention_deletes_schedule( + mock_project_cls, mock_schedule_cls, db_session_mock, +): + project = _mock_project(data_retention_enabled=True, data_retention_days=180) + mock_project_cls.get.return_value = project + mock_schedule_cls.get.return_value = _mock_schedule() + + from testgen.mcp.tools.projects import update_project + + with _patch_perms(): + update_project(project_code="demo", data_retention_enabled=False) + + project.save.assert_called_once() + mock_schedule_cls.delete_for_retention.assert_called_once_with("demo") + mock_schedule_cls.upsert_for_retention.assert_not_called() + assert project.data_retention_enabled is False + # Days cleared on disable (matches UI behaviour). + assert project.data_retention_days is None + + +@patch(f"{MODULE}.JobSchedule") +@patch(f"{MODULE}.Project") +def test_update_project_change_retention_days_upserts_schedule( + mock_project_cls, mock_schedule_cls, db_session_mock, +): + project = _mock_project(data_retention_enabled=True, data_retention_days=180) + mock_project_cls.get.return_value = project + mock_schedule_cls.get.return_value = _mock_schedule(cron_expr="0 1 * * *", cron_tz="UTC") + + from testgen.mcp.tools.projects import update_project + + with _patch_perms(): + update_project(project_code="demo", data_retention_days=30) + + project.save.assert_called_once() + mock_schedule_cls.upsert_for_retention.assert_called_once() + _, kwargs = mock_schedule_cls.upsert_for_retention.call_args + assert kwargs["project_code"] == "demo" + assert kwargs["retention_days"] == 30 + # Cron is preserved from existing schedule when not supplied. + assert kwargs["cron_expr"] == "0 1 * * *" + assert kwargs["cron_tz"] == "UTC" + + +@patch(f"{MODULE}.JobSchedule") +@patch(f"{MODULE}.Project") +def test_update_project_enable_retention_with_defaults(mock_project_cls, mock_schedule_cls, db_session_mock): + """Re-enabling retention without explicit days/cron falls back to the system defaults.""" + project = _mock_project(data_retention_enabled=False, data_retention_days=None) + mock_project_cls.get.return_value = project + mock_schedule_cls.get.return_value = None # no current schedule + + from testgen.mcp.tools.projects import update_project + + with _patch_perms(): + update_project(project_code="demo", data_retention_enabled=True) + + mock_schedule_cls.upsert_for_retention.assert_called_once() + _, kwargs = mock_schedule_cls.upsert_for_retention.call_args + assert kwargs["retention_days"] == 180 # default + assert kwargs["cron_expr"] == "0 1 * * *" # DEFAULT_DATA_CLEANUP_CRON + assert kwargs["cron_tz"] == "UTC" # DEFAULT_RETENTION_CRON_TZ + + +@patch(f"{MODULE}.JobSchedule") +@patch(f"{MODULE}.Project") +def test_update_project_change_cron_only(mock_project_cls, mock_schedule_cls, db_session_mock): + project = _mock_project(data_retention_enabled=True, data_retention_days=180) + mock_project_cls.get.return_value = project + mock_schedule_cls.get.return_value = _mock_schedule(cron_expr="0 1 * * *", cron_tz="UTC") + + from testgen.mcp.tools.projects import update_project + + with _patch_perms(): + out = update_project(project_code="demo", retention_cron_expr="0 2 * * *") + + mock_schedule_cls.upsert_for_retention.assert_called_once() + _, kwargs = mock_schedule_cls.upsert_for_retention.call_args + assert kwargs["cron_expr"] == "0 2 * * *" + # Days carried over from the project's current value. + assert kwargs["retention_days"] == 180 + assert "Retention cron expression" in out + + +@patch(f"{MODULE}.JobSchedule") +@patch(f"{MODULE}.Project") +def test_update_project_retention_days_rejected_when_disabled( + mock_project_cls, mock_schedule_cls, db_session_mock, +): + """Setting days while disabling retention is inconsistent — reject loudly instead of silently dropping.""" + project = _mock_project(data_retention_enabled=True) + mock_project_cls.get.return_value = project + mock_schedule_cls.get.return_value = _mock_schedule() + + from testgen.mcp.tools.projects import update_project + + with _patch_perms(), pytest.raises(MCPUserError) as exc: + update_project( + project_code="demo", + data_retention_enabled=False, + data_retention_days=90, + ) + msg = str(exc.value) + assert "data_retention_days: cannot be set when data_retention_enabled is False" in msg + project.save.assert_not_called() + mock_schedule_cls.upsert_for_retention.assert_not_called() + mock_schedule_cls.delete_for_retention.assert_not_called() + + +@patch(f"{MODULE}.JobSchedule") +@patch(f"{MODULE}.Project") +def test_update_project_cron_rejected_when_retention_currently_disabled( + mock_project_cls, mock_schedule_cls, db_session_mock, +): + project = _mock_project(data_retention_enabled=False, data_retention_days=None) + mock_project_cls.get.return_value = project + mock_schedule_cls.get.return_value = None + + from testgen.mcp.tools.projects import update_project + + with _patch_perms(), pytest.raises(MCPUserError) as exc: + update_project(project_code="demo", retention_cron_expr="0 2 * * *") + assert "retention_cron_expr: cannot be set when data_retention_enabled is False" in str(exc.value) + + +@patch(f"{MODULE}.JobSchedule") +@patch(f"{MODULE}.Project") +def test_update_project_retention_days_must_be_positive( + mock_project_cls, mock_schedule_cls, db_session_mock, +): + project = _mock_project(data_retention_enabled=True) + mock_project_cls.get.return_value = project + mock_schedule_cls.get.return_value = _mock_schedule() + + from testgen.mcp.tools.projects import update_project + + with _patch_perms(), pytest.raises(MCPUserError) as exc: + update_project(project_code="demo", data_retention_days=0) + assert "data_retention_days: must be a positive integer" in str(exc.value) + + +@patch(f"{MODULE}.JobSchedule") +@patch(f"{MODULE}.Project") +def test_update_project_name_only_does_not_touch_schedule( + mock_project_cls, mock_schedule_cls, db_session_mock, +): + """Renaming a project must NOT incidentally upsert / delete the retention schedule.""" + project = _mock_project(project_name="Old", data_retention_enabled=True) + mock_project_cls.get.return_value = project + mock_schedule_cls.get.return_value = _mock_schedule() + + from testgen.mcp.tools.projects import update_project + + with _patch_perms(): + update_project(project_code="demo", project_name="New") + + project.save.assert_called_once() + mock_schedule_cls.upsert_for_retention.assert_not_called() + mock_schedule_cls.delete_for_retention.assert_not_called() diff --git a/tests/unit/mcp/test_tools_test_suites.py b/tests/unit/mcp/test_tools_test_suites.py new file mode 100644 index 00000000..09907608 --- /dev/null +++ b/tests/unit/mcp/test_tools_test_suites.py @@ -0,0 +1,479 @@ +"""Tests for the MCP test-suite write tools — create / update for regular (non-monitor) suites.""" + +from unittest.mock import MagicMock, patch +from uuid import uuid4 + +import pytest + +from testgen.mcp.exceptions import MCPPermissionDenied, MCPResourceNotAccessible, MCPUserError +from testgen.mcp.permissions import ProjectPermissions + +pytestmark = pytest.mark.unit + +MODULE = "testgen.mcp.tools.test_suites" + + +def _patch_perms(allowed=("demo",), memberships=None, permission="edit", role="role_a"): + memberships = memberships or dict.fromkeys(allowed, role) + return patch( + "testgen.mcp.permissions._compute_project_permissions", + return_value=ProjectPermissions(memberships=memberships, permission=permission, username="test_user"), + ) + + +def _mock_table_group(**overrides) -> MagicMock: + tg_id = overrides.get("id", uuid4()) + tg = MagicMock() + tg.id = tg_id + tg.project_code = overrides.get("project_code", "demo") + tg.connection_id = overrides.get("connection_id", 42) + tg.table_groups_name = overrides.get("table_groups_name", "Sample TG") + return tg + + +def _mock_test_suite(**overrides) -> MagicMock: + suite = MagicMock() + suite.id = overrides.get("id", uuid4()) + suite.project_code = overrides.get("project_code", "demo") + suite.test_suite = overrides.get("test_suite", "Sample Suite") + suite.connection_id = overrides.get("connection_id", 42) + suite.table_groups_id = overrides.get("table_groups_id", uuid4()) + suite.test_suite_description = overrides.get("test_suite_description", None) + suite.severity = overrides.get("severity", None) + suite.export_to_observability = overrides.get("export_to_observability", False) + suite.dq_score_exclude = overrides.get("dq_score_exclude", False) + suite.component_key = overrides.get("component_key", None) + suite.component_type = overrides.get("component_type", None) + suite.component_name = overrides.get("component_name", None) + suite.is_monitor = overrides.get("is_monitor", False) + return suite + + +# --------------------------------------------------------------------------- +# create_test_suite +# --------------------------------------------------------------------------- + + +@patch(f"{MODULE}.TestSuite") +@patch(f"{MODULE}.resolve_table_group") +def test_create_test_suite_happy_path(mock_resolve, mock_suite_cls, db_session_mock): + tg = _mock_table_group() + mock_resolve.return_value = tg + instance = _mock_test_suite( + test_suite="Sales Tests", + project_code="demo", + table_groups_id=tg.id, + ) + mock_suite_cls.return_value = instance + + from testgen.mcp.tools.test_suites import create_test_suite + + with _patch_perms(): + out = create_test_suite(table_group_id=str(tg.id), test_suite_name="Sales Tests") + + # S2 review feedback: the tool no longer calls instance.save(); it relies on + # session.add + flush. The mock TestSuite constructor was invoked and the + # rendered response is returned — that's enough at the unit layer; the smoke + # test exercises the actual DB persistence. + mock_suite_cls.assert_called_once() + instance.save.assert_not_called() + assert "Test Suite `Sales Tests` created" in out + assert "**Project:** `demo`" in out + assert "Sample TG" in out + + +@patch(f"{MODULE}.TestSuite") +@patch(f"{MODULE}.resolve_table_group") +def test_create_test_suite_default_export_matches_ui(mock_resolve, mock_suite_cls, db_session_mock): + """S1 review feedback: export_to_observability defaults to False (matches the UI dialog), + not the model's legacy Y default.""" + tg = _mock_table_group() + mock_resolve.return_value = tg + mock_suite_cls.return_value = _mock_test_suite() + + from testgen.mcp.tools.test_suites import create_test_suite + + with _patch_perms(): + create_test_suite(table_group_id=str(tg.id), test_suite_name="Anything") + + _, kwargs = mock_suite_cls.call_args + assert kwargs["export_to_observability"] is False + # S3 review feedback: component_type defaults to "dataset" on create + assert kwargs["component_type"] == "dataset" + assert kwargs["dq_score_exclude"] is False + + +@patch(f"{MODULE}.TestSuite") +@patch(f"{MODULE}.resolve_table_group") +def test_create_test_suite_passes_all_new_args_to_model(mock_resolve, mock_suite_cls, db_session_mock): + """S9: create_test_suite accepts the full UI-editable surface (export, dq_score_exclude, component_*) + so the LLM can configure the suite in one round-trip.""" + tg = _mock_table_group() + mock_resolve.return_value = tg + mock_suite_cls.return_value = _mock_test_suite() + + from testgen.mcp.tools.test_suites import create_test_suite + + with _patch_perms(): + create_test_suite( + table_group_id=str(tg.id), + test_suite_name="Full", + description="d", + severity_default="Warning", + export_to_observability=True, + component_key="ck", + component_type="dataset", + component_name="cn", + dq_score_exclude=True, + ) + + _, kwargs = mock_suite_cls.call_args + assert kwargs["test_suite_description"] == "d" + assert kwargs["severity"] == "Warning" + assert kwargs["export_to_observability"] is True + assert kwargs["component_key"] == "ck" + assert kwargs["component_type"] == "dataset" + assert kwargs["component_name"] == "cn" + assert kwargs["dq_score_exclude"] is True + + +@patch(f"{MODULE}.TestSuite") +@patch(f"{MODULE}.resolve_table_group") +def test_create_test_suite_normalizes_empty_text_args_to_none(mock_resolve, mock_suite_cls, db_session_mock): + """S4: NullIfEmptyString writes "" -> NULL; pre-normalize on the way in so the model + receives the canonical value.""" + tg = _mock_table_group() + mock_resolve.return_value = tg + mock_suite_cls.return_value = _mock_test_suite() + + from testgen.mcp.tools.test_suites import create_test_suite + + with _patch_perms(): + create_test_suite( + table_group_id=str(tg.id), + test_suite_name="x", + description="", + component_key="", + component_type="", + component_name="", + ) + + _, kwargs = mock_suite_cls.call_args + assert kwargs["test_suite_description"] is None + assert kwargs["component_key"] is None + assert kwargs["component_type"] is None + assert kwargs["component_name"] is None + + +@patch(f"{MODULE}.TestSuite") +@patch(f"{MODULE}.resolve_table_group") +def test_create_test_suite_with_severity(mock_resolve, mock_suite_cls, db_session_mock): + tg = _mock_table_group() + mock_resolve.return_value = tg + instance = _mock_test_suite(test_suite="Sev Suite", severity="Fail") + mock_suite_cls.return_value = instance + + from testgen.mcp.tools.test_suites import create_test_suite + + with _patch_perms(): + out = create_test_suite( + table_group_id=str(tg.id), + test_suite_name="Sev Suite", + severity_default="Fail", + ) + + _, kwargs = mock_suite_cls.call_args + assert kwargs["severity"] == "Fail" + assert kwargs["is_monitor"] is False + assert "Fail" in out + + +@patch(f"{MODULE}.resolve_table_group") +def test_create_test_suite_invalid_severity(mock_resolve, db_session_mock): + mock_resolve.return_value = _mock_table_group() + + from testgen.mcp.tools.test_suites import create_test_suite + + with _patch_perms(), pytest.raises(MCPUserError) as exc: + create_test_suite( + table_group_id=str(uuid4()), + test_suite_name="Anything", + severity_default="Critical", + ) + msg = str(exc.value) + # S6: error is field-scoped and enumerates valid values without double-labeling. + assert "severity_default" in msg + assert "Fail" in msg and "Warning" in msg + assert "Critical" in msg + assert "Invalid severity" not in msg # the double-labeled phrasing is gone + + +def test_create_test_suite_empty_name_rejected(db_session_mock): + from testgen.mcp.tools.test_suites import create_test_suite + + with _patch_perms(), pytest.raises(MCPUserError) as exc: + create_test_suite(table_group_id=str(uuid4()), test_suite_name=" ") + assert "test_suite_name: must not be empty" in str(exc.value) + + +@patch(f"{MODULE}.TestSuite") +@patch(f"{MODULE}.resolve_table_group") +def test_create_test_suite_strips_whitespace(mock_resolve, mock_suite_cls, db_session_mock): + tg = _mock_table_group() + mock_resolve.return_value = tg + instance = _mock_test_suite(test_suite="Sales Tests") + mock_suite_cls.return_value = instance + + from testgen.mcp.tools.test_suites import create_test_suite + + with _patch_perms(): + create_test_suite(table_group_id=str(tg.id), test_suite_name=" Sales Tests ") + + _, kwargs = mock_suite_cls.call_args + assert kwargs["test_suite"] == "Sales Tests" + + +@patch(f"{MODULE}.resolve_table_group") +def test_create_test_suite_table_group_not_accessible(mock_resolve, db_session_mock): + mock_resolve.side_effect = MCPResourceNotAccessible("Table group", str(uuid4())) + + from testgen.mcp.tools.test_suites import create_test_suite + + with _patch_perms(allowed=("other",)), pytest.raises(MCPResourceNotAccessible): + create_test_suite(table_group_id=str(uuid4()), test_suite_name="Anything") + + +def test_create_test_suite_requires_edit(db_session_mock): + """role_c lacks edit (per conftest matrix).""" + from testgen.mcp.tools.test_suites import create_test_suite + + with _patch_perms(memberships={"demo": "role_c"}), pytest.raises(MCPPermissionDenied): + create_test_suite(table_group_id=str(uuid4()), test_suite_name="Anything") + + +@patch(f"{MODULE}.TestSuite") +@patch(f"{MODULE}.resolve_table_group") +def test_create_test_suite_allows_role_with_edit_but_not_administer( + mock_resolve, mock_suite_cls, db_session_mock, +): + """role_d has edit but NOT administer (per conftest matrix) — must still be allowed. + + Mirrors the production data_quality role: write access to suites without + project-level administer rights. + """ + tg = _mock_table_group() + mock_resolve.return_value = tg + instance = _mock_test_suite(test_suite="DQ Suite") + mock_suite_cls.return_value = instance + + from testgen.mcp.tools.test_suites import create_test_suite + + with _patch_perms(memberships={"demo": "role_d"}): + out = create_test_suite(table_group_id=str(tg.id), test_suite_name="DQ Suite") + + mock_suite_cls.assert_called_once() + instance.save.assert_not_called() + assert "Test Suite `DQ Suite` created" in out + + +# --------------------------------------------------------------------------- +# update_test_suite +# --------------------------------------------------------------------------- + + +def test_update_test_suite_no_fields_supplied(db_session_mock): + from testgen.mcp.tools.test_suites import update_test_suite + + with _patch_perms(), pytest.raises(MCPUserError) as exc: + update_test_suite(test_suite_id=str(uuid4())) + assert str(exc.value) == "No fields supplied to update." + + +@patch(f"{MODULE}.resolve_test_suite") +def test_update_test_suite_renames(mock_resolve, db_session_mock): + suite = _mock_test_suite(test_suite="Original Name") + mock_resolve.return_value = suite + + from testgen.mcp.tools.test_suites import update_test_suite + + with _patch_perms(): + out = update_test_suite(test_suite_id=str(suite.id), test_suite_name="New Name") + + suite.save.assert_called_once() + assert "Test Suite `New Name` updated" in out + assert "| Field | Before | After |" in out + assert "Original Name" in out + assert "New Name" in out + + +@patch(f"{MODULE}.resolve_test_suite") +def test_update_test_suite_multi_field_diff(mock_resolve, db_session_mock): + suite = _mock_test_suite( + test_suite="Suite", + severity=None, + export_to_observability=True, + dq_score_exclude=False, + ) + mock_resolve.return_value = suite + + from testgen.mcp.tools.test_suites import update_test_suite + + with _patch_perms(): + out = update_test_suite( + test_suite_id=str(suite.id), + severity_default="Warning", + export_to_observability=False, + dq_score_exclude=True, + ) + + suite.save.assert_called_once() + assert "Default severity" in out + assert "Warning" in out + # S13/S15: column header capitalises "Observability" (product name). + assert "Export to Observability" in out + assert "Exclude from quality scoring" in out + + +@patch(f"{MODULE}.resolve_test_suite") +def test_update_test_suite_no_op(mock_resolve, db_session_mock): + suite = _mock_test_suite(test_suite_description="same", severity="Fail") + mock_resolve.return_value = suite + + from testgen.mcp.tools.test_suites import update_test_suite + + with _patch_perms(): + out = update_test_suite( + test_suite_id=str(suite.id), + description="same", + severity_default="Fail", + ) + + suite.save.assert_not_called() + assert "No fields changed" in out + + +@patch(f"{MODULE}.resolve_test_suite") +def test_update_test_suite_empty_text_arg_on_null_field_is_noop(mock_resolve, db_session_mock): + """S4 phantom-diff regression: an "" arg on a currently-NULL NullIfEmptyString column must + NOT show up as a "changed" diff row (the DB would read back identical to before).""" + suite = _mock_test_suite( + test_suite_description=None, + component_key=None, + component_type=None, + component_name=None, + ) + mock_resolve.return_value = suite + + from testgen.mcp.tools.test_suites import update_test_suite + + with _patch_perms(): + out = update_test_suite( + test_suite_id=str(suite.id), + description="", + component_key="", + component_type="", + component_name="", + ) + + suite.save.assert_not_called() + assert "No fields changed" in out + + +@patch(f"{MODULE}.resolve_test_suite") +def test_update_test_suite_empty_text_arg_clears_populated_field(mock_resolve, db_session_mock): + """S4 complement: "" on a currently-populated field clears it to NULL and shows in the diff.""" + suite = _mock_test_suite(component_key="existing-key") + mock_resolve.return_value = suite + + from testgen.mcp.tools.test_suites import update_test_suite + + with _patch_perms(): + out = update_test_suite(test_suite_id=str(suite.id), component_key="") + + suite.save.assert_called_once() + assert "Component key" in out + assert "existing-key" in out + + +@patch(f"{MODULE}.resolve_test_suite") +def test_update_test_suite_empty_name_rejected(mock_resolve, db_session_mock): + suite = _mock_test_suite() + mock_resolve.return_value = suite + + from testgen.mcp.tools.test_suites import update_test_suite + + with _patch_perms(), pytest.raises(MCPUserError) as exc: + update_test_suite(test_suite_id=str(suite.id), test_suite_name=" ") + assert "test_suite_name: must not be empty" in str(exc.value) + suite.save.assert_not_called() + + +@patch(f"{MODULE}.resolve_test_suite") +def test_update_test_suite_invalid_severity_collected(mock_resolve, db_session_mock): + """Field-level validation collected, then raised together — name + severity errors in one message.""" + suite = _mock_test_suite() + mock_resolve.return_value = suite + + from testgen.mcp.tools.test_suites import update_test_suite + + with _patch_perms(), pytest.raises(MCPUserError) as exc: + update_test_suite( + test_suite_id=str(suite.id), + test_suite_name="", + severity_default="Critical", + ) + msg = str(exc.value) + assert "Update rejected" in msg + assert "test_suite_name: must not be empty" in msg + # S6: field-scoped phrasing, no double-labeled "Invalid severity" + assert "severity_default" in msg + assert "Fail" in msg and "Warning" in msg + assert "Invalid severity" not in msg + suite.save.assert_not_called() + + +def test_update_test_suite_monitor_suite_unified_wording(db_session_mock): + """Per TG-1053 review feedback: exercising an is_monitor=True suite must surface + the unified missing-or-inaccessible error, not a distinct rejection. The filter + side-effect is in ``resolve_test_suite``; this test exercises it end-to-end with + the real TestSuite.get patched to behave as it would in DB (filter applied → None).""" + with patch("testgen.mcp.tools.common.TestSuite") as mock_ts: + mock_ts.get.return_value = None # filter excluded the monitor suite + from testgen.mcp.tools.test_suites import update_test_suite + + with _patch_perms(), pytest.raises(MCPResourceNotAccessible) as exc: + update_test_suite(test_suite_id=str(uuid4()), test_suite_name="Anything") + assert "Test suite" in str(exc.value) + assert "not found or not accessible" in str(exc.value) + + +def test_update_test_suite_not_accessible(db_session_mock): + """Suite in a project the user can't access → unified wording.""" + with patch("testgen.mcp.tools.common.TestSuite") as mock_ts: + mock_ts.get.return_value = None + from testgen.mcp.tools.test_suites import update_test_suite + + with _patch_perms(allowed=("other",)), pytest.raises(MCPResourceNotAccessible): + update_test_suite(test_suite_id=str(uuid4()), test_suite_name="Anything") + + +def test_update_test_suite_requires_edit(db_session_mock): + from testgen.mcp.tools.test_suites import update_test_suite + + with _patch_perms(memberships={"demo": "role_c"}), pytest.raises(MCPPermissionDenied): + update_test_suite(test_suite_id=str(uuid4()), test_suite_name="Anything") + + +@patch(f"{MODULE}.resolve_test_suite") +def test_update_test_suite_allows_role_with_edit_but_not_administer(mock_resolve, db_session_mock): + """role_d has edit but NOT administer — must be allowed (mirrors production data_quality).""" + suite = _mock_test_suite(test_suite="DQ Suite") + mock_resolve.return_value = suite + + from testgen.mcp.tools.test_suites import update_test_suite + + with _patch_perms(memberships={"demo": "role_d"}): + out = update_test_suite(test_suite_id=str(suite.id), description="updated by data_quality") + + suite.save.assert_called_once() + assert "Test Suite `DQ Suite` updated" in out From 3a5bac849dd61a0e1d859c1cdd9f5df1febfb7d2 Mon Sep 17 00:00:00 2001 From: Diogo Basto Date: Fri, 26 Jun 2026 16:52:41 +0100 Subject: [PATCH 66/78] refactor(mcp): lift render_diff_table + resolve_project to common (TG-1070 review S5) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The diff-table render loop was duplicated three places (table_groups, projects, test_suites). The verify_access + Project.get + None-check trio in update_project also duplicated the resolve_* idiom already used for connections, table groups, and test suites. Both moves were flagged in the TG-1070 review. - render_diff_table(doc, before, after, *, attrs, labels, secret_attrs, value_renderer) in common.py owns the changed-attrs iteration, label lookup, secret redaction, and table emit. Each tool retains its own _DIFF_ORDER / _DIFF_LABELS and its own snapshot function (different shapes per entity). - resolve_project(project_code) joins the existing resolve_* family. Uses Project.get(code, Project.project_code.in_(allowed_codes)) so the unified not-found-or-not-accessible error path matches resolve_test_suite and resolve_table_group: out-of-scope projects never reveal whether they exist. update_connection keeps its own custom diff renderer — the "[secret] (rotated)" cue it emits for secret writes is semantically distinct from the generic "[secret]" redaction and isn't worth bending the helper to model both. Co-Authored-By: Claude Opus 4.7 --- testgen/mcp/tools/common.py | 81 ++++++++++++ testgen/mcp/tools/projects.py | 43 +++--- testgen/mcp/tools/table_groups.py | 20 +-- testgen/mcp/tools/test_suites.py | 20 +-- tests/unit/mcp/test_tools_common.py | 183 ++++++++++++++++++++++++++ tests/unit/mcp/test_tools_projects.py | 117 ++++++++-------- 6 files changed, 346 insertions(+), 118 deletions(-) diff --git a/testgen/mcp/tools/common.py b/testgen/mcp/tools/common.py index 581ef7e6..ffe127a6 100644 --- a/testgen/mcp/tools/common.py +++ b/testgen/mcp/tools/common.py @@ -1,3 +1,4 @@ +from collections.abc import Callable, Mapping, Sequence from dataclasses import dataclass, field from datetime import date, datetime from enum import StrEnum @@ -35,6 +36,7 @@ TestRunNotificationTrigger, ) from testgen.common.models.profiling_run import ProfilingRun +from testgen.common.models.project import Project from testgen.common.models.scheduler import SCHEDULABLE_JOB_KEYS, JobSchedule from testgen.common.models.scores import ScoreCategory, ScoreDefinition from testgen.common.models.table_group import TableGroup @@ -628,10 +630,89 @@ def format_page_footer(total: int, page: int, limit: int) -> str: return f"_Page {page} of {total_pages}. Use `page={page + 1}` for more._" +def _default_render_diff_value(value: object) -> str | None: + """Default formatter for before/after cells in :func:`render_diff_table`. + + * ``bool`` → ``"Yes"`` / ``"No"`` + * ``None`` or ``""`` → ``None`` (rendered as em-dash by the table cell) + * else → ``str(value)`` + """ + if isinstance(value, bool): + return "Yes" if value else "No" + if value is None or value == "": + return None + return str(value) + + +def render_diff_table( + doc: MdDoc, + before: Mapping[str, object], + after: Mapping[str, object], + *, + attrs: Sequence[str], + labels: Mapping[str, str], + secret_attrs: frozenset[str] = frozenset(), + value_renderer: Callable[[object], str | None] | None = None, +) -> bool: + """Emit a ``Field / Before / After`` table for attrs whose values changed. + + Update tools across the MCP surface share this shape: snapshot before, apply + field-by-field, snapshot after, render the delta. The shared helper keeps the + ordering, label lookup, and secret redaction consistent. + + Args: + doc: ``MdDoc`` to append the table to. + before / after: same-shape dicts of the snapshot keyed by attr name. + Only entries in ``attrs`` are inspected. + attrs: ordered tuple of attrs to consider; controls row order. + labels: attr → user-facing label for the first column. + secret_attrs: attrs whose values must not be echoed (API keys, passwords). + Rendered as ``[secret]`` when present, em-dash when absent — distinct + from a "rotated" cue. Tools that need to communicate rotation + specifically (e.g. ``update_connection``) keep their own renderer. + value_renderer: override for the non-secret cells. Defaults to + :func:`_default_render_diff_value`. + + Returns ``True`` if at least one row was rendered, ``False`` when nothing + changed (lets the caller emit a "No fields changed" message instead). + """ + changed = {attr for attr in attrs if before.get(attr) != after.get(attr)} + if not changed: + return False + render = value_renderer or _default_render_diff_value + rows: list[list[object]] = [] + for attr in attrs: + if attr not in changed: + continue + if attr in secret_attrs: + b = "[secret]" if before.get(attr) else None + a = "[secret]" if after.get(attr) else None + else: + b = render(before.get(attr)) + a = render(after.get(attr)) + rows.append([labels.get(attr, attr), b, a]) + doc.table(["Field", "Before", "After"], rows, code=[0]) + return True + + # Entity resolution helpers — see mcp-roadmap.md "Entity Resolution Helpers" guideline. # Extract a new resolve_ here when a second caller needs the same parse-uuid + # perm-scoped lookup + collapsed-error pattern. + +def resolve_project(project_code: str) -> Project: + """Resolve a project code to a Project, scoped to the user's allowed projects. + + Collapses missing-or-inaccessible into a single ``MCPResourceNotAccessible`` so + callers can't enumerate whether a project they can't see exists. + """ + perms = get_project_permissions() + project = Project.get(project_code, Project.project_code.in_(perms.allowed_codes)) + if project is None: + raise MCPResourceNotAccessible("Project", project_code) + return project + + def resolve_connection(connection_id: int) -> Connection: """Resolve a connection ID, collapsing missing-or-inaccessible into one error path.""" perms = get_project_permissions() diff --git a/testgen/mcp/tools/projects.py b/testgen/mcp/tools/projects.py index 2e60e3eb..325d95a6 100644 --- a/testgen/mcp/tools/projects.py +++ b/testgen/mcp/tools/projects.py @@ -18,9 +18,14 @@ DEFAULT_RETENTION_CRON_TZ, JobSchedule, ) -from testgen.mcp.exceptions import MCPResourceNotAccessible, MCPUserError -from testgen.mcp.permissions import get_project_permissions, mcp_permission -from testgen.mcp.tools.common import DocGroup, raise_validation_error +from testgen.mcp.exceptions import MCPUserError +from testgen.mcp.permissions import mcp_permission +from testgen.mcp.tools.common import ( + DocGroup, + raise_validation_error, + render_diff_table, + resolve_project, +) from testgen.mcp.tools.markdown import MdDoc _DOC_GROUP = DocGroup.MANAGE @@ -77,12 +82,7 @@ def update_project( if all(value is None for value in supplied.values()): raise MCPUserError("No fields supplied to update.") - perms = get_project_permissions() - perms.verify_access(project_code, not_found=MCPResourceNotAccessible("Project", project_code)) - - project = Project.get(project_code) - if project is None: - raise MCPResourceNotAccessible("Project", project_code) + project = resolve_project(project_code) schedule = JobSchedule.get( JobSchedule.project_code == project_code, @@ -164,12 +164,15 @@ def update_project( if effective_enabled else None, ) - changed = {attr for attr in before if before[attr] != after[attr]} doc = MdDoc() doc.heading(1, f"Project `{project_code}` updated") - if not changed: + rendered = render_diff_table( + doc, before, after, + attrs=_DIFF_ORDER, labels=_DIFF_LABELS, secret_attrs=_SECRET_ATTRS, + ) + if not rendered: doc.text("No fields changed — supplied values matched the current state.") return doc.render() @@ -196,12 +199,6 @@ def update_project( project_code=project_code, ) - rows: list[list[object]] = [] - for attr in _DIFF_ORDER: - if attr not in changed: - continue - rows.append([_DIFF_LABELS[attr], _render_field_value(attr, before[attr]), _render_field_value(attr, after[attr])]) - doc.table(["Field", "Before", "After"], rows, code=[0]) return doc.render() @@ -249,14 +246,4 @@ def _snapshot(project: Project, schedule: Any) -> dict[str, Any]: "retention_cron_tz": "Retention timezone", } - -def _render_field_value(attr: str, value: Any) -> str | None: - # Redact the API key — mcp-patterns "Secrets in inputs" rule. The diff still tells the LLM - # the field changed; it just never echoes the value (cleartext OR masked). - if attr == "observability_api_key" and value: - return "[secret]" - if isinstance(value, bool): - return "Yes" if value else "No" - if value is None or value == "": - return None - return str(value) +_SECRET_ATTRS: frozenset[str] = frozenset({"observability_api_key"}) diff --git a/testgen/mcp/tools/table_groups.py b/testgen/mcp/tools/table_groups.py index 109a670f..ca1524ae 100644 --- a/testgen/mcp/tools/table_groups.py +++ b/testgen/mcp/tools/table_groups.py @@ -25,6 +25,7 @@ format_flavor_label, format_page_footer, format_page_info, + render_diff_table, resolve_connection, resolve_table_group, validate_limit, @@ -424,13 +425,13 @@ def update_table_group( _raise_validation_error(errors, "Update rejected. No changes saved.") after = _snapshot(table_group) - changed = {attr for attr in before if before[attr] != after[attr]} doc = MdDoc() doc.heading(1, f"Table Group `{table_group.table_groups_name}` updated") doc.field("ID", str(table_group.id), code=True) - if not changed: + rendered = render_diff_table(doc, before, after, attrs=_DIFF_ATTRS, labels=_DIFF_LABELS) + if not rendered: doc.text("No fields changed — supplied values matched the current state.") return doc.render() @@ -440,13 +441,6 @@ def update_table_group( _maybe_raise_duplicate_name(err) raise - rows: list[list[object]] = [] - for attr in _DIFF_ATTRS: - if attr not in changed: - continue - label = _DIFF_LABELS.get(attr, attr) - rows.append([label, _render_field_value(before[attr]), _render_field_value(after[attr])]) - doc.table(["Field", "Before", "After"], rows, code=[0]) return doc.render() @@ -722,14 +716,6 @@ def _snapshot(table_group: TableGroup) -> dict[str, Any]: return {attr: getattr(table_group, attr, None) for attr in _DIFF_ATTRS} -def _render_field_value(value: Any) -> str | None: - if isinstance(value, bool): - return "Yes" if value else "No" - if value is None or value == "": - return None - return str(value) - - def _render_created_table_group(table_group: TableGroup, connection: Connection) -> str: doc = MdDoc() doc.heading(1, f"Table Group `{table_group.table_groups_name}` created") diff --git a/testgen/mcp/tools/test_suites.py b/testgen/mcp/tools/test_suites.py index a5773de3..56d37e86 100644 --- a/testgen/mcp/tools/test_suites.py +++ b/testgen/mcp/tools/test_suites.py @@ -19,6 +19,7 @@ DocGroup, parse_severity, raise_validation_error, + render_diff_table, resolve_table_group, resolve_test_suite, ) @@ -193,24 +194,17 @@ def update_test_suite( for attr, value in updates.items(): setattr(suite, attr, value) after = {attr: getattr(suite, attr, None) for attr in updates} - changed = {attr for attr in before if before[attr] != after[attr]} doc = MdDoc() doc.heading(1, f"Test Suite `{suite.test_suite}` updated") doc.field("ID", str(suite.id), code=True) - if not changed: + rendered = render_diff_table(doc, before, after, attrs=_DIFF_ORDER, labels=_DIFF_LABELS) + if not rendered: doc.text("No fields changed — supplied values matched the current state.") return doc.render() suite.save() - - rows: list[list[object]] = [] - for attr in _DIFF_ORDER: - if attr not in changed: - continue - rows.append([_DIFF_LABELS[attr], _render_field_value(before[attr]), _render_field_value(after[attr])]) - doc.table(["Field", "Before", "After"], rows, code=[0]) return doc.render() @@ -270,11 +264,3 @@ def _render_created_suite(suite: TestSuite, *, table_group_name: str) -> str: "component_name": "Component name", "dq_score_exclude": "Exclude from quality scoring", } - - -def _render_field_value(value: Any) -> str | None: - if isinstance(value, bool): - return "Yes" if value else "No" - if value is None or value == "": - return None - return str(value) diff --git a/tests/unit/mcp/test_tools_common.py b/tests/unit/mcp/test_tools_common.py index eb9ef43d..64de5f4e 100644 --- a/tests/unit/mcp/test_tools_common.py +++ b/tests/unit/mcp/test_tools_common.py @@ -1013,3 +1013,186 @@ def test_resolve_monitored_table_group_raises_when_tg_inaccessible(): pytest.raises(MCPResourceNotAccessible), ): resolve_monitored_table_group(bad_id) + + + +# --- resolve_project --- + + +def test_resolve_project_returns_project_when_in_scope(db_session_mock): + """Happy path: project_code in allowed_codes, Project.get returns the row.""" + from testgen.mcp.permissions import ProjectPermissions, _mcp_project_permissions + from testgen.mcp.tools.common import resolve_project + + project = MagicMock() + project.project_code = "demo" + + with patch("testgen.mcp.tools.common.Project") as mock_project_cls: + mock_project_cls.get.return_value = project + mock_project_cls.project_code = MagicMock() # for the .in_() filter clause + perms = ProjectPermissions(memberships={"demo": "admin"}, permission="administer", username="t") + token = _mcp_project_permissions.set(perms) + try: + with patch( + "testgen.mcp.permissions.PluginHook" + ) as mock_hook: + mock_hook.instance.return_value.rbac.get_roles_with_permission.return_value = ["admin"] + assert resolve_project("demo") is project + finally: + _mcp_project_permissions.reset(token) + + +def test_resolve_project_raises_unified_when_get_returns_none(db_session_mock): + """Project not in scope (or absent) → unified ``MCPResourceNotAccessible``.""" + from testgen.mcp.permissions import ProjectPermissions, _mcp_project_permissions + from testgen.mcp.tools.common import resolve_project + + with patch("testgen.mcp.tools.common.Project") as mock_project_cls: + mock_project_cls.get.return_value = None + mock_project_cls.project_code = MagicMock() + perms = ProjectPermissions(memberships={"demo": "admin"}, permission="administer", username="t") + token = _mcp_project_permissions.set(perms) + try: + with patch( + "testgen.mcp.permissions.PluginHook" + ) as mock_hook: + mock_hook.instance.return_value.rbac.get_roles_with_permission.return_value = ["admin"] + with pytest.raises(MCPResourceNotAccessible) as exc: + resolve_project("secret") + assert "Project `secret` not found or not accessible" in str(exc.value) + finally: + _mcp_project_permissions.reset(token) + + +# --- render_diff_table / _default_render_diff_value --- + + +def test_render_diff_table_emits_rows_for_changed_attrs(db_session_mock): + from testgen.mcp.tools.common import render_diff_table + from testgen.mcp.tools.markdown import MdDoc + + doc = MdDoc() + before = {"name": "Old", "active": True, "label": "x"} + after = {"name": "New", "active": True, "label": "y"} + + rendered = render_diff_table( + doc, before, after, + attrs=("name", "active", "label"), + labels={"name": "Name", "active": "Active", "label": "Label"}, + ) + + assert rendered is True + out = doc.render() + assert "| Field | Before | After |" in out + assert "Old" in out and "New" in out + assert "Label" in out and "x" in out and "y" in out + # Unchanged attr "active" must not appear + assert "Active" not in out + + +def test_render_diff_table_returns_false_when_nothing_changes(db_session_mock): + from testgen.mcp.tools.common import render_diff_table + from testgen.mcp.tools.markdown import MdDoc + + doc = MdDoc() + snap = {"name": "Same", "count": 3} + + rendered = render_diff_table( + doc, snap, snap, + attrs=("name", "count"), + labels={"name": "Name", "count": "Count"}, + ) + + assert rendered is False + assert doc.render() == "" # nothing appended + + +def test_render_diff_table_redacts_secret_attrs(db_session_mock): + """secret_attrs render as ``[secret]`` when present and em-dash when absent — value + is never echoed.""" + from testgen.mcp.tools.common import render_diff_table + from testgen.mcp.tools.markdown import MdDoc + + doc = MdDoc() + before = {"name": "Demo", "api_key": None} + after = {"name": "Demo", "api_key": "super-secret-value"} + + rendered = render_diff_table( + doc, before, after, + attrs=("name", "api_key"), + labels={"name": "Name", "api_key": "API key"}, + secret_attrs=frozenset({"api_key"}), + ) + + assert rendered is True + out = doc.render() + assert "API key" in out + assert "[secret]" in out + assert "super-secret-value" not in out + + +def test_render_diff_table_honors_attr_ordering(db_session_mock): + """Row order matches the supplied ``attrs`` tuple, not dict insertion order.""" + from testgen.mcp.tools.common import render_diff_table + from testgen.mcp.tools.markdown import MdDoc + + doc = MdDoc() + before = {"b": 1, "a": 1, "c": 1} + after = {"b": 2, "a": 2, "c": 2} + + render_diff_table( + doc, before, after, + attrs=("a", "b", "c"), + labels={"a": "A", "b": "B", "c": "C"}, + ) + + out = doc.render() + # Body rows render in attrs order. Label cells are code-wrapped (column 0 uses + # ``code=[0]`` in ``MdDoc.table``). + pos_a = out.find("| `A` |") + pos_b = out.find("| `B` |") + pos_c = out.find("| `C` |") + assert 0 < pos_a < pos_b < pos_c + + +def test_render_diff_table_custom_value_renderer(db_session_mock): + """``value_renderer`` overrides the default Yes/No/em-dash formatting.""" + from testgen.mcp.tools.common import render_diff_table + from testgen.mcp.tools.markdown import MdDoc + + doc = MdDoc() + before = {"n": 1} + after = {"n": 2} + + rendered = render_diff_table( + doc, before, after, + attrs=("n",), + labels={"n": "N"}, + value_renderer=lambda v: f"#{v}", + ) + + assert rendered is True + out = doc.render() + assert "#1" in out + assert "#2" in out + + +def test_default_render_diff_value_bool_yes_no(): + from testgen.mcp.tools.common import _default_render_diff_value + + assert _default_render_diff_value(True) == "Yes" + assert _default_render_diff_value(False) == "No" + + +def test_default_render_diff_value_none_and_empty(): + from testgen.mcp.tools.common import _default_render_diff_value + + assert _default_render_diff_value(None) is None + assert _default_render_diff_value("") is None + + +def test_default_render_diff_value_str(): + from testgen.mcp.tools.common import _default_render_diff_value + + assert _default_render_diff_value("hello") == "hello" + assert _default_render_diff_value(42) == "42" diff --git a/tests/unit/mcp/test_tools_projects.py b/tests/unit/mcp/test_tools_projects.py index 616d7890..7000e766 100644 --- a/tests/unit/mcp/test_tools_projects.py +++ b/tests/unit/mcp/test_tools_projects.py @@ -56,7 +56,12 @@ def test_update_project_no_fields_supplied(db_session_mock): assert str(exc.value) == "No fields supplied to update." -def test_update_project_inaccessible_uses_unified_wording(db_session_mock): +@patch(f"{MODULE}.resolve_project") +def test_update_project_inaccessible_uses_unified_wording(mock_resolve, db_session_mock): + """When resolve_project raises (out-of-scope or missing), the tool re-raises with the + unified wording. resolve_project's own behaviour is covered in test_tools_common.py.""" + mock_resolve.side_effect = MCPResourceNotAccessible("Project", "secret") + from testgen.mcp.tools.projects import update_project with _patch_perms(allowed=("demo",)), pytest.raises(MCPResourceNotAccessible) as exc: @@ -65,10 +70,10 @@ def test_update_project_inaccessible_uses_unified_wording(db_session_mock): @patch(f"{MODULE}.JobSchedule") -@patch(f"{MODULE}.Project") -def test_update_project_not_found_uses_unified_wording(mock_project_cls, mock_schedule_cls, db_session_mock): - """User has the project code in scope, but Project.get returns None — still the unified error.""" - mock_project_cls.get.return_value = None +@patch(f"{MODULE}.resolve_project") +def test_update_project_not_found_uses_unified_wording(mock_resolve, mock_schedule_cls, db_session_mock): + """resolve_project collapses 'no row' into the unified error.""" + mock_resolve.side_effect = MCPResourceNotAccessible("Project", "demo") mock_schedule_cls.get.return_value = None from testgen.mcp.tools.projects import update_project @@ -92,10 +97,10 @@ def test_update_project_requires_administer(db_session_mock): @patch(f"{MODULE}.JobSchedule") -@patch(f"{MODULE}.Project") -def test_update_project_renames(mock_project_cls, mock_schedule_cls, db_session_mock): +@patch(f"{MODULE}.resolve_project") +def test_update_project_renames(mock_resolve, mock_schedule_cls, db_session_mock): project = _mock_project(project_name="Demo Project") - mock_project_cls.get.return_value = project + mock_resolve.return_value = project mock_schedule_cls.get.return_value = _mock_schedule() from testgen.mcp.tools.projects import update_project @@ -111,10 +116,10 @@ def test_update_project_renames(mock_project_cls, mock_schedule_cls, db_session_ @patch(f"{MODULE}.JobSchedule") -@patch(f"{MODULE}.Project") -def test_update_project_no_op_when_value_unchanged(mock_project_cls, mock_schedule_cls, db_session_mock): +@patch(f"{MODULE}.resolve_project") +def test_update_project_no_op_when_value_unchanged(mock_resolve, mock_schedule_cls, db_session_mock): project = _mock_project(project_name="Demo Project") - mock_project_cls.get.return_value = project + mock_resolve.return_value = project mock_schedule_cls.get.return_value = _mock_schedule() from testgen.mcp.tools.projects import update_project @@ -127,10 +132,10 @@ def test_update_project_no_op_when_value_unchanged(mock_project_cls, mock_schedu @patch(f"{MODULE}.JobSchedule") -@patch(f"{MODULE}.Project") -def test_update_project_empty_name_rejected(mock_project_cls, mock_schedule_cls, db_session_mock): +@patch(f"{MODULE}.resolve_project") +def test_update_project_empty_name_rejected(mock_resolve, mock_schedule_cls, db_session_mock): project = _mock_project(project_name="Demo") - mock_project_cls.get.return_value = project + mock_resolve.return_value = project mock_schedule_cls.get.return_value = _mock_schedule() from testgen.mcp.tools.projects import update_project @@ -142,10 +147,10 @@ def test_update_project_empty_name_rejected(mock_project_cls, mock_schedule_cls, @patch(f"{MODULE}.JobSchedule") -@patch(f"{MODULE}.Project") -def test_update_project_strips_whitespace(mock_project_cls, mock_schedule_cls, db_session_mock): +@patch(f"{MODULE}.resolve_project") +def test_update_project_strips_whitespace(mock_resolve, mock_schedule_cls, db_session_mock): project = _mock_project(project_name="Demo Project") - mock_project_cls.get.return_value = project + mock_resolve.return_value = project mock_schedule_cls.get.return_value = _mock_schedule() from testgen.mcp.tools.projects import update_project @@ -163,12 +168,12 @@ def test_update_project_strips_whitespace(mock_project_cls, mock_schedule_cls, d @patch(f"{MODULE}.JobExecution") @patch(f"{MODULE}.JobSchedule") -@patch(f"{MODULE}.Project") +@patch(f"{MODULE}.resolve_project") def test_update_project_weights_toggle_submits_recalc( - mock_project_cls, mock_schedule_cls, mock_job_exec_cls, db_session_mock, + mock_resolve, mock_schedule_cls, mock_job_exec_cls, db_session_mock, ): project = _mock_project(use_dq_score_weights=True) - mock_project_cls.get.return_value = project + mock_resolve.return_value = project mock_schedule_cls.get.return_value = _mock_schedule() from testgen.mcp.tools.projects import update_project @@ -185,12 +190,12 @@ def test_update_project_weights_toggle_submits_recalc( @patch(f"{MODULE}.JobExecution") @patch(f"{MODULE}.JobSchedule") -@patch(f"{MODULE}.Project") +@patch(f"{MODULE}.resolve_project") def test_update_project_weights_unchanged_no_recalc( - mock_project_cls, mock_schedule_cls, mock_job_exec_cls, db_session_mock, + mock_resolve, mock_schedule_cls, mock_job_exec_cls, db_session_mock, ): project = _mock_project(use_dq_score_weights=True) - mock_project_cls.get.return_value = project + mock_resolve.return_value = project mock_schedule_cls.get.return_value = _mock_schedule() from testgen.mcp.tools.projects import update_project @@ -208,10 +213,10 @@ def test_update_project_weights_unchanged_no_recalc( @patch(f"{MODULE}.JobSchedule") -@patch(f"{MODULE}.Project") -def test_update_project_observability_url_diff_visible(mock_project_cls, mock_schedule_cls, db_session_mock): +@patch(f"{MODULE}.resolve_project") +def test_update_project_observability_url_diff_visible(mock_resolve, mock_schedule_cls, db_session_mock): project = _mock_project(observability_api_url=None) - mock_project_cls.get.return_value = project + mock_resolve.return_value = project mock_schedule_cls.get.return_value = _mock_schedule() from testgen.mcp.tools.projects import update_project @@ -225,11 +230,11 @@ def test_update_project_observability_url_diff_visible(mock_project_cls, mock_sc @patch(f"{MODULE}.JobSchedule") -@patch(f"{MODULE}.Project") -def test_update_project_observability_key_redacted_in_diff(mock_project_cls, mock_schedule_cls, db_session_mock): +@patch(f"{MODULE}.resolve_project") +def test_update_project_observability_key_redacted_in_diff(mock_resolve, mock_schedule_cls, db_session_mock): """Per mcp-patterns 'Secrets in inputs': the key is consumed but never echoed back.""" project = _mock_project(observability_api_key=None) - mock_project_cls.get.return_value = project + mock_resolve.return_value = project mock_schedule_cls.get.return_value = _mock_schedule() from testgen.mcp.tools.projects import update_project @@ -244,13 +249,13 @@ def test_update_project_observability_key_redacted_in_diff(mock_project_cls, moc @patch(f"{MODULE}.JobSchedule") -@patch(f"{MODULE}.Project") +@patch(f"{MODULE}.resolve_project") def test_update_project_observability_url_empty_string_clears( - mock_project_cls, mock_schedule_cls, db_session_mock, + mock_resolve, mock_schedule_cls, db_session_mock, ): """Empty string clears the field (NullIfEmptyString column).""" project = _mock_project(observability_api_url="https://existing.example") - mock_project_cls.get.return_value = project + mock_resolve.return_value = project mock_schedule_cls.get.return_value = _mock_schedule() from testgen.mcp.tools.projects import update_project @@ -268,12 +273,12 @@ def test_update_project_observability_url_empty_string_clears( @patch(f"{MODULE}.JobSchedule") -@patch(f"{MODULE}.Project") +@patch(f"{MODULE}.resolve_project") def test_update_project_disable_retention_deletes_schedule( - mock_project_cls, mock_schedule_cls, db_session_mock, + mock_resolve, mock_schedule_cls, db_session_mock, ): project = _mock_project(data_retention_enabled=True, data_retention_days=180) - mock_project_cls.get.return_value = project + mock_resolve.return_value = project mock_schedule_cls.get.return_value = _mock_schedule() from testgen.mcp.tools.projects import update_project @@ -290,12 +295,12 @@ def test_update_project_disable_retention_deletes_schedule( @patch(f"{MODULE}.JobSchedule") -@patch(f"{MODULE}.Project") +@patch(f"{MODULE}.resolve_project") def test_update_project_change_retention_days_upserts_schedule( - mock_project_cls, mock_schedule_cls, db_session_mock, + mock_resolve, mock_schedule_cls, db_session_mock, ): project = _mock_project(data_retention_enabled=True, data_retention_days=180) - mock_project_cls.get.return_value = project + mock_resolve.return_value = project mock_schedule_cls.get.return_value = _mock_schedule(cron_expr="0 1 * * *", cron_tz="UTC") from testgen.mcp.tools.projects import update_project @@ -314,11 +319,11 @@ def test_update_project_change_retention_days_upserts_schedule( @patch(f"{MODULE}.JobSchedule") -@patch(f"{MODULE}.Project") -def test_update_project_enable_retention_with_defaults(mock_project_cls, mock_schedule_cls, db_session_mock): +@patch(f"{MODULE}.resolve_project") +def test_update_project_enable_retention_with_defaults(mock_resolve, mock_schedule_cls, db_session_mock): """Re-enabling retention without explicit days/cron falls back to the system defaults.""" project = _mock_project(data_retention_enabled=False, data_retention_days=None) - mock_project_cls.get.return_value = project + mock_resolve.return_value = project mock_schedule_cls.get.return_value = None # no current schedule from testgen.mcp.tools.projects import update_project @@ -334,10 +339,10 @@ def test_update_project_enable_retention_with_defaults(mock_project_cls, mock_sc @patch(f"{MODULE}.JobSchedule") -@patch(f"{MODULE}.Project") -def test_update_project_change_cron_only(mock_project_cls, mock_schedule_cls, db_session_mock): +@patch(f"{MODULE}.resolve_project") +def test_update_project_change_cron_only(mock_resolve, mock_schedule_cls, db_session_mock): project = _mock_project(data_retention_enabled=True, data_retention_days=180) - mock_project_cls.get.return_value = project + mock_resolve.return_value = project mock_schedule_cls.get.return_value = _mock_schedule(cron_expr="0 1 * * *", cron_tz="UTC") from testgen.mcp.tools.projects import update_project @@ -354,13 +359,13 @@ def test_update_project_change_cron_only(mock_project_cls, mock_schedule_cls, db @patch(f"{MODULE}.JobSchedule") -@patch(f"{MODULE}.Project") +@patch(f"{MODULE}.resolve_project") def test_update_project_retention_days_rejected_when_disabled( - mock_project_cls, mock_schedule_cls, db_session_mock, + mock_resolve, mock_schedule_cls, db_session_mock, ): """Setting days while disabling retention is inconsistent — reject loudly instead of silently dropping.""" project = _mock_project(data_retention_enabled=True) - mock_project_cls.get.return_value = project + mock_resolve.return_value = project mock_schedule_cls.get.return_value = _mock_schedule() from testgen.mcp.tools.projects import update_project @@ -379,12 +384,12 @@ def test_update_project_retention_days_rejected_when_disabled( @patch(f"{MODULE}.JobSchedule") -@patch(f"{MODULE}.Project") +@patch(f"{MODULE}.resolve_project") def test_update_project_cron_rejected_when_retention_currently_disabled( - mock_project_cls, mock_schedule_cls, db_session_mock, + mock_resolve, mock_schedule_cls, db_session_mock, ): project = _mock_project(data_retention_enabled=False, data_retention_days=None) - mock_project_cls.get.return_value = project + mock_resolve.return_value = project mock_schedule_cls.get.return_value = None from testgen.mcp.tools.projects import update_project @@ -395,12 +400,12 @@ def test_update_project_cron_rejected_when_retention_currently_disabled( @patch(f"{MODULE}.JobSchedule") -@patch(f"{MODULE}.Project") +@patch(f"{MODULE}.resolve_project") def test_update_project_retention_days_must_be_positive( - mock_project_cls, mock_schedule_cls, db_session_mock, + mock_resolve, mock_schedule_cls, db_session_mock, ): project = _mock_project(data_retention_enabled=True) - mock_project_cls.get.return_value = project + mock_resolve.return_value = project mock_schedule_cls.get.return_value = _mock_schedule() from testgen.mcp.tools.projects import update_project @@ -411,13 +416,13 @@ def test_update_project_retention_days_must_be_positive( @patch(f"{MODULE}.JobSchedule") -@patch(f"{MODULE}.Project") +@patch(f"{MODULE}.resolve_project") def test_update_project_name_only_does_not_touch_schedule( - mock_project_cls, mock_schedule_cls, db_session_mock, + mock_resolve, mock_schedule_cls, db_session_mock, ): """Renaming a project must NOT incidentally upsert / delete the retention schedule.""" project = _mock_project(project_name="Old", data_retention_enabled=True) - mock_project_cls.get.return_value = project + mock_resolve.return_value = project mock_schedule_cls.get.return_value = _mock_schedule() from testgen.mcp.tools.projects import update_project From e5ec5c44b678a524c93758f33e0568e485aea522 Mon Sep 17 00:00:00 2001 From: Aarthy Adityan Date: Mon, 29 Jun 2026 11:08:41 -0400 Subject: [PATCH 67/78] fix(mcp): minor tweaks to description and arg order --- testgen/mcp/tools/projects.py | 2 +- testgen/mcp/tools/test_suites.py | 30 +++++++++++++++--------------- 2 files changed, 16 insertions(+), 16 deletions(-) diff --git a/testgen/mcp/tools/projects.py b/testgen/mcp/tools/projects.py index 325d95a6..be075233 100644 --- a/testgen/mcp/tools/projects.py +++ b/testgen/mcp/tools/projects.py @@ -63,7 +63,7 @@ def update_project( Never echoed back in tool output. data_retention_enabled: Whether old profiling and test history is automatically deleted. data_retention_days: How many days of history to keep (only when retention is enabled). - Defaults to 180 when retention is being enabled without specifying days. + Defaults to 180 when retention is enabled without specifying days. retention_cron_expr: Cron expression for the retention cleanup job (only when retention is enabled). Defaults to daily at 01:00. retention_cron_tz: Timezone for the retention cleanup cron diff --git a/testgen/mcp/tools/test_suites.py b/testgen/mcp/tools/test_suites.py index 56d37e86..dc1ab7bf 100644 --- a/testgen/mcp/tools/test_suites.py +++ b/testgen/mcp/tools/test_suites.py @@ -51,11 +51,11 @@ def create_test_suite( *, description: str | None = None, severity_default: str | None = None, + dq_score_exclude: bool = False, export_to_observability: bool = False, component_key: str | None = None, component_type: str | None = "dataset", component_name: str | None = None, - dq_score_exclude: bool = False, ) -> str: """Create a test suite under a table group. @@ -67,14 +67,14 @@ def create_test_suite( description: Optional free-text description. severity_default: Optional default severity applied to tests that do not set their own. Accepts `Fail` or `Warning`. + dq_score_exclude: Whether to exclude this suite's results from data + quality scoring. Defaults to False. export_to_observability: Whether to export test results to the - configured DataOps Observability backend. Defaults to False. + configured DataOps Observability API. Defaults to False. component_key: Component identifier in DataOps Observability. component_type: Component type in DataOps Observability (e.g. `dataset`). Defaults to `dataset`. component_name: Component display name in DataOps Observability. - dq_score_exclude: Whether to exclude this suite's results from data - quality scoring. Defaults to False. """ name = test_suite_name.strip() if test_suite_name else "" errors: list[str] = [] @@ -93,11 +93,11 @@ def create_test_suite( test_suite=name, test_suite_description=_empty_to_none(description), severity=parsed_severity.value if parsed_severity is not None else None, + dq_score_exclude=dq_score_exclude, export_to_observability=export_to_observability, component_key=_empty_to_none(component_key), component_type=_empty_to_none(component_type), component_name=_empty_to_none(component_name), - dq_score_exclude=dq_score_exclude, is_monitor=False, ) session = get_current_session() @@ -116,11 +116,11 @@ def update_test_suite( test_suite_name: str | None = None, description: str | None = None, severity_default: str | None = None, + dq_score_exclude: bool | None = None, export_to_observability: bool | None = None, component_key: str | None = None, component_type: str | None = None, component_name: str | None = None, - dq_score_exclude: bool | None = None, ) -> str: """Update fields on a test suite. Atomic — nothing saved unless every supplied value is valid. @@ -130,26 +130,26 @@ def update_test_suite( description: New free-text description. Pass an empty string to clear. severity_default: New default severity for tests that do not set their own. Accepts `Fail` or `Warning`. + dq_score_exclude: Whether this suite's results are excluded from data + quality scoring. export_to_observability: Whether test results are exported to the - configured DataOps Observability backend. + configured DataOps Observability API. component_key: Component identifier in DataOps Observability. Pass an empty string to clear. component_type: Component type in DataOps Observability (e.g. `dataset`). Pass an empty string to clear. component_name: Component display name in DataOps Observability. Pass an empty string to clear. - dq_score_exclude: Whether this suite's results are excluded from data - quality scoring. """ supplied = { "test_suite_name": test_suite_name, "description": description, "severity_default": severity_default, + "dq_score_exclude": dq_score_exclude, "export_to_observability": export_to_observability, "component_key": component_key, "component_type": component_type, "component_name": component_name, - "dq_score_exclude": dq_score_exclude, } if all(value is None for value in supplied.values()): raise MCPUserError("No fields supplied to update.") @@ -175,6 +175,8 @@ def update_test_suite( if parsed is not None: updates["severity"] = parsed.value + if dq_score_exclude is not None: + updates["dq_score_exclude"] = dq_score_exclude if export_to_observability is not None: updates["export_to_observability"] = export_to_observability if component_key is not None: @@ -183,8 +185,6 @@ def update_test_suite( updates["component_type"] = _empty_to_none(component_type) if component_name is not None: updates["component_name"] = _empty_to_none(component_name) - if dq_score_exclude is not None: - updates["dq_score_exclude"] = dq_score_exclude if errors: raise_validation_error(errors, "Update rejected. No changes saved.") @@ -230,6 +230,7 @@ def _render_created_suite(suite: TestSuite, *, table_group_name: str) -> str: doc.field("Description", suite.test_suite_description) if suite.severity: doc.field("Default severity", suite.severity) + doc.field("Exclude from quality scoring", suite.dq_score_exclude) doc.field("Export to Observability", suite.export_to_observability) if suite.component_key: doc.field("Component key", suite.component_key) @@ -237,7 +238,6 @@ def _render_created_suite(suite: TestSuite, *, table_group_name: str) -> str: doc.field("Component type", suite.component_type) if suite.component_name: doc.field("Component name", suite.component_name) - doc.field("Exclude from quality scoring", suite.dq_score_exclude) return doc.render() @@ -247,20 +247,20 @@ def _render_created_suite(suite: TestSuite, *, table_group_name: str) -> str: "test_suite", "test_suite_description", "severity", + "dq_score_exclude", "export_to_observability", "component_key", "component_type", "component_name", - "dq_score_exclude", ) _DIFF_LABELS: dict[str, str] = { "test_suite": "Name", "test_suite_description": "Description", "severity": "Default severity", + "dq_score_exclude": "Exclude from quality scoring", "export_to_observability": "Export to Observability", "component_key": "Component key", "component_type": "Component type", "component_name": "Component name", - "dq_score_exclude": "Exclude from quality scoring", } From 4cee4f6158462a059124cb4869c14063c072a7f0 Mon Sep 17 00:00:00 2001 From: Astor Date: Mon, 8 Jun 2026 11:44:33 -0300 Subject: [PATCH 68/78] feat(TG-1085): add data_classification tag field MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Adds a data_classification tag field to the data model, supporting three-level inheritance (column → table → table_group) throughout the application, consistent with existing tag fields. - Migration 0193: ADD COLUMN data_classification to table_groups, data_table_chars, data_column_chars, score_definition_results_breakdown; recreate 8 scoring views with COALESCE inheritance - Schema setup: ADD COLUMN and COALESCE in all 8 standard views - ORM: data_classification added to TableGroup model and ScoreDefinitionBreakdownItem / ScoreCategory / SCORE_CATEGORIES - Queries: COALESCE inheritance in profiling anomaly, test result, and scoring queries; added to TAG_FIELDS and scoring categories list - Data Catalog: fix tag inheritance in list view (join table_groups in get_table_group_columns); fix CSV export to COALESCE from table_groups; add seeded filter defaults (Public/Internal/Confidential/Restricted) via merge_tag_defaults; Excel export column - Table Group Form: Data Classification input in TaggingForm - JS: TAG_KEYS/TAG_HELP in metadata_tags.js; three-level fallback chain in data_catalog.js - Reports: data_classification in PDF hygiene and test result report metadata - Import: "data classification" header mapped in import_metadata_dialog - Tests: 5 unit tests for merge_tag_defaults dedupe logic Co-Authored-By: Claude Sonnet 4.6 test: update TAG_FIELDS count to 9 after adding data_classification Co-Authored-By: Claude Sonnet 4.6 fix(TG-1085): address round-2 review comments --- testgen/common/data_catalog_service.py | 1 + testgen/common/models/data_column.py | 1 + testgen/common/models/data_table.py | 1 + testgen/common/models/scores.py | 4 ++ testgen/common/models/table_group.py | 1 + testgen/mcp/tools/common.py | 6 +++ testgen/mcp/tools/data_catalog.py | 3 +- testgen/mcp/tools/quality_scores.py | 14 ++++--- testgen/mcp/tools/table_groups.py | 14 +++++++ .../030_initialize_new_schema_structure.sql | 4 ++ .../dbsetup/060_create_standard_views.sql | 15 +++++++ .../dbupgrade/0199_incremental_upgrade.sql | 6 +++ .../js/data_profiling/data_profiling_utils.js | 5 +++ .../js/data_profiling/metadata_tags.js | 2 + .../frontend/js/pages/data_catalog.js | 14 ++++++- .../frontend/js/pages/score_explorer.js | 2 + testgen/ui/pdf/hygiene_issue_report.py | 1 + testgen/ui/pdf/test_result_report.py | 1 + testgen/ui/queries/profiling_queries.py | 6 ++- testgen/ui/queries/scoring_queries.py | 6 ++- testgen/ui/queries/test_result_queries.py | 6 ++- .../static/js/components/score_breakdown.js | 1 + .../static/js/components/table_group_form.js | 15 +++++++ testgen/ui/views/data_catalog.py | 32 +++++++++++++-- .../views/dialogs/import_metadata_dialog.py | 3 +- testgen/utils/__init__.py | 1 + .../unit/common/test_data_catalog_service.py | 5 ++- tests/unit/mcp/test_model_data_column.py | 1 + tests/unit/mcp/test_model_data_table.py | 1 + tests/unit/mcp/test_tools_common.py | 2 + tests/unit/mcp/test_tools_table_groups.py | 1 + tests/unit/ui/conftest.py | 1 + tests/unit/ui/test_data_catalog.py | 39 +++++++++++++++++++ 33 files changed, 194 insertions(+), 21 deletions(-) create mode 100644 testgen/template/dbupgrade/0199_incremental_upgrade.sql create mode 100644 tests/unit/ui/test_data_catalog.py diff --git a/testgen/common/data_catalog_service.py b/testgen/common/data_catalog_service.py index d31a0c55..ba385411 100644 --- a/testgen/common/data_catalog_service.py +++ b/testgen/common/data_catalog_service.py @@ -36,6 +36,7 @@ "transform_level", "aggregation_level", "data_product", + "data_classification", ] # Metadata fields settable per target type, keyed by their data_*_chars column names. diff --git a/testgen/common/models/data_column.py b/testgen/common/models/data_column.py index 75734d01..12a7473b 100644 --- a/testgen/common/models/data_column.py +++ b/testgen/common/models/data_column.py @@ -270,6 +270,7 @@ class DataColumnChars(Entity): transform_level: str | None = Column(String(40)) aggregation_level: str | None = Column(String(40)) data_product: str | None = Column(String(40)) + data_classification: str | None = Column(String(40)) drop_date: datetime | None = Column(postgresql.TIMESTAMP) last_complete_profile_run_id: UUID | None = Column(postgresql.UUID(as_uuid=True)) dq_score_profiling: float | None = Column(Float) diff --git a/testgen/common/models/data_table.py b/testgen/common/models/data_table.py index 51cccf1c..d3d3c7d0 100644 --- a/testgen/common/models/data_table.py +++ b/testgen/common/models/data_table.py @@ -77,6 +77,7 @@ class DataTable(Entity): transform_level: str | None = Column(String(40)) aggregation_level: str | None = Column(String(40)) data_product: str | None = Column(String(40)) + data_classification: str | None = Column(String(40)) drop_date: datetime | None = Column(postgresql.TIMESTAMP) last_complete_profile_run_id: UUID | None = Column(postgresql.UUID(as_uuid=True)) dq_score_profiling: float | None = Column(Float) diff --git a/testgen/common/models/scores.py b/testgen/common/models/scores.py index 4dfb4c10..7b2e8d40 100644 --- a/testgen/common/models/scores.py +++ b/testgen/common/models/scores.py @@ -52,6 +52,7 @@ "stakeholder_group", "transform_level", "data_product", + "data_classification", ] Categories = Literal[ "column_name", @@ -68,6 +69,7 @@ "stakeholder_group", "transform_level", "data_product", + "data_classification", ] ScoreTypes = Literal["score", "cde_score"] @@ -88,6 +90,7 @@ class ScoreCategory(enum.Enum): dq_dimension = "dq_dimension" impact_dimension = "impact_dimension" data_product = "data_product" + data_classification = "data_classification" class ScoreDefinition(Base): @@ -826,6 +829,7 @@ class ScoreDefinitionBreakdownItem(Base): stakeholder_group: str = Column(String, nullable=True) transform_level: str = Column(String, nullable=True) data_product: str = Column(String, nullable=True) + data_classification: str = Column(String, nullable=True) impact: float = Column(Float) score: float = Column(Float) issue_ct: int = Column(Integer) diff --git a/testgen/common/models/table_group.py b/testgen/common/models/table_group.py index b107e1d3..5314d4ec 100644 --- a/testgen/common/models/table_group.py +++ b/testgen/common/models/table_group.py @@ -261,6 +261,7 @@ class TableGroup(Entity): stakeholder_group: str = Column(NullIfEmptyString) transform_level: str = Column(NullIfEmptyString) data_product: str = Column(NullIfEmptyString) + data_classification: str = Column(NullIfEmptyString) last_complete_profile_run_id: UUID = Column(postgresql.UUID(as_uuid=True)) dq_score_profiling: float = Column(Float) dq_score_testing: float = Column(Float) diff --git a/testgen/mcp/tools/common.py b/testgen/mcp/tools/common.py index ffe127a6..f90038df 100644 --- a/testgen/mcp/tools/common.py +++ b/testgen/mcp/tools/common.py @@ -158,6 +158,7 @@ class ScoreGroupBy(StrEnum): STAKEHOLDER_GROUP = "Stakeholder Group" TRANSFORM_LEVEL = "Transform Level" DATA_PRODUCT = "Data Product" + DATA_CLASSIFICATION = "Data Classification" # Translates the user-facing label to the internal DB column name used by @@ -175,6 +176,7 @@ class ScoreGroupBy(StrEnum): ScoreGroupBy.STAKEHOLDER_GROUP: "stakeholder_group", ScoreGroupBy.TRANSFORM_LEVEL: "transform_level", ScoreGroupBy.DATA_PRODUCT: "data_product", + ScoreGroupBy.DATA_CLASSIFICATION: "data_classification", } @@ -197,6 +199,7 @@ class ScoreFilterField(StrEnum): STAKEHOLDER_GROUP = "Stakeholder Group" TRANSFORM_LEVEL = "Transform Level" DATA_PRODUCT = "Data Product" + DATA_CLASSIFICATION = "Data Classification" SCORE_FILTER_FIELD_TO_COLUMN: dict[ScoreFilterField, str] = { @@ -210,6 +213,7 @@ class ScoreFilterField(StrEnum): ScoreFilterField.STAKEHOLDER_GROUP: "stakeholder_group", ScoreFilterField.TRANSFORM_LEVEL: "transform_level", ScoreFilterField.DATA_PRODUCT: "data_product", + ScoreFilterField.DATA_CLASSIFICATION: "data_classification", } @@ -233,6 +237,7 @@ class ScoreCategoryArg(StrEnum): QUALITY_DIMENSION = "Quality Dimension" IMPACT_DIMENSION = "Impact Dimension" DATA_PRODUCT = "Data Product" + DATA_CLASSIFICATION = "Data Classification" SCORE_CATEGORY_ARG_TO_COLUMN: dict[ScoreCategoryArg, str] = { @@ -247,6 +252,7 @@ class ScoreCategoryArg(StrEnum): ScoreCategoryArg.QUALITY_DIMENSION: "dq_dimension", ScoreCategoryArg.IMPACT_DIMENSION: "impact_dimension", ScoreCategoryArg.DATA_PRODUCT: "data_product", + ScoreCategoryArg.DATA_CLASSIFICATION: "data_classification", } diff --git a/testgen/mcp/tools/data_catalog.py b/testgen/mcp/tools/data_catalog.py index 334c65bb..645d160a 100644 --- a/testgen/mcp/tools/data_catalog.py +++ b/testgen/mcp/tools/data_catalog.py @@ -48,7 +48,8 @@ def update_catalog_metadata(updates: list[dict]) -> str: pii: Whether the column contains personally identifiable information (true/false). Columns only. data_source, source_system, source_process, business_domain, - stakeholder_group, transform_level, aggregation_level, data_product: + stakeholder_group, transform_level, aggregation_level, data_product, + data_classification: Catalog tags (max 40 characters each). """ if not updates: diff --git a/testgen/mcp/tools/quality_scores.py b/testgen/mcp/tools/quality_scores.py index 1545d7f4..692b9995 100644 --- a/testgen/mcp/tools/quality_scores.py +++ b/testgen/mcp/tools/quality_scores.py @@ -87,7 +87,7 @@ def get_quality_scores( ``"Table Group"``, ``"Data Location"``, ``"Data Source"``, ``"Source System"``, ``"Source Process"``, ``"Business Domain"``, ``"Stakeholder Group"``, ``"Transform Level"``, ``"Semantic Data Type"``, - ``"Data Product"``. To target specific tables or columns, chain a + ``"Data Product"``, ``"Data Classification"``. To target specific tables or columns, chain a ``"Table Group"`` filter via ``others`` into ``"Table"`` (optionally then ``"Column"``); sibling chains OR. ``"Impact Dimension"`` and ``"Quality Dimension"`` are valid as ``group_by`` only, not as filter @@ -106,7 +106,7 @@ def get_quality_scores( ``"Table Group"``, ``"Data Location"``, ``"Data Source"``, ``"Source System"``, ``"Source Process"``, ``"Business Domain"``, ``"Stakeholder Group"``, ``"Transform Level"``, - ``"Data Product"``. + ``"Data Product"``, ``"Data Classification"``. score_type: Narrow returned scores. Omit to show all four (Total, CDE, Profiling, Testing); pass ``"Total"`` for Total + Profiling + Testing, or ``"CDE"`` for CDE alone. @@ -551,7 +551,7 @@ def create_scorecard( ``"Table Group"``, ``"Data Location"``, ``"Data Source"``, ``"Source System"``, ``"Source Process"``, ``"Business Domain"``, ``"Stakeholder Group"``, ``"Transform Level"``, ``"Semantic Data Type"``, - ``"Data Product"``. To target specific tables or columns, chain a + ``"Data Product"``, ``"Data Classification"``. To target specific tables or columns, chain a ``"Table Group"`` filter via ``others`` into ``"Table"`` (optionally then ``"Column"``); sibling chains OR. @@ -563,7 +563,8 @@ def create_scorecard( ``"Quality Dimension"``, ``"Impact Dimension"``, ``"Data Source"``, ``"Business Domain"``, ``"Stakeholder Group"``, ``"Table Group"``, ``"Transform Level"``, ``"Data Location"``, - ``"Source System"``, ``"Source Process"``, ``"Data Product"``. + ``"Source System"``, ``"Source Process"``, ``"Data Product"``, + ``"Data Classification"``. show_total_score: Whether the scorecard exposes the Total Score. show_cde_score: Whether the scorecard exposes the CDE Score. """ @@ -624,7 +625,7 @@ def update_scorecard( ``"Table Group"``, ``"Data Location"``, ``"Data Source"``, ``"Source System"``, ``"Source Process"``, ``"Business Domain"``, ``"Stakeholder Group"``, ``"Transform Level"``, ``"Semantic Data Type"``, - ``"Data Product"``. To target specific tables or columns, chain a + ``"Data Product"``, ``"Data Classification"``. To target specific tables or columns, chain a ``"Table Group"`` filter via ``others`` into ``"Table"`` (optionally then ``"Column"``); sibling chains OR. @@ -638,7 +639,8 @@ def update_scorecard( ``"Quality Dimension"``, ``"Impact Dimension"``, ``"Data Source"``, ``"Business Domain"``, ``"Stakeholder Group"``, ``"Table Group"``, ``"Transform Level"``, ``"Data Location"``, - ``"Source System"``, ``"Source Process"``, ``"Data Product"``. + ``"Source System"``, ``"Source Process"``, ``"Data Product"``, + ``"Data Classification"``. Pass ``""`` to clear an existing category. filters: List of filter entries. See **Filters** above for shape. """ diff --git a/testgen/mcp/tools/table_groups.py b/testgen/mcp/tools/table_groups.py index ca1524ae..9f4fa259 100644 --- a/testgen/mcp/tools/table_groups.py +++ b/testgen/mcp/tools/table_groups.py @@ -214,6 +214,7 @@ def create_table_group( stakeholder_group: str | None = None, transform_level: str | None = None, data_product: str | None = None, + data_classification: str | None = None, ) -> str: """Create a table group on an existing connection. @@ -254,6 +255,8 @@ def create_table_group( ``Conformed``, ``Processed``, ``Reporting``) or Medallion level (``bronze``, ``silver``, ``gold``). data_product: Catalog tag — data domain that comprises the dataset. + data_classification: Catalog tag — information classification level of the dataset + (e.g. ``Public``, ``Internal``, ``Confidential``, ``Restricted``). """ connection = resolve_connection(connection_id) @@ -288,6 +291,7 @@ def create_table_group( stakeholder_group=stakeholder_group, transform_level=transform_level, data_product=data_product, + data_classification=data_classification, ) errors = validate_table_group_fields(table_group) @@ -332,6 +336,7 @@ def update_table_group( stakeholder_group: str | None = None, transform_level: str | None = None, data_product: str | None = None, + data_classification: str | None = None, ) -> str: """Update fields on an existing table group. Atomic — no partial save. @@ -371,6 +376,8 @@ def update_table_group( ``Conformed``, ``Processed``, ``Reporting``) or Medallion level (``bronze``, ``silver``, ``gold``). data_product: Catalog tag — data domain that comprises the dataset. + data_classification: Catalog tag — information classification level of the dataset + (e.g. ``Public``, ``Internal``, ``Confidential``, ``Restricted``). """ supplied = { "table_group_name": table_group_name, @@ -397,6 +404,7 @@ def update_table_group( "stakeholder_group": stakeholder_group, "transform_level": transform_level, "data_product": data_product, + "data_classification": data_classification, } if all(value is None for value in supplied.values()): raise MCPUserError("No fields supplied to update.") @@ -545,6 +553,7 @@ def preview_table_group( "stakeholder_group", "transform_level", "data_product", + "data_classification", ) _DIFF_LABELS: dict[str, str] = { @@ -572,6 +581,7 @@ def preview_table_group( "stakeholder_group": "Stakeholder group", "transform_level": "Transform level", "data_product": "Data product", + "data_classification": "Data classification", } _CATALOG_ATTRS: tuple[str, ...] = ( @@ -583,6 +593,7 @@ def preview_table_group( "stakeholder_group", "transform_level", "data_product", + "data_classification", ) @@ -646,6 +657,7 @@ def _apply_args_to_table_group( stakeholder_group: str | None = None, transform_level: str | None = None, data_product: str | None = None, + data_classification: str | None = None, ) -> None: """Apply every non-None arg to its model field. @@ -700,6 +712,8 @@ def _apply_args_to_table_group( table_group.transform_level = transform_level if data_product is not None: table_group.data_product = data_product + if data_classification is not None: + table_group.data_classification = data_classification def _raise_validation_error(errors: list[str], header: str) -> None: diff --git a/testgen/template/dbsetup/030_initialize_new_schema_structure.sql b/testgen/template/dbsetup/030_initialize_new_schema_structure.sql index 75f37b15..32e6b569 100644 --- a/testgen/template/dbsetup/030_initialize_new_schema_structure.sql +++ b/testgen/template/dbsetup/030_initialize_new_schema_structure.sql @@ -131,6 +131,7 @@ CREATE TABLE table_groups stakeholder_group VARCHAR(40), transform_level VARCHAR(40), data_product VARCHAR(40), + data_classification VARCHAR(40), last_complete_profile_run_id UUID, dq_score_profiling FLOAT, dq_score_testing FLOAT @@ -429,6 +430,7 @@ CREATE TABLE data_table_chars ( transform_level VARCHAR(40), aggregation_level VARCHAR(40), data_product VARCHAR(40), + data_classification VARCHAR(40), add_date TIMESTAMP, drop_date TIMESTAMP, last_refresh_date TIMESTAMP, @@ -468,6 +470,7 @@ CREATE TABLE data_column_chars ( transform_level VARCHAR(40), aggregation_level VARCHAR(40), data_product VARCHAR(40), + data_classification VARCHAR(40), add_date TIMESTAMP, last_mod_date TIMESTAMP, drop_date TIMESTAMP, @@ -851,6 +854,7 @@ CREATE TABLE IF NOT EXISTS score_definition_results_breakdown ( stakeholder_group TEXT DEFAULT NULL, transform_level TEXT DEFAULT NULL, data_product TEXT DEFAULT NULL, + data_classification TEXT DEFAULT NULL, impact DOUBLE PRECISION DEFAULT 0.0, score DOUBLE PRECISION DEFAULT 0.0, issue_ct INTEGER DEFAULT 0 diff --git a/testgen/template/dbsetup/060_create_standard_views.sql b/testgen/template/dbsetup/060_create_standard_views.sql index fffaa445..5b99ae67 100644 --- a/testgen/template/dbsetup/060_create_standard_views.sql +++ b/testgen/template/dbsetup/060_create_standard_views.sql @@ -122,6 +122,7 @@ SELECT COALESCE(dcc.transform_level, dtc.transform_level, tg.transform_level) as transform_level, COALESCE(dcc.critical_data_element, dtc.critical_data_element) as critical_data_element, COALESCE(dcc.data_product, dtc.data_product, tg.data_product) as data_product, + COALESCE(dcc.data_classification, dtc.data_classification, tg.data_classification) as data_classification, dcc.functional_data_type as semantic_data_type, dtc.table_name, dcc.column_name, pr.profiling_starttime as profiling_run_date, @@ -162,6 +163,7 @@ SELECT tg.project_code, COALESCE(dcc.transform_level, dtc.transform_level, tg.transform_level) as transform_level, COALESCE(dcc.critical_data_element, dtc.critical_data_element) as critical_data_element, COALESCE(dcc.data_product, dtc.data_product, tg.data_product) as data_product, + COALESCE(dcc.data_classification, dtc.data_classification, tg.data_classification) as data_classification, dcc.functional_data_type as semantic_data_type, t.dq_dimension, pr.table_name, @@ -205,6 +207,7 @@ GROUP BY pr.profile_run_id, pr.table_groups_id, COALESCE(dcc.transform_level, dtc.transform_level, tg.transform_level), COALESCE(dcc.critical_data_element, dtc.critical_data_element), COALESCE(dcc.data_product, dtc.data_product, tg.data_product), + COALESCE(dcc.data_classification, dtc.data_classification, tg.data_classification), dcc.functional_data_type, t.dq_dimension, pr.run_date, tg.project_code ; @@ -232,6 +235,7 @@ SELECT COALESCE(dcc.transform_level, dtc.transform_level, tg.transform_level) as transform_level, COALESCE(dcc.critical_data_element, dtc.critical_data_element) as critical_data_element, COALESCE(dcc.data_product, dtc.data_product, tg.data_product) as data_product, + COALESCE(dcc.data_classification, dtc.data_classification, tg.data_classification) as data_classification, dcc.functional_data_type as semantic_data_type, r.test_time, r.table_name, r.column_names as column_name, COUNT(*) as test_ct, @@ -271,6 +275,7 @@ GROUP BY r.table_groups_id, r.table_name, r.column_names, tg.stakeholder_group, dcc.transform_level, dtc.transform_level, tg.transform_level, dcc.critical_data_element, dtc.critical_data_element, dcc.data_product, dtc.data_product, tg.data_product, + dcc.data_classification, dtc.data_classification, tg.data_classification, dcc.functional_data_type, r.test_time, tg.project_code; @@ -312,6 +317,7 @@ SELECT COALESCE(dcc.transform_level, dtc.transform_level, tg.transform_level) as transform_level, COALESCE(dcc.critical_data_element, dtc.critical_data_element) as critical_data_element, COALESCE(dcc.data_product, dtc.data_product, tg.data_product) as data_product, + COALESCE(dcc.data_classification, dtc.data_classification, tg.data_classification) as data_classification, dcc.functional_data_type as semantic_data_type, r.dq_dimension, r.test_time, r.table_name, dcc.column_name, @@ -347,6 +353,7 @@ GROUP BY r.table_groups_id, r.test_run_id, r.test_suite_id, tg.stakeholder_group, dcc.transform_level, dtc.transform_level, tg.transform_level, dcc.critical_data_element, dtc.critical_data_element, dcc.data_product, dtc.data_product, tg.data_product, + dcc.data_classification, dtc.data_classification, tg.data_classification, dcc.functional_data_type, r.dq_dimension, r.test_time, r.table_name, dcc.column_name, tg.project_code; @@ -368,6 +375,7 @@ SELECT tg.project_code, COALESCE(dcc.transform_level, dtc.transform_level, tg.transform_level) as transform_level, COALESCE(dcc.critical_data_element, dtc.critical_data_element) as critical_data_element, COALESCE(dcc.data_product, dtc.data_product, tg.data_product) as data_product, + COALESCE(dcc.data_classification, dtc.data_classification, tg.data_classification) as data_classification, dcc.functional_data_type as semantic_data_type, t.impact_dimension, pr.table_name, @@ -411,6 +419,7 @@ GROUP BY pr.profile_run_id, pr.table_groups_id, COALESCE(dcc.transform_level, dtc.transform_level, tg.transform_level), COALESCE(dcc.critical_data_element, dtc.critical_data_element), COALESCE(dcc.data_product, dtc.data_product, tg.data_product), + COALESCE(dcc.data_classification, dtc.data_classification, tg.data_classification), dcc.functional_data_type, t.impact_dimension, pr.run_date, tg.project_code; @@ -450,6 +459,7 @@ SELECT COALESCE(dcc.transform_level, dtc.transform_level, tg.transform_level) as transform_level, COALESCE(dcc.critical_data_element, dtc.critical_data_element) as critical_data_element, COALESCE(dcc.data_product, dtc.data_product, tg.data_product) as data_product, + COALESCE(dcc.data_classification, dtc.data_classification, tg.data_classification) as data_classification, dcc.functional_data_type as semantic_data_type, r.impact_dimension, r.test_time, r.table_name, dcc.column_name, @@ -485,6 +495,7 @@ GROUP BY r.table_groups_id, r.test_run_id, r.test_suite_id, tg.stakeholder_group, dcc.transform_level, dtc.transform_level, tg.transform_level, dcc.critical_data_element, dtc.critical_data_element, dcc.data_product, dtc.data_product, tg.data_product, + dcc.data_classification, dtc.data_classification, tg.data_classification, dcc.functional_data_type, r.impact_dimension, r.test_time, r.table_name, dcc.column_name, tg.project_code; @@ -511,6 +522,7 @@ SELECT tg.project_code, COALESCE(dcc.transform_level, dtc.transform_level, tg.transform_level) as transform_level, COALESCE(dcc.critical_data_element, dtc.critical_data_element) as critical_data_element, COALESCE(dcc.data_product, dtc.data_product, tg.data_product) as data_product, + COALESCE(dcc.data_classification, dtc.data_classification, tg.data_classification) as data_classification, dcc.functional_data_type as semantic_data_type, pr.table_name, pr.column_name, @@ -557,6 +569,7 @@ GROUP BY pr.profile_run_id, COALESCE(dcc.transform_level, dtc.transform_level, tg.transform_level), COALESCE(dcc.critical_data_element, dtc.critical_data_element), COALESCE(dcc.data_product, dtc.data_product, tg.data_product), + COALESCE(dcc.data_classification, dtc.data_classification, tg.data_classification), dcc.functional_data_type, pr.run_date, tg.project_code ; @@ -581,6 +594,7 @@ SELECT COALESCE(dcc.transform_level, dtc.transform_level, tg.transform_level) as transform_level, COALESCE(dcc.critical_data_element, dtc.critical_data_element) as critical_data_element, COALESCE(dcc.data_product, dtc.data_product, tg.data_product) as data_product, + COALESCE(dcc.data_classification, dtc.data_classification, tg.data_classification) as data_classification, dcc.functional_data_type as semantic_data_type, r.test_time, r.table_name, r.column_names as column_name, COUNT(*) as test_ct, @@ -623,5 +637,6 @@ GROUP BY sr.definition_id, tg.stakeholder_group, dcc.transform_level, dtc.transform_level, tg.transform_level, dcc.critical_data_element, dtc.critical_data_element, dcc.data_product, dtc.data_product, tg.data_product, + dcc.data_classification, dtc.data_classification, tg.data_classification, dcc.functional_data_type, r.test_time, tg.project_code; diff --git a/testgen/template/dbupgrade/0199_incremental_upgrade.sql b/testgen/template/dbupgrade/0199_incremental_upgrade.sql new file mode 100644 index 00000000..43aaf9c9 --- /dev/null +++ b/testgen/template/dbupgrade/0199_incremental_upgrade.sql @@ -0,0 +1,6 @@ +SET SEARCH_PATH TO {SCHEMA_NAME}; + +ALTER TABLE table_groups ADD COLUMN IF NOT EXISTS data_classification VARCHAR(40); +ALTER TABLE data_table_chars ADD COLUMN IF NOT EXISTS data_classification VARCHAR(40); +ALTER TABLE data_column_chars ADD COLUMN IF NOT EXISTS data_classification VARCHAR(40); +ALTER TABLE score_definition_results_breakdown ADD COLUMN IF NOT EXISTS data_classification TEXT DEFAULT NULL; diff --git a/testgen/ui/components/frontend/js/data_profiling/data_profiling_utils.js b/testgen/ui/components/frontend/js/data_profiling/data_profiling_utils.js index 5089bea5..d44255ec 100644 --- a/testgen/ui/components/frontend/js/data_profiling/data_profiling_utils.js +++ b/testgen/ui/components/frontend/js/data_profiling/data_profiling_utils.js @@ -55,6 +55,7 @@ * @property {string?} transform_level * @property {string?} aggregation_level * @property {string?} data_product + * @property {string?} data_classification * * Table Tags * @property {boolean?} table_critical_data_element * @property {string?} table_data_source @@ -65,6 +66,7 @@ * @property {string?} table_transform_level * @property {string?} table_aggregation_level * @property {string?} table_data_product + * @property {string?} table_data_classification * * Table Group Tags * @property {string} table_group_data_source * @property {string} table_group_source_system @@ -73,6 +75,7 @@ * @property {string} table_group_stakeholder_group * @property {string} table_group_transform_level * @property {string} table_group_data_product + * @property {string} table_group_data_classification * * Profile & Test Runs * @property {string?} profile_run_id * @property {number?} profile_run_date @@ -169,6 +172,7 @@ * @property {string} transform_level * @property {string} aggregation_level * @property {string} data_product + * @property {string} data_classification * * Table Group Tags * @property {string} table_group_data_source * @property {string} table_group_source_system @@ -177,6 +181,7 @@ * @property {string} table_group_stakeholder_group * @property {string} table_group_transform_level * @property {string} table_group_data_product + * @property {string} table_group_data_classification * * Profile & Test Runs * @property {string} profile_run_id * @property {number} profile_run_date diff --git a/testgen/ui/components/frontend/js/data_profiling/metadata_tags.js b/testgen/ui/components/frontend/js/data_profiling/metadata_tags.js index fbb9ad8f..5ae1092a 100644 --- a/testgen/ui/components/frontend/js/data_profiling/metadata_tags.js +++ b/testgen/ui/components/frontend/js/data_profiling/metadata_tags.js @@ -61,6 +61,7 @@ const TAG_KEYS = [ 'transform_level', 'aggregation_level', 'data_product', + 'data_classification', ]; const TAG_HELP = { data_source: 'Original source of the dataset', @@ -71,6 +72,7 @@ const TAG_HELP = { transform_level: 'Data warehouse processing stage, e.g., Raw, Conformed, Processed, Reporting, or Medallion level (bronze, silver, gold)', aggregation_level: 'Data granularity of the dataset, e.g. atomic, historical, snapshot, aggregated, time-rollup, rolling, summary', data_product: 'Data domain that comprises the dataset', + data_classification: 'Information sensitivity level of the dataset, e.g., Public, Internal, Confidential, Restricted', }; /** diff --git a/testgen/ui/components/frontend/js/pages/data_catalog.js b/testgen/ui/components/frontend/js/pages/data_catalog.js index 3ed1c10c..92857742 100644 --- a/testgen/ui/components/frontend/js/pages/data_catalog.js +++ b/testgen/ui/components/frontend/js/pages/data_catalog.js @@ -29,6 +29,7 @@ * @property {string} transform_level * @property {string} aggregation_level * @property {string} data_product + * @property {string} data_classification * @property {string} table_data_source * @property {string} table_source_system * @property {string} table_source_process @@ -37,6 +38,15 @@ * @property {string} table_transform_level * @property {string} table_aggregation_level * @property {string} table_data_product + * @property {string} table_data_classification + * @property {string} table_group_data_source + * @property {string} table_group_source_system + * @property {string} table_group_source_process + * @property {string} table_group_business_domain + * @property {string} table_group_stakeholder_group + * @property {string} table_group_transform_level + * @property {string} table_group_data_product + * @property {string} table_group_data_classification * * @typedef Permissions * @type {object} @@ -130,7 +140,7 @@ const DataCatalog = (/** @type Properties */ props) => { criticalDataElement: !!item.table_critical_data_element, children: [], }; - TAG_KEYS.forEach(key => tables[table_id][key] = item[`table_${key}`]); + TAG_KEYS.forEach(key => tables[table_id][key] = item[`table_${key}`] ?? item[`table_group_${key}`]); } const columnNode = { id: column_id, @@ -156,7 +166,7 @@ const DataCatalog = (/** @type Properties */ props) => { excludedDataElement: !!item.excluded_data_element, piiFlag: !!item.pii_flag, }; - TAG_KEYS.forEach(key => columnNode[key] = item[key] ?? item[`table_${key}`]); + TAG_KEYS.forEach(key => columnNode[key] = item[key] ?? item[`table_${key}`] ?? item[`table_group_${key}`]); tables[table_id].children.push(columnNode); }); return Object.values(tables); diff --git a/testgen/ui/components/frontend/js/pages/score_explorer.js b/testgen/ui/components/frontend/js/pages/score_explorer.js index addd28b3..7f8d8d12 100644 --- a/testgen/ui/components/frontend/js/pages/score_explorer.js +++ b/testgen/ui/components/frontend/js/pages/score_explorer.js @@ -82,6 +82,7 @@ const TRANSLATIONS = { dq_dimension: 'Quality Dimension', impact_dimension: 'Impact Dimension', data_product: 'Data Product', + data_classification: 'Data Classification', }; const ScoreExplorer = (/** @type {Properties} */ props) => { @@ -188,6 +189,7 @@ const Toolbar = ( 'dq_dimension', 'impact_dimension', 'data_product', + 'data_classification', ]; const filterableFields = categories.filter((c) => c !== 'dq_dimension' && c !== 'impact_dimension'); const filters = van.state(definition.filters.map((f, idx) => ({key: `${f.field}-${idx}-${getRandomId()}`, field: f.field, value: van.state(f.value), others: f.others ?? [] }))); diff --git a/testgen/ui/pdf/hygiene_issue_report.py b/testgen/ui/pdf/hygiene_issue_report.py index a6d24fe3..b754ab19 100644 --- a/testgen/ui/pdf/hygiene_issue_report.py +++ b/testgen/ui/pdf/hygiene_issue_report.py @@ -138,6 +138,7 @@ def build_summary_table(document, hi_data): "transform_level", "aggregation_level", "data_product", + "data_classification", ] if hi_data[tag] ], diff --git a/testgen/ui/pdf/test_result_report.py b/testgen/ui/pdf/test_result_report.py index 9b57ff73..455b2bdf 100644 --- a/testgen/ui/pdf/test_result_report.py +++ b/testgen/ui/pdf/test_result_report.py @@ -154,6 +154,7 @@ def build_summary_table(document, tr_data): "transform_level", "aggregation_level", "data_product", + "data_classification", ] if tr_data[tag] ], diff --git a/testgen/ui/queries/profiling_queries.py b/testgen/ui/queries/profiling_queries.py index 5a64d238..49087d13 100644 --- a/testgen/ui/queries/profiling_queries.py +++ b/testgen/ui/queries/profiling_queries.py @@ -613,7 +613,8 @@ def get_profiling_anomalies( COALESCE(dcc.stakeholder_group, dtc.stakeholder_group, tg.stakeholder_group) as stakeholder_group, COALESCE(dcc.transform_level, dtc.transform_level, tg.transform_level) as transform_level, COALESCE(dcc.aggregation_level, dtc.aggregation_level) as aggregation_level, - COALESCE(dcc.data_product, dtc.data_product, tg.data_product) as data_product + COALESCE(dcc.data_product, dtc.data_product, tg.data_product) as data_product, + COALESCE(dcc.data_classification, dtc.data_classification, tg.data_classification) as data_classification FROM profile_anomaly_results r INNER JOIN profile_anomaly_types t ON r.anomaly_id = t.id @@ -682,7 +683,8 @@ def get_profiling_anomalies_by_ids(anomaly_ids: list[str]) -> pd.DataFrame: COALESCE(dcc.stakeholder_group, dtc.stakeholder_group, tg.stakeholder_group) as stakeholder_group, COALESCE(dcc.transform_level, dtc.transform_level, tg.transform_level) as transform_level, COALESCE(dcc.aggregation_level, dtc.aggregation_level) as aggregation_level, - COALESCE(dcc.data_product, dtc.data_product, tg.data_product) as data_product + COALESCE(dcc.data_product, dtc.data_product, tg.data_product) as data_product, + COALESCE(dcc.data_classification, dtc.data_classification, tg.data_classification) as data_classification FROM profile_anomaly_results r INNER JOIN profile_anomaly_types t ON r.anomaly_id = t.id diff --git a/testgen/ui/queries/scoring_queries.py b/testgen/ui/queries/scoring_queries.py index 019c9bb7..23f9e00b 100644 --- a/testgen/ui/queries/scoring_queries.py +++ b/testgen/ui/queries/scoring_queries.py @@ -59,6 +59,7 @@ def get_score_card_issue_reports(selected_issues: list["SelectedIssue"]) -> list COALESCE(column_chars.transform_level, table_chars.transform_level, groups.transform_level) as transform_level, COALESCE(column_chars.aggregation_level, table_chars.aggregation_level) as aggregation_level, COALESCE(column_chars.data_product, table_chars.data_product, groups.data_product) as data_product, + COALESCE(column_chars.data_classification, table_chars.data_classification, groups.data_classification) as data_classification, types.impact_dimension, types.dq_dimension FROM profile_anomaly_results results @@ -128,7 +129,9 @@ def get_score_card_issue_reports(selected_issues: list["SelectedIssue"]) -> list COALESCE(column_chars.transform_level, table_chars.transform_level, groups.transform_level) as transform_level, COALESCE(column_chars.aggregation_level, table_chars.aggregation_level) as aggregation_level, COALESCE(column_chars.data_product, table_chars.data_product, groups.data_product) as data_product, - COALESCE(results.impact_dimension, types.impact_dimension) as impact_dimension FROM test_results results + COALESCE(column_chars.data_classification, table_chars.data_classification, groups.data_classification) as data_classification, + COALESCE(results.impact_dimension, types.impact_dimension) as impact_dimension + FROM test_results results INNER JOIN test_types types ON (results.test_type = types.test_type) INNER JOIN test_suites suites @@ -179,6 +182,7 @@ def get_score_category_values(project_code: str) -> dict[ScoreCategory, list[str "stakeholder_group", "transform_level", "data_product", + "data_classification", ] quote = lambda v: f"'{v}'" diff --git a/testgen/ui/queries/test_result_queries.py b/testgen/ui/queries/test_result_queries.py index 92f38673..ed3c520c 100644 --- a/testgen/ui/queries/test_result_queries.py +++ b/testgen/ui/queries/test_result_queries.py @@ -131,7 +131,8 @@ def get_test_results( COALESCE(dcc.stakeholder_group, dtc.stakeholder_group, tg.stakeholder_group) as stakeholder_group, COALESCE(dcc.transform_level, dtc.transform_level, tg.transform_level) as transform_level, COALESCE(dcc.aggregation_level, dtc.aggregation_level) as aggregation_level, - COALESCE(dcc.data_product, dtc.data_product, tg.data_product) as data_product + COALESCE(dcc.data_product, dtc.data_product, tg.data_product) as data_product, + COALESCE(dcc.data_classification, dtc.data_classification, tg.data_classification) as data_classification FROM run_results r INNER JOIN test_types tt ON (r.test_type = tt.test_type) @@ -218,7 +219,8 @@ def get_test_results_by_ids(test_result_ids: list[str]) -> pd.DataFrame: COALESCE(dcc.stakeholder_group, dtc.stakeholder_group, tg.stakeholder_group) as stakeholder_group, COALESCE(dcc.transform_level, dtc.transform_level, tg.transform_level) as transform_level, COALESCE(dcc.aggregation_level, dtc.aggregation_level) as aggregation_level, - COALESCE(dcc.data_product, dtc.data_product, tg.data_product) as data_product + COALESCE(dcc.data_product, dtc.data_product, tg.data_product) as data_product, + COALESCE(dcc.data_classification, dtc.data_classification, tg.data_classification) as data_classification FROM test_results r INNER JOIN test_runs tr ON (r.test_run_id = tr.id) diff --git a/testgen/ui/static/js/components/score_breakdown.js b/testgen/ui/static/js/components/score_breakdown.js index 0b843405..cb069ee5 100644 --- a/testgen/ui/static/js/components/score_breakdown.js +++ b/testgen/ui/static/js/components/score_breakdown.js @@ -194,6 +194,7 @@ const CATEGORIES = { stakeholder_group: 'Stakeholder Group', transform_level: 'Transform Level', data_product: 'Data Product', + data_classification: 'Data Classification', }; const BREAKDOWN_COLUMN_LABEL = { diff --git a/testgen/ui/static/js/components/table_group_form.js b/testgen/ui/static/js/components/table_group_form.js index 85f53b66..739f9210 100644 --- a/testgen/ui/static/js/components/table_group_form.js +++ b/testgen/ui/static/js/components/table_group_form.js @@ -29,6 +29,7 @@ * @property {string?} stakeholder_group * @property {string?} transform_level * @property {string?} data_product + * @property {string?} data_classification * * @typedef FormState * @type {object} @@ -104,6 +105,7 @@ const TableGroupForm = (props) => { const stakeholderGroup = van.state(tableGroup.stakeholder_group); const transformLevel = van.state(tableGroup.transform_level); const dataProduct = van.state(tableGroup.data_product); + const dataClassification = van.state(tableGroup.data_classification); const connectionOptions = van.derive(() => { const connections = getValue(props.connections) ?? []; @@ -159,6 +161,7 @@ const TableGroupForm = (props) => { stakeholder_group: stakeholderGroup.val, transform_level: transformLevel.val, data_product: dataProduct.val, + data_classification: dataClassification.val, }; }); const dirty = van.derive(() => !isEqual(updatedTableGroup.val, tableGroup)); @@ -233,6 +236,7 @@ const TableGroupForm = (props) => { stakeholderGroup, transformLevel, dataProduct, + dataClassification, ), ); }; @@ -480,6 +484,7 @@ const TaggingForm = ( stakeholderGroup, transformLevel, dataProduct, + dataClassification, ) => { return ExpansionPanel( { title: 'Table Group Tags' }, @@ -575,6 +580,16 @@ const TaggingForm = ( options.setValidity?.('data_product', state.valid); }, }), + Input({ + name: 'data_classification', + label: 'Data Classification', + value: dataClassification, + help: 'Information sensitivity level of the dataset, e.g., Public, Internal, Confidential, Restricted', + onChange: (value, state) => { + dataClassification.val = value; + options.setValidity?.('data_classification', state.valid); + }, + }), ), ); }; diff --git a/testgen/ui/views/data_catalog.py b/testgen/ui/views/data_catalog.py index ae8dcdd2..ca068ebd 100644 --- a/testgen/ui/views/data_catalog.py +++ b/testgen/ui/views/data_catalog.py @@ -496,7 +496,7 @@ def get_excel_report_data(update_progress: PROGRESS_UPDATE_TYPE, table_group: Ta lambda val: date_service.format_friendly_datetime(val, "%b %-d %Y, %-I:%M %p") if not pd.isna(val) and not isinstance(val, str) else val ) - for key in ["data_source", "source_system", "source_process", "business_domain", "stakeholder_group", "transform_level", "aggregation_level", "data_product"]: + for key in ["data_source", "source_system", "source_process", "business_domain", "stakeholder_group", "transform_level", "aggregation_level", "data_product", "data_classification"]: data[key] = data.apply( lambda row: row[key] or row[f"table_{key}"] or row.get(f"table_group_{key}"), axis=1, @@ -587,6 +587,7 @@ def get_excel_report_data(update_progress: PROGRESS_UPDATE_TYPE, table_group: Ta "transform_level": {}, "aggregation_level": {}, "data_product": {}, + "data_classification": {}, } return get_excel_file_data( data, @@ -686,7 +687,8 @@ def get_table_group_columns(table_group_id: str) -> list[dict]: column_chars.pii_flag, column_chars.excluded_data_element, {", ".join([ f"column_chars.{tag}" for tag in TAG_FIELDS ])}, - {", ".join([ f"table_chars.{tag} AS table_{tag}" for tag in TAG_FIELDS ])} + {", ".join([ f"table_chars.{tag} AS table_{tag}" for tag in TAG_FIELDS ])}, + {", ".join([ f"table_groups.{tag} AS table_group_{tag}" for tag in TAG_FIELDS if tag != "aggregation_level" ])} FROM data_column_chars column_chars LEFT JOIN data_table_chars table_chars ON ( column_chars.table_id = table_chars.table_id @@ -696,6 +698,9 @@ def get_table_group_columns(table_group_id: str) -> list[dict]: AND column_chars.table_name = profile_results.table_name AND column_chars.column_name = profile_results.column_name ) + LEFT JOIN table_groups ON ( + column_chars.table_groups_id = table_groups.id + ) WHERE column_chars.table_groups_id = :table_group_id ORDER BY LOWER(table_chars.table_name), ordinal_position; """ @@ -901,6 +906,25 @@ def get_preview_data( } +TAG_FIELD_DEFAULTS: dict[str, list[str]] = { + "data_classification": ["Public", "Internal", "Confidential", "Restricted"], +} + + +def merge_tag_defaults( + values: dict[str, list[str]], + defaults: dict[str, list[str]], +) -> dict[str, list[str]]: + result: dict[str, list[str]] = defaultdict(list, {k: list(v) for k, v in values.items()}) + for tag, tag_defaults in defaults.items(): + existing_lower = {v.lower() for v in result[tag]} + for default in tag_defaults: + if default.lower() not in existing_lower: + result[tag].append(default) + result[tag] = sorted(result[tag], key=str.lower) + return result + + @st.cache_data(show_spinner=False) def get_tag_values() -> dict[str, list[str]]: quote = lambda v: f"'{v}'" @@ -926,8 +950,8 @@ def get_tag_values() -> dict[str, list[str]]: """ results = fetch_all_from_db(query) - values = defaultdict(list) + values: dict[str, list[str]] = defaultdict(list) for row in results: if row.tag and row.value: values[row.tag].append(row.value) - return values + return merge_tag_defaults(values, TAG_FIELD_DEFAULTS) diff --git a/testgen/ui/views/dialogs/import_metadata_dialog.py b/testgen/ui/views/dialogs/import_metadata_dialog.py index 09b0b91b..b7147240 100644 --- a/testgen/ui/views/dialogs/import_metadata_dialog.py +++ b/testgen/ui/views/dialogs/import_metadata_dialog.py @@ -32,7 +32,8 @@ "stakeholder group": "stakeholder_group", "transform level": "transform_level", "aggregation level": "aggregation_level", - "data product": "data_product", + "data product": "data_product", + "data classification": "data_classification", } METADATA_COLUMNS = ["description", "critical_data_element", "excluded_data_element", "pii_flag", *TAG_FIELDS] diff --git a/testgen/utils/__init__.py b/testgen/utils/__init__.py index 5218d3d0..75760994 100644 --- a/testgen/utils/__init__.py +++ b/testgen/utils/__init__.py @@ -181,6 +181,7 @@ def format_score_card(score_card: ScoreCard | None) -> ScoreCard: "dq_dimension": "Quality Dimension", "impact_dimension": "Impact Dimension", "data_product": "Data Product", + "data_classification": "Data Classification", } if not score_card: return { diff --git a/tests/unit/common/test_data_catalog_service.py b/tests/unit/common/test_data_catalog_service.py index 17445bf9..2d0c7f9a 100644 --- a/tests/unit/common/test_data_catalog_service.py +++ b/tests/unit/common/test_data_catalog_service.py @@ -298,7 +298,8 @@ def test_disable_autoflags_disables_both(): assert disabled == ["profile_flag_cdes", "profile_flag_pii"] -def test_tag_fields_has_eight_entries(): - assert len(TAG_FIELDS) == 8 +def test_tag_fields_has_nine_entries(): + assert len(TAG_FIELDS) == 9 assert "data_source" in TAG_FIELDS assert "aggregation_level" in TAG_FIELDS + assert "data_classification" in TAG_FIELDS diff --git a/tests/unit/mcp/test_model_data_column.py b/tests/unit/mcp/test_model_data_column.py index dffdb21e..abdcae6f 100644 --- a/tests/unit/mcp/test_model_data_column.py +++ b/tests/unit/mcp/test_model_data_column.py @@ -14,6 +14,7 @@ "transform_level", "aggregation_level", "data_product", + "data_classification", } diff --git a/tests/unit/mcp/test_model_data_table.py b/tests/unit/mcp/test_model_data_table.py index f5424963..2f04386f 100644 --- a/tests/unit/mcp/test_model_data_table.py +++ b/tests/unit/mcp/test_model_data_table.py @@ -26,6 +26,7 @@ def test_default_order_by_compiles_for_entity_select(): "transform_level", "aggregation_level", "data_product", + "data_classification", } diff --git a/tests/unit/mcp/test_tools_common.py b/tests/unit/mcp/test_tools_common.py index 64de5f4e..142946f8 100644 --- a/tests/unit/mcp/test_tools_common.py +++ b/tests/unit/mcp/test_tools_common.py @@ -710,6 +710,7 @@ def test_parse_score_type_invalid_lists_valid_values(): ("Stakeholder Group", ScoreCategory.stakeholder_group), ("Transform Level", ScoreCategory.transform_level), ("Data Product", ScoreCategory.data_product), + ("Data Classification", ScoreCategory.data_classification), ], ) def test_parse_category_display_form_returns_column_form_enum(display_value, expected): @@ -738,6 +739,7 @@ def test_parse_category_translation_dict_covers_all_args(): "stakeholder_group", "transform_level", "data_product", + "data_classification", ], ) def test_parse_category_rejects_column_form_input(internal): diff --git a/tests/unit/mcp/test_tools_table_groups.py b/tests/unit/mcp/test_tools_table_groups.py index c6c63415..febbb853 100644 --- a/tests/unit/mcp/test_tools_table_groups.py +++ b/tests/unit/mcp/test_tools_table_groups.py @@ -71,6 +71,7 @@ def _mock_table_group(**overrides) -> MagicMock: tg.stakeholder_group = overrides.get("stakeholder_group", None) tg.transform_level = overrides.get("transform_level", None) tg.data_product = overrides.get("data_product", None) + tg.data_classification = overrides.get("data_classification", None) return tg diff --git a/tests/unit/ui/conftest.py b/tests/unit/ui/conftest.py index 0e70ee57..cd74174c 100644 --- a/tests/unit/ui/conftest.py +++ b/tests/unit/ui/conftest.py @@ -9,3 +9,4 @@ # widgets/__init__.py calls components_v2.component() at module level; stub it out so # views that import widgets don't fail collection without a Streamlit runtime. sys.modules.setdefault("testgen.ui.components.widgets", MagicMock()) +sys.modules.setdefault("testgen.ui.components.widgets.download_dialog", MagicMock()) diff --git a/tests/unit/ui/test_data_catalog.py b/tests/unit/ui/test_data_catalog.py new file mode 100644 index 00000000..51e15e93 --- /dev/null +++ b/tests/unit/ui/test_data_catalog.py @@ -0,0 +1,39 @@ +import pytest + +from testgen.ui.views.data_catalog import merge_tag_defaults + +pytestmark = pytest.mark.unit + + +class Test_MergeTagDefaults: + def test_db_value_beats_matching_seed_case_insensitively(self): + values = {"data_classification": ["confidential", "public"]} + defaults = {"data_classification": ["Confidential", "Internal", "Public", "Restricted"]} + result = merge_tag_defaults(values, defaults) + assert result["data_classification"] == ["confidential", "Internal", "public", "Restricted"] + + def test_seed_added_when_no_db_match(self): + values = {"data_classification": []} + defaults = {"data_classification": ["Confidential", "Internal", "Public", "Restricted"]} + result = merge_tag_defaults(values, defaults) + assert result["data_classification"] == ["Confidential", "Internal", "Public", "Restricted"] + + def test_result_sorted_case_insensitively(self): + values = {"data_classification": ["RESTRICTED"]} + defaults = {"data_classification": ["Confidential", "Internal", "Public", "Restricted"]} + result = merge_tag_defaults(values, defaults) + assert result["data_classification"] == ["Confidential", "Internal", "Public", "RESTRICTED"] + + def test_non_classified_tags_unaffected(self): + values = {"data_product": ["prod-a"], "data_classification": ["internal"]} + defaults = {"data_classification": ["Confidential", "Internal", "Public", "Restricted"]} + result = merge_tag_defaults(values, defaults) + assert result["data_product"] == ["prod-a"] + assert "Internal" not in result["data_classification"] + assert "internal" in result["data_classification"] + + def test_empty_db_values_for_tagged_field(self): + values = {} + defaults = {"data_classification": ["Confidential"]} + result = merge_tag_defaults(values, defaults) + assert result["data_classification"] == ["Confidential"] From 78d1c29faad722361a6d1229777a130674f5c6a9 Mon Sep 17 00:00:00 2001 From: Diogo Basto Date: Thu, 25 Jun 2026 19:51:21 +0100 Subject: [PATCH 69/78] =?UTF-8?q?feat(mcp):=20add=20monitor=20L2=20read=20?= =?UTF-8?q?tools=20=E2=80=94=20per-table=20events,=20configs,=20and=20sche?= =?UTF-8?q?ma-change=20audit=20(TG-1092)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Three drill-down tools on top of the L1 group / table inventory: - list_monitor_events(table_group_id, table_name, monitor_type, monitor_id?, include_predictions?, limit?, page?) — per-table event history scoped to one monitor type. monitor_id is required when monitor_type="metric" (Metric is the only multi-instance type — singletons reject monitor_id with a clear pointer). Per-type column shapes: * Volume: Time | Status | Row count | Lower bound | Upper bound * Freshness: Time | Status | Update detected | Detail (parsed from the SQL template's structured message) * Schema: Time | Status | Table change | Columns added | Columns dropped | Columns modified * Metric: Time | Status | Value | Lower bound | Upper bound (metric name in the heading, not as a column) include_predictions=True appends a separate `## Forecast` section listing future timestamps with predicted bounds — never interleaved into the historical events. Schema short-circuits to "not applicable"; non- Prediction-Model monitors render a "not available" note. - list_monitors(table_group_id, table_name) — configured monitors for a table. Per row: monitor_id, type (Title Case, reuses _MONITOR_LABEL), metric_name, threshold mode (derived from history_calculation: "PREDICT" → Prediction; non-empty otherwise and not Freshness → Historical; empty → Static), bounds, and the metric expression for Metric monitors. Sensitivity is surfaced as a top-level "Prediction model sensitivity" field. - list_monitor_schema_changes(table_group_id, table_name?, since?, limit?, page?) — newest-first audit log of column-level schema events (added / dropped / modified), independent of monitor-run results. All three follow the L1 contract: gated on view, resolve_monitored_table_group returns the literal "This table group is not monitored." when the group has no linked monitor suite, and never expose internal test_type codes, test_definition_id, or test_suite_id on the surface. Backing model methods: - DataStructureLog ORM wraps the existing data_structure_log audit table (no ORM mapping until now) and exposes list_for_table_group(*, table_name, since, until, page, limit). The dashboard's get_data_structure_logs delegates here. - TestDefinition.list_monitor_configs_for_table returns MonitorConfig with threshold_mode derived per row. - TestDefinition.get_singleton_monitor(suite, table, type) for the non- metric forecast lookup path. - forecast_points_from_prediction(prediction, sensitivity) — standalone helper that reads forecast rows from the prediction JSONB using epoch-ms keys with numeric comparison (matches the dashboard's reader). - TestResult.list_monitor_events_for_table runs the dashboard's CTE under the ORM with optional monitor_type filtering. ORDER BY includes results.id NULLS LAST, active_runs.id as a stable tiebreaker so the Python-side pagination doesn't duplicate or skip rows on test_time ties. - TestResult.list_metric_monitor_events — separate, simpler query path for Metric (no run-by-type CROSS JOIN, no synthesized pending rows). The dashboard's per-type events transform stays in monitors_dashboard.py because its shape is bespoke to the chart payload; centralizing just the SQL would still require a non-trivial adapter on the dashboard side. monitor_id joins the Followable IDs table as an alias for the underlying test_definition.id UUID — the abstraction stays stable if monitors move to a dedicated table later. Status helpers (_summary_status, _format_monitor_cell, _format_schema_cell, _event_status) return Title Case ("Error", "Pending", "Training", "Anomaly", "Ok"), matching the rest of the MCP layer. parse_monitor_type accepts either case so values round-trip. _parse_kv_pairs splits input_parameters on ";" (matches the writer side at execute_tests_query.py:321 and the dashboard's dict_from_kv). Closes TG-1092 Co-Authored-By: Claude Opus 4.7 --- testgen/common/models/data_structure_log.py | 93 ++++ testgen/common/models/test_definition.py | 180 ++++++ testgen/common/models/test_result.py | 265 ++++++++- testgen/mcp/server.py | 6 + testgen/mcp/tools/common.py | 10 +- testgen/mcp/tools/monitors.py | 424 +++++++++++++- testgen/ui/views/monitors_dashboard.py | 48 +- tests/unit/mcp/test_model_test_definition.py | 158 +++++- tests/unit/mcp/test_model_test_result.py | 102 +++- tests/unit/mcp/test_tools_monitors.py | 555 ++++++++++++++++++- 10 files changed, 1785 insertions(+), 56 deletions(-) create mode 100644 testgen/common/models/data_structure_log.py diff --git a/testgen/common/models/data_structure_log.py b/testgen/common/models/data_structure_log.py new file mode 100644 index 00000000..b9a34942 --- /dev/null +++ b/testgen/common/models/data_structure_log.py @@ -0,0 +1,93 @@ +from dataclasses import dataclass +from datetime import date, datetime +from uuid import UUID, uuid4 + +from sqlalchemy import Column, String, asc, desc, select +from sqlalchemy.dialects import postgresql + +from testgen.common.models import get_current_session +from testgen.common.models.entity import Entity, EntityMinimal + +# Schema-change codes stored in ``data_structure_log.change``. ``M`` only ever +# appears on column-level rows (column-name set, old/new data type populated). +SCHEMA_CHANGE_ADDED = "A" +SCHEMA_CHANGE_DROPPED = "D" +SCHEMA_CHANGE_MODIFIED = "M" + + +@dataclass +class DataStructureLogEntry(EntityMinimal): + """One schema-change event for a table or column in a table group. + + ``column_name`` is ``None`` for table-level events (add / drop of an entire + table); set for column-level events. ``change`` is the stored single-letter + code (``A`` / ``D`` / ``M``); callers translate to user-facing words. + """ + log_id: UUID + table_groups_id: UUID + table_name: str + column_name: str | None + change_date: datetime + change: str + old_data_type: str | None + new_data_type: str | None + + +class DataStructureLog(Entity): + __tablename__ = "data_structure_log" + + log_id: UUID = Column(postgresql.UUID(as_uuid=True), primary_key=True, default=uuid4) + table_groups_id: UUID = Column(postgresql.UUID(as_uuid=True)) + table_id: UUID = Column(postgresql.UUID(as_uuid=True)) + column_id: UUID = Column(postgresql.UUID(as_uuid=True)) + table_name: str = Column(String) + column_name: str = Column(String) + change_date: datetime = Column(postgresql.TIMESTAMP) + change: str = Column(String) + old_data_type: str = Column(String) + new_data_type: str = Column(String) + + _get_by = "log_id" + _default_order_by = (desc(change_date), asc(table_name), asc(column_name)) + + @classmethod + def list_for_table_group( + cls, + table_group_id: str | UUID, + *, + table_name: str | None = None, + since: date | datetime | None = None, + until: date | datetime | None = None, + page: int = 1, + limit: int | None = 20, + ) -> tuple[list[DataStructureLogEntry], int]: + """Paginated schema-change audit log for one table group, newest first. + + Filters by ``table_name`` (exact match; the audit log stores names + case-sensitively as the source emitted them), ``since`` (lower-bound on + ``change_date``), and ``until`` (upper-bound). ``limit=None`` skips + pagination — the caller gets every matching row in one shot. + """ + query = select( + cls.log_id, + cls.table_groups_id, + cls.table_name, + cls.column_name, + cls.change_date, + cls.change, + cls.old_data_type, + cls.new_data_type, + ).where(cls.table_groups_id == table_group_id) + if table_name is not None: + query = query.where(cls.table_name == table_name) + if since is not None: + query = query.where(cls.change_date >= since) + if until is not None: + query = query.where(cls.change_date <= until) + query = query.order_by(*cls._default_order_by) + + if limit is None: + rows = get_current_session().execute(query).mappings().all() + entries = [DataStructureLogEntry(**row) for row in rows] + return entries, len(entries) + return cls._paginate(query, page=page, limit=limit, data_class=DataStructureLogEntry) diff --git a/testgen/common/models/test_definition.py b/testgen/common/models/test_definition.py index 878bdc8c..7220017b 100644 --- a/testgen/common/models/test_definition.py +++ b/testgen/common/models/test_definition.py @@ -27,6 +27,7 @@ from sqlalchemy.orm import InstrumentedAttribute from sqlalchemy.sql.expression import case, literal +from testgen.common.enums import MonitorType from testgen.common.models import Base, get_current_session from testgen.common.models.custom_types import NullIfEmptyString, YNString, ZeroIfEmptyInteger from testgen.common.models.entity import Entity, EntityMinimal @@ -270,6 +271,96 @@ class TestDefinitionMinimal(EntityMinimal): test_name_short: str +# Threshold-mode labels for monitor TestDefinitions — derived from which fields +# on the definition are populated. See ``TestDefinition._derive_threshold_mode``. +THRESHOLD_MODE_PREDICTION = "Prediction Model" +THRESHOLD_MODE_HISTORICAL = "Historical Calculation" +THRESHOLD_MODE_STATIC = "Static" +THRESHOLD_MODE_NONE = "N/A" + + +@dataclass +class MonitorForecastPoint(EntityMinimal): + """One future forecast point read from a Prediction-Model monitor's stored + ``prediction`` JSONB. Each point is a ``(test_time, lower_bound, upper_bound)`` + triple at a specific upcoming timestamp; collectively they extend the + historical event series forward under the suite's active prediction + sensitivity. Surfaced as a separate forecast section on + ``list_monitor_events``, never as an event.""" + test_time: datetime + lower_bound: float | None + upper_bound: float | None + + +def forecast_points_from_prediction( + prediction: dict | None, + sensitivity: str, +) -> list[MonitorForecastPoint]: + """Extract forecast points for a sensitivity from a monitor's stored + ``prediction`` JSONB, sorted by time ascending (nearest future point + first). + + The JSONB is keyed as ``"lower_tolerance|" → {epoch_ms: value}`` + and ``"upper_tolerance|" → {epoch_ms: value}`` — matches the + format the dashboard reads at ``monitors_dashboard.py`` via + ``datetime.fromtimestamp(int(timestamp) / 1000.0, UTC)``. Returns ``[]`` + when the monitor isn't in Prediction Model mode (no JSONB) or when the + sensitivity has no stored series. Standalone so it works against either + the ``TestDefinition`` ORM row or the ``TestDefinitionSummary`` dataclass + — both carry the same ``prediction`` field shape. + """ + if not prediction: + return [] + lower_series = prediction.get(f"lower_tolerance|{sensitivity}") or {} + upper_series = prediction.get(f"upper_tolerance|{sensitivity}") or {} + if not lower_series and not upper_series: + return [] + + all_keys = sorted(set(lower_series) | set(upper_series), key=int) + points: list[MonitorForecastPoint] = [] + for k in all_keys: + ts = datetime.fromtimestamp(int(k) / 1000.0, UTC) + lower = lower_series.get(k) + upper = upper_series.get(k) + points.append(MonitorForecastPoint( + test_time=ts, + lower_bound=float(lower) if lower is not None else None, + upper_bound=float(upper) if upper is not None else None, + )) + return points + + +@dataclass +class MonitorConfig(EntityMinimal): + """One configured monitor — produced by ``TestDefinition.list_monitor_configs_for_table``. + + ``threshold_mode`` is derived from the underlying definition's + ``history_calculation``: ``"PREDICT"`` flags Prediction Model; any other + non-empty value flags Historical Calculation (not available for Freshness); + empty falls through to Static — the default for Freshness, Volume, and + Metric. Schema_Drift is presence-only and reports N/A. + + ``threshold_lower`` / ``threshold_upper`` carry the bounds active under + the current mode — static tolerances for Static, history-calc expressions + (e.g. ``"Minimum"`` / ``"Maximum"``) for Historical, ``None`` for + Prediction (the runtime bands live on each event, not the configuration). + For Freshness in Static mode only the upper bound applies. + + ``metric_name`` is the user-defined name for a Metric monitor (stored on + the underlying definition's ``column_name`` column, but it is the metric's + name rather than a column reference). ``custom_query`` is the metric's SQL + expression. Both are only set for ``Metric_Trend``. + """ + monitor_id: UUID + test_type: str + table_name: str + metric_name: str | None + threshold_mode: str + threshold_lower: str | None + threshold_upper: str | None + custom_query: str | None + + class QueryString(TypeDecorator): impl = String cache_ok = True @@ -557,6 +648,95 @@ def list_for_suite( query = query.order_by(*cls._default_order_by) return cls._paginate(query, page=page, limit=limit, data_class=TestDefinitionSummary) + @classmethod + def get_singleton_monitor( + cls, + test_suite_id: str | UUID, + table_name: str, + test_type: str, + ) -> "TestDefinition | None": + """Return the single ``TestDefinition`` row for a singleton monitor + type (Freshness / Volume / Schema) on a given table. Metric is + multi-instance and should be looked up by ``id`` instead — this + helper would silently pick the first match. Returns ``None`` when no + monitor is configured.""" + session = get_current_session() + return session.execute( + select(cls) + .where(cls.test_suite_id == test_suite_id) + .where(cls.table_name == table_name) + .where(cls.test_type == test_type) + .limit(1) + ).scalars().first() + + @classmethod + def list_monitor_configs_for_table( + cls, + test_suite_id: str | UUID, + table_name: str, + ) -> list[MonitorConfig]: + """List configured monitors for a single table within a monitor suite. + + Returns one entry per ``test_definition`` row whose ``test_type`` is one + of the four monitor types. Typically 3-4 rows per table; Metric_Trend + contributes one entry per metric, so a table may have more. + """ + monitor_codes = [m.value for m in MonitorType] + defs = cls.select_where( + cls.test_suite_id == test_suite_id, + cls.table_name == table_name, + cls.test_type.in_(monitor_codes), + ) + return [cls._build_monitor_config(td) for td in defs] + + @classmethod + def _build_monitor_config(cls, td: "TestDefinition") -> MonitorConfig: + mode, threshold_lower, threshold_upper = cls._derive_threshold_mode(td) + return MonitorConfig( + monitor_id=td.id, + test_type=td.test_type, + table_name=td.table_name, + metric_name=td.column_name or None if td.test_type == MonitorType.METRIC.value else None, + threshold_mode=mode, + threshold_lower=threshold_lower, + threshold_upper=threshold_upper, + custom_query=td.custom_query if td.test_type == MonitorType.METRIC.value else None, + ) + + @classmethod + def _derive_threshold_mode( + cls, td: "TestDefinition", + ) -> tuple[str, str | None, str | None]: + """Pick a mode and the bounds tuple that applies under that mode. + + Detection mirrors the UI form (``test_definition_form.js``): a + ``history_calculation`` of exactly ``"PREDICT"`` flags Prediction + mode; any other non-empty value flags Historical (not available for + Freshness); empty falls through to Static — the default for + Freshness / Volume / Metric. Schema never has thresholds. + + Bounds returned per mode: + + * Prediction: ``None, None``. Runtime bounds live in the per-run + prediction JSONB; they are not configuration and do not surface + here. + * Historical: ``(history_calculation, history_calculation_upper)`` — + stored as expressions like ``"Minimum"`` / ``"Maximum"`` that the + execution layer evaluates against the lookback window. + * Static for Freshness: ``(None, upper_tolerance)``. Only an upper + bound applies. + * Static (Volume / Metric): ``(lower_tolerance, upper_tolerance)``. + """ + if td.test_type == MonitorType.SCHEMA.value: + return THRESHOLD_MODE_NONE, None, None + if td.history_calculation == "PREDICT": + return THRESHOLD_MODE_PREDICTION, None, None + if td.history_calculation and td.test_type != MonitorType.FRESHNESS.value: + return THRESHOLD_MODE_HISTORICAL, td.history_calculation, td.history_calculation_upper + if td.test_type == MonitorType.FRESHNESS.value: + return THRESHOLD_MODE_STATIC, None, td.upper_tolerance + return THRESHOLD_MODE_STATIC, td.lower_tolerance, td.upper_tolerance + @classmethod def select_page( cls, diff --git a/testgen/common/models/test_result.py b/testgen/common/models/test_result.py index 4ffc3782..6c5f9719 100644 --- a/testgen/common/models/test_result.py +++ b/testgen/common/models/test_result.py @@ -5,11 +5,12 @@ from typing import Self from uuid import UUID, uuid4 -from sqlalchemy import Boolean, Column, Enum, ForeignKey, Integer, String, desc, func, or_, select +from sqlalchemy import Boolean, Column, Enum, ForeignKey, Integer, String, desc, func, or_, select, text from sqlalchemy.dialects import postgresql from sqlalchemy.orm import aliased from sqlalchemy.sql.expression import case +from testgen.common.enums import MonitorType from testgen.common.models import get_current_session from testgen.common.models.entity import Entity from testgen.common.models.test_definition import TestType @@ -126,6 +127,118 @@ class RunDiff: removed_tests: list[DiffRow] = field(default_factory=list) +@dataclass +class MonitorEvent: + """One monitor result for a (table, monitor_type) pair within the lookback window. + + Type-specific fields are populated by ``test_type``: + + * ``Freshness_Trend``: ``signal`` carries minutes-since-last-update (string + like ``"120"`` / ``"Unknown"``); ``message`` carries the human-readable + detection text. + * ``Volume_Trend``: ``signal`` carries the row count (string of integer); + ``lower_bound`` / ``upper_bound`` carry the tolerance bounds the run was + evaluated against (sourced from ``input_parameters`` at run time). + * ``Schema_Drift``: ``schema_change_kind`` is the table-level code + (``A`` / ``D`` / ``M``); ``column_adds`` / ``column_drops`` / + ``column_mods`` carry per-event column-change counts. + * ``Metric_Trend``: same as Volume_Trend plus ``metric_name`` (the + user-defined name for the metric — the underlying SQL expression lives + on ``MonitorConfig``, not on the event). + + Pending rows have no underlying ``test_results`` row — ``monitor_id``, + ``test_time``, and result flags are ``None``; ``is_pending`` is True. + Forecast points (future timestamps with predicted bounds) are NOT + events and surface separately via ``TestDefinition.get_forecast_points``. + """ + monitor_id: UUID | None + test_type: str + test_time: datetime | None + is_anomaly: bool | None + is_training: bool | None + is_pending: bool + is_error: bool + message: str | None + signal: str | None + lower_bound: str | None = None + upper_bound: str | None = None + schema_change_kind: str | None = None + column_adds: int | None = None + column_drops: int | None = None + column_mods: int | None = None + metric_name: str | None = None + + +def _parse_kv_pairs(raw: str | None) -> dict[str, str]: + """Parse an ``input_parameters`` blob (``key=value; key=value; ...``) into a dict. + + ``input_parameters`` is built with ``"; ".join(...)`` in ``execute_tests_query`` + and read by the dashboard via ``dict_from_kv`` (default separator ``;``). + Tolerant of missing values, trailing/leading whitespace, empty strings. + Returns ``{}`` on empty input. + """ + if not raw: + return {} + pairs: dict[str, str] = {} + for entry in raw.split(";"): + if "=" not in entry: + continue + key, _, value = entry.partition("=") + pairs[key.strip()] = value.strip() + return pairs + + +def _build_monitor_event(row) -> MonitorEvent: + """Translate one CTE row into a ``MonitorEvent``, populating type-specific extras.""" + is_pending = row["result_id"] is None + result_code = row["result_code"] + is_anomaly = (result_code == 0) if result_code is not None else None + is_training = (result_code == -1) if result_code is not None else None + is_error = (row["result_status"] == "Error") + + event = MonitorEvent( + monitor_id=row["test_definition_id"], + test_type=row["test_type"], + test_time=row["test_time"] if not is_pending else None, + is_anomaly=is_anomaly, + is_training=is_training, + is_pending=is_pending, + is_error=is_error, + message=row["result_message"], + signal=row["result_signal"], + ) + + params = _parse_kv_pairs(row["input_parameters"]) + if event.test_type in (MonitorType.VOLUME.value, MonitorType.METRIC.value): + event.lower_bound = params.get("lower_tolerance") or None + event.upper_bound = params.get("upper_tolerance") or None + if event.test_type == MonitorType.METRIC.value: + event.metric_name = row["column_names"] or None + elif event.test_type == MonitorType.SCHEMA.value: + signal = row["result_signal"] + if signal: + parts = signal.split("|") + event.schema_change_kind = parts[0] or None + event.column_adds = _int_or_none(parts, 1) + event.column_drops = _int_or_none(parts, 2) + event.column_mods = _int_or_none(parts, 3) + + return event + + +def _int_or_none(parts: list[str], index: int) -> int | None: + try: + value = parts[index] + except IndexError: + return None + if not value: + return None + try: + return int(value) + except ValueError: + return None + + class TestResult(Entity): __tablename__ = "test_results" @@ -540,3 +653,153 @@ def _row(tid: UUID, baseline_info: dict | None, target_info: dict | None) -> Dif diff.removed_tests.append(_row(tid, baseline_results[tid], None)) return diff + + @classmethod + def list_monitor_events_for_table( + cls, + test_suite_id: str | UUID, + table_name: str, + *, + monitor_type: str | None = None, + lookback_multiplier: int = 1, + page: int = 1, + limit: int | None = None, + ) -> tuple[list[MonitorEvent], int]: + """Per-table monitor events within the suite's lookback window, newest first. + + ``monitor_type`` is the internal ``test_type`` value; when ``None`` events + for all four monitor types are returned (used by the dashboard, which + groups by type on the JS side). ``lookback_multiplier`` extends the + active window for the "show more history" toggle. ``limit=None`` + skips pagination entirely (the caller wants every event in the window). + + Forecast points for Prediction-Model monitors are NOT included here — + events are only past, observed runs. Read forecasts separately via + ``TestDefinition.get_forecast_points(sensitivity)``. + """ + monitor_codes = ( + [monitor_type] if monitor_type is not None + else [m.value for m in MonitorType] + ) + type_filter_sql = "AND results.test_type = :monitor_type" if monitor_type else "" + + query = f""" + WITH ranked_test_runs AS ( + SELECT + test_runs.id, + test_runs.test_starttime, + COALESCE(test_suites.monitor_lookback, 1) * :lookback_multiplier AS lookback, + ROW_NUMBER() OVER (PARTITION BY test_runs.test_suite_id ORDER BY test_runs.test_starttime DESC) AS position + FROM test_suites + INNER JOIN test_runs ON (test_suites.id = test_runs.test_suite_id) + WHERE test_suites.id = :test_suite_id + ), + active_runs AS ( + SELECT id, test_starttime FROM ranked_test_runs WHERE position <= lookback + ), + target_types AS ( + SELECT UNNEST(CAST(:monitor_codes AS TEXT[])) AS test_type + ) + SELECT + COALESCE(results.test_time, active_runs.test_starttime) AS test_time, + tt.test_type, + results.id AS result_id, + results.test_definition_id, + results.result_code, + COALESCE(results.result_status, 'Log') AS result_status, + results.result_signal, + results.result_message, + COALESCE(results.input_parameters, '') AS input_parameters, + results.column_names + FROM active_runs + CROSS JOIN target_types tt + LEFT JOIN test_results AS results + ON results.test_run_id = active_runs.id + AND results.test_type = tt.test_type + AND results.table_name = :table_name + {type_filter_sql} + ORDER BY test_time DESC, tt.test_type, results.id NULLS LAST, active_runs.id + """ + + params: dict = { + "test_suite_id": str(test_suite_id), + "table_name": table_name, + "lookback_multiplier": lookback_multiplier, + "monitor_codes": monitor_codes, + } + if monitor_type is not None: + params["monitor_type"] = monitor_type + + session = get_current_session() + rows = session.execute(text(query), params).mappings().all() + events = [_build_monitor_event(row) for row in rows] + + # Paginate in Python — the CTE is bounded by lookback x |monitor_types| + # (typically <= ~120 rows for a single table). Revisit if either grows. + total = len(events) + if limit is not None: + start = (page - 1) * limit + events = events[start:start + limit] + return events, total + + @classmethod + def list_metric_monitor_events( + cls, + test_suite_id: str | UUID, + test_definition_id: str | UUID, + *, + lookback_multiplier: int = 1, + page: int = 1, + limit: int | None = None, + ) -> tuple[list[MonitorEvent], int]: + """Per-metric event history within the suite's lookback window, newest + first. Scoped to one ``test_definition_id`` since Metric_Trend is the + only multi-instance monitor type — table + type alone would interleave + events across every metric on the table. + + Distinct from ``list_monitor_events_for_table`` in two ways: (1) no + cross join with target_types — only one monitor type to query; (2) no + synthesized pending rows — a pending result for a specific + ``test_definition_id`` can't be distinguished from "no run yet" + without a results row to anchor on, so the model returns only rows + that actually ran. ``limit=None`` skips pagination. + """ + query = """ + WITH suite_window AS ( + SELECT COALESCE(monitor_lookback, 1) * :lookback_multiplier AS lookback + FROM test_suites + WHERE id = :test_suite_id + ) + SELECT + results.test_time, + results.test_type, + results.id AS result_id, + results.test_definition_id, + results.result_code, + COALESCE(results.result_status, 'Log') AS result_status, + results.result_signal, + results.result_message, + COALESCE(results.input_parameters, '') AS input_parameters, + results.column_names + FROM test_results AS results + WHERE results.test_suite_id = :test_suite_id + AND results.test_definition_id = :test_definition_id + ORDER BY results.test_time DESC, results.id + LIMIT (SELECT lookback FROM suite_window) + """ + + params: dict = { + "test_suite_id": str(test_suite_id), + "test_definition_id": str(test_definition_id), + "lookback_multiplier": lookback_multiplier, + } + + session = get_current_session() + rows = session.execute(text(query), params).mappings().all() + events = [_build_monitor_event(row) for row in rows] + + total = len(events) + if limit is not None: + start = (page - 1) * limit + events = events[start:start + limit] + return events, total diff --git a/testgen/mcp/server.py b/testgen/mcp/server.py index cf8ac6ff..c74fe5e7 100644 --- a/testgen/mcp/server.py +++ b/testgen/mcp/server.py @@ -177,7 +177,10 @@ def build_mcp_server( enable_monitors, get_monitor_settings, get_monitor_summary, + list_monitor_events, + list_monitor_schema_changes, list_monitored_tables, + list_monitors, update_monitor_settings, ) from testgen.mcp.tools.notifications import ( @@ -330,6 +333,9 @@ def safe_prompt(fn): safe_tool(get_schema_history) safe_tool(get_monitor_summary) safe_tool(list_monitored_tables) + safe_tool(list_monitor_events) + safe_tool(list_monitors) + safe_tool(list_monitor_schema_changes) safe_tool(enable_monitors) safe_tool(get_monitor_settings) safe_tool(update_monitor_settings) diff --git a/testgen/mcp/tools/common.py b/testgen/mcp/tools/common.py index f90038df..25b302fa 100644 --- a/testgen/mcp/tools/common.py +++ b/testgen/mcp/tools/common.py @@ -392,11 +392,13 @@ def next_scheduled_run( def parse_monitor_type(value: str, label: str = "monitor_type") -> MonitorType: """Validate a user-facing monitor type label and return the stored ``MonitorType``. - Accepts ``freshness`` / ``volume`` / ``schema`` / ``metric``. ``label`` names the - caller's argument in the error message — pass ``"anomaly_type"`` when the - public arg is named differently from ``monitor_type``. + Accepts ``Freshness`` / ``Volume`` / ``Schema`` / ``Metric`` (Title Case + as rendered in tool output) or the equivalent lowercase forms. ``label`` + names the caller's argument in the error message — pass + ``"anomaly_type"`` when the public arg is named differently from + ``monitor_type``. """ - db_value = _MONITOR_TYPE_USER_TO_DB.get(value) + db_value = _MONITOR_TYPE_USER_TO_DB.get(value.lower()) if isinstance(value, str) else None if db_value is None: valid = ", ".join(_MONITOR_TYPE_USER_TO_DB) raise MCPUserError(f"Invalid {label} `{value}`. Valid values: {valid}") diff --git a/testgen/mcp/tools/monitors.py b/testgen/mcp/tools/monitors.py index f1493585..11c10104 100644 --- a/testgen/mcp/tools/monitors.py +++ b/testgen/mcp/tools/monitors.py @@ -1,3 +1,4 @@ +import re from datetime import UTC, datetime from typing import Any, cast from uuid import UUID @@ -8,9 +9,20 @@ from testgen.common.enums import MonitorType from testgen.common.holiday_service import is_supported_holiday_code from testgen.common.models import get_current_session, with_database_session +from testgen.common.models.data_structure_log import ( + SCHEMA_CHANGE_ADDED, + SCHEMA_CHANGE_DROPPED, + SCHEMA_CHANGE_MODIFIED, + DataStructureLog, +) from testgen.common.models.job_execution import JobExecution from testgen.common.models.scheduler import RUN_MONITORS_JOB_KEY, JobSchedule from testgen.common.models.table_group import MonitorTableSummary, TableGroup +from testgen.common.models.test_definition import ( + TestDefinition, + forecast_points_from_prediction, +) +from testgen.common.models.test_result import MonitorEvent, TestResult from testgen.common.models.test_suite import PredictSensitivity, TestSuite from testgen.common.monitor_service import disable_monitoring, enable_monitoring, update_monitoring from testgen.mcp.exceptions import MCPUserError @@ -23,6 +35,7 @@ next_scheduled_run, parse_monitor_table_sort, parse_monitor_type, + parse_since_arg, resolve_monitored_table_group, validate_limit, validate_page, @@ -239,12 +252,12 @@ def list_monitored_tables( def _summary_status(*, is_pending: bool, is_training: bool, has_errors: bool) -> str: """Single status word for the group-level summary cell.""" if has_errors: - return "error" + return "Error" if is_pending: - return "no results yet or not configured" + return "No results yet or not configured" if is_training: - return "training" - return "ok" + return "Training" + return "Ok" def _format_monitor_cell( @@ -262,13 +275,13 @@ def _format_monitor_cell( shows the X count in its column. """ if has_error: - return "error" + return "Error" if count > 0: return str(count) if is_pending: - return "pending" + return "Pending" if is_training: - return "training" + return "Training" return "0" @@ -276,9 +289,9 @@ def _format_schema_cell(row: MonitorTableSummary) -> str: """Schema anomaly count, or pending / error status. Verbose detail is in the sibling ``Schema change`` column.""" if row.schema_error_message is not None: - return "error" + return "Error" if row.schema_is_pending: - return "pending" + return "Pending" return str(row.schema_anomalies) @@ -573,3 +586,396 @@ def _render_monitor_settings(doc: MdDoc, monitor_suite: TestSuite, schedule: Job if (last_run := _last_monitor_run(schedule.id)) is not None: # job_executions timestamps are tz-aware UTC; convert so the rendered " UTC" suffix is accurate. doc.field("Last run", last_run.astimezone(UTC)) +# list_monitor_events +# --------------------------------------------------------------------------- + + +@with_database_session +@mcp_permission("view") +def list_monitor_events( + table_group_id: str, + table_name: str, + monitor_type: str, + monitor_id: str | None = None, + include_predictions: bool = False, + limit: int = 20, + page: int = 1, +) -> str: + """List per-table monitor events for one monitor type within the lookback window, newest first. + + Args: + table_group_id: UUID of the table group, e.g. from ``list_table_groups``. + table_name: Table name exactly as stored in TestGen (case-sensitive). + monitor_type: One of ``freshness`` / ``volume`` / ``schema`` / ``metric``. + monitor_id: Required when ``monitor_type="metric"``; rejected otherwise. Metric monitors are user-defined (many per table), so events must be scoped to one ``monitor_id`` to avoid interleaving every metric's history. Find the id via ``list_monitors(table_group_id, table_name)``. Singleton types (``freshness`` / ``volume`` / ``schema``) are uniquely identified by table + type. + include_predictions: When True, append the monitor's forecasted future events (predicted bounds for upcoming time points) after the historical events. Only Volume / Metric / Freshness monitors in Prediction Model mode have a forecast; Schema and other modes render a note explaining why no forecast follows. + limit: Page size (default 20, max 100). + page: Page number starting at 1 (default 1). + """ + validate_page(page) + validate_limit(limit, 100) + parsed_type = parse_monitor_type(monitor_type) + + if parsed_type == MonitorType.METRIC and monitor_id is None: + raise MCPUserError( + '`monitor_id` is required for `monitor_type="metric"`. ' + "Metric monitors are user-defined and many can apply to one table — " + "use `list_monitors(table_group_id, table_name)` to find the metric's " + "`monitor_id` first, then call this tool with it." + ) + if monitor_id is not None and parsed_type != MonitorType.METRIC: + raise MCPUserError( + '`monitor_id` only applies when `monitor_type="metric"`. ' + "For singleton monitor types (`freshness`, `volume`, `schema`), " + "the monitor is uniquely identified by table + type." + ) + + tg, suite = resolve_monitored_table_group(table_group_id) + if suite is None: + return _NOT_MONITORED_OUTPUT + + if parsed_type == MonitorType.METRIC: + events, total = TestResult.list_metric_monitor_events( + suite.id, + monitor_id, + page=page, + limit=limit, + ) + monitor_def = TestDefinition.get(monitor_id) + metric_name = ( + monitor_def.column_name or None + if monitor_def is not None and monitor_def.test_type == MonitorType.METRIC.value + else None + ) + else: + events, total = TestResult.list_monitor_events_for_table( + suite.id, + table_name, + monitor_type=parsed_type.value, + page=page, + limit=limit, + ) + # Singleton lookup only needed when we will render a forecast — Schema + # short-circuits to "not applicable" without touching the definition. + monitor_def = ( + TestDefinition.get_singleton_monitor(suite.id, table_name, parsed_type.value) + if include_predictions and parsed_type != MonitorType.SCHEMA + else None + ) + metric_name = None + + doc = MdDoc() + if parsed_type == MonitorType.METRIC and metric_name: + doc.heading( + 1, + f"Monitor events: `{metric_name}` on `{table_name}` in `{tg.table_groups_name}`", + ) + else: + doc.heading( + 1, + f"Monitor events: `{table_name}` — `{monitor_type}` in `{tg.table_groups_name}`", + ) + + page_info = format_page_info(total, page, limit) + if page_info: + doc.text(page_info) + + if not events: + if page > 1: + doc.text(f"No events on page {page} (total: {total}).") + else: + doc.text("_No monitor events in the active lookback window._") + if include_predictions: + _render_forecast_section(doc, parsed_type, monitor_def, suite) + return doc.render() + + if parsed_type == MonitorType.SCHEMA: + doc.table( + ["Time", "Status", "Table change", "Columns added", "Columns dropped", "Columns modified"], + [ + [ + event.test_time, + _event_status(event), + _format_schema_change_kind(event.schema_change_kind), + event.column_adds, + event.column_drops, + event.column_mods, + ] + for event in events + ], + ) + elif parsed_type == MonitorType.FRESHNESS: + doc.table( + ["Time", "Status", "Update detected", "Detail"], + [ + [event.test_time, _event_status(event), *_parse_freshness_message(event.message)] + for event in events + ], + ) + elif parsed_type == MonitorType.METRIC: + doc.table( + ["Time", "Status", "Value", "Lower bound", "Upper bound"], + [ + [ + event.test_time, + _event_status(event), + event.signal, + event.lower_bound, + event.upper_bound, + ] + for event in events + ], + ) + else: # VOLUME + doc.table( + ["Time", "Status", "Row count", "Lower bound", "Upper bound"], + [ + [ + event.test_time, + _event_status(event), + event.signal, + event.lower_bound, + event.upper_bound, + ] + for event in events + ], + ) + + if footer := format_page_footer(total, page, limit): + doc.text(footer) + + if include_predictions: + _render_forecast_section(doc, parsed_type, monitor_def, suite) + + return doc.render() + + +def _render_forecast_section(doc, parsed_type, monitor_def, suite) -> None: + """Append the forecast section (or an explanatory note) under the + historical events table. Predictions are NEVER mixed into the events + table — they live on their own under a `## Forecast` heading so the LLM + consumer can tell observed-past from predicted-future at a glance.""" + doc.heading(2, "Forecast") + + if parsed_type == MonitorType.SCHEMA: + doc.text( + "_Predictions not applicable — Schema monitors are presence-only and " + "do not have tolerance bands._" + ) + return + + if monitor_def is None: + doc.text("_No monitor configured for this table._") + return + + if monitor_def.history_calculation != "PREDICT": + doc.text( + "_Predictions not available — this monitor's threshold mode is " + "Static or Historical Calculation, not Prediction Model._" + ) + return + + sensitivity = ( + suite.predict_sensitivity.value + if suite.predict_sensitivity is not None + else "medium" + ) + points = forecast_points_from_prediction(monitor_def.prediction, sensitivity) + if not points: + doc.text( + "_No forecast stored yet — Prediction Model monitors populate their " + "forecast after the first successful run._" + ) + return + + doc.field("Sensitivity", sensitivity) + doc.table( + ["Time", "Predicted lower", "Predicted upper"], + [[p.test_time, p.lower_bound, p.upper_bound] for p in points], + ) + + +def _event_status(event: MonitorEvent) -> str: + """Single status word for a monitor-event row.""" + if event.is_error: + return "Error" + if event.is_pending: + return "Pending" + if event.is_training: + return "Training" + if event.is_anomaly: + return "Anomaly" + return "Ok" + + +# Freshness events carry a structured message of the form +# ``"Table update detected: {Yes|No}[. {Detail}]"`` — see the SQL templates at +# ``template/dbsetup_test_types/test_types_Freshness_Trend.yaml``. Parse it into +# separate columns so the LLM doesn't have to re-derive the same fields. +_FRESHNESS_MESSAGE_RE = re.compile( + r"^Table update detected:\s*(Yes|No)(?:\.\s*(.+?))?\.?\s*$", + re.IGNORECASE, +) + + +def _parse_freshness_message(message: str | None) -> tuple[str, str]: + """Return ``(update_detected, detail)``: ``"Yes"``/``"No"``/``"—"`` and the + descriptive tail (``"On time"`` / ``"Earlier than expected"`` / + ``"Later than expected"`` / ``"Late"`` / ``"—"``). Falls back to the raw + message text when the format doesn't match (e.g. error rows).""" + if not message: + return "—", "—" + match = _FRESHNESS_MESSAGE_RE.match(message.strip()) + if not match: + return "—", message + detected = match.group(1).capitalize() + detail = match.group(2) or "—" + return detected, detail + + +# --------------------------------------------------------------------------- +# list_monitors +# --------------------------------------------------------------------------- + + +_SCHEMA_CHANGE_LABEL: dict[str, str] = { + SCHEMA_CHANGE_ADDED: "added", + SCHEMA_CHANGE_DROPPED: "dropped", + SCHEMA_CHANGE_MODIFIED: "modified", +} + + +def _format_schema_change_kind(code: str | None) -> str | None: + """Map an audit-log change code (``A`` / ``D`` / ``M``) to its user-facing word.""" + if code is None: + return None + return _SCHEMA_CHANGE_LABEL.get(code, code) + + +@with_database_session +@mcp_permission("view") +def list_monitors(table_group_id: str, table_name: str) -> str: + """List configured monitors for a table: monitor IDs, types, threshold modes, and bounds. + + Args: + table_group_id: UUID of the table group, e.g. from ``list_table_groups``. + table_name: Table name exactly as stored in TestGen (case-sensitive). + """ + tg, suite = resolve_monitored_table_group(table_group_id) + if suite is None: + return _NOT_MONITORED_OUTPUT + + configs = TestDefinition.list_monitor_configs_for_table(suite.id, table_name) + + doc = MdDoc() + doc.heading(1, f"Monitors on `{table_name}` in `{tg.table_groups_name}`") + sensitivity = ( + suite.predict_sensitivity.value + if suite.predict_sensitivity is not None + else "medium" + ) + doc.field("Prediction model sensitivity", sensitivity) + + if not configs: + doc.text("_No monitors configured for this table._") + return doc.render() + + doc.table( + ["Monitor ID", "Type", "Metric name", "Threshold mode", "Lower", "Upper", "Metric expression"], + [ + [ + str(c.monitor_id), + _MONITOR_LABEL[MonitorType(c.test_type)], + c.metric_name, + c.threshold_mode, + c.threshold_lower, + c.threshold_upper, + c.custom_query, + ] + for c in configs + ], + code=[0, 6], + ) + + return doc.render() + + +# --------------------------------------------------------------------------- +# list_monitor_schema_changes +# --------------------------------------------------------------------------- + + +@with_database_session +@mcp_permission("view") +def list_monitor_schema_changes( + table_group_id: str, + table_name: str | None = None, + since: str | None = None, + limit: int = 20, + page: int = 1, +) -> str: + """List schema changes detected for a table group, newest first. + + Args: + table_group_id: UUID of the table group, e.g. from ``list_table_groups``. + table_name: Filter to one table (exact, case-sensitive). Omit to list changes across every table in the group. + since: Lower-bound date (e.g. ``'7 days'``, ``'2 weeks'``, ``'2026-04-01'``). Omit to include the entire stored history. + limit: Page size (default 20, max 100). + page: Page number starting at 1 (default 1). + """ + validate_page(page) + validate_limit(limit, 100) + since_date = parse_since_arg(since) if since is not None else None + + tg, suite = resolve_monitored_table_group(table_group_id) + if suite is None: + return _NOT_MONITORED_OUTPUT + + entries, total = DataStructureLog.list_for_table_group( + tg.id, + table_name=table_name, + since=since_date, + page=page, + limit=limit, + ) + + doc = MdDoc() + scope_parts: list[str] = [] + if table_name: + scope_parts.append(f"table `{table_name}`") + if since: + scope_parts.append(f"since `{since}`") + scope = f" — {' — '.join(scope_parts)}" if scope_parts else "" + doc.heading(1, f"Schema changes in `{tg.table_groups_name}`{scope}") + + page_info = format_page_info(total, page, limit) + if page_info: + doc.text(page_info) + + if not entries: + if page > 1: + doc.text(f"No schema changes on page {page} (total: {total}).") + else: + doc.text("_No schema changes recorded for this scope._") + return doc.render() + + doc.table( + ["Time", "Change", "Table", "Column", "Old type", "New type"], + [ + [ + entry.change_date, + _format_schema_change_kind(entry.change), + entry.table_name, + entry.column_name, + entry.old_data_type, + entry.new_data_type, + ] + for entry in entries + ], + code=[2, 3], + ) + + if footer := format_page_footer(total, page, limit): + doc.text(footer) + + return doc.render() diff --git a/testgen/ui/views/monitors_dashboard.py b/testgen/ui/views/monitors_dashboard.py index 9bc4e687..5c7baa7e 100644 --- a/testgen/ui/views/monitors_dashboard.py +++ b/testgen/ui/views/monitors_dashboard.py @@ -10,6 +10,7 @@ from testgen.common.cron_service import get_cron_sample from testgen.common.freshness_service import add_business_minutes, get_schedule_params, resolve_holiday_dates from testgen.common.models import with_database_session +from testgen.common.models.data_structure_log import DataStructureLog from testgen.common.models.notification_settings import ( MonitorNotificationSettings, MonitorNotificationTrigger, @@ -894,29 +895,32 @@ def get_monitor_events_for_table(test_suite_id: str, table_name: str, lookback_m @st.cache_data(show_spinner=False) -def get_data_structure_logs(table_group_id: str, table_name: str, start_time: str, end_time: str): - query = """ - SELECT - change_date, - change, - old_data_type, - new_data_type, - column_name - FROM data_structure_log - WHERE table_groups_id = :table_group_id - AND table_name = :table_name - AND change_date > :start_time ::TIMESTAMP - AND change_date <= :end_time ::TIMESTAMP; - """ - params = { - "table_group_id": str(table_group_id), - "table_name": table_name, - "start_time": datetime.fromtimestamp(start_time, UTC), - "end_time": datetime.fromtimestamp(end_time, UTC), - } +def get_data_structure_logs(table_group_id: str, table_name: str, start_time: float, end_time: float): + """Schema-change rows for a (table-group, table) inside an epoch-seconds window. - results = fetch_all_from_db(query, params) - return [ dict(row) for row in results ] + The dashboard's chart layer sends timestamps as epoch seconds (the JS-side + convention); we convert to UTC datetimes for the model method, which expects + real date/datetime bounds. + """ + since = datetime.fromtimestamp(start_time, UTC) + until = datetime.fromtimestamp(end_time, UTC) + entries, _total = DataStructureLog.list_for_table_group( + table_group_id, + table_name=table_name, + since=since, + until=until, + limit=None, + ) + return [ + { + "change_date": entry.change_date, + "change": entry.change, + "old_data_type": entry.old_data_type, + "new_data_type": entry.new_data_type, + "column_name": entry.column_name, + } + for entry in entries + ] @with_database_session diff --git a/tests/unit/mcp/test_model_test_definition.py b/tests/unit/mcp/test_model_test_definition.py index bc77630a..a987d07f 100644 --- a/tests/unit/mcp/test_model_test_definition.py +++ b/tests/unit/mcp/test_model_test_definition.py @@ -1,10 +1,18 @@ +from datetime import UTC, datetime +from types import SimpleNamespace from unittest.mock import patch from uuid import uuid4 import pytest from sqlalchemy.dialects import postgresql -from testgen.common.models.test_definition import TestDefinition +from testgen.common.models.test_definition import ( + THRESHOLD_MODE_HISTORICAL, + THRESHOLD_MODE_NONE, + THRESHOLD_MODE_PREDICTION, + THRESHOLD_MODE_STATIC, + TestDefinition, +) @pytest.fixture @@ -54,3 +62,151 @@ def test_list_for_suite_excludes_monitor_suites(session_mock): sql_joined = "\n".join(_compiled_sql(q) for q in queries) assert "test_suites.is_monitor IS NOT true" in sql_joined assert "JOIN test_suites" in sql_joined + + +# --------------------------------------------------------------------------- +# _derive_threshold_mode — must match the UI form's detection at +# test_definition_form.js:274 (history_calculation == "PREDICT" → Prediction; +# any other non-empty value → Historical; empty → Static; schema → N/A). +# ``threshold_value`` must never leak into the returned bounds. +# --------------------------------------------------------------------------- + + +def _td(**overrides): + defaults = { + "test_type": "Volume_Trend", + "history_calculation": None, + "history_calculation_upper": None, + "lower_tolerance": None, + "upper_tolerance": None, + "threshold_value": None, + "prediction": None, + } + defaults.update(overrides) + return SimpleNamespace(**defaults) + + +def test_derive_threshold_mode_schema_returns_none(): + td = _td(test_type="Schema_Drift", history_calculation="PREDICT", lower_tolerance="5") + mode, lower, upper = TestDefinition._derive_threshold_mode(td) + assert (mode, lower, upper) == (THRESHOLD_MODE_NONE, None, None) + + +def test_derive_threshold_mode_prediction_keys_on_history_calculation_not_prediction_jsonb(): + """``history_calculation == "PREDICT"`` is the canonical signal — the + ``prediction`` JSONB is empty during training or after a failed run, so + keying off it would misreport prediction monitors as Static.""" + td = _td(test_type="Volume_Trend", history_calculation="PREDICT", prediction=None) + mode, lower, upper = TestDefinition._derive_threshold_mode(td) + assert mode == THRESHOLD_MODE_PREDICTION + assert lower is None and upper is None + + +def test_derive_threshold_mode_historical_for_non_freshness(): + td = _td( + test_type="Volume_Trend", + history_calculation="Minimum", + history_calculation_upper="Maximum", + ) + mode, lower, upper = TestDefinition._derive_threshold_mode(td) + assert (mode, lower, upper) == (THRESHOLD_MODE_HISTORICAL, "Minimum", "Maximum") + + +def test_derive_threshold_mode_historical_blocked_for_freshness(): + """Freshness does not support Historical mode — a stray + ``history_calculation`` value falls through to Static.""" + td = _td( + test_type="Freshness_Trend", + history_calculation="Minimum", + upper_tolerance="720", + ) + mode, lower, upper = TestDefinition._derive_threshold_mode(td) + assert mode == THRESHOLD_MODE_STATIC + assert lower is None # Freshness has no lower bound + assert upper == "720" + + +def test_derive_threshold_mode_static_freshness_uses_upper_only(): + td = _td(test_type="Freshness_Trend", lower_tolerance="60", upper_tolerance="720") + mode, lower, upper = TestDefinition._derive_threshold_mode(td) + assert mode == THRESHOLD_MODE_STATIC + assert lower is None # lower is silently dropped for Freshness even if populated + assert upper == "720" + + +def test_derive_threshold_mode_static_volume_uses_both_tolerances(): + td = _td(test_type="Volume_Trend", lower_tolerance="900", upper_tolerance="1100") + mode, lower, upper = TestDefinition._derive_threshold_mode(td) + assert (mode, lower, upper) == (THRESHOLD_MODE_STATIC, "900", "1100") + + +def test_derive_threshold_mode_threshold_value_never_returned(): + """``threshold_value`` is not a monitor configuration — even when no + tolerances are set it must not leak into the bounds.""" + td = _td(test_type="Volume_Trend", threshold_value="42") + mode, lower, upper = TestDefinition._derive_threshold_mode(td) + assert mode == THRESHOLD_MODE_STATIC + assert lower is None + assert upper is None + + +# --------------------------------------------------------------------------- +# forecast_points_from_prediction — extracts forecast rows from a monitor's +# stored prediction JSONB. Keys are epoch-millisecond integer strings; the +# helper must agree with the dashboard's reader at ``monitors_dashboard.py``. +# Standalone (not bound to the ORM class) so it works against either the +# ORM row or the ``TestDefinitionSummary`` dataclass. +# --------------------------------------------------------------------------- + + +def _epoch_ms(ts: datetime) -> str: + return str(int(ts.timestamp() * 1000)) + + +def test_forecast_points_returns_sorted_per_sensitivity_band(): + from testgen.common.models.test_definition import forecast_points_from_prediction + + t1 = datetime(2026, 7, 1, 12, 0, tzinfo=UTC) + t2 = datetime(2026, 7, 8, 12, 0, tzinfo=UTC) + t3 = datetime(2026, 7, 15, 12, 0, tzinfo=UTC) + prediction = { + "lower_tolerance|medium": {_epoch_ms(t2): 90.0, _epoch_ms(t1): 100.0, _epoch_ms(t3): 80.0}, + "upper_tolerance|medium": {_epoch_ms(t1): 110.0, _epoch_ms(t2): 120.0, _epoch_ms(t3): 130.0}, + # Other sensitivities present but not selected: + "lower_tolerance|low": {_epoch_ms(t1): 200.0}, + } + + points = forecast_points_from_prediction(prediction, "medium") + + assert [p.test_time for p in points] == [t1, t2, t3] + assert [p.lower_bound for p in points] == [100.0, 90.0, 80.0] + assert [p.upper_bound for p in points] == [110.0, 120.0, 130.0] + + +def test_forecast_points_returns_empty_when_no_prediction(): + from testgen.common.models.test_definition import forecast_points_from_prediction + assert forecast_points_from_prediction(None, "medium") == [] + assert forecast_points_from_prediction({}, "medium") == [] + + +def test_forecast_points_returns_empty_for_unknown_sensitivity(): + from testgen.common.models.test_definition import forecast_points_from_prediction + prediction = {"lower_tolerance|medium": {_epoch_ms(datetime(2026, 7, 1, tzinfo=UTC)): 1.0}} + assert forecast_points_from_prediction(prediction, "high") == [] + + +def test_forecast_points_handles_one_sided_bounds(): + """When a sensitivity has only one of lower/upper stored, the other + bound on each point should come through as ``None`` rather than skipping + the point entirely.""" + from testgen.common.models.test_definition import forecast_points_from_prediction + + t1 = datetime(2026, 7, 1, 12, 0, tzinfo=UTC) + prediction = {"upper_tolerance|medium": {_epoch_ms(t1): 130.0}} + + points = forecast_points_from_prediction(prediction, "medium") + + assert len(points) == 1 + assert points[0].test_time == t1 + assert points[0].lower_bound is None + assert points[0].upper_bound == 130.0 diff --git a/tests/unit/mcp/test_model_test_result.py b/tests/unit/mcp/test_model_test_result.py index 4b34bac8..8abd9018 100644 --- a/tests/unit/mcp/test_model_test_result.py +++ b/tests/unit/mcp/test_model_test_result.py @@ -4,7 +4,12 @@ import pytest from sqlalchemy.dialects import postgresql -from testgen.common.models.test_result import TestResult, TestResultStatus, TestRunResultRow +from testgen.common.models.test_result import ( + TestResult, + TestResultStatus, + TestRunResultRow, + _parse_kv_pairs, +) @pytest.fixture @@ -166,3 +171,98 @@ def test_select_history_excludes_monitor_suites(session_mock): sql = _compiled_sql(session_mock.scalars.call_args[0][0]) assert "test_suites.is_monitor IS NOT true" in sql assert "JOIN test_suites" in sql + + +# --------------------------------------------------------------------------- +# _parse_kv_pairs — contract with the producer of ``input_parameters`` +# --------------------------------------------------------------------------- + + +def test_parse_kv_pairs_splits_on_semicolon_for_multi_field_input(): + """``input_parameters`` is built with ``"; ".join(...)`` on the writer side + and consumed by ``dict_from_kv`` with ``;`` as the default separator. The + parser must agree, otherwise a row carrying both ``lower_tolerance`` and + ``upper_tolerance`` returns one entry with the second pair embedded in the + first value, and the upper bound never surfaces.""" + raw = "lower_tolerance=5; upper_tolerance=10; baseline_value=42" + assert _parse_kv_pairs(raw) == { + "lower_tolerance": "5", + "upper_tolerance": "10", + "baseline_value": "42", + } + + +def test_parse_kv_pairs_returns_empty_for_empty_input(): + assert _parse_kv_pairs(None) == {} + assert _parse_kv_pairs("") == {} + + +def test_parse_kv_pairs_tolerates_whitespace_and_skips_unkeyed_entries(): + raw = " k1 = v1 ;k2=v2; just_text ; k3=v3" + assert _parse_kv_pairs(raw) == {"k1": "v1", "k2": "v2", "k3": "v3"} + + +# --------------------------------------------------------------------------- +# list_monitor_events_for_table — ORDER BY needs a stable tiebreaker because +# the result is then paginated in Python; without it a row sharing ``test_time`` +# with another can duplicate or skip across pages. +# --------------------------------------------------------------------------- + + +def test_list_monitor_events_for_table_orders_by_stable_tiebreaker(session_mock): + session_mock.execute.return_value.mappings.return_value.all.return_value = [] + + TestResult.list_monitor_events_for_table(uuid4(), "orders") + + raw_sql = str(session_mock.execute.call_args[0][0]) + order_clause = raw_sql.split("ORDER BY", 1)[1] + assert "results.id" in order_clause + assert "active_runs.id" in order_clause + + +# --------------------------------------------------------------------------- +# list_metric_monitor_events — scoped to one test_definition_id, separate +# query path from the run-by-type CTE (no synthesized pending rows). +# --------------------------------------------------------------------------- + + +def test_list_metric_monitor_events_scopes_to_test_definition_id(session_mock): + session_mock.execute.return_value.mappings.return_value.all.return_value = [] + + suite_id = uuid4() + monitor_id = uuid4() + TestResult.list_metric_monitor_events(suite_id, monitor_id) + + raw_sql = str(session_mock.execute.call_args[0][0]) + params = session_mock.execute.call_args[0][1] + assert "results.test_definition_id = :test_definition_id" in raw_sql + assert "results.test_suite_id = :test_suite_id" in raw_sql + assert params["test_definition_id"] == str(monitor_id) + assert params["test_suite_id"] == str(suite_id) + + +def test_list_metric_monitor_events_has_stable_order_with_tiebreaker(session_mock): + """Like the multi-type CTE, this path paginates in Python — the ORDER BY + must include a tiebreaker so rows sharing ``test_time`` don't shuffle + between calls.""" + session_mock.execute.return_value.mappings.return_value.all.return_value = [] + + TestResult.list_metric_monitor_events(uuid4(), uuid4()) + + raw_sql = str(session_mock.execute.call_args[0][0]) + order_clause = raw_sql.split("ORDER BY", 1)[1] + assert "results.test_time DESC" in order_clause + assert "results.id" in order_clause + + +def test_list_metric_monitor_events_bounds_by_suite_lookback(session_mock): + """The suite's ``monitor_lookback`` (defaulting to 1) caps how many rows + come back. The query baking the LIMIT directly via a CTE means the model + doesn't need a second round-trip to read the lookback value.""" + session_mock.execute.return_value.mappings.return_value.all.return_value = [] + + TestResult.list_metric_monitor_events(uuid4(), uuid4()) + + raw_sql = str(session_mock.execute.call_args[0][0]) + assert "monitor_lookback" in raw_sql + assert "lookback_multiplier" in raw_sql diff --git a/tests/unit/mcp/test_tools_monitors.py b/tests/unit/mcp/test_tools_monitors.py index 553d8c88..964fe50e 100644 --- a/tests/unit/mcp/test_tools_monitors.py +++ b/tests/unit/mcp/test_tools_monitors.py @@ -1,8 +1,10 @@ -"""Tests for the MCP monitor tools — read (``get_monitor_summary`` / ``list_monitored_tables``) -and lifecycle/settings (``enable_monitors`` / ``get_monitor_settings`` / ``update_monitor_settings`` +"""Tests for the MCP monitor tools — read (``get_monitor_summary`` / ``list_monitored_tables`` / +``list_monitor_events`` / ``list_monitors`` / ``list_monitor_schema_changes``) and +lifecycle/settings (``enable_monitors`` / ``get_monitor_settings`` / ``update_monitor_settings`` / ``disable_monitors``).""" -from datetime import UTC, datetime +from datetime import UTC, date, datetime +from types import SimpleNamespace from unittest.mock import MagicMock, patch from uuid import uuid4 @@ -43,10 +45,13 @@ def _mock_table_group(**overrides) -> MagicMock: def _mock_monitor_suite(**overrides) -> MagicMock: + from testgen.common.models.test_suite import PredictSensitivity + suite = MagicMock() suite.id = overrides.get("id", uuid4()) suite.is_monitor = True suite.monitor_lookback = overrides.get("monitor_lookback", 7) + suite.predict_sensitivity = overrides.get("predict_sensitivity", PredictSensitivity.medium) return suite @@ -135,10 +140,10 @@ def test_get_monitor_summary_happy_path(mock_resolve, mock_tg_cls, mock_next, db assert "(override)" not in out assert "**Next scheduled run:** not scheduled" in out # Per-type rows - assert "| Freshness | 2 | ok |" in out - assert "| Volume | 0 | ok |" in out - assert "| Schema | 1 | ok |" in out - assert "| Metric | 0 | training |" in out + assert "| Freshness | 2 | Ok |" in out + assert "| Volume | 0 | Ok |" in out + assert "| Schema | 1 | Ok |" in out + assert "| Metric | 0 | Training |" in out @patch(f"{MODULE}.next_scheduled_run", return_value=datetime(2026, 6, 2, 18, 0, tzinfo=UTC)) @@ -246,7 +251,7 @@ def test_get_monitor_summary_empty_state_lookback_zero( assert "**Window start:**" not in out assert "**Window end:**" not in out # All per-type status cells reflect "no results" - assert out.count("no results yet or not configured") == 4 + assert out.count("No results yet or not configured") == 4 @patch(f"{MODULE}.next_scheduled_run", return_value=None) @@ -269,10 +274,10 @@ def test_get_monitor_summary_renders_error_and_pending_states( with _patch_perms(): out = get_monitor_summary(str(tg.id)) - assert "| Freshness | 2 | error |" in out - assert "| Volume | 0 | no results yet or not configured |" in out - assert "| Schema | 1 | no results yet or not configured |" in out - assert "| Metric | 0 | no results yet or not configured |" in out + assert "| Freshness | 2 | Error |" in out + assert "| Volume | 0 | No results yet or not configured |" in out + assert "| Schema | 1 | No results yet or not configured |" in out + assert "| Metric | 0 | No results yet or not configured |" in out # --------------------------------------------------------------------------- @@ -473,9 +478,9 @@ def test_list_monitored_tables_training_and_pending_cells( out = list_monitored_tables(str(tg.id)) row = next(line for line in out.splitlines() if "`t1`" in line) - assert "training" in row - assert "pending" in row - assert "error" in row + assert "Training" in row + assert "Pending" in row + assert "Error" in row @patch(f"{MODULE}.TableGroup") @@ -518,11 +523,11 @@ def test_list_monitored_tables_count_wins_over_training_and_pending( assert " 5 " in row_count, "freshness count should render despite is_training" assert " 3 " in row_count, "volume count should render despite is_pending" assert " 2 " in row_count, "metric count should render despite is_training" - assert "training" not in row_count - assert "pending" not in row_count + assert "Training" not in row_count + assert "Pending" not in row_count # Row 2: error still wins over count row_error = next(line for line in out.splitlines() if "`t_error_with_count`" in line) - assert "error" in row_error + assert "Error" in row_error assert " 4 " not in row_error, "error must win over count (measurement is suspect)" @@ -677,6 +682,304 @@ def test_get_monitor_settings_not_monitored(mock_resolve, db_session_mock): with _patch_perms(): out = get_monitor_settings(str(tg.id)) +# list_monitor_events (TG-1092) +# --------------------------------------------------------------------------- + + +def _monitor_event(test_type="Volume_Trend", **overrides): + from testgen.common.models.test_result import MonitorEvent + + defaults: dict = { + "monitor_id": uuid4(), + "test_type": test_type, + "test_time": datetime(2026, 6, 1, 12, 0, tzinfo=UTC), + "is_anomaly": False, + "is_training": False, + "is_pending": False, + "is_error": False, + "message": None, + "signal": "1000", + "lower_bound": "900", + "upper_bound": "1100", + "schema_change_kind": None, + "column_adds": None, + "column_drops": None, + "column_mods": None, + "metric_name": None, + } + defaults.update(overrides) + return MonitorEvent(**defaults) + + +@patch(f"{MODULE}.TestResult") +@patch(f"{MODULE}.resolve_monitored_table_group") +def test_list_monitor_events_happy_volume(mock_resolve, mock_tr_cls, db_session_mock): + tg = _mock_table_group() + mock_resolve.return_value = (tg, _mock_monitor_suite()) + mock_tr_cls.list_monitor_events_for_table.return_value = ( + [ + _monitor_event(signal="1000", is_anomaly=True), + _monitor_event(signal="800", is_training=True), + ], + 2, + ) + + from testgen.mcp.tools.monitors import list_monitor_events + + with _patch_perms(): + out = list_monitor_events(str(tg.id), "orders", "volume") + + assert "Monitor events: `orders` — `volume` in `Sales`" in out + assert "| Time | Status | Row count | Lower bound | Upper bound |" in out + assert "Anomaly" in out + assert "Training" in out + # Volume rendering must not pull in Metric's name header + assert "Metric name" not in out + mock_tr_cls.list_monitor_events_for_table.assert_called_once_with( + mock_resolve.return_value[1].id, + "orders", + monitor_type="Volume_Trend", + page=1, + limit=20, + ) + + +@patch(f"{MODULE}.TestResult") +@patch(f"{MODULE}.resolve_monitored_table_group") +def test_list_monitor_events_freshness_parses_message_into_columns( + mock_resolve, mock_tr_cls, db_session_mock, +): + """Freshness events carry a structured message from the SQL template: + ``"Table update detected: {Yes|No}[. {Detail}]"``. Render it as two + distinct columns so the LLM consumer doesn't have to re-parse the text.""" + tg = _mock_table_group() + mock_resolve.return_value = (tg, _mock_monitor_suite()) + mock_tr_cls.list_monitor_events_for_table.return_value = ( + [ + _monitor_event(test_type="Freshness_Trend", message="Table update detected: Yes. On time."), + _monitor_event(test_type="Freshness_Trend", message="Table update detected: No. Late.", is_anomaly=True), + _monitor_event(test_type="Freshness_Trend", message="Table update detected: Yes. Later than expected.", is_anomaly=True), + _monitor_event(test_type="Freshness_Trend", message="Table update detected: No"), + ], + 4, + ) + + from testgen.mcp.tools.monitors import list_monitor_events + + with _patch_perms(): + out = list_monitor_events(str(tg.id), "orders", "freshness") + + assert "| Time | Status | Update detected | Detail |" in out + # Freshness rendering should not surface raw bound columns (they map to + # the static-mode upper_tolerance, which is configuration, not per-event data). + assert "Lower bound" not in out + assert "Row count" not in out + # Each row's structured fields land in their own columns. + assert "| Yes | On time |" in out + assert "| No | Late |" in out + assert "| Yes | Later than expected |" in out + # Missing detail renders as em-dash. + assert "| No | — |" in out + + +@patch(f"{MODULE}.TestDefinition") +@patch(f"{MODULE}.TestResult") +@patch(f"{MODULE}.resolve_monitored_table_group") +def test_list_monitor_events_metric_renders_name_in_heading( + mock_resolve, mock_tr_cls, mock_td_cls, db_session_mock, +): + """For Metric monitors the metric name belongs in the heading (one metric + per call — see the ``monitor_id`` requirement). The data table itself + only carries Time / Status / Value / bounds.""" + monitor_id = str(uuid4()) + tg = _mock_table_group() + mock_resolve.return_value = (tg, _mock_monitor_suite()) + mock_tr_cls.list_metric_monitor_events.return_value = ( + [_monitor_event(test_type="Metric_Trend", metric_name="total_amount")], + 1, + ) + mock_td_cls.get.return_value = SimpleNamespace( + test_type="Metric_Trend", + column_name="total_amount", + ) + + from testgen.mcp.tools.monitors import list_monitor_events + + with _patch_perms(): + out = list_monitor_events(str(tg.id), "orders", "metric", monitor_id=monitor_id) + + assert "Monitor events: `total_amount` on `orders` in `Sales`" in out + assert "| Time | Status | Value | Lower bound | Upper bound |" in out + # metric_name lives in the heading, not as a column + assert "Metric name" not in out + + +@patch(f"{MODULE}.resolve_monitored_table_group") +def test_list_monitor_events_metric_requires_monitor_id(mock_resolve, db_session_mock): + """Metric is the only multi-instance monitor type — without ``monitor_id`` + the query would interleave every metric on the table.""" + from testgen.mcp.exceptions import MCPUserError + from testgen.mcp.tools.monitors import list_monitor_events + + tg = _mock_table_group() + mock_resolve.return_value = (tg, _mock_monitor_suite()) + + with _patch_perms(), pytest.raises(MCPUserError, match="monitor_id.*required.*metric"): + list_monitor_events(str(tg.id), "orders", "metric") + + +@patch(f"{MODULE}.resolve_monitored_table_group") +def test_list_monitor_events_monitor_id_rejected_for_non_metric(mock_resolve, db_session_mock): + """Singleton monitor types (freshness/volume/schema) are uniquely + identified by table + type; ``monitor_id`` is meaningless for them and + must be rejected so the caller gets a clear discovery path.""" + from testgen.mcp.exceptions import MCPUserError + from testgen.mcp.tools.monitors import list_monitor_events + + tg = _mock_table_group() + mock_resolve.return_value = (tg, _mock_monitor_suite()) + + with _patch_perms(), pytest.raises(MCPUserError, match="monitor_id.*only applies.*metric"): + list_monitor_events(str(tg.id), "orders", "volume", monitor_id=str(uuid4())) + + +@patch(f"{MODULE}.TestResult") +@patch(f"{MODULE}.resolve_monitored_table_group") +def test_list_monitor_events_schema_uses_dedicated_columns(mock_resolve, mock_tr_cls, db_session_mock): + tg = _mock_table_group() + mock_resolve.return_value = (tg, _mock_monitor_suite()) + mock_tr_cls.list_monitor_events_for_table.return_value = ( + [ + _monitor_event( + test_type="Schema_Drift", + signal="A|3|1|0|2026-05-01T00:00:00", + schema_change_kind="A", + column_adds=3, column_drops=1, column_mods=0, + is_anomaly=True, + ), + ], + 1, + ) + + from testgen.mcp.tools.monitors import list_monitor_events + + with _patch_perms(): + out = list_monitor_events(str(tg.id), "orders", "schema") + + assert "| Time | Status | Table change | Columns added | Columns dropped | Columns modified |" in out + assert "added" in out + # Internal codes never leak + assert "Schema_Drift" not in out + assert "| A |" not in out + + +@patch(f"{MODULE}.TestDefinition") +@patch(f"{MODULE}.TestResult") +@patch(f"{MODULE}.resolve_monitored_table_group") +def test_list_monitor_events_forecast_renders_separate_section( + mock_resolve, mock_tr_cls, mock_td_cls, db_session_mock, +): + """``include_predictions=True`` against a Prediction-Model monitor appends + a ``## Forecast`` section listing forecast points (future timestamps with + predicted bounds). Predictions never bleed into the historical events + table — the LLM consumer can distinguish observed from predicted at a + glance.""" + tg = _mock_table_group() + mock_resolve.return_value = (tg, _mock_monitor_suite()) + mock_tr_cls.list_monitor_events_for_table.return_value = ( + [_monitor_event()], + 1, + ) + # Build a prediction JSONB the helper will read at render time. Keys are + # epoch-ms strings — same format the dashboard uses. + def _epoch_ms(ts): + return str(int(ts.timestamp() * 1000)) + + forecast_t1 = datetime(2026, 6, 22, 12, 0, tzinfo=UTC) + forecast_t2 = datetime(2026, 6, 29, 12, 0, tzinfo=UTC) + monitor_def = MagicMock() + monitor_def.history_calculation = "PREDICT" + monitor_def.prediction = { + "lower_tolerance|medium": {_epoch_ms(forecast_t1): 500.4, _epoch_ms(forecast_t2): 500.3}, + "upper_tolerance|medium": {_epoch_ms(forecast_t1): 503.6, _epoch_ms(forecast_t2): 503.7}, + } + mock_td_cls.get_singleton_monitor.return_value = monitor_def + + from testgen.mcp.tools.monitors import list_monitor_events + + with _patch_perms(): + out = list_monitor_events(str(tg.id), "orders", "volume", include_predictions=True) + + # Forecast section is separate from the historical table + assert "## Forecast" in out + assert "**Sensitivity:** medium" in out + assert "| Time | Predicted lower | Predicted upper |" in out + assert "500.4" in out and "503.6" in out + # Singleton lookup ran for the non-metric forecast path + mock_td_cls.get_singleton_monitor.assert_called_once() + + +@patch(f"{MODULE}.TestDefinition") +@patch(f"{MODULE}.TestResult") +@patch(f"{MODULE}.resolve_monitored_table_group") +def test_list_monitor_events_predictions_unavailable_note( + mock_resolve, mock_tr_cls, mock_td_cls, db_session_mock, +): + """include_predictions=True against a non-Prediction-Model monitor surfaces + a note under the forecast heading explaining why no forecast follows. The + historical events table itself stays clean.""" + tg = _mock_table_group() + mock_resolve.return_value = (tg, _mock_monitor_suite()) + mock_tr_cls.list_monitor_events_for_table.return_value = ([_monitor_event()], 1) + monitor_def = MagicMock() + monitor_def.history_calculation = None # Static or Historical, not Prediction + mock_td_cls.get_singleton_monitor.return_value = monitor_def + + from testgen.mcp.tools.monitors import list_monitor_events + + with _patch_perms(): + out = list_monitor_events(str(tg.id), "orders", "volume", include_predictions=True) + + assert "## Forecast" in out + assert "Predictions not available" in out + assert "Static or Historical Calculation" in out + + +@patch(f"{MODULE}.TestDefinition") +@patch(f"{MODULE}.TestResult") +@patch(f"{MODULE}.resolve_monitored_table_group") +def test_list_monitor_events_predictions_not_applicable_for_schema( + mock_resolve, mock_tr_cls, mock_td_cls, db_session_mock, +): + tg = _mock_table_group() + mock_resolve.return_value = (tg, _mock_monitor_suite()) + mock_tr_cls.list_monitor_events_for_table.return_value = ( + [_monitor_event(test_type="Schema_Drift", schema_change_kind="M", column_mods=1)], + 1, + ) + + from testgen.mcp.tools.monitors import list_monitor_events + + with _patch_perms(): + out = list_monitor_events(str(tg.id), "orders", "schema", include_predictions=True) + + assert "## Forecast" in out + assert "Predictions not applicable" in out + assert "Schema monitors are presence-only" in out + # Schema short-circuits before any TestDefinition lookup — predictions + # never apply, so don't waste a DB round-trip. + mock_td_cls.get_singleton_monitor.assert_not_called() + + +@patch(f"{MODULE}.resolve_monitored_table_group") +def test_list_monitor_events_not_monitored(mock_resolve, db_session_mock): + tg = _mock_table_group(monitor_test_suite_id=None) + mock_resolve.return_value = (tg, None) + + from testgen.mcp.tools.monitors import list_monitor_events + + with _patch_perms(): + out = list_monitor_events(str(tg.id), "orders", "volume") assert out == "This table group is not monitored." @@ -854,3 +1157,219 @@ def test_disable_monitors_not_enabled(mock_resolve, mock_disable, db_session_moc disable_monitors(str(tg.id)) mock_disable.assert_not_called() +def test_list_monitor_events_invalid_monitor_type(db_session_mock): + from testgen.mcp.tools.monitors import list_monitor_events + + with _patch_perms(), pytest.raises(MCPUserError) as exc: + list_monitor_events(str(uuid4()), "orders", "metrics") + assert "Invalid monitor_type" in str(exc.value) + + +# --------------------------------------------------------------------------- +# list_monitors (TG-1092) +# --------------------------------------------------------------------------- + + +def _monitor_config(**overrides): + from testgen.common.models.test_definition import ( + THRESHOLD_MODE_STATIC, + MonitorConfig, + ) + + defaults: dict = { + "monitor_id": uuid4(), + "test_type": "Volume_Trend", + "table_name": "orders", + "metric_name": None, + "threshold_mode": THRESHOLD_MODE_STATIC, + "threshold_lower": "900", + "threshold_upper": "1100", + "custom_query": None, + } + defaults.update(overrides) + return MonitorConfig(**defaults) + + +@patch(f"{MODULE}.TestDefinition") +@patch(f"{MODULE}.resolve_monitored_table_group") +def test_list_monitors_happy_path(mock_resolve, mock_td_cls, db_session_mock): + tg = _mock_table_group() + suite = _mock_monitor_suite() + mock_resolve.return_value = (tg, suite) + mock_td_cls.list_monitor_configs_for_table.return_value = [ + _monitor_config(test_type="Freshness_Trend", threshold_lower=None, threshold_upper="60"), + _monitor_config( + test_type="Volume_Trend", + threshold_mode="Prediction Model", + threshold_lower=None, threshold_upper=None, + ), + _monitor_config( + test_type="Metric_Trend", + metric_name="total_amount", + custom_query="SUM(total_amount)", + ), + ] + + from testgen.mcp.tools.monitors import list_monitors + + with _patch_perms(): + out = list_monitors(str(tg.id), "orders") + + assert "Monitors on `orders` in `Sales`" in out + assert "**Prediction model sensitivity:** medium" in out + # Type column uses Title Case (reuses _MONITOR_LABEL — no duplicate + # lowercase dict needed; parse_monitor_type accepts either casing). + assert "Freshness" in out + assert "Volume" in out + assert "Metric" in out + # Metric expression column only populated for Metric + assert "SUM(total_amount)" in out + # Configuration tool must not surface runtime prediction bands + assert "Prediction bands" not in out + # Internal codes never leak + assert "Volume_Trend" not in out + assert "Freshness_Trend" not in out + + +@patch(f"{MODULE}.TestDefinition") +@patch(f"{MODULE}.resolve_monitored_table_group") +def test_list_monitors_empty(mock_resolve, mock_td_cls, db_session_mock): + tg = _mock_table_group() + mock_resolve.return_value = (tg, _mock_monitor_suite()) + mock_td_cls.list_monitor_configs_for_table.return_value = [] + + from testgen.mcp.tools.monitors import list_monitors + + with _patch_perms(): + out = list_monitors(str(tg.id), "orders") + + assert "_No monitors configured for this table._" in out + + +@patch(f"{MODULE}.resolve_monitored_table_group") +def test_list_monitors_not_monitored(mock_resolve, db_session_mock): + tg = _mock_table_group(monitor_test_suite_id=None) + mock_resolve.return_value = (tg, None) + + from testgen.mcp.tools.monitors import list_monitors + + with _patch_perms(): + out = list_monitors(str(tg.id), "orders") + + assert out == "This table group is not monitored." + + +# --------------------------------------------------------------------------- +# list_monitor_schema_changes (TG-1092) +# --------------------------------------------------------------------------- + + +def _schema_log_entry(**overrides): + from testgen.common.models.data_structure_log import DataStructureLogEntry + + defaults: dict = { + "log_id": uuid4(), + "table_groups_id": uuid4(), + "table_name": "orders", + "column_name": "customer_id", + "change_date": datetime(2026, 6, 1, 12, 0, tzinfo=UTC), + "change": "A", + "old_data_type": None, + "new_data_type": "INTEGER", + } + defaults.update(overrides) + return DataStructureLogEntry(**defaults) + + +@patch(f"{MODULE}.DataStructureLog") +@patch(f"{MODULE}.resolve_monitored_table_group") +def test_list_monitor_schema_changes_happy_path(mock_resolve, mock_dsl_cls, db_session_mock): + tg = _mock_table_group() + mock_resolve.return_value = (tg, _mock_monitor_suite()) + mock_dsl_cls.list_for_table_group.return_value = ( + [ + _schema_log_entry(change="A", column_name="new_col", old_data_type=None, new_data_type="TEXT"), + _schema_log_entry(change="D", column_name="old_col", old_data_type="INTEGER", new_data_type=None), + _schema_log_entry(change="M", column_name="total", old_data_type="NUMERIC", new_data_type="DECIMAL(10,2)"), + ], + 3, + ) + + from testgen.mcp.tools.monitors import list_monitor_schema_changes + + with _patch_perms(): + out = list_monitor_schema_changes(str(tg.id)) + + assert "Schema changes in `Sales`" in out + assert "| Time | Change | Table | Column | Old type | New type |" in out + # Codes are mapped to user-facing words; raw codes don't leak + assert "added" in out + assert "dropped" in out + assert "modified" in out + out_row_section = out.split("| --- |", 1)[1] + assert "| A |" not in out_row_section + assert "| D |" not in out_row_section + assert "| M |" not in out_row_section + + +@patch(f"{MODULE}.DataStructureLog") +@patch(f"{MODULE}.resolve_monitored_table_group") +def test_list_monitor_schema_changes_table_filter_in_heading( + mock_resolve, mock_dsl_cls, db_session_mock, +): + tg = _mock_table_group() + mock_resolve.return_value = (tg, _mock_monitor_suite()) + mock_dsl_cls.list_for_table_group.return_value = ([], 0) + + from testgen.mcp.tools.monitors import list_monitor_schema_changes + + with _patch_perms(): + out = list_monitor_schema_changes(str(tg.id), table_name="orders", since="7 days") + + assert "table `orders`" in out + assert "since `7 days`" in out + + +@patch(f"{MODULE}.DataStructureLog") +@patch(f"{MODULE}.resolve_monitored_table_group") +def test_list_monitor_schema_changes_since_passed_to_model( + mock_resolve, mock_dsl_cls, db_session_mock, +): + tg = _mock_table_group() + mock_resolve.return_value = (tg, _mock_monitor_suite()) + mock_dsl_cls.list_for_table_group.return_value = ([], 0) + + from testgen.mcp.tools.monitors import list_monitor_schema_changes + + with _patch_perms(): + list_monitor_schema_changes(str(tg.id), since="2026-05-01") + + call = mock_dsl_cls.list_for_table_group.call_args + assert call.kwargs["since"] == date(2026, 5, 1) + + +@patch(f"{MODULE}.resolve_monitored_table_group") +def test_list_monitor_schema_changes_not_monitored(mock_resolve, db_session_mock): + tg = _mock_table_group(monitor_test_suite_id=None) + mock_resolve.return_value = (tg, None) + + from testgen.mcp.tools.monitors import list_monitor_schema_changes + + with _patch_perms(): + out = list_monitor_schema_changes(str(tg.id)) + + assert out == "This table group is not monitored." + + +def test_list_monitor_schema_changes_invalid_since(db_session_mock): + from testgen.mcp.tools.monitors import list_monitor_schema_changes + + with _patch_perms(), pytest.raises(MCPUserError): + list_monitor_schema_changes(str(uuid4()), since="bogus") + + +def test_list_monitor_schema_changes_limit_out_of_range(db_session_mock): + from testgen.mcp.tools.monitors import list_monitor_schema_changes + + with _patch_perms(), pytest.raises(MCPUserError): + list_monitor_schema_changes(str(uuid4()), limit=500) From c9de0cc004fc688e1c0065323087747713410237 Mon Sep 17 00:00:00 2001 From: Aarthy Adityan Date: Mon, 29 Jun 2026 16:59:36 -0400 Subject: [PATCH 70/78] fix(mcp): align monitor L2 forecasts with the UI and harden the read tools (TG-1092) Review follow-ups on the monitor L2 read tools: - list_monitor_events: validate the metric monitor_id as a UUID and scope the lookup to the suite + table + Metric_Trend, so a monitor_id belonging to a metric on another table or suite can't be rendered under the wrong heading. - Forecasts now mirror the dashboard exactly. The forecast logic (predicted next-update window, coupled baseline-then-refresh band) moved into the UI-agnostic common.monitor_forecast module shared by the dashboard and the MCP tool. Volume/Metric monitors coupled to a Freshness monitor render the same band the UI plots; Freshness renders its predicted next-update window. No internal coupling terminology is exposed. - Threshold-mode labels are a ThresholdMode StrEnum; DataStructureLog .list_for_table_group takes *clauses like its siblings (the dashboard keeps its exclusive lower-bound schema-change window). - Type the forecast renderer, fix stale docstrings, and make the forecast-pending note accurate (a forecast appears once the Prediction Model has trained on enough history, not after the first run). Co-Authored-By: Claude Opus 4.8 (1M context) --- testgen/common/models/data_structure_log.py | 24 +-- testgen/common/models/test_definition.py | 27 +-- testgen/common/models/test_result.py | 4 +- testgen/common/monitor_forecast.py | 150 ++++++++++++++ testgen/mcp/tools/monitors.py | 199 +++++++++++++++---- testgen/ui/views/monitors_dashboard.py | 116 +++-------- tests/unit/mcp/test_model_test_definition.py | 19 +- tests/unit/mcp/test_tools_monitors.py | 181 +++++++++++++++-- 8 files changed, 521 insertions(+), 199 deletions(-) create mode 100644 testgen/common/monitor_forecast.py diff --git a/testgen/common/models/data_structure_log.py b/testgen/common/models/data_structure_log.py index b9a34942..552c932d 100644 --- a/testgen/common/models/data_structure_log.py +++ b/testgen/common/models/data_structure_log.py @@ -1,5 +1,5 @@ from dataclasses import dataclass -from datetime import date, datetime +from datetime import datetime from uuid import UUID, uuid4 from sqlalchemy import Column, String, asc, desc, select @@ -54,19 +54,17 @@ class DataStructureLog(Entity): def list_for_table_group( cls, table_group_id: str | UUID, - *, - table_name: str | None = None, - since: date | datetime | None = None, - until: date | datetime | None = None, + *clauses, page: int = 1, limit: int | None = 20, ) -> tuple[list[DataStructureLogEntry], int]: """Paginated schema-change audit log for one table group, newest first. - Filters by ``table_name`` (exact match; the audit log stores names - case-sensitively as the source emitted them), ``since`` (lower-bound on - ``change_date``), and ``until`` (upper-bound). ``limit=None`` skips - pagination — the caller gets every matching row in one shot. + Caller-supplied ``*clauses`` are WHERE expressions on this model — e.g. + ``cls.table_name == name`` (exact match; the audit log stores names + case-sensitively as the source emitted them) or a ``cls.change_date`` + bound. ``limit=None`` skips pagination — the caller gets every matching + row in one shot. """ query = select( cls.log_id, @@ -77,13 +75,7 @@ def list_for_table_group( cls.change, cls.old_data_type, cls.new_data_type, - ).where(cls.table_groups_id == table_group_id) - if table_name is not None: - query = query.where(cls.table_name == table_name) - if since is not None: - query = query.where(cls.change_date >= since) - if until is not None: - query = query.where(cls.change_date <= until) + ).where(cls.table_groups_id == table_group_id, *clauses) query = query.order_by(*cls._default_order_by) if limit is None: diff --git a/testgen/common/models/test_definition.py b/testgen/common/models/test_definition.py index 7220017b..1479ef2e 100644 --- a/testgen/common/models/test_definition.py +++ b/testgen/common/models/test_definition.py @@ -271,12 +271,13 @@ class TestDefinitionMinimal(EntityMinimal): test_name_short: str -# Threshold-mode labels for monitor TestDefinitions — derived from which fields -# on the definition are populated. See ``TestDefinition._derive_threshold_mode``. -THRESHOLD_MODE_PREDICTION = "Prediction Model" -THRESHOLD_MODE_HISTORICAL = "Historical Calculation" -THRESHOLD_MODE_STATIC = "Static" -THRESHOLD_MODE_NONE = "N/A" +class ThresholdMode(StrEnum): + """How a monitor's bounds are determined — derived from which fields on the + definition are populated. See ``TestDefinition._derive_threshold_mode``.""" + PREDICTION = "Prediction Model" + HISTORICAL = "Historical Calculation" + STATIC = "Static" + NONE = "N/A" @dataclass @@ -355,7 +356,7 @@ class MonitorConfig(EntityMinimal): test_type: str table_name: str metric_name: str | None - threshold_mode: str + threshold_mode: ThresholdMode threshold_lower: str | None threshold_upper: str | None custom_query: str | None @@ -706,7 +707,7 @@ def _build_monitor_config(cls, td: "TestDefinition") -> MonitorConfig: @classmethod def _derive_threshold_mode( cls, td: "TestDefinition", - ) -> tuple[str, str | None, str | None]: + ) -> tuple[ThresholdMode, str | None, str | None]: """Pick a mode and the bounds tuple that applies under that mode. Detection mirrors the UI form (``test_definition_form.js``): a @@ -728,14 +729,14 @@ def _derive_threshold_mode( * Static (Volume / Metric): ``(lower_tolerance, upper_tolerance)``. """ if td.test_type == MonitorType.SCHEMA.value: - return THRESHOLD_MODE_NONE, None, None + return ThresholdMode.NONE, None, None if td.history_calculation == "PREDICT": - return THRESHOLD_MODE_PREDICTION, None, None + return ThresholdMode.PREDICTION, None, None if td.history_calculation and td.test_type != MonitorType.FRESHNESS.value: - return THRESHOLD_MODE_HISTORICAL, td.history_calculation, td.history_calculation_upper + return ThresholdMode.HISTORICAL, td.history_calculation, td.history_calculation_upper if td.test_type == MonitorType.FRESHNESS.value: - return THRESHOLD_MODE_STATIC, None, td.upper_tolerance - return THRESHOLD_MODE_STATIC, td.lower_tolerance, td.upper_tolerance + return ThresholdMode.STATIC, None, td.upper_tolerance + return ThresholdMode.STATIC, td.lower_tolerance, td.upper_tolerance @classmethod def select_page( diff --git a/testgen/common/models/test_result.py b/testgen/common/models/test_result.py index 6c5f9719..a8ed2753 100644 --- a/testgen/common/models/test_result.py +++ b/testgen/common/models/test_result.py @@ -149,7 +149,7 @@ class MonitorEvent: Pending rows have no underlying ``test_results`` row — ``monitor_id``, ``test_time``, and result flags are ``None``; ``is_pending`` is True. Forecast points (future timestamps with predicted bounds) are NOT - events and surface separately via ``TestDefinition.get_forecast_points``. + events and surface separately via ``forecast_points_from_prediction``. """ monitor_id: UUID | None test_type: str @@ -675,7 +675,7 @@ def list_monitor_events_for_table( Forecast points for Prediction-Model monitors are NOT included here — events are only past, observed runs. Read forecasts separately via - ``TestDefinition.get_forecast_points(sensitivity)``. + ``forecast_points_from_prediction(prediction, sensitivity)``. """ monitor_codes = ( [monitor_type] if monitor_type is not None diff --git a/testgen/common/monitor_forecast.py b/testgen/common/monitor_forecast.py new file mode 100644 index 00000000..6c292e15 --- /dev/null +++ b/testgen/common/monitor_forecast.py @@ -0,0 +1,150 @@ +"""Monitor forecast computation shared by the monitors dashboard and the MCP +monitor tools. + +The forecast a monitor displays depends on its threshold mode and type: + +* Volume / Metric in Prediction Model mode: a value band extending forward. For + monitors coupled to a Freshness monitor the band holds at the baseline until + the next expected refresh, then steps to the forecast value (``gated_forecast_prediction``); + otherwise it is the raw per-step prediction band. +* Freshness in Prediction Model mode: a predicted next-update *time window* + rather than a value band (``next_update_window``). + +These functions are intentionally free of any UI dependency so both surfaces +render the same forecast from the same source of truth. +""" + +from datetime import UTC, date, datetime + +import pandas as pd + +from testgen.common.freshness_service import add_business_minutes, get_schedule_params, resolve_holiday_dates +from testgen.common.models.test_definition import MonitorForecastPoint, TestDefinition, TestDefinitionSummary +from testgen.common.models.test_suite import TestSuite + +# A monitor definition carrying the forecast-relevant fields — satisfied by both +# the ORM row and the read-only summary dataclass. +MonitorDefinition = TestDefinition | TestDefinitionSummary + + +def resolve_suite_holiday_dates(test_suite: TestSuite) -> set[date] | None: + """Holiday dates excluded from the suite's schedule, over a window spanning + recent history through the forecast horizon. ``None`` when no calendars are set.""" + if not test_suite.holiday_codes_list: + return None + now = pd.Timestamp.now("UTC") + idx = pd.DatetimeIndex([now - pd.Timedelta(days=7), now + pd.Timedelta(days=30)]) + return resolve_holiday_dates(test_suite.holiday_codes_list, idx) + + +def next_update_window( + freshness_definition: MonitorDefinition | None, + last_detection_time: datetime | None, + *, + exclude_weekends: bool, + holiday_dates: set[date] | None, + cron_tz: str | None, +) -> dict | None: + """Predicted next-update window as ``{"start", "end"}`` epoch-ms, or ``None``. + + The schedule-derived business-time interval from the last detected update out + to the lower/upper staleness tolerance. ``start`` is ``None`` when only an + upper tolerance is configured (the update is expected by ``end``). Returns + ``None`` unless the Freshness monitor is in Prediction Model mode with a + trained schedule and at least one observed update. + """ + if ( + freshness_definition is None + or freshness_definition.history_calculation != "PREDICT" + or (freshness_definition.prediction and not freshness_definition.prediction.get("schedule_stage")) + or freshness_definition.upper_tolerance is None + or last_detection_time is None + ): + return None + + tz = cron_tz or "UTC" + sched = get_schedule_params(freshness_definition.prediction) + + window_end = add_business_minutes( + pd.Timestamp(last_detection_time), + float(freshness_definition.upper_tolerance), + exclude_weekends, + holiday_dates, tz, + excluded_days=sched.excluded_days, + ) + window_start = None + if lower_minutes := (float(freshness_definition.lower_tolerance) if freshness_definition.lower_tolerance else None): + window_start = add_business_minutes( + pd.Timestamp(last_detection_time), + lower_minutes, + exclude_weekends, + holiday_dates, tz, + excluded_days=sched.excluded_days, + ) + + return { + "start": int(window_start.timestamp() * 1000) if window_start else None, + "end": int(window_end.timestamp() * 1000), + } + + +def gated_forecast_prediction( + definition: MonitorDefinition, + freshness_window: dict | None, + last_run_time: datetime | None, +) -> dict | None: + """Coupled forecast band for a Volume/Metric monitor whose value holds at its + baseline between refreshes. + + A flat baseline line from the latest run up to the predicted next-update + window, then a step to the forecast's next-refresh value with the band + opening to its tolerance. The anchor is never earlier than the latest run, so + the forecast extends forward rather than back over history. Returns ``None`` + when there is no usable forward window — no predicted window, the window has + already elapsed, or no stored baseline — keyed as ``lower_tolerance`` / + ``upper_tolerance`` (no sensitivity suffix); read with ``forecast_band_points``. + """ + baseline = definition.prediction.get("baseline_value") if definition.prediction else None + now_ms = int(pd.Timestamp(last_run_time).timestamp() * 1000) if last_run_time is not None else None + if freshness_window is None or now_ms is None or baseline is None: + return None + window_end = freshness_window.get("end") + if window_end is None or window_end <= now_ms: + return None + + forecast_means = (definition.prediction.get("mean") if definition.prediction else None) or {} + next_refresh_mean = forecast_means[min(forecast_means, key=lambda k: int(k))] if forecast_means else baseline + flat_anchor = max(freshness_window.get("start") or now_ms, now_ms) + # lower/upper_tolerance are VARCHAR columns — coerce to float so the band dicts are numerically + # typed throughout (baseline is already a float). + lower_tol = float(definition.lower_tolerance) if definition.lower_tolerance is not None else None + upper_tol = float(definition.upper_tolerance) if definition.upper_tolerance is not None else None + return { + "method": "predict", + "mean": {flat_anchor: baseline, window_end: next_refresh_mean}, + "lower_tolerance": {flat_anchor: baseline, window_end: lower_tol}, + "upper_tolerance": {flat_anchor: baseline, window_end: upper_tol}, + } + + +def forecast_band_points(prediction: dict | None) -> list[MonitorForecastPoint]: + """Convert a forecast band dict with plain ``lower_tolerance`` / ``upper_tolerance`` + epoch-ms series (as built by ``gated_forecast_prediction``) into time-ordered + forecast points. For the raw per-sensitivity prediction JSONB use + ``forecast_points_from_prediction`` instead.""" + if not prediction: + return [] + lower_series = prediction.get("lower_tolerance") or {} + upper_series = prediction.get("upper_tolerance") or {} + keys = sorted(set(lower_series) | set(upper_series), key=lambda k: int(k)) + points: list[MonitorForecastPoint] = [] + for k in keys: + ts = datetime.fromtimestamp(int(k) / 1000.0, UTC) + lower = lower_series.get(k) + upper = upper_series.get(k) + points.append(MonitorForecastPoint( + test_time=ts, + lower_bound=float(lower) if lower is not None else None, + upper_bound=float(upper) if upper is not None else None, + )) + return points diff --git a/testgen/mcp/tools/monitors.py b/testgen/mcp/tools/monitors.py index 11c10104..919d9bca 100644 --- a/testgen/mcp/tools/monitors.py +++ b/testgen/mcp/tools/monitors.py @@ -1,4 +1,5 @@ import re +from dataclasses import dataclass, field from datetime import UTC, datetime from typing import Any, cast from uuid import UUID @@ -19,11 +20,19 @@ from testgen.common.models.scheduler import RUN_MONITORS_JOB_KEY, JobSchedule from testgen.common.models.table_group import MonitorTableSummary, TableGroup from testgen.common.models.test_definition import ( + MonitorForecastPoint, TestDefinition, + TestDefinitionSummary, forecast_points_from_prediction, ) from testgen.common.models.test_result import MonitorEvent, TestResult from testgen.common.models.test_suite import PredictSensitivity, TestSuite +from testgen.common.monitor_forecast import ( + forecast_band_points, + gated_forecast_prediction, + next_update_window, + resolve_suite_holiday_dates, +) from testgen.common.monitor_service import disable_monitoring, enable_monitoring, update_monitoring from testgen.mcp.exceptions import MCPUserError from testgen.mcp.permissions import mcp_permission @@ -36,6 +45,7 @@ parse_monitor_table_sort, parse_monitor_type, parse_since_arg, + parse_uuid, resolve_monitored_table_group, validate_limit, validate_page, @@ -46,6 +56,11 @@ _NOT_MONITORED_OUTPUT = "This table group is not monitored." +_FORECAST_PENDING_NOTE = ( + "_No forecast yet — monitors in Prediction Model mode produce a forecast once they " + "have trained on enough history._" +) + _MONITOR_LABEL: dict[MonitorType, str] = { MonitorType.FRESHNESS: "Freshness", MonitorType.VOLUME: "Volume", @@ -607,8 +622,8 @@ def list_monitor_events( table_group_id: UUID of the table group, e.g. from ``list_table_groups``. table_name: Table name exactly as stored in TestGen (case-sensitive). monitor_type: One of ``freshness`` / ``volume`` / ``schema`` / ``metric``. - monitor_id: Required when ``monitor_type="metric"``; rejected otherwise. Metric monitors are user-defined (many per table), so events must be scoped to one ``monitor_id`` to avoid interleaving every metric's history. Find the id via ``list_monitors(table_group_id, table_name)``. Singleton types (``freshness`` / ``volume`` / ``schema``) are uniquely identified by table + type. - include_predictions: When True, append the monitor's forecasted future events (predicted bounds for upcoming time points) after the historical events. Only Volume / Metric / Freshness monitors in Prediction Model mode have a forecast; Schema and other modes render a note explaining why no forecast follows. + monitor_id: Required for ``monitor_type="metric"``, rejected for other types. Get it from ``list_monitors``. + include_predictions: When True, append the monitor's forecast after the historical events, if available. limit: Page size (default 20, max 100). page: Page number starting at 1 (default 1). """ @@ -634,19 +649,35 @@ def list_monitor_events( if suite is None: return _NOT_MONITORED_OUTPUT + monitor_def: TestDefinition | TestDefinitionSummary | None if parsed_type == MonitorType.METRIC: + monitor_uuid = parse_uuid(monitor_id, "monitor_id") + # Scope the lookup to this suite + table + Metric_Trend so a monitor_id + # for a metric on another table (or another suite) can't be rendered + # under this table's heading. + monitor_def = next( + iter( + TestDefinition.select_where( + TestDefinition.id == monitor_uuid, + TestDefinition.test_suite_id == suite.id, + TestDefinition.table_name == table_name, + TestDefinition.test_type == MonitorType.METRIC.value, + ) + ), + None, + ) + if monitor_def is None: + raise MCPUserError( + f"No metric monitor `{monitor_id}` found on table `{table_name}`. " + "Use `list_monitors(table_group_id, table_name)` to find a metric's `monitor_id`." + ) events, total = TestResult.list_metric_monitor_events( suite.id, - monitor_id, + monitor_uuid, page=page, limit=limit, ) - monitor_def = TestDefinition.get(monitor_id) - metric_name = ( - monitor_def.column_name or None - if monitor_def is not None and monitor_def.test_type == MonitorType.METRIC.value - else None - ) + metric_name = monitor_def.column_name or None else: events, total = TestResult.list_monitor_events_for_table( suite.id, @@ -668,7 +699,7 @@ def list_monitor_events( if parsed_type == MonitorType.METRIC and metric_name: doc.heading( 1, - f"Monitor events: `{metric_name}` on `{table_name}` in `{tg.table_groups_name}`", + f"Monitor events: Metric `{metric_name}` on `{table_name}` in `{tg.table_groups_name}`", ) else: doc.heading( @@ -686,7 +717,7 @@ def list_monitor_events( else: doc.text("_No monitor events in the active lookback window._") if include_predictions: - _render_forecast_section(doc, parsed_type, monitor_def, suite) + _render_forecast_section(doc, _compute_forecast(suite, table_name, parsed_type, monitor_def, events)) return doc.render() if parsed_type == MonitorType.SCHEMA: @@ -745,54 +776,129 @@ def list_monitor_events( doc.text(footer) if include_predictions: - _render_forecast_section(doc, parsed_type, monitor_def, suite) + _render_forecast_section(doc, _compute_forecast(suite, table_name, parsed_type, monitor_def, events)) return doc.render() -def _render_forecast_section(doc, parsed_type, monitor_def, suite) -> None: - """Append the forecast section (or an explanatory note) under the - historical events table. Predictions are NEVER mixed into the events - table — they live on their own under a `## Forecast` heading so the LLM - consumer can tell observed-past from predicted-future at a glance.""" - doc.heading(2, "Forecast") +@dataclass +class _Forecast: + """The forecast to render under a monitor's events. Exactly one of + ``window`` / ``points`` / ``note`` is meaningful — see ``_compute_forecast``.""" + note: str | None = None + sensitivity: str | None = None + points: list[MonitorForecastPoint] = field(default_factory=list) + window: dict | None = None # {"start": epoch_ms | None, "end": epoch_ms} + +def _compute_forecast( + suite: TestSuite, + table_name: str, + parsed_type: MonitorType, + monitor_def: TestDefinition | TestDefinitionSummary | None, + events: list[MonitorEvent], +) -> _Forecast: + """Compute the same forecast the monitors dashboard plots for this monitor. + + Volume / Metric in Prediction Model mode get a forward value band — the + coupled baseline-then-refresh band when the monitor is tied to a Freshness + monitor, otherwise the per-step prediction band. Freshness gets its predicted + next-update window. Everything else gets an explanatory note.""" if parsed_type == MonitorType.SCHEMA: - doc.text( - "_Predictions not applicable — Schema monitors are presence-only and " - "do not have tolerance bands._" + return _Forecast( + note="_Predictions not applicable to Schema monitors._" ) - return - if monitor_def is None: - doc.text("_No monitor configured for this table._") - return - + return _Forecast(note="_No monitor configured for this table._") if monitor_def.history_calculation != "PREDICT": - doc.text( - "_Predictions not available — this monitor's threshold mode is " + return _Forecast( + note="_Predictions not available — this monitor's threshold mode is " "Static or Historical Calculation, not Prediction Model._" ) - return - sensitivity = ( - suite.predict_sensitivity.value - if suite.predict_sensitivity is not None - else "medium" - ) + if parsed_type == MonitorType.FRESHNESS: + window = _next_update_window_for_table(suite, monitor_def, events) + if window is None: + return _Forecast(note=_FORECAST_PENDING_NOTE) + return _Forecast(window=window) + + last_run_time = max((e.test_time for e in events if e.test_time is not None), default=None) + + # A monitor coupled to a Freshness monitor holds at its baseline until the + # next expected refresh, so its band is the coupled baseline-then-refresh + # shape keyed off the freshness next-update window — not the raw per-step + # prediction series. + if monitor_def.prediction and monitor_def.prediction.get("freshness_gated"): + freshness_def = TestDefinition.get_singleton_monitor( + suite.id, table_name, MonitorType.FRESHNESS.value + ) + freshness_events, _ = TestResult.list_monitor_events_for_table( + suite.id, table_name, monitor_type=MonitorType.FRESHNESS.value, limit=None + ) + window = _next_update_window_for_table(suite, freshness_def, freshness_events) + points = forecast_band_points(gated_forecast_prediction(monitor_def, window, last_run_time)) + if not points: + return _Forecast(note=_FORECAST_PENDING_NOTE) + return _Forecast(points=points) + + sensitivity = suite.predict_sensitivity.value if suite.predict_sensitivity is not None else "medium" points = forecast_points_from_prediction(monitor_def.prediction, sensitivity) if not points: - doc.text( - "_No forecast stored yet — Prediction Model monitors populate their " - "forecast after the first successful run._" + return _Forecast(note=_FORECAST_PENDING_NOTE) + return _Forecast(sensitivity=sensitivity, points=points) + + +def _next_update_window_for_table( + suite: TestSuite, + freshness_def: TestDefinition | TestDefinitionSummary | None, + freshness_events: list[MonitorEvent], +) -> dict | None: + """Predicted next-update window from the table's Freshness monitor, computed + the same way the dashboard does (last detected update + business-time tolerance).""" + last_detection = max( + ( + e.test_time for e in freshness_events + if e.test_time is not None and not e.is_training and not e.is_pending + and _parse_freshness_message(e.message)[0] == "Yes" + ), + default=None, + ) + schedule = JobSchedule.get_for_monitor_suite(suite.id) + return next_update_window( + freshness_def, + last_detection, + exclude_weekends=suite.predict_exclude_weekends, + holiday_dates=resolve_suite_holiday_dates(suite), + cron_tz=schedule.cron_tz if schedule else None, + ) + + +def _render_forecast_section(doc: MdDoc, forecast: _Forecast) -> None: + """Append the forecast under its own `## Forecast` heading. Predictions are + NEVER mixed into the events table so the consumer can tell observed-past from + predicted-future at a glance.""" + doc.heading(2, "Forecast") + + if forecast.window is not None: + end = datetime.fromtimestamp(forecast.window["end"] / 1000.0, UTC) + start_ms = forecast.window.get("start") + if start_ms is not None: + start = datetime.fromtimestamp(start_ms / 1000.0, UTC) + doc.field("Next update expected", f"{start:%Y-%m-%d %H:%M} to {end:%Y-%m-%d %H:%M} UTC") + else: + doc.field("Next update expected by", f"{end:%Y-%m-%d %H:%M} UTC") + return + + if forecast.points: + if forecast.sensitivity: + doc.field("Sensitivity", forecast.sensitivity) + doc.table( + ["Time", "Predicted lower", "Predicted upper"], + [[p.test_time, p.lower_bound, p.upper_bound] for p in forecast.points], ) return - doc.field("Sensitivity", sensitivity) - doc.table( - ["Time", "Predicted lower", "Predicted upper"], - [[p.test_time, p.lower_bound, p.upper_bound] for p in points], - ) + doc.text(forecast.note or "_No forecast available._") def _event_status(event: MonitorEvent) -> str: @@ -931,10 +1037,15 @@ def list_monitor_schema_changes( if suite is None: return _NOT_MONITORED_OUTPUT + clauses = [] + if table_name is not None: + clauses.append(DataStructureLog.table_name == table_name) + if since_date is not None: + clauses.append(DataStructureLog.change_date >= since_date) + entries, total = DataStructureLog.list_for_table_group( tg.id, - table_name=table_name, - since=since_date, + *clauses, page=page, limit=limit, ) diff --git a/testgen/ui/views/monitors_dashboard.py b/testgen/ui/views/monitors_dashboard.py index 5c7baa7e..845aa0f1 100644 --- a/testgen/ui/views/monitors_dashboard.py +++ b/testgen/ui/views/monitors_dashboard.py @@ -1,14 +1,12 @@ import dataclasses import logging -from datetime import UTC, date, datetime +from datetime import UTC, datetime from math import ceil from typing import Any, ClassVar, Literal, cast -import pandas as pd import streamlit as st from testgen.common.cron_service import get_cron_sample -from testgen.common.freshness_service import add_business_minutes, get_schedule_params, resolve_holiday_dates from testgen.common.models import with_database_session from testgen.common.models.data_structure_log import DataStructureLog from testgen.common.models.notification_settings import ( @@ -20,6 +18,11 @@ from testgen.common.models.table_group import TableGroup, TableGroupMinimal from testgen.common.models.test_definition import TestDefinition, TestDefinitionSummary from testgen.common.models.test_suite import PredictSensitivity, TestSuite +from testgen.common.monitor_forecast import ( + gated_forecast_prediction, + next_update_window, + resolve_suite_holiday_dates, +) from testgen.common.monitor_service import disable_monitoring, enable_monitoring, update_monitoring from testgen.ui.components import widgets as testgen from testgen.ui.navigation.menu import MenuItem @@ -509,102 +512,28 @@ def build_schema_changes_data(table_group: TableGroupMinimal, payload: dict) -> return data, handlers -def _resolve_holiday_dates(test_suite: TestSuite) -> set[date] | None: - if not test_suite.holiday_codes_list: - return None - now = pd.Timestamp.now("UTC") - idx = pd.DatetimeIndex([now - pd.Timedelta(days=7), now + pd.Timedelta(days=30)]) - return resolve_holiday_dates(test_suite.holiday_codes_list, idx) - - def _freshness_next_update_window( freshness_definition: TestDefinition | None, events: dict, test_suite: TestSuite, monitor_schedule, ) -> dict | None: - """Predicted next freshness-update window as {"start", "end"} epoch-ms, or None. - - The schedule-derived business-time interval from the last detected update out to the - lower/upper staleness tolerance. Drives the Freshness_Trend display window and couples the - freshness-gated Volume/Metric forecast to the expected next refresh. - """ - if ( - freshness_definition is None - or freshness_definition.history_calculation != "PREDICT" - or (freshness_definition.prediction and not freshness_definition.prediction.get("schedule_stage")) - or freshness_definition.upper_tolerance is None - ): - return None - + """Predicted next freshness-update window for the table — drives the + Freshness_Trend display window and couples the Volume/Metric forecast to the + expected next refresh. Extracts the last detected update from the chart's + event payload and delegates the business-time math to the shared service.""" last_update_events = [ e for e in events["freshness_events"] if e["changed"] and not e["is_training"] and not e["is_pending"] ] - if not last_update_events: - return None - - last_detection_time = max(e["time"] for e in last_update_events) - holiday_dates = _resolve_holiday_dates(test_suite) - tz = monitor_schedule.cron_tz or "UTC" if monitor_schedule else None - sched = get_schedule_params(freshness_definition.prediction) - - window_end = add_business_minutes( - pd.Timestamp(last_detection_time), - float(freshness_definition.upper_tolerance), - test_suite.predict_exclude_weekends, - holiday_dates, tz, - excluded_days=sched.excluded_days, + last_detection_time = max((e["time"] for e in last_update_events), default=None) + return next_update_window( + freshness_definition, + last_detection_time, + exclude_weekends=test_suite.predict_exclude_weekends, + holiday_dates=resolve_suite_holiday_dates(test_suite), + cron_tz=monitor_schedule.cron_tz if monitor_schedule else None, ) - window_start = None - if lower_minutes := (float(freshness_definition.lower_tolerance) if freshness_definition.lower_tolerance else None): - window_start = add_business_minutes( - pd.Timestamp(last_detection_time), - lower_minutes, - test_suite.predict_exclude_weekends, - holiday_dates, tz, - excluded_days=sched.excluded_days, - ) - - return { - "start": int(window_start.timestamp() * 1000) if window_start else None, - "end": int(window_end.timestamp() * 1000), - } - - -def _build_gated_forecast_prediction( - definition: TestDefinition, - freshness_window: dict | None, - last_run_time: datetime | None, -) -> dict | None: - """Coupled forecast payload for a freshness-gated Volume/Metric monitor. - - Holds a flat baseline line from the latest run up to the predicted next-update window, then - steps to the forecast's next-refresh value with the band opening to its tolerance. The anchor - is never earlier than the latest run, so the forecast extends forward rather than back over - history. Returns None when there is no usable forward window — the caller then falls back to a - flat band — i.e. when freshness has no predicted window, the window has already elapsed, or the - gated prediction carries no baseline. - """ - baseline = definition.prediction.get("baseline_value") if definition.prediction else None - window_end = freshness_window.get("end") if freshness_window else None - now_ms = int(pd.Timestamp(last_run_time).timestamp() * 1000) if last_run_time is not None else None - if window_end is None or now_ms is None or window_end <= now_ms or baseline is None: - return None - - forecast_means = (definition.prediction.get("mean") if definition.prediction else None) or {} - next_refresh_mean = forecast_means[min(forecast_means, key=lambda k: int(k))] if forecast_means else baseline - flat_anchor = max(freshness_window.get("start") or now_ms, now_ms) - # lower/upper_tolerance are VARCHAR columns — coerce to float so the band dicts are numerically - # typed throughout (baseline is already a float). - lower_tol = float(definition.lower_tolerance) if definition.lower_tolerance is not None else None - upper_tol = float(definition.upper_tolerance) if definition.upper_tolerance is not None else None - return { - "method": "predict", - "mean": {flat_anchor: baseline, window_end: next_refresh_mean}, - "lower_tolerance": {flat_anchor: baseline, window_end: lower_tol}, - "upper_tolerance": {flat_anchor: baseline, window_end: upper_tol}, - } @with_database_session @@ -702,7 +631,7 @@ def on_close_trends(_payload=None): ): # A freshness-gated monitor holds at its baseline between refreshes (the stale-period # check is value == baseline), so it must never render the rising forecast cone. - gated_prediction = _build_gated_forecast_prediction( + gated_prediction = gated_forecast_prediction( definition, freshness_window, last_run_time_per_test_key.get(test_key), ) if gated_prediction is not None: @@ -904,11 +833,14 @@ def get_data_structure_logs(table_group_id: str, table_name: str, start_time: fl """ since = datetime.fromtimestamp(start_time, UTC) until = datetime.fromtimestamp(end_time, UTC) + # Half-open window (exclusive lower, inclusive upper) so a change at the + # window-start boundary (the previous run's timestamp) is attributed to the + # earlier interval, matching the chart's data-point selection. entries, _total = DataStructureLog.list_for_table_group( table_group_id, - table_name=table_name, - since=since, - until=until, + DataStructureLog.table_name == table_name, + DataStructureLog.change_date > since, + DataStructureLog.change_date <= until, limit=None, ) return [ diff --git a/tests/unit/mcp/test_model_test_definition.py b/tests/unit/mcp/test_model_test_definition.py index a987d07f..1cb30f84 100644 --- a/tests/unit/mcp/test_model_test_definition.py +++ b/tests/unit/mcp/test_model_test_definition.py @@ -7,11 +7,8 @@ from sqlalchemy.dialects import postgresql from testgen.common.models.test_definition import ( - THRESHOLD_MODE_HISTORICAL, - THRESHOLD_MODE_NONE, - THRESHOLD_MODE_PREDICTION, - THRESHOLD_MODE_STATIC, TestDefinition, + ThresholdMode, ) @@ -89,7 +86,7 @@ def _td(**overrides): def test_derive_threshold_mode_schema_returns_none(): td = _td(test_type="Schema_Drift", history_calculation="PREDICT", lower_tolerance="5") mode, lower, upper = TestDefinition._derive_threshold_mode(td) - assert (mode, lower, upper) == (THRESHOLD_MODE_NONE, None, None) + assert (mode, lower, upper) == (ThresholdMode.NONE, None, None) def test_derive_threshold_mode_prediction_keys_on_history_calculation_not_prediction_jsonb(): @@ -98,7 +95,7 @@ def test_derive_threshold_mode_prediction_keys_on_history_calculation_not_predic keying off it would misreport prediction monitors as Static.""" td = _td(test_type="Volume_Trend", history_calculation="PREDICT", prediction=None) mode, lower, upper = TestDefinition._derive_threshold_mode(td) - assert mode == THRESHOLD_MODE_PREDICTION + assert mode == ThresholdMode.PREDICTION assert lower is None and upper is None @@ -109,7 +106,7 @@ def test_derive_threshold_mode_historical_for_non_freshness(): history_calculation_upper="Maximum", ) mode, lower, upper = TestDefinition._derive_threshold_mode(td) - assert (mode, lower, upper) == (THRESHOLD_MODE_HISTORICAL, "Minimum", "Maximum") + assert (mode, lower, upper) == (ThresholdMode.HISTORICAL, "Minimum", "Maximum") def test_derive_threshold_mode_historical_blocked_for_freshness(): @@ -121,7 +118,7 @@ def test_derive_threshold_mode_historical_blocked_for_freshness(): upper_tolerance="720", ) mode, lower, upper = TestDefinition._derive_threshold_mode(td) - assert mode == THRESHOLD_MODE_STATIC + assert mode == ThresholdMode.STATIC assert lower is None # Freshness has no lower bound assert upper == "720" @@ -129,7 +126,7 @@ def test_derive_threshold_mode_historical_blocked_for_freshness(): def test_derive_threshold_mode_static_freshness_uses_upper_only(): td = _td(test_type="Freshness_Trend", lower_tolerance="60", upper_tolerance="720") mode, lower, upper = TestDefinition._derive_threshold_mode(td) - assert mode == THRESHOLD_MODE_STATIC + assert mode == ThresholdMode.STATIC assert lower is None # lower is silently dropped for Freshness even if populated assert upper == "720" @@ -137,7 +134,7 @@ def test_derive_threshold_mode_static_freshness_uses_upper_only(): def test_derive_threshold_mode_static_volume_uses_both_tolerances(): td = _td(test_type="Volume_Trend", lower_tolerance="900", upper_tolerance="1100") mode, lower, upper = TestDefinition._derive_threshold_mode(td) - assert (mode, lower, upper) == (THRESHOLD_MODE_STATIC, "900", "1100") + assert (mode, lower, upper) == (ThresholdMode.STATIC, "900", "1100") def test_derive_threshold_mode_threshold_value_never_returned(): @@ -145,7 +142,7 @@ def test_derive_threshold_mode_threshold_value_never_returned(): tolerances are set it must not leak into the bounds.""" td = _td(test_type="Volume_Trend", threshold_value="42") mode, lower, upper = TestDefinition._derive_threshold_mode(td) - assert mode == THRESHOLD_MODE_STATIC + assert mode == ThresholdMode.STATIC assert lower is None assert upper is None diff --git a/tests/unit/mcp/test_tools_monitors.py b/tests/unit/mcp/test_tools_monitors.py index 964fe50e..60f8e473 100644 --- a/tests/unit/mcp/test_tools_monitors.py +++ b/tests/unit/mcp/test_tools_monitors.py @@ -3,13 +3,14 @@ lifecycle/settings (``enable_monitors`` / ``get_monitor_settings`` / ``update_monitor_settings`` / ``disable_monitors``).""" -from datetime import UTC, date, datetime +from datetime import UTC, datetime from types import SimpleNamespace from unittest.mock import MagicMock, patch from uuid import uuid4 import pytest +from testgen.common.models.data_structure_log import DataStructureLog from testgen.common.models.table_group import MonitorGroupSummary, MonitorTableSummary from testgen.common.models.test_suite import PredictSensitivity from testgen.mcp.exceptions import MCPResourceNotAccessible, MCPUserError @@ -798,17 +799,17 @@ def test_list_monitor_events_metric_renders_name_in_heading( [_monitor_event(test_type="Metric_Trend", metric_name="total_amount")], 1, ) - mock_td_cls.get.return_value = SimpleNamespace( - test_type="Metric_Trend", - column_name="total_amount", - ) + # Scoped lookup returns the metric's summary (select_where, not the bare get). + mock_td_cls.select_where.return_value = [ + SimpleNamespace(test_type="Metric_Trend", column_name="total_amount"), + ] from testgen.mcp.tools.monitors import list_monitor_events with _patch_perms(): out = list_monitor_events(str(tg.id), "orders", "metric", monitor_id=monitor_id) - assert "Monitor events: `total_amount` on `orders` in `Sales`" in out + assert "Monitor events: Metric `total_amount` on `orders` in `Sales`" in out assert "| Time | Status | Value | Lower bound | Upper bound |" in out # metric_name lives in the heading, not as a column assert "Metric name" not in out @@ -843,6 +844,45 @@ def test_list_monitor_events_monitor_id_rejected_for_non_metric(mock_resolve, db list_monitor_events(str(tg.id), "orders", "volume", monitor_id=str(uuid4())) +@patch(f"{MODULE}.resolve_monitored_table_group") +def test_list_monitor_events_metric_invalid_monitor_id(mock_resolve, db_session_mock): + """A malformed ``monitor_id`` is rejected with a clean MCPUserError before + it reaches the query — not surfaced as a raw Postgres UUID-cast error.""" + from testgen.mcp.tools.monitors import list_monitor_events + + tg = _mock_table_group() + mock_resolve.return_value = (tg, _mock_monitor_suite()) + + with _patch_perms(), pytest.raises(MCPUserError, match="Invalid monitor_id"): + list_monitor_events(str(tg.id), "orders", "metric", monitor_id="not-a-uuid") + + +@patch(f"{MODULE}.TestDefinition") +@patch(f"{MODULE}.TestResult") +@patch(f"{MODULE}.resolve_monitored_table_group") +def test_list_monitor_events_metric_monitor_id_scoped_and_not_found( + mock_resolve, mock_tr_cls, mock_td_cls, db_session_mock, +): + """The metric lookup is scoped to this suite + table + Metric_Trend, so a + ``monitor_id`` that doesn't resolve within scope (wrong table, wrong suite, + or not a metric) yields a discovery-hint error rather than another table's + events under this table's heading.""" + from testgen.mcp.tools.monitors import list_monitor_events + + tg = _mock_table_group() + suite = _mock_monitor_suite() + mock_resolve.return_value = (tg, suite) + mock_td_cls.select_where.return_value = [] # no metric matches the scoped clauses + + with _patch_perms(), pytest.raises(MCPUserError, match=r"No metric monitor.*list_monitors"): + list_monitor_events(str(tg.id), "orders", "metric", monitor_id=str(uuid4())) + + # The events query never runs when the scoped definition lookup fails. + mock_tr_cls.list_metric_monitor_events.assert_not_called() + # Lookup was scoped (id + suite + table + type clauses), not a bare PK fetch. + assert len(mock_td_cls.select_where.call_args[0]) == 4 + + @patch(f"{MODULE}.TestResult") @patch(f"{MODULE}.resolve_monitored_table_group") def test_list_monitor_events_schema_uses_dedicated_columns(mock_resolve, mock_tr_cls, db_session_mock): @@ -964,13 +1004,99 @@ def test_list_monitor_events_predictions_not_applicable_for_schema( out = list_monitor_events(str(tg.id), "orders", "schema", include_predictions=True) assert "## Forecast" in out - assert "Predictions not applicable" in out - assert "Schema monitors are presence-only" in out + assert "Predictions not applicable to Schema monitors" in out # Schema short-circuits before any TestDefinition lookup — predictions # never apply, so don't waste a DB round-trip. mock_td_cls.get_singleton_monitor.assert_not_called() +@patch(f"{MODULE}.next_update_window") +@patch(f"{MODULE}.resolve_suite_holiday_dates") +@patch(f"{MODULE}.JobSchedule") +@patch(f"{MODULE}.TestDefinition") +@patch(f"{MODULE}.TestResult") +@patch(f"{MODULE}.resolve_monitored_table_group") +def test_list_monitor_events_freshness_prediction_shows_window( + mock_resolve, mock_tr_cls, mock_td_cls, mock_sched, mock_holidays, mock_window, db_session_mock, +): + """A Freshness monitor in Prediction Model mode forecasts a next-update time + window (the same one the dashboard computes), not a value band — so the + forecast section presents that window rather than a band table or a note.""" + tg = _mock_table_group() + mock_resolve.return_value = (tg, _mock_monitor_suite()) + mock_tr_cls.list_monitor_events_for_table.return_value = ( + [_monitor_event(test_type="Freshness_Trend", message="Table update detected: Yes. On time.")], + 1, + ) + monitor_def = MagicMock() + monitor_def.history_calculation = "PREDICT" + mock_td_cls.get_singleton_monitor.return_value = monitor_def + mock_holidays.return_value = None + mock_sched.get_for_monitor_suite.return_value = SimpleNamespace(cron_tz="UTC") + start_ms = int(datetime(2026, 7, 1, 9, 0, tzinfo=UTC).timestamp() * 1000) + end_ms = int(datetime(2026, 7, 1, 17, 0, tzinfo=UTC).timestamp() * 1000) + mock_window.return_value = {"start": start_ms, "end": end_ms} + + from testgen.mcp.tools.monitors import list_monitor_events + + with _patch_perms(): + out = list_monitor_events(str(tg.id), "orders", "freshness", include_predictions=True) + + assert "## Forecast" in out + assert "Next update expected" in out + assert "2026-07-01 09:00 to 2026-07-01 17:00 UTC" in out + # A window, not a value-band table. + assert "| Time | Predicted lower | Predicted upper |" not in out + + +@patch(f"{MODULE}.next_update_window") +@patch(f"{MODULE}.resolve_suite_holiday_dates") +@patch(f"{MODULE}.JobSchedule") +@patch(f"{MODULE}.TestDefinition") +@patch(f"{MODULE}.TestResult") +@patch(f"{MODULE}.resolve_monitored_table_group") +def test_list_monitor_events_freshness_coupled_volume_shows_gated_band( + mock_resolve, mock_tr_cls, mock_td_cls, mock_sched, mock_holidays, mock_window, db_session_mock, +): + """A Volume/Metric monitor coupled to a Freshness monitor holds at its + baseline until the next expected refresh — the forecast renders that coupled + band (the UI's), and the internal coupling is never named in the output.""" + tg = _mock_table_group() + mock_resolve.return_value = (tg, _mock_monitor_suite()) + # First call: the volume monitor's own events; second: the freshness events + # the window is derived from. + mock_tr_cls.list_monitor_events_for_table.side_effect = [ + ([_monitor_event()], 1), + ([_monitor_event(test_type="Freshness_Trend", message="Table update detected: Yes.")], 1), + ] + volume_def = MagicMock() + volume_def.history_calculation = "PREDICT" + end_ms = int(datetime(2026, 6, 2, 12, 0, tzinfo=UTC).timestamp() * 1000) + volume_def.prediction = {"freshness_gated": True, "baseline_value": 500.0, "mean": {str(end_ms): 480.0}} + volume_def.lower_tolerance = "450" + volume_def.upper_tolerance = "550" + freshness_def = MagicMock() + # list_monitor_events looks up the volume singleton; _compute_forecast then + # looks up the table's freshness definition for the window. + mock_td_cls.get_singleton_monitor.side_effect = [volume_def, freshness_def] + mock_holidays.return_value = None + mock_sched.get_for_monitor_suite.return_value = SimpleNamespace(cron_tz="UTC") + start_ms = int(datetime(2026, 6, 1, 18, 0, tzinfo=UTC).timestamp() * 1000) + mock_window.return_value = {"start": start_ms, "end": end_ms} + + from testgen.mcp.tools.monitors import list_monitor_events + + with _patch_perms(): + out = list_monitor_events(str(tg.id), "orders", "volume", include_predictions=True) + + assert "## Forecast" in out + # The coupled band IS rendered (the step to the next-refresh tolerances). + assert "| Time | Predicted lower | Predicted upper |" in out + assert "450" in out and "550" in out + # Internal coupling terminology must never reach the client. + assert "gated" not in out.lower() + + @patch(f"{MODULE}.resolve_monitored_table_group") def test_list_monitor_events_not_monitored(mock_resolve, db_session_mock): tg = _mock_table_group(monitor_test_suite_id=None) @@ -1172,8 +1298,8 @@ def test_list_monitor_events_invalid_monitor_type(db_session_mock): def _monitor_config(**overrides): from testgen.common.models.test_definition import ( - THRESHOLD_MODE_STATIC, MonitorConfig, + ThresholdMode, ) defaults: dict = { @@ -1181,7 +1307,7 @@ def _monitor_config(**overrides): "test_type": "Volume_Trend", "table_name": "orders", "metric_name": None, - "threshold_mode": THRESHOLD_MODE_STATIC, + "threshold_mode": ThresholdMode.STATIC, "threshold_lower": "900", "threshold_upper": "1100", "custom_query": None, @@ -1281,12 +1407,12 @@ def _schema_log_entry(**overrides): return DataStructureLogEntry(**defaults) -@patch(f"{MODULE}.DataStructureLog") +@patch.object(DataStructureLog, "list_for_table_group") @patch(f"{MODULE}.resolve_monitored_table_group") -def test_list_monitor_schema_changes_happy_path(mock_resolve, mock_dsl_cls, db_session_mock): +def test_list_monitor_schema_changes_happy_path(mock_resolve, mock_list, db_session_mock): tg = _mock_table_group() mock_resolve.return_value = (tg, _mock_monitor_suite()) - mock_dsl_cls.list_for_table_group.return_value = ( + mock_list.return_value = ( [ _schema_log_entry(change="A", column_name="new_col", old_data_type=None, new_data_type="TEXT"), _schema_log_entry(change="D", column_name="old_col", old_data_type="INTEGER", new_data_type=None), @@ -1312,14 +1438,14 @@ def test_list_monitor_schema_changes_happy_path(mock_resolve, mock_dsl_cls, db_s assert "| M |" not in out_row_section -@patch(f"{MODULE}.DataStructureLog") +@patch.object(DataStructureLog, "list_for_table_group") @patch(f"{MODULE}.resolve_monitored_table_group") def test_list_monitor_schema_changes_table_filter_in_heading( - mock_resolve, mock_dsl_cls, db_session_mock, + mock_resolve, mock_list, db_session_mock, ): tg = _mock_table_group() mock_resolve.return_value = (tg, _mock_monitor_suite()) - mock_dsl_cls.list_for_table_group.return_value = ([], 0) + mock_list.return_value = ([], 0) from testgen.mcp.tools.monitors import list_monitor_schema_changes @@ -1328,24 +1454,37 @@ def test_list_monitor_schema_changes_table_filter_in_heading( assert "table `orders`" in out assert "since `7 days`" in out + # table_name + since reach the model as WHERE clauses, not named kwargs. + sql = _compile_clauses(mock_list) + assert "data_structure_log.table_name = 'orders'" in sql + assert "data_structure_log.change_date >=" in sql -@patch(f"{MODULE}.DataStructureLog") +@patch.object(DataStructureLog, "list_for_table_group") @patch(f"{MODULE}.resolve_monitored_table_group") def test_list_monitor_schema_changes_since_passed_to_model( - mock_resolve, mock_dsl_cls, db_session_mock, + mock_resolve, mock_list, db_session_mock, ): tg = _mock_table_group() mock_resolve.return_value = (tg, _mock_monitor_suite()) - mock_dsl_cls.list_for_table_group.return_value = ([], 0) + mock_list.return_value = ([], 0) from testgen.mcp.tools.monitors import list_monitor_schema_changes with _patch_perms(): list_monitor_schema_changes(str(tg.id), since="2026-05-01") - call = mock_dsl_cls.list_for_table_group.call_args - assert call.kwargs["since"] == date(2026, 5, 1) + # ``since`` is passed as a ``change_date >= `` WHERE clause, not a kwarg. + sql = _compile_clauses(mock_list) + assert "data_structure_log.change_date >=" in sql + assert "2026-05-01" in sql + + +def _compile_clauses(mock_method) -> str: + """Compile the ``*clauses`` positional args of a captured + ``list_for_table_group`` call into one SQL string.""" + clauses = mock_method.call_args[0][1:] # drop table_group_id (first positional) + return " ".join(str(c.compile(compile_kwargs={"literal_binds": True})) for c in clauses) @patch(f"{MODULE}.resolve_monitored_table_group") From 53310c5063b841ec527e03e18851ad3b3ac82bbe Mon Sep 17 00:00:00 2001 From: Aarthy Adityan Date: Mon, 29 Jun 2026 17:42:14 -0400 Subject: [PATCH 71/78] fix(mcp): correct monitor forecast pagination, coupling parity, and notes (TG-1092) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Addresses review findings on the forecast work: - Render the forecast only on the first page — it is forward-looking and page-independent, but its anchor (the latest event) was read from the current page, producing a stale forecast on page >= 2. - Match the dashboard's coupled-forecast precondition: a freshness-coupled Volume/Metric monitor with no configured tolerance has no band on either surface (it was entering the coupled path on the coupling flag alone). - Resolve schedule holidays only after the next-update-window guards pass, restoring the dashboard's short-circuit (no calendar work for non-Prediction monitors); next_update_window now takes the test suite. - Exclude Error-status freshness events from the next-update anchor, matching the dashboard. - Use an accurate note when a forecast is absent for a non-training reason (elapsed window / no baseline) instead of telling a trained monitor to keep training. - Fix the dashboard unit test that imported the relocated forecast helper (this broke the Python Tests CI job), and add coverage for the coupled-without-tolerance case. Co-Authored-By: Claude Opus 4.8 (1M context) --- testgen/common/monitor_forecast.py | 5 +-- testgen/mcp/tools/monitors.py | 22 ++++++++----- testgen/ui/views/monitors_dashboard.py | 4 +-- tests/unit/mcp/test_tools_monitors.py | 41 ++++++++++++++++++++---- tests/unit/ui/test_monitors_dashboard.py | 26 +++++++-------- 5 files changed, 66 insertions(+), 32 deletions(-) diff --git a/testgen/common/monitor_forecast.py b/testgen/common/monitor_forecast.py index 6c292e15..a88e52f2 100644 --- a/testgen/common/monitor_forecast.py +++ b/testgen/common/monitor_forecast.py @@ -41,8 +41,7 @@ def next_update_window( freshness_definition: MonitorDefinition | None, last_detection_time: datetime | None, *, - exclude_weekends: bool, - holiday_dates: set[date] | None, + test_suite: TestSuite, cron_tz: str | None, ) -> dict | None: """Predicted next-update window as ``{"start", "end"}`` epoch-ms, or ``None``. @@ -63,6 +62,8 @@ def next_update_window( return None tz = cron_tz or "UTC" + exclude_weekends = test_suite.predict_exclude_weekends + holiday_dates = resolve_suite_holiday_dates(test_suite) sched = get_schedule_params(freshness_definition.prediction) window_end = add_business_minutes( diff --git a/testgen/mcp/tools/monitors.py b/testgen/mcp/tools/monitors.py index 919d9bca..426190be 100644 --- a/testgen/mcp/tools/monitors.py +++ b/testgen/mcp/tools/monitors.py @@ -31,7 +31,6 @@ forecast_band_points, gated_forecast_prediction, next_update_window, - resolve_suite_holiday_dates, ) from testgen.common.monitor_service import disable_monitoring, enable_monitoring, update_monitoring from testgen.mcp.exceptions import MCPUserError @@ -61,6 +60,11 @@ "have trained on enough history._" ) +# Used where a forecast can be absent for reasons other than training (e.g. the +# predicted next-update window has already passed, or no baseline is stored), so +# the message must not claim the model still needs to train. +_FORECAST_UNAVAILABLE_NOTE = "_No forecast available for this monitor right now._" + _MONITOR_LABEL: dict[MonitorType, str] = { MonitorType.FRESHNESS: "Freshness", MonitorType.VOLUME: "Volume", @@ -623,7 +627,7 @@ def list_monitor_events( table_name: Table name exactly as stored in TestGen (case-sensitive). monitor_type: One of ``freshness`` / ``volume`` / ``schema`` / ``metric``. monitor_id: Required for ``monitor_type="metric"``, rejected for other types. Get it from ``list_monitors``. - include_predictions: When True, append the monitor's forecast after the historical events, if available. + include_predictions: When True, append the monitor's forecast, if available (shown on the first page only). limit: Page size (default 20, max 100). page: Page number starting at 1 (default 1). """ @@ -716,7 +720,7 @@ def list_monitor_events( doc.text(f"No events on page {page} (total: {total}).") else: doc.text("_No monitor events in the active lookback window._") - if include_predictions: + if include_predictions and page == 1: _render_forecast_section(doc, _compute_forecast(suite, table_name, parsed_type, monitor_def, events)) return doc.render() @@ -827,8 +831,11 @@ def _compute_forecast( # A monitor coupled to a Freshness monitor holds at its baseline until the # next expected refresh, so its band is the coupled baseline-then-refresh # shape keyed off the freshness next-update window — not the raw per-step - # prediction series. + # prediction series. The tolerance precondition mirrors the dashboard: a + # coupled monitor with no configured tolerance has no band on either surface. if monitor_def.prediction and monitor_def.prediction.get("freshness_gated"): + if monitor_def.lower_tolerance is None and monitor_def.upper_tolerance is None: + return _Forecast(note=_FORECAST_UNAVAILABLE_NOTE) freshness_def = TestDefinition.get_singleton_monitor( suite.id, table_name, MonitorType.FRESHNESS.value ) @@ -838,7 +845,7 @@ def _compute_forecast( window = _next_update_window_for_table(suite, freshness_def, freshness_events) points = forecast_band_points(gated_forecast_prediction(monitor_def, window, last_run_time)) if not points: - return _Forecast(note=_FORECAST_PENDING_NOTE) + return _Forecast(note=_FORECAST_UNAVAILABLE_NOTE) return _Forecast(points=points) sensitivity = suite.predict_sensitivity.value if suite.predict_sensitivity is not None else "medium" @@ -858,7 +865,7 @@ def _next_update_window_for_table( last_detection = max( ( e.test_time for e in freshness_events - if e.test_time is not None and not e.is_training and not e.is_pending + if e.test_time is not None and not e.is_training and not e.is_pending and not e.is_error and _parse_freshness_message(e.message)[0] == "Yes" ), default=None, @@ -867,8 +874,7 @@ def _next_update_window_for_table( return next_update_window( freshness_def, last_detection, - exclude_weekends=suite.predict_exclude_weekends, - holiday_dates=resolve_suite_holiday_dates(suite), + test_suite=suite, cron_tz=schedule.cron_tz if schedule else None, ) diff --git a/testgen/ui/views/monitors_dashboard.py b/testgen/ui/views/monitors_dashboard.py index 845aa0f1..08838ba6 100644 --- a/testgen/ui/views/monitors_dashboard.py +++ b/testgen/ui/views/monitors_dashboard.py @@ -21,7 +21,6 @@ from testgen.common.monitor_forecast import ( gated_forecast_prediction, next_update_window, - resolve_suite_holiday_dates, ) from testgen.common.monitor_service import disable_monitoring, enable_monitoring, update_monitoring from testgen.ui.components import widgets as testgen @@ -530,8 +529,7 @@ def _freshness_next_update_window( return next_update_window( freshness_definition, last_detection_time, - exclude_weekends=test_suite.predict_exclude_weekends, - holiday_dates=resolve_suite_holiday_dates(test_suite), + test_suite=test_suite, cron_tz=monitor_schedule.cron_tz if monitor_schedule else None, ) diff --git a/tests/unit/mcp/test_tools_monitors.py b/tests/unit/mcp/test_tools_monitors.py index 60f8e473..591bd74f 100644 --- a/tests/unit/mcp/test_tools_monitors.py +++ b/tests/unit/mcp/test_tools_monitors.py @@ -1011,13 +1011,12 @@ def test_list_monitor_events_predictions_not_applicable_for_schema( @patch(f"{MODULE}.next_update_window") -@patch(f"{MODULE}.resolve_suite_holiday_dates") @patch(f"{MODULE}.JobSchedule") @patch(f"{MODULE}.TestDefinition") @patch(f"{MODULE}.TestResult") @patch(f"{MODULE}.resolve_monitored_table_group") def test_list_monitor_events_freshness_prediction_shows_window( - mock_resolve, mock_tr_cls, mock_td_cls, mock_sched, mock_holidays, mock_window, db_session_mock, + mock_resolve, mock_tr_cls, mock_td_cls, mock_sched, mock_window, db_session_mock, ): """A Freshness monitor in Prediction Model mode forecasts a next-update time window (the same one the dashboard computes), not a value band — so the @@ -1031,7 +1030,6 @@ def test_list_monitor_events_freshness_prediction_shows_window( monitor_def = MagicMock() monitor_def.history_calculation = "PREDICT" mock_td_cls.get_singleton_monitor.return_value = monitor_def - mock_holidays.return_value = None mock_sched.get_for_monitor_suite.return_value = SimpleNamespace(cron_tz="UTC") start_ms = int(datetime(2026, 7, 1, 9, 0, tzinfo=UTC).timestamp() * 1000) end_ms = int(datetime(2026, 7, 1, 17, 0, tzinfo=UTC).timestamp() * 1000) @@ -1050,13 +1048,12 @@ def test_list_monitor_events_freshness_prediction_shows_window( @patch(f"{MODULE}.next_update_window") -@patch(f"{MODULE}.resolve_suite_holiday_dates") @patch(f"{MODULE}.JobSchedule") @patch(f"{MODULE}.TestDefinition") @patch(f"{MODULE}.TestResult") @patch(f"{MODULE}.resolve_monitored_table_group") def test_list_monitor_events_freshness_coupled_volume_shows_gated_band( - mock_resolve, mock_tr_cls, mock_td_cls, mock_sched, mock_holidays, mock_window, db_session_mock, + mock_resolve, mock_tr_cls, mock_td_cls, mock_sched, mock_window, db_session_mock, ): """A Volume/Metric monitor coupled to a Freshness monitor holds at its baseline until the next expected refresh — the forecast renders that coupled @@ -1079,7 +1076,6 @@ def test_list_monitor_events_freshness_coupled_volume_shows_gated_band( # list_monitor_events looks up the volume singleton; _compute_forecast then # looks up the table's freshness definition for the window. mock_td_cls.get_singleton_monitor.side_effect = [volume_def, freshness_def] - mock_holidays.return_value = None mock_sched.get_for_monitor_suite.return_value = SimpleNamespace(cron_tz="UTC") start_ms = int(datetime(2026, 6, 1, 18, 0, tzinfo=UTC).timestamp() * 1000) mock_window.return_value = {"start": start_ms, "end": end_ms} @@ -1097,6 +1093,39 @@ def test_list_monitor_events_freshness_coupled_volume_shows_gated_band( assert "gated" not in out.lower() +@patch(f"{MODULE}.TestDefinition") +@patch(f"{MODULE}.TestResult") +@patch(f"{MODULE}.resolve_monitored_table_group") +def test_list_monitor_events_freshness_coupled_volume_without_tolerances_shows_note( + mock_resolve, mock_tr_cls, mock_td_cls, db_session_mock, +): + """A coupled Volume/Metric monitor with no configured tolerance has no band on + the dashboard either, so the MCP shows a note (and skips the window queries), + rather than diverging by computing a band the UI never plots.""" + tg = _mock_table_group() + mock_resolve.return_value = (tg, _mock_monitor_suite()) + mock_tr_cls.list_monitor_events_for_table.return_value = ([_monitor_event()], 1) + volume_def = MagicMock() + volume_def.history_calculation = "PREDICT" + volume_def.prediction = {"freshness_gated": True, "baseline_value": 500.0} + volume_def.lower_tolerance = None + volume_def.upper_tolerance = None + mock_td_cls.get_singleton_monitor.return_value = volume_def + + from testgen.mcp.tools.monitors import list_monitor_events + + with _patch_perms(): + out = list_monitor_events(str(tg.id), "orders", "volume", include_predictions=True) + + assert "## Forecast" in out + assert "No forecast available for this monitor right now" in out + # No band, and no terminology leak. + assert "| Time | Predicted lower | Predicted upper |" not in out + assert "gated" not in out.lower() + # The freshness window queries are skipped when there's no tolerance to band. + assert mock_td_cls.get_singleton_monitor.call_count == 1 # only the volume singleton lookup + + @patch(f"{MODULE}.resolve_monitored_table_group") def test_list_monitor_events_not_monitored(mock_resolve, db_session_mock): tg = _mock_table_group(monitor_test_suite_id=None) diff --git a/tests/unit/ui/test_monitors_dashboard.py b/tests/unit/ui/test_monitors_dashboard.py index 84e2fd01..877a0ffc 100644 --- a/tests/unit/ui/test_monitors_dashboard.py +++ b/tests/unit/ui/test_monitors_dashboard.py @@ -4,14 +4,14 @@ import pandas as pd import pytest -from testgen.ui.views.monitors_dashboard import ( - _build_gated_forecast_prediction, - _freshness_next_update_window, -) +from testgen.common.monitor_forecast import gated_forecast_prediction +from testgen.ui.views.monitors_dashboard import _freshness_next_update_window pytestmark = pytest.mark.unit -MODULE = "testgen.ui.views.monitors_dashboard" +# The business-time helpers live in (and are called from) the shared forecast +# module, so patch them there even when exercising the dashboard wrapper. +MODULE = "testgen.common.monitor_forecast" def _freshness_def(history_calculation="PREDICT", prediction=None, upper="1000", lower="500"): @@ -93,7 +93,7 @@ def test_window_start_is_none_when_no_lower_tolerance(mock_abm, _mock_sched): assert window["end"] == int(pd.Timestamp("2026-06-23 19:00").timestamp() * 1000) -# --- _build_gated_forecast_prediction --- +# --- gated_forecast_prediction --- LAST_RUN = datetime(2026, 6, 23, 16, 0) NOW_MS = int(pd.Timestamp(LAST_RUN).timestamp() * 1000) @@ -111,29 +111,29 @@ def _gated_def(baseline=1000.0, mean=None, lower="950", upper="1400"): def test_gated_prediction_none_without_window(): - assert _build_gated_forecast_prediction(_gated_def(), None, LAST_RUN) is None + assert gated_forecast_prediction(_gated_def(), None, LAST_RUN) is None def test_gated_prediction_none_when_window_elapsed(): window = {"start": NOW_MS - 12 * HOUR, "end": NOW_MS - HOUR} # window_end already in the past - assert _build_gated_forecast_prediction(_gated_def(), window, LAST_RUN) is None + assert gated_forecast_prediction(_gated_def(), window, LAST_RUN) is None def test_gated_prediction_none_without_baseline(): window = {"start": NOW_MS - HOUR, "end": NOW_MS + 3 * HOUR} - assert _build_gated_forecast_prediction(_gated_def(baseline=None), window, LAST_RUN) is None + assert gated_forecast_prediction(_gated_def(baseline=None), window, LAST_RUN) is None def test_gated_prediction_none_without_last_run(): window = {"start": NOW_MS - HOUR, "end": NOW_MS + 3 * HOUR} - assert _build_gated_forecast_prediction(_gated_def(), window, None) is None + assert gated_forecast_prediction(_gated_def(), window, None) is None def test_gated_prediction_anchors_at_now_when_window_started(): # window_start is before the latest run → anchor clamps to now (forecast never draws backward) window = {"start": NOW_MS - 12 * HOUR, "end": NOW_MS + 3 * HOUR} mean = {str(NOW_MS + 3 * HOUR): 1200.0, str(NOW_MS + 27 * HOUR): 1300.0} - result = _build_gated_forecast_prediction(_gated_def(mean=mean), window, LAST_RUN) + result = gated_forecast_prediction(_gated_def(mean=mean), window, LAST_RUN) assert result["method"] == "predict" assert result["mean"] == {NOW_MS: 1000.0, NOW_MS + 3 * HOUR: 1200.0} # tolerances coerced from VARCHAR to float @@ -145,11 +145,11 @@ def test_gated_prediction_anchors_at_window_start_when_future(): # window opens after the latest run → flat segment runs out to window_start window = {"start": NOW_MS + HOUR, "end": NOW_MS + 3 * HOUR} mean = {str(NOW_MS + 3 * HOUR): 1200.0} - result = _build_gated_forecast_prediction(_gated_def(mean=mean), window, LAST_RUN) + result = gated_forecast_prediction(_gated_def(mean=mean), window, LAST_RUN) assert set(result["mean"].keys()) == {NOW_MS + HOUR, NOW_MS + 3 * HOUR} def test_gated_prediction_next_mean_falls_back_to_baseline_without_forecast(): window = {"start": NOW_MS - HOUR, "end": NOW_MS + 3 * HOUR} - result = _build_gated_forecast_prediction(_gated_def(mean=None), window, LAST_RUN) + result = gated_forecast_prediction(_gated_def(mean=None), window, LAST_RUN) assert result["mean"][NOW_MS + 3 * HOUR] == 1000.0 # baseline used as the step value From 73f00a44fbc32e187221867d02b64468c8d2e7f3 Mon Sep 17 00:00:00 2001 From: Luis Date: Thu, 18 Jun 2026 13:00:48 -0400 Subject: [PATCH 72/78] feat(telemetry): async batched Mixpanel queue + MCP tool auto-instrumentation send_event enqueues to a bounded queue drained by a daemon worker batching up to 50 events; every MCP tool call auto-emits an mcp-tool-call event, with bounded shutdown drain. Co-Authored-By: Claude Opus 4.8 (1M context) --- testgen/common/mixpanel_service.py | 108 +++++++- testgen/mcp/permissions.py | 5 + testgen/mcp/server.py | 255 +++++++++++------- testgen/mcp/tools/discovery.py | 2 - testgen/server/__init__.py | 12 +- tests/unit/common/test_mixpanel_service.py | 146 ++++++++++ tests/unit/mcp/test_server_instrumentation.py | 125 +++++++++ 7 files changed, 539 insertions(+), 114 deletions(-) create mode 100644 tests/unit/common/test_mixpanel_service.py create mode 100644 tests/unit/mcp/test_server_instrumentation.py diff --git a/testgen/common/mixpanel_service.py b/testgen/common/mixpanel_service.py index 1af2ae4c..eb26e57e 100644 --- a/testgen/common/mixpanel_service.py +++ b/testgen/common/mixpanel_service.py @@ -1,7 +1,11 @@ +import atexit import functools import json import logging +import queue import ssl +import threading +import time import uuid from base64 import b64encode from functools import cached_property, wraps @@ -18,6 +22,13 @@ LOG = logging.getLogger("testgen") +_BATCH_SIZE = 50 # Mixpanel /track array cap +_FLUSH_INTERVAL_SEC = 10 # Time-based flush +_QUEUE_MAX_SIZE = 1000 # Memory cap before drop-on-overflow +_DRAIN_TIMEOUT_SEC = 5 # Bounded shutdown drain + +_SHUTDOWN = object() # sentinel enqueued by drain() to stop the worker + def safe_method(method): @wraps(method) @@ -33,6 +44,14 @@ def wrapped(*args, **kwargs): class MixpanelService(Singleton): + def __init__(self) -> None: + self._queue: queue.Queue | None = None + self._worker: threading.Thread | None = None + self._worker_lock = threading.Lock() + self._started = False + self._stopped = False + atexit.register(self.drain) + @cached_property @with_database_session def instance_id(self): @@ -54,33 +73,106 @@ def _hash_value(self, value: bytes | str, digest_size: int = 8) -> str: @safe_method def send_event(self, event_name, include_usage=False, **properties): - self._track(event_name, include_usage=include_usage, **properties) + self._enqueue(self._build_event(event_name, include_usage=include_usage, **properties)) def send_feedback(self, **properties): - # User-submitted feedback is content the user explicitly chose to share - # so it is not gated by the TG_ANALYTICS opt-out. + # User-submitted feedback is content the user explicitly chose to share, + # so it is not gated by the TG_ANALYTICS opt-out. It is a foreground action + # posted synchronously — never enqueued — so it never starts the worker. try: - self._track("feedback", **properties) + self.send_mp_request("track?ip=1", self._build_event("feedback", **properties)) except Exception: LOG.exception("Error sending feedback") - def _track(self, event_name, include_usage=False, **properties): + def _build_event(self, event_name, include_usage=False, **properties) -> dict: properties.setdefault("instance_id", self.instance_id) properties.setdefault("edition", settings.DOCKER_HUB_REPOSITORY) properties.setdefault("version", settings.VERSION) properties.setdefault("username", session.auth.user_display if session.auth else None) properties.setdefault("distinct_id", self.get_distinct_id(properties["username"])) + properties.setdefault("time", int(time.time())) if include_usage: properties.update(self.get_usage()) - track_payload = { + return { "event": event_name, "properties": { "token": settings.MIXPANEL_TOKEN, **properties, - } + }, } - self.send_mp_request("track?ip=1", track_payload) + + def _ensure_worker(self) -> None: + if self._started: + return + with self._worker_lock: + if self._started: + return + self._queue = queue.Queue(maxsize=_QUEUE_MAX_SIZE) + self._worker = threading.Thread(target=self._worker_loop, name="mixpanel-flush", daemon=True) + self._worker.start() + self._started = True + + def _enqueue(self, event: dict) -> None: + if self._stopped: + LOG.warning("analytics worker stopped; dropping event") + return + self._ensure_worker() + try: + self._queue.put_nowait(event) + except queue.Full: + LOG.warning("analytics queue full; dropping event") + + def _worker_loop(self) -> None: + while True: + batch, stop = self._next_batch() + if batch: + self._flush(batch) + if stop: + return + + def _next_batch(self) -> tuple[list[dict], bool]: + """Block up to _FLUSH_INTERVAL_SEC for the first event, then drain up to + _BATCH_SIZE without blocking. Returns (batch, stop).""" + try: + first = self._queue.get(timeout=_FLUSH_INTERVAL_SEC) + except queue.Empty: + return [], False + if first is _SHUTDOWN: + return [], True + batch = [first] + while len(batch) < _BATCH_SIZE: + try: + event = self._queue.get_nowait() + except queue.Empty: + break + if event is _SHUTDOWN: + return batch, True + batch.append(event) + return batch, False + + def _flush(self, events: list[dict]) -> None: + for start in range(0, len(events), _BATCH_SIZE): + chunk = events[start:start + _BATCH_SIZE] + try: + self.send_mp_request("track?ip=1", chunk) + except Exception: + LOG.exception("Failed to flush analytics batch") + + def drain(self) -> None: + """Flush queued events and stop the worker. Idempotent; bounded by + _DRAIN_TIMEOUT_SEC. Called by the atexit hook and the server lifespan.""" + if not self._started or self._stopped: + return + self._stopped = True + try: + self._queue.put_nowait(_SHUTDOWN) + except queue.Full: + LOG.warning("analytics queue full at shutdown; in-flight events may be dropped") + if self._worker is not None: + self._worker.join(timeout=_DRAIN_TIMEOUT_SEC) + if self._worker.is_alive(): + LOG.warning("analytics drain timed out after %ss", _DRAIN_TIMEOUT_SEC) def get_ssl_context(self): ssl_context = ssl.create_default_context() diff --git a/testgen/mcp/permissions.py b/testgen/mcp/permissions.py index d65483a1..6042d5dd 100644 --- a/testgen/mcp/permissions.py +++ b/testgen/mcp/permissions.py @@ -80,6 +80,11 @@ def set_mcp_token(token: str | None) -> None: _mcp_token.set(token) +def get_mcp_username() -> str | None: + """Return the authenticated username for the current MCP request (or None).""" + return _mcp_username.get() + + def get_authorized_mcp_user() -> User: """Get the authenticated and authorized User for the current MCP request. diff --git a/testgen/mcp/server.py b/testgen/mcp/server.py index c74fe5e7..3d5c4287 100644 --- a/testgen/mcp/server.py +++ b/testgen/mcp/server.py @@ -1,4 +1,7 @@ +import functools import logging +import time +from enum import StrEnum from urllib.parse import urlparse from mcp.server.auth.provider import AccessToken @@ -10,10 +13,26 @@ from testgen import settings from testgen.common.auth import AuthError, decode_jwt_token -from testgen.mcp.permissions import set_mcp_token, set_mcp_username +from testgen.common.mixpanel_service import MixpanelService +from testgen.mcp.exceptions import MCPPermissionDenied, MCPUserError +from testgen.mcp.permissions import get_mcp_username, set_mcp_token, set_mcp_username LOG = logging.getLogger("testgen") + +class MCPCallStatus(StrEnum): + SUCCESS = "success" + ERROR = "error" + PERMISSION_DENIED = "permission_denied" + USER_ERROR = "user_error" + + +class HandlerKind(StrEnum): + TOOL = "tool" + RESOURCE = "resource" + PROMPT = "prompt" + + SERVER_INSTRUCTIONS = """\ TestGen is a data quality platform that profiles databases, generates tests, and monitors tables. @@ -126,6 +145,42 @@ def _build_transport_security() -> TransportSecuritySettings: ) +def _instrument(fn, kind: HandlerKind): + """Wrap an MCP handler to emit one ``mcp-call`` event per invocation. + + Sits inside ``mcp_error_handler`` so it observes the raised exception type + before it is converted to a text response. ``kind`` distinguishes tools, + resources, and prompts in the emitted event. + """ + + @functools.wraps(fn) + def wrapper(*args, **kwargs): + start = time.monotonic() + status = MCPCallStatus.SUCCESS + try: + return fn(*args, **kwargs) + except MCPPermissionDenied: + status = MCPCallStatus.PERMISSION_DENIED + raise + except MCPUserError: + status = MCPCallStatus.USER_ERROR + raise + except Exception: + status = MCPCallStatus.ERROR + raise + finally: + MixpanelService().send_event( + "mcp-call", + kind=kind, + handler_name=fn.__name__, + username=get_mcp_username(), + latency_ms=int((time.monotonic() - start) * 1000), + status=status, + ) + + return wrapper + + def build_mcp_server( api_base_url: str, server_url: str | None = None, @@ -284,109 +339,109 @@ def build_mcp_server( ) _configure_mcp_logging() - def safe_tool(fn): - mcp.tool()(mcp_error_handler(fn)) + def tool_wrapper(fn): + mcp.tool()(mcp_error_handler(_instrument(fn, HandlerKind.TOOL))) def safe_resource(uri, fn): - mcp.resource(uri)(mcp_error_handler(fn)) + mcp.resource(uri)(mcp_error_handler(_instrument(fn, HandlerKind.RESOURCE))) def safe_prompt(fn): - mcp.prompt()(mcp_error_handler(fn)) + mcp.prompt()(mcp_error_handler(_instrument(fn, HandlerKind.PROMPT))) # Tools - safe_tool(get_data_inventory) - safe_tool(list_projects) - safe_tool(get_project) - safe_tool(list_tables) - safe_tool(list_test_suites) - safe_tool(get_test_suite) - safe_tool(list_test_runs) - safe_tool(get_test_run) - safe_tool(list_test_results) - safe_tool(list_test_result_history) - safe_tool(get_failure_summary) - safe_tool(search_test_results) - safe_tool(get_failure_trend) - safe_tool(compare_test_runs) - safe_tool(update_test_result) - safe_tool(bulk_update_test_results) - safe_tool(get_test_type) - safe_tool(get_source_data) - safe_tool(get_source_data_query) - safe_tool(list_tests) - safe_tool(get_test) - safe_tool(list_test_notes) - safe_tool(list_test_types) - safe_tool(get_table) - safe_tool(generate_create_table_script) - safe_tool(get_table_sample) - safe_tool(list_column_profiles) - safe_tool(list_profiling_summaries) - safe_tool(list_profiling_runs) - safe_tool(get_profiling_run) - safe_tool(get_column_profile_detail) - safe_tool(get_column_frequent_values) - safe_tool(get_column_patterns) - safe_tool(search_columns) - safe_tool(compare_profiling_runs) - safe_tool(get_profiling_trends) - safe_tool(get_schema_history) - safe_tool(get_monitor_summary) - safe_tool(list_monitored_tables) - safe_tool(list_monitor_events) - safe_tool(list_monitors) - safe_tool(list_monitor_schema_changes) - safe_tool(enable_monitors) - safe_tool(get_monitor_settings) - safe_tool(update_monitor_settings) - safe_tool(disable_monitors) - safe_tool(run_tests) - safe_tool(run_profiling) - safe_tool(cancel_test_run) - safe_tool(cancel_profiling_run) - safe_tool(generate_tests) - safe_tool(create_test) - safe_tool(update_test) - safe_tool(validate_custom_test) - safe_tool(bulk_update_tests) - safe_tool(export_tests) - safe_tool(import_tests) - safe_tool(create_test_note) - safe_tool(update_test_note) - safe_tool(delete_test_note) - safe_tool(list_hygiene_issues) - safe_tool(get_hygiene_issue) - safe_tool(search_hygiene_issues) - safe_tool(update_hygiene_issue) - safe_tool(create_profiling_schedule) - safe_tool(create_test_run_schedule) - safe_tool(list_schedules) - safe_tool(get_schedule) - safe_tool(update_schedule) - safe_tool(delete_schedule) - safe_tool(get_quality_scores) - safe_tool(list_scorecards) - safe_tool(get_scorecard) - safe_tool(create_scorecard) - safe_tool(update_scorecard) - safe_tool(delete_scorecard) - safe_tool(list_notifications) - safe_tool(get_notification) - safe_tool(create_notification) - safe_tool(update_notification) - safe_tool(delete_notification) - safe_tool(list_connections) - safe_tool(get_connection) - safe_tool(test_connection) - safe_tool(list_table_groups) - safe_tool(get_table_group) - safe_tool(create_table_group) - safe_tool(update_table_group) - safe_tool(preview_table_group) - safe_tool(update_project) - safe_tool(create_test_suite) - safe_tool(update_test_suite) - safe_tool(update_catalog_metadata) + tool_wrapper(get_data_inventory) + tool_wrapper(list_projects) + tool_wrapper(get_project) + tool_wrapper(list_tables) + tool_wrapper(list_test_suites) + tool_wrapper(get_test_suite) + tool_wrapper(list_test_runs) + tool_wrapper(get_test_run) + tool_wrapper(list_test_results) + tool_wrapper(list_test_result_history) + tool_wrapper(get_failure_summary) + tool_wrapper(search_test_results) + tool_wrapper(get_failure_trend) + tool_wrapper(compare_test_runs) + tool_wrapper(update_test_result) + tool_wrapper(bulk_update_test_results) + tool_wrapper(get_test_type) + tool_wrapper(get_source_data) + tool_wrapper(get_source_data_query) + tool_wrapper(list_tests) + tool_wrapper(get_test) + tool_wrapper(list_test_notes) + tool_wrapper(list_test_types) + tool_wrapper(get_table) + tool_wrapper(generate_create_table_script) + tool_wrapper(get_table_sample) + tool_wrapper(list_column_profiles) + tool_wrapper(list_profiling_summaries) + tool_wrapper(list_profiling_runs) + tool_wrapper(get_profiling_run) + tool_wrapper(get_column_profile_detail) + tool_wrapper(get_column_frequent_values) + tool_wrapper(get_column_patterns) + tool_wrapper(search_columns) + tool_wrapper(compare_profiling_runs) + tool_wrapper(get_profiling_trends) + tool_wrapper(get_schema_history) + tool_wrapper(get_monitor_summary) + tool_wrapper(list_monitored_tables) + tool_wrapper(list_monitor_events) + tool_wrapper(list_monitors) + tool_wrapper(list_monitor_schema_changes) + tool_wrapper(enable_monitors) + tool_wrapper(get_monitor_settings) + tool_wrapper(update_monitor_settings) + tool_wrapper(disable_monitors) + tool_wrapper(run_tests) + tool_wrapper(run_profiling) + tool_wrapper(cancel_test_run) + tool_wrapper(cancel_profiling_run) + tool_wrapper(generate_tests) + tool_wrapper(create_test) + tool_wrapper(update_test) + tool_wrapper(validate_custom_test) + tool_wrapper(bulk_update_tests) + tool_wrapper(export_tests) + tool_wrapper(import_tests) + tool_wrapper(create_test_note) + tool_wrapper(update_test_note) + tool_wrapper(delete_test_note) + tool_wrapper(list_hygiene_issues) + tool_wrapper(get_hygiene_issue) + tool_wrapper(search_hygiene_issues) + tool_wrapper(update_hygiene_issue) + tool_wrapper(create_profiling_schedule) + tool_wrapper(create_test_run_schedule) + tool_wrapper(list_schedules) + tool_wrapper(get_schedule) + tool_wrapper(update_schedule) + tool_wrapper(delete_schedule) + tool_wrapper(get_quality_scores) + tool_wrapper(list_scorecards) + tool_wrapper(get_scorecard) + tool_wrapper(create_scorecard) + tool_wrapper(update_scorecard) + tool_wrapper(delete_scorecard) + tool_wrapper(list_notifications) + tool_wrapper(get_notification) + tool_wrapper(create_notification) + tool_wrapper(update_notification) + tool_wrapper(delete_notification) + tool_wrapper(list_connections) + tool_wrapper(get_connection) + tool_wrapper(test_connection) + tool_wrapper(list_table_groups) + tool_wrapper(get_table_group) + tool_wrapper(create_table_group) + tool_wrapper(update_table_group) + tool_wrapper(preview_table_group) + tool_wrapper(update_project) + tool_wrapper(create_test_suite) + tool_wrapper(update_test_suite) + tool_wrapper(update_catalog_metadata) # Resources safe_resource("testgen://test-types", test_types_resource) @@ -404,14 +459,14 @@ 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. + # Register plugin-provided tools through the same 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) + tool_wrapper(tool) except Exception: LOG.warning("Plugin %s failed to load; skipping its MCP tools", plugin.package) diff --git a/testgen/mcp/tools/discovery.py b/testgen/mcp/tools/discovery.py index af82bb46..84eb03a1 100644 --- a/testgen/mcp/tools/discovery.py +++ b/testgen/mcp/tools/discovery.py @@ -1,4 +1,3 @@ -from testgen.common.mixpanel_service import MixpanelService from testgen.common.models import with_database_session from testgen.common.models.data_table import DataTable from testgen.common.models.project import Project @@ -31,7 +30,6 @@ def get_data_inventory() -> str: from testgen.mcp.services.inventory_service import get_inventory perms = get_project_permissions() - MixpanelService().send_event("mcp-get-data-inventory", username=perms.username) return get_inventory( project_codes=perms.allowed_codes, view_project_codes=perms.codes_allowed_to("view"), diff --git a/testgen/server/__init__.py b/testgen/server/__init__.py index 9d8dd1d8..ba8dd5b3 100644 --- a/testgen/server/__init__.py +++ b/testgen/server/__init__.py @@ -23,6 +23,7 @@ from testgen.api.oauth.routes import router as oauth_router from testgen.api.oauth.server import create_authorization_server from testgen.common import version_service +from testgen.common.mixpanel_service import MixpanelService from testgen.common.models import with_database_session from testgen.server.middleware import BodySizeLimitMiddleware, SecurityHeadersMiddleware @@ -85,11 +86,14 @@ def create_app(version: str | None = None) -> FastAPI: @asynccontextmanager async def lifespan(_app: FastAPI) -> AsyncIterator[None]: - if mcp_session_manager is not None: - async with mcp_session_manager.run(): + try: + if mcp_session_manager is not None: + async with mcp_session_manager.run(): + yield + else: yield - else: - yield + finally: + MixpanelService().drain() tags_metadata = [ {"name": "Jobs", "description": "Submit, poll, cancel, and list job executions (profiling, tests, generation)."}, diff --git a/tests/unit/common/test_mixpanel_service.py b/tests/unit/common/test_mixpanel_service.py new file mode 100644 index 00000000..58bc598a --- /dev/null +++ b/tests/unit/common/test_mixpanel_service.py @@ -0,0 +1,146 @@ +import queue +from unittest.mock import MagicMock + +import pytest + +from testgen import settings +from testgen.common import mixpanel_service as mp_module +from testgen.common.mixpanel_service import MixpanelService +from testgen.utils.singleton import SingletonType + + +@pytest.fixture +def service(monkeypatch): + """Fresh, isolated MixpanelService with analytics on and no real network/DB.""" + SingletonType._instances.pop(MixpanelService, None) + monkeypatch.setattr(settings, "ANALYTICS_ENABLED", True) + monkeypatch.setattr(settings, "MIXPANEL_TOKEN", "tok") + # Avoid Streamlit + DB: session.auth falsy, instance_id pre-seeded. + monkeypatch.setattr(mp_module, "session", MagicMock(auth=None)) + svc = MixpanelService() + svc.__dict__["instance_id"] = "iid" + yield svc + svc.drain() + SingletonType._instances.pop(MixpanelService, None) + + +def test_send_event_enqueues_and_worker_flushes_one_event(service, monkeypatch): + captured = [] + monkeypatch.setattr(service, "send_mp_request", lambda endpoint, payload: captured.append((endpoint, payload))) + + service.send_event("nav-home") + service.drain() + + assert len(captured) == 1 + endpoint, payload = captured[0] + assert endpoint == "track?ip=1" + assert isinstance(payload, list) + assert len(payload) == 1 + assert payload[0]["event"] == "nav-home" + assert payload[0]["properties"]["token"] == "tok" # noqa: S105 (test token, not a secret) + assert payload[0]["properties"]["instance_id"] == "iid" + assert "time" in payload[0]["properties"] + + +def test_next_batch_caps_at_batch_size(service, monkeypatch): + monkeypatch.setattr(mp_module, "_BATCH_SIZE", 2) + service._queue = queue.Queue() + for i in range(5): + service._queue.put({"event": f"e{i}", "properties": {}}) + + batch, stop = service._next_batch() + + assert [e["event"] for e in batch] == ["e0", "e1"] + assert stop is False + + +def test_next_batch_stops_on_sentinel_and_returns_pending(service): + service._queue = queue.Queue() + service._queue.put({"event": "a", "properties": {}}) + service._queue.put({"event": "b", "properties": {}}) + service._queue.put(mp_module._SHUTDOWN) + + batch, stop = service._next_batch() + + assert [e["event"] for e in batch] == ["a", "b"] + assert stop is True + + +def test_flush_chunks_over_batch_size(service, monkeypatch): + monkeypatch.setattr(mp_module, "_BATCH_SIZE", 2) + posts = [] + monkeypatch.setattr(service, "send_mp_request", lambda _endpoint, payload: posts.append(len(payload))) + + service._flush([{"event": f"e{i}", "properties": {}} for i in range(3)]) + + assert posts == [2, 1] + + +def test_queue_overflow_drops_and_warns(service, monkeypatch, caplog): + # Pretend the worker is running so _enqueue does not start a real one, + # and give it a size-1 queue that nothing drains. + service._started = True + service._queue = queue.Queue(maxsize=1) + + service._enqueue({"event": "a", "properties": {}}) + with caplog.at_level("WARNING"): + service._enqueue({"event": "b", "properties": {}}) + + assert service._queue.qsize() == 1 + assert "analytics queue full" in caplog.text + + +def test_drain_stops_worker_and_is_idempotent(service, monkeypatch): + monkeypatch.setattr(service, "send_mp_request", lambda *_: None) + service.send_event("x") # starts worker + + service.drain() + assert service._worker is not None + assert not service._worker.is_alive() + + service.drain() # second call is a no-op, must not raise + + +def test_analytics_disabled_never_starts_worker(service, monkeypatch): + monkeypatch.setattr(settings, "ANALYTICS_ENABLED", False) + posts = [] + monkeypatch.setattr(service, "send_mp_request", lambda _endpoint, payload: posts.append(payload)) + + service.send_event("nav-home") + service.drain() + + assert service._started is False + assert service._worker is None + assert posts == [] + + +def test_send_event_after_drain_drops_and_warns(service, monkeypatch, caplog): + posts = [] + monkeypatch.setattr(service, "send_mp_request", lambda _endpoint, payload: posts.append(payload)) + + service.send_event("pre-drain") + service.drain() + + posts.clear() + with caplog.at_level("WARNING"): + service.send_event("late") + + assert "stopped" in caplog.text + assert posts == [] + + +def test_send_feedback_posts_synchronously_without_worker(service, monkeypatch): + posts = [] + monkeypatch.setattr(service, "send_mp_request", lambda endpoint, payload: posts.append((endpoint, payload))) + + service.send_feedback(comment="great tool") + + # Posted immediately (no drain needed), as a single dict (not a list), worker never started. + assert len(posts) == 1 + endpoint, payload = posts[0] + assert endpoint == "track?ip=1" + assert isinstance(payload, dict) + assert payload["event"] == "feedback" + assert payload["properties"]["comment"] == "great tool" + assert service._started is False + assert service._worker is None diff --git a/tests/unit/mcp/test_server_instrumentation.py b/tests/unit/mcp/test_server_instrumentation.py new file mode 100644 index 00000000..38a968bb --- /dev/null +++ b/tests/unit/mcp/test_server_instrumentation.py @@ -0,0 +1,125 @@ +from unittest.mock import patch + +import pytest + +from testgen.mcp.exceptions import MCPPermissionDenied, MCPResourceNotAccessible, MCPUserError +from testgen.mcp.permissions import set_mcp_username +from testgen.mcp.server import HandlerKind, MCPCallStatus, _instrument + + +def _emit_capture(): + calls = [] + + def fake_send_event(self, event_name, include_usage=False, **properties): + calls.append((event_name, properties)) + + return calls, fake_send_event + + +def test_instrument_emits_success_event(): + set_mcp_username("alice") + calls, fake = _emit_capture() + + def my_tool(): + return "ok" + + with patch("testgen.common.mixpanel_service.MixpanelService.send_event", fake): + wrapped = _instrument(my_tool, HandlerKind.TOOL) + assert wrapped() == "ok" + + assert len(calls) == 1 + event_name, props = calls[0] + assert event_name == "mcp-call" + assert props["kind"] == HandlerKind.TOOL + assert props["handler_name"] == "my_tool" + assert props["username"] == "alice" + assert props["status"] == MCPCallStatus.SUCCESS + assert isinstance(props["latency_ms"], int) + + +def test_instrument_resource_emits_event(): + set_mcp_username("alice") + calls, fake = _emit_capture() + + def my_resource(): + return "doc" + + with patch("testgen.common.mixpanel_service.MixpanelService.send_event", fake): + wrapped = _instrument(my_resource, HandlerKind.RESOURCE) + assert wrapped() == "doc" + + event_name, props = calls[0] + assert event_name == "mcp-call" + assert props["kind"] == HandlerKind.RESOURCE + assert props["handler_name"] == "my_resource" + assert props["status"] == MCPCallStatus.SUCCESS + + +def test_instrument_prompt_emits_event(): + set_mcp_username("alice") + calls, fake = _emit_capture() + + def my_prompt(): + return "template" + + with patch("testgen.common.mixpanel_service.MixpanelService.send_event", fake): + wrapped = _instrument(my_prompt, HandlerKind.PROMPT) + assert wrapped() == "template" + + event_name, props = calls[0] + assert event_name == "mcp-call" + assert props["kind"] == HandlerKind.PROMPT + assert props["handler_name"] == "my_prompt" + assert props["status"] == MCPCallStatus.SUCCESS + + +def test_instrument_permission_denied_status(): + set_mcp_username("alice") + calls, fake = _emit_capture() + + def denied_tool(): + raise MCPResourceNotAccessible("Project", "demo") + + with patch("testgen.common.mixpanel_service.MixpanelService.send_event", fake): + wrapped = _instrument(denied_tool, HandlerKind.TOOL) + with pytest.raises(MCPPermissionDenied): + wrapped() + + assert calls[0][1]["status"] == MCPCallStatus.PERMISSION_DENIED + + +def test_instrument_user_error_status(): + set_mcp_username("alice") + calls, fake = _emit_capture() + + def bad_input(): + raise MCPUserError("invalid uuid") + + with patch("testgen.common.mixpanel_service.MixpanelService.send_event", fake): + wrapped = _instrument(bad_input, HandlerKind.TOOL) + with pytest.raises(MCPUserError): + wrapped() + + assert calls[0][1]["status"] == MCPCallStatus.USER_ERROR + + +def test_instrument_error_status(): + set_mcp_username("alice") + calls, fake = _emit_capture() + + def boom(): + raise ValueError("nope") + + with patch("testgen.common.mixpanel_service.MixpanelService.send_event", fake): + wrapped = _instrument(boom, HandlerKind.TOOL) + with pytest.raises(ValueError): + wrapped() + + assert calls[0][1]["status"] == MCPCallStatus.ERROR + + +def test_instrument_preserves_name(): + def my_tool(): + return "ok" + + assert _instrument(my_tool, HandlerKind.TOOL).__name__ == "my_tool" From bef569312ada0829fd142019e2a9444951c8955e Mon Sep 17 00:00:00 2001 From: Aarthy Adityan Date: Tue, 30 Jun 2026 02:02:33 -0400 Subject: [PATCH 73/78] fix(monitors): keep gated forecast band continuous (no zero-width pinch) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The freshness-gated forecast band collapsed to a zero-width point at the flat anchor (lower=upper=baseline), rendering as an hourglass for a future next-update window and as a band detached from history for an imminent one. Carry the next-refresh tolerance across the whole forecast so the band stays a continuous width that connects to the historical band; the mean line still holds flat at baseline then steps. Display-only — no effect on anomaly evaluation. TG-1131 Co-Authored-By: Claude Opus 4.8 (1M context) --- testgen/common/monitor_forecast.py | 9 +++++++-- tests/unit/ui/test_monitors_dashboard.py | 8 +++++--- 2 files changed, 12 insertions(+), 5 deletions(-) diff --git a/testgen/common/monitor_forecast.py b/testgen/common/monitor_forecast.py index a88e52f2..a3c65a46 100644 --- a/testgen/common/monitor_forecast.py +++ b/testgen/common/monitor_forecast.py @@ -120,11 +120,16 @@ def gated_forecast_prediction( # typed throughout (baseline is already a float). lower_tol = float(definition.lower_tolerance) if definition.lower_tolerance is not None else None upper_tol = float(definition.upper_tolerance) if definition.upper_tolerance is not None else None + # The band carries the next-refresh tolerance across the whole forecast, flat anchor included. + # Anchoring the band to the tolerance — rather than collapsing it to the baseline at the flat + # anchor — keeps it a continuous width that connects to the historical band and avoids a + # zero-width pinch (which renders as an hourglass for a future window, or a band detached from + # history for an imminent one). The mean line still holds flat at baseline, then steps. return { "method": "predict", "mean": {flat_anchor: baseline, window_end: next_refresh_mean}, - "lower_tolerance": {flat_anchor: baseline, window_end: lower_tol}, - "upper_tolerance": {flat_anchor: baseline, window_end: upper_tol}, + "lower_tolerance": {flat_anchor: lower_tol, window_end: lower_tol}, + "upper_tolerance": {flat_anchor: upper_tol, window_end: upper_tol}, } diff --git a/tests/unit/ui/test_monitors_dashboard.py b/tests/unit/ui/test_monitors_dashboard.py index 877a0ffc..75f06cf0 100644 --- a/tests/unit/ui/test_monitors_dashboard.py +++ b/tests/unit/ui/test_monitors_dashboard.py @@ -135,10 +135,12 @@ def test_gated_prediction_anchors_at_now_when_window_started(): mean = {str(NOW_MS + 3 * HOUR): 1200.0, str(NOW_MS + 27 * HOUR): 1300.0} result = gated_forecast_prediction(_gated_def(mean=mean), window, LAST_RUN) assert result["method"] == "predict" + # mean holds flat at baseline then steps to the next-refresh forecast assert result["mean"] == {NOW_MS: 1000.0, NOW_MS + 3 * HOUR: 1200.0} - # tolerances coerced from VARCHAR to float - assert result["lower_tolerance"] == {NOW_MS: 1000.0, NOW_MS + 3 * HOUR: 950.0} - assert result["upper_tolerance"] == {NOW_MS: 1000.0, NOW_MS + 3 * HOUR: 1400.0} + # band carries the next-refresh tolerance across the whole forecast (continuous width, no + # zero-width pinch at the flat anchor); coerced from VARCHAR to float + assert result["lower_tolerance"] == {NOW_MS: 950.0, NOW_MS + 3 * HOUR: 950.0} + assert result["upper_tolerance"] == {NOW_MS: 1400.0, NOW_MS + 3 * HOUR: 1400.0} def test_gated_prediction_anchors_at_window_start_when_future(): From 00e6efb0e81e1a4a49976f686db0335af7603a1b Mon Sep 17 00:00:00 2001 From: Aarthy Adityan Date: Tue, 30 Jun 2026 02:24:37 -0400 Subject: [PATCH 74/78] fix(data-catalog): bugs in editing tags --- .../frontend/js/data_profiling/metadata_tags.js | 15 +++++++++------ 1 file changed, 9 insertions(+), 6 deletions(-) diff --git a/testgen/ui/components/frontend/js/data_profiling/metadata_tags.js b/testgen/ui/components/frontend/js/data_profiling/metadata_tags.js index 5ae1092a..c469b0e6 100644 --- a/testgen/ui/components/frontend/js/data_profiling/metadata_tags.js +++ b/testgen/ui/components/frontend/js/data_profiling/metadata_tags.js @@ -100,6 +100,7 @@ const MetadataTagsCard = (props, item) => { help: TAG_HELP[key], label: key === 'pii_flag' ? 'PII Data' : capitalize(key.replaceAll('_', ' ')), state: van.state(value), + initialValue: value, inheritTableGroup: item[`table_group_${key}`] ?? null, // Table group values inherited by table or column inheritTable: item[`table_${key}`] ?? null, // Table values inherited by column }; @@ -198,13 +199,15 @@ const MetadataTagsCard = (props, item) => { content, editingContent, onSave: () => { const items = [{ type: item.type, id: item.id }]; - const tags = attributes.reduce((object, { key, state }) => { - object[key] = state.rawVal; + const tags = attributes.reduce((object, { key, state, initialValue }) => { + if (state.rawVal !== initialValue) { + object[key] = state.rawVal; + } return object; }, {}); - warnCde.val = props.autoflagSettings.profile_flag_cdes && tags.critical_data_element !== item.critical_data_element; - warnPii.val = props.autoflagSettings.profile_flag_pii && tags.pii_flag !== item.pii_flag; + warnCde.val = props.autoflagSettings.profile_flag_cdes && 'critical_data_element' in tags; + warnPii.val = props.autoflagSettings.profile_flag_pii && 'pii_flag' in tags; if (warnCde.val || warnPii.val) { const disableFlags = []; @@ -221,8 +224,8 @@ const MetadataTagsCard = (props, item) => { } }, // Reset states to original values on cancel - onCancel: () => attributes.forEach(({ key, state }) => state.val = item[key]), - hasChanges: () => attributes.some(({ key, state }) => state.val !== item[key]), + onCancel: () => attributes.forEach(({ state, initialValue }) => state.val = initialValue), + hasChanges: () => attributes.some(({ state, initialValue }) => state.val !== initialValue), }), WarningDialog(warningDialogOpen, pendingSaveAction, warnCde, warnPii), ); From d83af2f51f06375518a1a08f3a67bbd3aa6312d1 Mon Sep 17 00:00:00 2001 From: Aarthy Adityan Date: Tue, 30 Jun 2026 02:40:53 -0400 Subject: [PATCH 75/78] fix(test definitions): display labels from test type yaml --- .../dbsetup_test_types/test_types_Valid_Characters.yaml | 3 ++- .../dbsetup_test_types/test_types_Valid_US_Zip.yaml | 3 ++- .../dbsetup_test_types/test_types_Valid_US_Zip3.yaml | 3 ++- .../ui/components/frontend/js/pages/test_definitions.js | 9 ++++++++- testgen/ui/static/js/components/test_definition_form.js | 3 ++- 5 files changed, 16 insertions(+), 5 deletions(-) diff --git a/testgen/template/dbsetup_test_types/test_types_Valid_Characters.yaml b/testgen/template/dbsetup_test_types/test_types_Valid_Characters.yaml index c3437dcf..d4f5d3e9 100644 --- a/testgen/template/dbsetup_test_types/test_types_Valid_Characters.yaml +++ b/testgen/template/dbsetup_test_types/test_types_Valid_Characters.yaml @@ -21,7 +21,8 @@ test_types: default_parm_columns: threshold_value default_parm_values: |- 0 - default_parm_prompts: null + default_parm_prompts: |- + Threshold Invalid Value Count default_parm_help: |- The acceptable number of records with invalid character values present. default_severity: Warning diff --git a/testgen/template/dbsetup_test_types/test_types_Valid_US_Zip.yaml b/testgen/template/dbsetup_test_types/test_types_Valid_US_Zip.yaml index 2aad44e1..b60bf716 100644 --- a/testgen/template/dbsetup_test_types/test_types_Valid_US_Zip.yaml +++ b/testgen/template/dbsetup_test_types/test_types_Valid_US_Zip.yaml @@ -21,7 +21,8 @@ test_types: default_parm_columns: threshold_value default_parm_values: |- 0 - default_parm_prompts: null + default_parm_prompts: |- + Threshold Invalid Value Count default_parm_help: null default_severity: Warning run_type: CAT diff --git a/testgen/template/dbsetup_test_types/test_types_Valid_US_Zip3.yaml b/testgen/template/dbsetup_test_types/test_types_Valid_US_Zip3.yaml index af148a48..756aefe1 100644 --- a/testgen/template/dbsetup_test_types/test_types_Valid_US_Zip3.yaml +++ b/testgen/template/dbsetup_test_types/test_types_Valid_US_Zip3.yaml @@ -21,7 +21,8 @@ test_types: default_parm_columns: threshold_value default_parm_values: |- 0 - default_parm_prompts: null + default_parm_prompts: |- + Threshold Invalid Value Count default_parm_help: null default_severity: Warning run_type: CAT diff --git a/testgen/ui/components/frontend/js/pages/test_definitions.js b/testgen/ui/components/frontend/js/pages/test_definitions.js index af8d602d..cbd09824 100644 --- a/testgen/ui/components/frontend/js/pages/test_definitions.js +++ b/testgen/ui/components/frontend/js/pages/test_definitions.js @@ -21,6 +21,7 @@ import { ProfilingResultsDialog } from '../shared/profiling_results_dialog.js'; import { AXES, FACET_AXES, GROUP_BY_AXES, EMPTY, appliesToSelectedColumn } from '/app/static/js/components/test_picker_taxonomy.js'; import { enterPage, exitPage, getPageSignal } from '/app/static/js/page_lifecycle.js'; import { jsonObject, maxLength } from '/app/static/js/form_validators.js'; +import { capitalize } from '/app/static/js/display_utils.js'; const { button: btn, div, i: icon, span, strong } = van.tags; @@ -820,6 +821,8 @@ const DetailPanel = (row) => { const paramCols = row.default_parm_columns ? row.default_parm_columns.split(',').map(c => c.trim()).filter(Boolean) : []; + const paramLabels = (row.default_parm_prompts || '').split(',').map(v => v.trim()); + const paramHelp = (row.default_parm_help || '').split('|').map(v => v.trim()); return div( { class: 'flex-column fx-gap-3 border border-radius-1 p-4 mt-2' }, @@ -836,7 +839,11 @@ const DetailPanel = (row) => { Attribute({ label: 'Lock Refresh', value: row.lock_refresh_display }), Attribute({ label: 'Urgency', value: row.urgency }), Attribute({ label: 'Export to Observability', value: row.export_to_observability_display }), - ...paramCols.map(col => Attribute({ label: col, value: String(row[col] ?? '') })), + ...paramCols.map((col, index) => Attribute({ + label: paramLabels[index] || capitalize(col.replaceAll('_', ' ')), + help: paramHelp[index] || null, + value: String(row[col] ?? ''), + })), ), div( { class: 'flex-column fx-flex fx-gap-3' }, diff --git a/testgen/ui/static/js/components/test_definition_form.js b/testgen/ui/static/js/components/test_definition_form.js index ec3265b3..88dae2ba 100644 --- a/testgen/ui/static/js/components/test_definition_form.js +++ b/testgen/ui/static/js/components/test_definition_form.js @@ -73,6 +73,7 @@ import { Textarea } from './textarea.js'; import { RadioGroup } from './radio_group.js'; import { Caption } from './caption.js'; import { numberBetween, required } from '../form_validators.js'; +import { capitalize } from '../display_utils.js'; const { div, span } = van.tags; @@ -108,7 +109,7 @@ const TestDefinitionForm = (/** @type Properties */ props) => { .map((column, index) => ({ ...(PARAMETER_CONFIG[column] || { type: 'text' }), column, - label: paramLabels[index] || column.replaceAll('_', ' '), + label: paramLabels[index] || capitalize(column.replaceAll('_', ' ')), help: paramHelp[index] || null, validators: paramRequired[index] ? [required] : undefined, })) From 26454333ab2661bfc9d62f7dada32446596a3b69 Mon Sep 17 00:00:00 2001 From: Aarthy Adityan Date: Tue, 30 Jun 2026 09:54:43 -0400 Subject: [PATCH 76/78] fix(mcp): correct list_table_groups join, profiling-run PII counts, schedule id display (TG-1135) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - list_table_groups: join job_executions on profiling_runs.id (the shared-PK FK) instead of a non-existent profiling_runs.job_execution_id column, which raised a 500. - get_profiling_run: render the hygiene breakdown from HygieneIssue.count_for_run so Potential PII stays separate from the "possible" likelihood bucket, matching the REST profiling-run issue_counts. - create_test_run_schedule / create_profiling_schedule: flush after save so the default-generated JobSchedule id is populated before rendering (was "Schedule ID: —"). Co-Authored-By: Claude Opus 4.8 (1M context) Claude-Session: https://claude.ai/code/session_01HE8KakxW43wFn3pW5cqV5b --- testgen/common/models/table_group.py | 2 +- testgen/mcp/tools/profiling.py | 21 ++++++++++++++------- testgen/mcp/tools/schedules.py | 2 ++ 3 files changed, 17 insertions(+), 8 deletions(-) diff --git a/testgen/common/models/table_group.py b/testgen/common/models/table_group.py index 5314d4ec..095f54e3 100644 --- a/testgen/common/models/table_group.py +++ b/testgen/common/models/table_group.py @@ -583,7 +583,7 @@ def _list_with_activity( latest_profile AS ( SELECT pr.table_groups_id, MAX(je.started_at) AS started_at FROM profiling_runs pr - LEFT JOIN job_executions je ON je.id = pr.job_execution_id + LEFT JOIN job_executions je ON je.id = pr.id GROUP BY pr.table_groups_id ), latest_test AS ( diff --git a/testgen/mcp/tools/profiling.py b/testgen/mcp/tools/profiling.py index 9e566434..5ef11af9 100644 --- a/testgen/mcp/tools/profiling.py +++ b/testgen/mcp/tools/profiling.py @@ -14,6 +14,7 @@ DataColumnChars, ) from testgen.common.models.data_table import DataTable +from testgen.common.models.hygiene_issue import HygieneIssue from testgen.common.models.job_execution import JobExecution from testgen.common.models.profile_result import ProfileResult from testgen.common.models.profiling_run import ProfilingRun, ProfilingRunSummary @@ -543,13 +544,19 @@ def get_profiling_run(job_execution_id: str) -> str: doc.field("Columns profiled", summary.column_ct or 0) if summary.record_ct is not None: doc.field("Records", summary.record_ct) - doc.field( - "Hygiene issues (confirmed)", - f"{(summary.anomalies_definite_ct or 0) + (summary.anomalies_likely_ct or 0) + (summary.anomalies_possible_ct or 0)} total " - f"— {summary.anomalies_definite_ct or 0} definite, " - f"{summary.anomalies_likely_ct or 0} likely, " - f"{summary.anomalies_possible_ct or 0} possible", - ) + if summary.profiling_run_id: + # Count from the canonical source so likelihood buckets and Potential PII + # stay separate (matches the REST profiling-run issue_counts). + counts = HygieneIssue.count_for_run(summary.profiling_run_id) + hygiene = counts.hygiene_issues + doc.field( + "Hygiene issues (confirmed)", + f"{hygiene.definite + hygiene.likely + hygiene.possible} total " + f"— {hygiene.definite} definite, {hygiene.likely} likely, {hygiene.possible} possible", + ) + pii = counts.potential_pii + if pii.high or pii.moderate: + doc.field("Potential PII", f"{pii.high} high, {pii.moderate} moderate") if summary.dq_score_profiling is not None: doc.field("Profiling Score", friendly_score(summary.dq_score_profiling)) diff --git a/testgen/mcp/tools/schedules.py b/testgen/mcp/tools/schedules.py index 3aa0a021..d1db736e 100644 --- a/testgen/mcp/tools/schedules.py +++ b/testgen/mcp/tools/schedules.py @@ -205,6 +205,7 @@ def create_profiling_schedule( active=active, ) sched.save() + get_current_session().flush() # populate the default-generated id before rendering doc = MdDoc() doc.heading(1, f"Profiling schedule created for `{table_group.table_groups_name}`") @@ -239,6 +240,7 @@ def create_test_run_schedule( active=active, ) sched.save() + get_current_session().flush() # populate the default-generated id before rendering doc = MdDoc() doc.heading(1, f"Test run schedule created for `{suite.test_suite}`") From 098ded036edcf9abdc72e869c9b5b33aaa7c51f3 Mon Sep 17 00:00:00 2001 From: Aarthy Adityan Date: Tue, 30 Jun 2026 10:11:15 -0400 Subject: [PATCH 77/78] test(mcp): update get_profiling_run test for count_for_run-based hygiene rendering (TG-1135) get_profiling_run now reads the hygiene breakdown from HygieneIssue.count_for_run instead of the select_summary buckets, so the unit test mocks count_for_run and asserts the separated likelihood / Potential PII rendering. Co-Authored-By: Claude Opus 4.8 (1M context) Claude-Session: https://claude.ai/code/session_01HE8KakxW43wFn3pW5cqV5b --- tests/unit/mcp/test_tools_profiling.py | 13 ++++++++++++- 1 file changed, 12 insertions(+), 1 deletion(-) diff --git a/tests/unit/mcp/test_tools_profiling.py b/tests/unit/mcp/test_tools_profiling.py index 924cef38..b0136ba1 100644 --- a/tests/unit/mcp/test_tools_profiling.py +++ b/tests/unit/mcp/test_tools_profiling.py @@ -5,6 +5,7 @@ import pytest from testgen.common.models.data_column import ColumnProfileDetail, ColumnProfileSummary, DataColumnChars +from testgen.common.models.hygiene_issue import HygieneIssueCounts, IssueCounts, PotentialPiiCounts from testgen.common.pii_masking import PII_REDACTED from testgen.mcp.exceptions import MCPResourceNotAccessible, MCPUserError from testgen.mcp.permissions import ProjectPermissions @@ -635,8 +636,9 @@ def test_list_profiling_runs_malformed_schedule_raises(mock_tg_cls, mock_run_cls # ---------------------------------------------------------------------- +@patch("testgen.mcp.tools.profiling.HygieneIssue") @patch("testgen.mcp.tools.profiling.ProfilingRun") -def test_get_profiling_run_returns_detail(mock_run_cls, db_session_mock): +def test_get_profiling_run_returns_detail(mock_run_cls, mock_hygiene_cls, db_session_mock): summary = _mock_profiling_run() mock_run_cls.select_summary.return_value = ([summary], 1) mock_run = MagicMock(project_code="demo") @@ -644,6 +646,11 @@ def test_get_profiling_run_returns_detail(mock_run_cls, db_session_mock): mock_run_cls.select_table_breakdown.return_value = [ MagicMock(schema_name="demo", table_name="orders", record_ct=1000, column_ct=5, anomaly_ct=2), ] + mock_hygiene_cls.count_for_run.return_value = IssueCounts( + hygiene_issues=HygieneIssueCounts(definite=2, likely=2, possible=4), + potential_pii=PotentialPiiCounts(high=0, moderate=8), + dismissed=0, + ) with patch("testgen.mcp.permissions._compute_project_permissions") as mock_compute: mock_compute.return_value = ProjectPermissions( @@ -659,6 +666,10 @@ def test_get_profiling_run_returns_detail(mock_run_cls, db_session_mock): assert "Completed" in result assert "Per-table breakdown" in result assert "orders" in result + # Hygiene breakdown comes from count_for_run, keeping Potential PII separate from "possible". + mock_hygiene_cls.count_for_run.assert_called_once_with(summary.profiling_run_id) + assert "8 total — 2 definite, 2 likely, 4 possible" in result + assert "Potential PII:** 0 high, 8 moderate" in result @patch("testgen.mcp.tools.profiling.ProfilingRun") From 0bbbd2e51b56716e86346ee947af47a07777b0db Mon Sep 17 00:00:00 2001 From: Aarthy Adityan Date: Tue, 30 Jun 2026 15:13:25 -0400 Subject: [PATCH 78/78] release: 5.48.0 -> 5.70.2 --- pyproject.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index 92f75a7f..b64fa8fa 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -8,7 +8,7 @@ build-backend = "setuptools.build_meta" [project] name = "dataops-testgen" -version = "5.48.0" +version = "5.70.2" description = "DataKitchen's Data Quality DataOps TestGen" authors = [ { "name" = "DataKitchen, Inc.", "email" = "info@datakitchen.io" },
{{profiling_run.log_message}}
View {{format_number issue_count}} issues > @@ -241,6 +242,8 @@ def send_profiling_run_notifications(profiling_run: ProfilingRun, result_list_ct if not notifications: return + job_execution = profiling_run.job_execution + previous_run = profiling_run.get_previous() issues = list( HygieneIssue.select_with_diff( @@ -251,7 +254,7 @@ def send_profiling_run_notifications(profiling_run: ProfilingRun, result_list_ct ) triggers = {ProfilingRunNotificationTrigger.always} - if profiling_run.status in ("Error", "Cancelled") or {None, True} & {is_new for _, is_new in issues}: + if job_execution.status in (JobStatus.ERROR, JobStatus.CANCELED) or {None, True} & {is_new for _, is_new in issues}: triggers.add(ProfilingRunNotificationTrigger.on_changes) notifications = [ns for ns in notifications if ns.trigger in triggers] @@ -320,10 +323,11 @@ def send_profiling_run_notifications(profiling_run: ProfilingRun, result_list_ct "&source=email" ) ), - "start_time": profiling_run.profiling_starttime, - "end_time": profiling_run.profiling_endtime, - "status": profiling_run.status, - "log_message": profiling_run.log_message, + "start_time": job_execution.started_at, + "end_time": job_execution.completed_at, + "status": job_execution.status, + "status_label": JOB_STATUS_LABEL.get(job_execution.status, job_execution.status), + "log_message": job_execution.error_message, "table_ct": profiling_run.table_ct, "column_ct": profiling_run.column_ct, }, diff --git a/testgen/common/notifications/test_run.py b/testgen/common/notifications/test_run.py index 4f5985a8..9ad2e08e 100644 --- a/testgen/common/notifications/test_run.py +++ b/testgen/common/notifications/test_run.py @@ -3,6 +3,7 @@ from sqlalchemy import case, literal, select from testgen import settings +from testgen.common.enums import JobStatus from testgen.common.models import get_current_session, with_database_session from testgen.common.models.notification_settings import ( TestRunNotificationSettings, @@ -261,7 +262,7 @@ def send_test_run_notifications(test_run: TestRun, result_list_ct=20, result_sta changed_td_id_list.extend(td_id_list) triggers = {TestRunNotificationTrigger.always} - if test_run.status in ("Error", "Cancelled"): + if test_run.job_execution.status in (JobStatus.ERROR, JobStatus.CANCELED): triggers.update(TestRunNotificationTrigger) else: if test_run.error_ct + test_run.failed_ct: diff --git a/tests/unit/common/notifications/test_monitor_run_notifications.py b/tests/unit/common/notifications/test_monitor_run_notifications.py index 248fcd41..9e3c9266 100644 --- a/tests/unit/common/notifications/test_monitor_run_notifications.py +++ b/tests/unit/common/notifications/test_monitor_run_notifications.py @@ -29,6 +29,12 @@ def create_test_result(table_name, test_type, message, result_code=0): return mock +def make_monitor_run(completed_at="2024-01-15T10:30:00Z"): + run = TestRun(id="monitor-run-id", test_suite_id="monitor-suite-id") + run.job_execution = Mock(completed_at=completed_at) + return run + + @pytest.fixture def ns_select_result(): return [ @@ -110,11 +116,7 @@ def test_send_monitor_notifications( send_mock, persisted_setting_mock, ): - test_run = TestRun( - id="monitor-run-id", - test_suite_id="monitor-suite-id", - test_endtime="2024-01-15T10:30:00Z", - ) + test_run = make_monitor_run() table_group = Mock(spec=TableGroup) table_group.id = "tg-id" @@ -193,11 +195,7 @@ def test_send_monitor_notifications_early_exit( test_result_select_where_mock, send_mock, ): - test_run = TestRun( - id="monitor-run-id", - test_suite_id="monitor-suite-id", - test_endtime="2024-01-15T10:30:00Z", - ) + test_run = make_monitor_run() if not has_notifications: ns_select_patched.return_value = [] @@ -219,11 +217,7 @@ def test_send_monitor_notifications_anomaly_counts( send_mock, persisted_setting_mock, ): - test_run = TestRun( - id="monitor-run-id", - test_suite_id="monitor-suite-id", - test_endtime="2024-01-15T10:30:00Z", - ) + test_run = make_monitor_run() table_group = Mock(spec=TableGroup) table_group.id = "tg-id" @@ -270,11 +264,7 @@ def test_send_monitor_notifications_url_construction( send_mock, persisted_setting_mock, ): - test_run = TestRun( - id="monitor-run-id", - test_suite_id="monitor-suite-id", - test_endtime="2024-01-15T10:30:00Z", - ) + test_run = make_monitor_run() table_group = Mock(spec=TableGroup) table_group.id = "tg-123" diff --git a/tests/unit/common/notifications/test_profiling_run_notifications.py b/tests/unit/common/notifications/test_profiling_run_notifications.py index 15cd4e8b..5e985cd2 100644 --- a/tests/unit/common/notifications/test_profiling_run_notifications.py +++ b/tests/unit/common/notifications/test_profiling_run_notifications.py @@ -4,6 +4,7 @@ import pytest +from testgen.common.enums import JOB_STATUS_LABEL, JobStatus from testgen.common.models.hygiene_issue import IssueCount from testgen.common.models.notification_settings import ( ProfilingRunNotificationSettings, @@ -15,6 +16,10 @@ pytestmark = pytest.mark.unit +def make_job_execution(status, started_at=None, completed_at=None, error_message=None): + return Mock(status=status, started_at=started_at, completed_at=completed_at, error_message=error_message) + + def create_ns(**kwargs): with patch("testgen.common.notifications.profiling_run.ProfilingRunNotificationSettings.save"): return ProfilingRunNotificationSettings.create("proj", None, **kwargs) @@ -79,12 +84,12 @@ def get_prev_mock(): @pytest.mark.parametrize( ("profiling_run_status", "has_prev_run", "issue_count", "new_issue_count", "expected_triggers"), ( - ("Error", True, 25, 0, ("always", "on_changes")), - ("Error", True, 0, 0, ("always", "on_changes")), - ("Cancelled", True, 50, 10, ("always", "on_changes")), - ("Complete", True, 50, 10, ("always", "on_changes")), - ("Complete", True, 15, 0, ("always",)), - ("Complete", False, 15, 15, ("always", "on_changes")), + (JobStatus.ERROR, True, 25, 0, ("always", "on_changes")), + (JobStatus.ERROR, True, 0, 0, ("always", "on_changes")), + (JobStatus.CANCELED, True, 50, 10, ("always", "on_changes")), + (JobStatus.COMPLETED, True, 50, 10, ("always", "on_changes")), + (JobStatus.COMPLETED, True, 15, 0, ("always",)), + (JobStatus.COMPLETED, False, 15, 15, ("always", "on_changes")), ), ) def test_send_profiling_run_notification( @@ -104,9 +109,9 @@ def test_send_profiling_run_notification( id="pr-id", job_execution_id="pr-id", table_groups_id="tg-id", - status=profiling_run_status, project_code="proj", ) + profiling_run.job_execution = make_job_execution(profiling_run_status) get_prev_mock.return_value = ProfilingRun(id="pr-prev-id") if has_prev_run else None new_count = iter(count()) priorities = ("Definite", "Likely", "Possible", "High", "Moderate") @@ -144,6 +149,7 @@ def test_send_profiling_run_notification( "start_time": None, "end_time": None, "status": profiling_run_status, + "status_label": JOB_STATUS_LABEL[profiling_run_status], "log_message": None, "table_ct": None, "column_ct": None, diff --git a/tests/unit/common/notifications/test_test_run_notifications.py b/tests/unit/common/notifications/test_test_run_notifications.py index beb7a2e8..3b4a035b 100644 --- a/tests/unit/common/notifications/test_test_run_notifications.py +++ b/tests/unit/common/notifications/test_test_run_notifications.py @@ -3,6 +3,7 @@ import pytest +from testgen.common.enums import JobStatus from testgen.common.models.notification_settings import TestRunNotificationSettings, TestRunNotificationTrigger from testgen.common.models.test_result import TestResultStatus from testgen.common.models.test_run import TestRun @@ -11,6 +12,10 @@ pytestmark = pytest.mark.unit +def make_job_execution(status): + return Mock(status=status) + + def create_ns(**kwargs): with patch("testgen.common.notifications.test_run.TestRunNotificationSettings.save"): return TestRunNotificationSettings.create("proj", None, **kwargs) @@ -98,19 +103,19 @@ def select_summary_mock(): "failed_expected", "warning_expected", "error_expected", "expected_triggers" ), [ - ("Complete", 0, 0, 0, {}, 0, 0, 0, ["always"]), - ("Complete", 0, 5, 0, {}, 0, 5, 0, ["always", "on_warnings"]), - ("Complete", 1, 1, 1, {}, 1, 1, 1, ["always", "on_failures", "on_warnings"]), - ("Complete", 50, 50, 50, {"failed": 2, "warning": 3}, 10, 5, 5, [ + (JobStatus.COMPLETED, 0, 0, 0, {}, 0, 0, 0, ["always"]), + (JobStatus.COMPLETED, 0, 5, 0, {}, 0, 5, 0, ["always", "on_warnings"]), + (JobStatus.COMPLETED, 1, 1, 1, {}, 1, 1, 1, ["always", "on_failures", "on_warnings"]), + (JobStatus.COMPLETED, 50, 50, 50, {"failed": 2, "warning": 3}, 10, 5, 5, [ "always", "on_failures", "on_warnings", "on_changes", ]), - ("Complete", 0, 0, 50, {"error": 50}, 0, 0, 20, ["always", "on_failures", "on_warnings", "on_changes"]), - ("Complete", 50, 0, 0, None, 20, 0, 0, ["always", "on_failures", "on_warnings"]), - ("Complete", 50, 0, 10, {"failed": 5}, 15, 0, 5, ["always", "on_failures", "on_warnings", "on_changes"]), - ("Error", 0, 0, 0, {}, 0, 0, 0, ["always", "on_failures", "on_warnings", "on_changes"]), - ("Error", 20, 10, 0, None, 15, 5, 0, ["always", "on_failures", "on_warnings", "on_changes"]), - ("Cancelled", 0, 0, 0, {}, 0, 0, 0, ["always", "on_failures", "on_warnings", "on_changes"]), - ("Cancelled", 30, 20, 0, {}, 15, 5, 0, ["always", "on_failures", "on_warnings", "on_changes"]), + (JobStatus.COMPLETED, 0, 0, 50, {"error": 50}, 0, 0, 20, ["always", "on_failures", "on_warnings", "on_changes"]), + (JobStatus.COMPLETED, 50, 0, 0, None, 20, 0, 0, ["always", "on_failures", "on_warnings"]), + (JobStatus.COMPLETED, 50, 0, 10, {"failed": 5}, 15, 0, 5, ["always", "on_failures", "on_warnings", "on_changes"]), + (JobStatus.ERROR, 0, 0, 0, {}, 0, 0, 0, ["always", "on_failures", "on_warnings", "on_changes"]), + (JobStatus.ERROR, 20, 10, 0, None, 15, 5, 0, ["always", "on_failures", "on_warnings", "on_changes"]), + (JobStatus.CANCELED, 0, 0, 0, {}, 0, 0, 0, ["always", "on_failures", "on_warnings", "on_changes"]), + (JobStatus.CANCELED, 30, 20, 0, {}, 15, 5, 0, ["always", "on_failures", "on_warnings", "on_changes"]), ] ) def test_send_test_run_notification( @@ -135,12 +140,12 @@ def test_send_test_run_notification( test_run = TestRun( id="tr-id", job_execution_id="tr-id", - status=test_run_status, test_suite_id="ts-id", failed_ct=failed_ct, warning_ct=warning_ct, error_ct=error_ct, ) + test_run.job_execution = make_job_execution(test_run_status) # SA 2.0 Row objects expose ._mapping; mock them accordingly def _make_row(): diff --git a/tests/unit/common/test_enums.py b/tests/unit/common/test_enums.py new file mode 100644 index 00000000..2525839d --- /dev/null +++ b/tests/unit/common/test_enums.py @@ -0,0 +1,18 @@ +import pytest + +from testgen.common.enums import JOB_STATUS_LABEL, JobStatus +from testgen.common.models.profiling_run import ProfilingRunSummary +from testgen.common.models.test_run import TestRunSummary + +pytestmark = pytest.mark.unit + + +def test_job_status_label_covers_every_status(): + # A missing entry would surface a raw lowercase status code (e.g. "cancel_requested") + # in run lists, emails, and MCP output instead of a display label. + assert set(JOB_STATUS_LABEL) == set(JobStatus) + + +def test_run_summaries_share_the_job_status_label_map(): + assert TestRunSummary.STATUS_LABEL is JOB_STATUS_LABEL + assert ProfilingRunSummary.STATUS_LABEL is JOB_STATUS_LABEL From ec7e4ab19c4bc46e7e4de5f72045e7269ac3ecc0 Mon Sep 17 00:00:00 2001 From: Diogo Basto Date: Fri, 22 May 2026 15:48:35 +0100 Subject: [PATCH 09/78] feat(test-definitions): support copy/move across projects (TG-1075) Add a Target Project dropdown to the Copy/Move Test dialog so tests can be copied or moved into table groups in another project. Switching the project re-scopes the Target Table Group and Test Suite options and resets downstream selections. - Adds Authentication.get_projects_with_permission as the hook that populates the dropdown. Default returns all projects. - on_copy_confirmed and on_move_confirmed re-resolve the target table group's project_code and revalidate the user's permission server-side before mutating, since the frontend selection is not trustworthy. Co-Authored-By: Claude Opus 4.7 --- testgen/ui/auth.py | 4 ++ .../frontend/js/pages/test_definitions.js | 27 ++++++++++- testgen/ui/views/test_definitions.py | 48 ++++++++++++++++--- 3 files changed, 70 insertions(+), 9 deletions(-) diff --git a/testgen/ui/auth.py b/testgen/ui/auth.py index 8bb5b788..d86fd154 100644 --- a/testgen/ui/auth.py +++ b/testgen/ui/auth.py @@ -12,6 +12,7 @@ from testgen.ui.services.query_cache import ( get_membership_by_user_and_project, get_user, + select_projects_where, select_users_where, ) from testgen.ui.session import session @@ -55,6 +56,9 @@ def user_has_permission(self, permission: Permission, /, project_code: str | Non def user_has_project_access(self, project_code: str) -> bool: # noqa: ARG002 return True + def get_projects_with_permission(self, permission: Permission, /) -> list[str]: # noqa: ARG002 + return [p.project_code for p in select_projects_where()] + def get_jwt_hashing_key(self) -> bytes: try: return get_jwt_signing_key() diff --git a/testgen/ui/components/frontend/js/pages/test_definitions.js b/testgen/ui/components/frontend/js/pages/test_definitions.js index 47bf21d2..7e83fbcc 100644 --- a/testgen/ui/components/frontend/js/pages/test_definitions.js +++ b/testgen/ui/components/frontend/js/pages/test_definitions.js @@ -1313,6 +1313,7 @@ const CopyMoveDialogComponent = ({ open, info, onClose }, emit) => { const dialogInfo = van.derive(() => getValue(info) ?? null); const collision = van.derive(() => dialogInfo.val?.collision ?? null); + const targetProjectCode = van.state(null); const targetTgId = van.state(null); const targetTsId = van.state(null); const targetTableName = van.state(null); @@ -1324,6 +1325,7 @@ const CopyMoveDialogComponent = ({ open, info, onClose }, emit) => { const isOpen = open.val; if (isOpen && !wasOpen.val) { const di = dialogInfo.val; + targetProjectCode.val = di?.current_project_code ?? null; targetTgId.val = di?.current_table_group_id ?? null; targetTsId.val = null; targetTableName.val = null; @@ -1334,10 +1336,16 @@ const CopyMoveDialogComponent = ({ open, info, onClose }, emit) => { } }); - const tableGroupOptions = van.derive(() => - (dialogInfo.val?.table_groups ?? []).map(tg => ({ label: tg.table_groups_name, value: tg.id })) + const projectOptions = van.derive(() => + (dialogInfo.val?.projects ?? []).map(p => ({ label: p.project_name, value: p.project_code })) ); + const tableGroupOptions = van.derive(() => { + const project = targetProjectCode.val; + const tgs = dialogInfo.val?.table_groups_by_project?.[project] ?? []; + return tgs.map(tg => ({ label: tg.table_groups_name, value: tg.id })); + }); + const testSuiteOptions = van.derive(() => { const tg = targetTgId.val; const suites = dialogInfo.val?.test_suites_by_table_group?.[tg] ?? []; @@ -1414,6 +1422,21 @@ const CopyMoveDialogComponent = ({ open, info, onClose }, emit) => { { class: 'flex-column fx-gap-4 td-form-dialog' }, () => div({ class: 'text-caption' }, `Selected tests: ${(dialogInfo.val?.selected ?? []).length}`), + () => Select({ + label: 'Target Project', + value: targetProjectCode.val, + options: projectOptions.val, + required: true, + filterable: true, + onChange: (value) => { + targetProjectCode.val = value; + targetTgId.val = null; + targetTsId.val = null; + targetTableName.val = null; + targetColumnName.val = null; + }, + }), + () => Select({ label: 'Target Table Group', value: targetTgId.val, diff --git a/testgen/ui/views/test_definitions.py b/testgen/ui/views/test_definitions.py index f460d873..c2d5843c 100644 --- a/testgen/ui/views/test_definitions.py +++ b/testgen/ui/views/test_definitions.py @@ -12,6 +12,7 @@ from testgen.common.enums import JobSource from testgen.common.models import with_database_session from testgen.common.models.job_execution import JobExecution +from testgen.common.models.project import Project from testgen.common.models.table_group import TableGroup, TableGroupMinimal from testgen.common.models.test_definition import ( TestDefinition, @@ -37,6 +38,7 @@ get_connection, get_table_group_minimal, get_test_suite, + select_projects_where, select_table_groups_minimal_where, select_test_definitions_minimal_where, select_test_definitions_page, @@ -157,11 +159,6 @@ def render( test_types = run_test_type_lookup_query().to_dict("records") table_columns = get_columns(str(table_group.id)) filter_columns_df = get_test_suite_columns(test_suite_id) - table_groups = select_table_groups_minimal_where(TableGroup.project_code == project_code) - all_test_suites = select_test_suites_minimal_where( - TestSuite.table_groups_id.in_([str(tg.id) for tg in table_groups]), - TestSuite.is_monitor.isnot(True), - ) # Build filter options table_options = sorted(filter_columns_df["table_name"].dropna().unique().tolist(), key=str.lower) @@ -230,15 +227,38 @@ def render( copy_move_dialog = None if selected := st.session_state.get(TD_COPY_MOVE_DIALOG_KEY): + editable_project_codes = session.auth.get_projects_with_permission("edit") + editable_projects = ( + select_projects_where(Project.project_code.in_(editable_project_codes)) + if editable_project_codes + else [] + ) + editable_table_groups = select_table_groups_minimal_where( + TableGroup.project_code.in_(editable_project_codes) + ) + editable_test_suites = select_test_suites_minimal_where( + TestSuite.table_groups_id.in_([str(tg.id) for tg in editable_table_groups]), + TestSuite.is_monitor.isnot(True), + ) + tgs_by_project: dict[str, list] = {} + for tg in editable_table_groups: + tgs_by_project.setdefault(tg.project_code, []).append( + {"id": str(tg.id), "table_groups_name": tg.table_groups_name} + ) suites_by_tg: dict[str, list] = {} - for ts in all_test_suites: + for ts in editable_test_suites: suites_by_tg.setdefault(str(ts.table_groups_id), []).append( {"id": str(ts.id), "test_suite": ts.test_suite} ) copy_move_dialog = { "open": True, "selected": selected, - "table_groups": [{"id": str(tg.id), "table_groups_name": tg.table_groups_name} for tg in table_groups], + "projects": [ + {"project_code": p.project_code, "project_name": p.project_name} + for p in editable_projects + ], + "current_project_code": project_code, + "table_groups_by_project": tgs_by_project, "current_table_group_id": str(table_group.id), "current_test_suite_id": str(test_suite.id), "test_suites_by_table_group": suites_by_tg, @@ -402,6 +422,10 @@ def on_update_attribute_all(payload: dict) -> None: TestDefinition.set_status_attribute(attribute, all_ids, value) st.cache_data.clear() + def _resolve_target_project(target_tg_id: str) -> str | None: + target_tg = TableGroup.get_minimal(target_tg_id) + return target_tg.project_code if target_tg else None + @with_database_session def on_copy_confirmed(payload: dict) -> None: ids = payload["ids"] @@ -409,6 +433,11 @@ def on_copy_confirmed(payload: dict) -> None: target_ts_id = payload["target_test_suite_id"] target_table = payload.get("target_table_name") target_col = payload.get("target_column_name") + target_project = _resolve_target_project(target_tg_id) + if not target_project or not session.auth.user_has_permission("edit", target_project): + LOG.warning("Refusing copy to table group %s — user lacks edit permission", target_tg_id) + st.toast("You don't have edit permission for the target project.", icon=":material/error:") + return overwrite_ids = st.session_state.pop(TD_COPY_MOVE_OVERWRITE_KEY, []) if overwrite_ids: TestDefinition.delete_where(TestDefinition.id.in_(overwrite_ids)) @@ -426,6 +455,11 @@ def on_move_confirmed(payload: dict) -> None: target_ts_id = payload["target_test_suite_id"] target_table = payload.get("target_table_name") target_col = payload.get("target_column_name") + target_project = _resolve_target_project(target_tg_id) + if not target_project or not session.auth.user_has_permission("edit", target_project): + LOG.warning("Refusing move to table group %s — user lacks edit permission", target_tg_id) + st.toast("You don't have edit permission for the target project.", icon=":material/error:") + return overwrite_ids = st.session_state.pop(TD_COPY_MOVE_OVERWRITE_KEY, []) if overwrite_ids: TestDefinition.delete_where(TestDefinition.id.in_(overwrite_ids)) From f2ec74264032b2ad5a7a867d818bdd1ba3649912 Mon Sep 17 00:00:00 2001 From: Ricardo Boni Date: Thu, 4 Jun 2026 17:21:32 -0400 Subject: [PATCH 10/78] fix(mcp): consistent error surfacing and scope validation - get_failure_summary: validate group_by against the allowed set instead of falling through to a generic error - list_test_suites: use the standard not-found-or-not-accessible wording on access denial; keep the specific message for the access-granted empty case - get_failure_summary/get_failure_trend: reject a cross-project or inaccessible test suite / table group instead of silently returning empty; brings get_failure_trend up to the same validation bar - MCP auth boundary: surface a clean authentication error when the token's user no longer exists or the token was revoked, instead of a generic "unexpected error" plus an error-level traceback. Introduce an AuthError domain exception raised by decode_jwt_token/authorize_token; REST's 401 path is unchanged TG-1119 Co-Authored-By: Claude Opus 4.8 (1M context) --- testgen/api/deps.py | 6 +- testgen/common/auth.py | 12 ++- testgen/mcp/exceptions.py | 8 ++ testgen/mcp/permissions.py | 15 +++- testgen/mcp/server.py | 4 +- testgen/mcp/tools/common.py | 62 ++++++++++++++++ testgen/mcp/tools/discovery.py | 3 +- testgen/mcp/tools/test_results.py | 46 ++++++------ tests/unit/common/test_auth.py | 9 ++- tests/unit/mcp/test_auth.py | 5 +- tests/unit/mcp/test_permissions.py | 17 +++-- tests/unit/mcp/test_tools_discovery.py | 2 +- tests/unit/mcp/test_tools_test_results.py | 89 +++++++++++++++++++++-- 13 files changed, 219 insertions(+), 59 deletions(-) diff --git a/testgen/api/deps.py b/testgen/api/deps.py index e1f834ec..e59d783d 100644 --- a/testgen/api/deps.py +++ b/testgen/api/deps.py @@ -6,7 +6,7 @@ from fastapi.security import HTTPAuthorizationCredentials, HTTPBearer from sqlalchemy import select -from testgen.common.auth import authorize_token, decode_jwt_token +from testgen.common.auth import AuthError, authorize_token, decode_jwt_token from testgen.common.enums import PublicJobKey from testgen.common.models import Session, _current_session_wrapper, get_current_session from testgen.common.models.job_execution import JobExecution @@ -48,7 +48,7 @@ def get_authorized_user(credentials: HTTPAuthorizationCredentials = _bearer_secu try: payload = decode_jwt_token(credentials.credentials) - except ValueError: + except AuthError: raise _invalid from None username = payload.get("username") @@ -58,7 +58,7 @@ def get_authorized_user(credentials: HTTPAuthorizationCredentials = _bearer_secu session = get_current_session() try: return authorize_token(credentials.credentials, username, session) - except ValueError: + except AuthError: raise _invalid from None diff --git a/testgen/common/auth.py b/testgen/common/auth.py index 1dcdc479..141b032e 100644 --- a/testgen/common/auth.py +++ b/testgen/common/auth.py @@ -10,6 +10,10 @@ LOG = logging.getLogger("testgen") +class AuthError(Exception): + """Token authentication failed — invalid/expired token, unknown user, or revoked token.""" + + def get_jwt_signing_key() -> bytes: """Decode the base64-encoded JWT signing key from settings.""" return base64.b64decode(settings.JWT_HASHING_KEY_B64.encode("ascii")) @@ -27,13 +31,13 @@ def create_jwt_token(username: str, expiry_seconds: int = 86400) -> str: def decode_jwt_token(token_str: str) -> dict: """Decode and validate a JWT token. Returns the payload dict. - Raises ValueError if the token is invalid or expired. + Raises ``AuthError`` if the token is invalid or expired. PyJWT auto-validates the standard ``exp`` claim during decode. """ try: return jwt.decode(token_str, get_jwt_signing_key(), algorithms=["HS256"]) except jwt.InvalidTokenError as e: - raise ValueError(f"Invalid token: {e}") from e + raise AuthError(f"Invalid token: {e}") from e def authorize_token(token_str: str, username: str, session): @@ -48,11 +52,11 @@ def authorize_token(token_str: str, username: str, session): user = session.scalars(select(User).where(func.lower(User.username) == func.lower(username))).first() if user is None: - raise ValueError("User not found") + raise AuthError("User not found") token_record = session.scalars(select(OAuth2Token).where(OAuth2Token.access_token == token_str)).first() if token_record and token_record.access_token_revoked_at: - raise ValueError("Token has been revoked") + raise AuthError("Token has been revoked") return user diff --git a/testgen/mcp/exceptions.py b/testgen/mcp/exceptions.py index fc89f98a..713f5b40 100644 --- a/testgen/mcp/exceptions.py +++ b/testgen/mcp/exceptions.py @@ -20,6 +20,14 @@ class MCPUserError(Exception): """ +class MCPAuthenticationError(MCPUserError): + """Token authenticated at the transport but failed authorization at the tool boundary. + + Raised when the token's user no longer exists or the token was revoked — distinct from + per-project permission denial. Carries a uniform, re-authenticate message for the client. + """ + + class MCPPermissionDenied(MCPUserError): """Raised when access is denied due to insufficient project permissions.""" diff --git a/testgen/mcp/permissions.py b/testgen/mcp/permissions.py index 0850c753..77a597ba 100644 --- a/testgen/mcp/permissions.py +++ b/testgen/mcp/permissions.py @@ -2,14 +2,17 @@ import contextvars import functools +import logging from collections.abc import Callable from dataclasses import dataclass from testgen.common.models.project_membership import ProjectMembership from testgen.common.models.user import User -from testgen.mcp.exceptions import MCPPermissionDenied +from testgen.mcp.exceptions import MCPAuthenticationError, MCPPermissionDenied from testgen.utils.plugins import PluginHook +LOG = logging.getLogger("testgen") + _NOT_SET = object() _mcp_username: contextvars.ContextVar[str | None] = contextvars.ContextVar("mcp_username", default=None) @@ -83,7 +86,7 @@ def get_authorized_mcp_user() -> User: Checks user existence and token revocation status. Must be called within @with_database_session scope. """ - from testgen.common.auth import authorize_token + from testgen.common.auth import AuthError, authorize_token from testgen.common.models import get_current_session username = _mcp_username.get() @@ -92,7 +95,13 @@ def get_authorized_mcp_user() -> User: token_str = _mcp_token.get() session = get_current_session() - return authorize_token(token_str or "", username, session) + try: + return authorize_token(token_str or "", username, session) + except AuthError as err: + LOG.warning("MCP token authorization failed: %s", err) + raise MCPAuthenticationError( + "Authentication failed: your access token is no longer valid. Please sign in again." + ) from err def _compute_project_permissions(user: User, permission: str) -> ProjectPermissions: diff --git a/testgen/mcp/server.py b/testgen/mcp/server.py index 17f555d3..f7f474e3 100644 --- a/testgen/mcp/server.py +++ b/testgen/mcp/server.py @@ -9,7 +9,7 @@ from starlette.applications import Starlette from testgen import settings -from testgen.common.auth import decode_jwt_token +from testgen.common.auth import AuthError, decode_jwt_token from testgen.mcp.permissions import set_mcp_token, set_mcp_username LOG = logging.getLogger("testgen") @@ -65,7 +65,7 @@ async def verify_token(self, token: str) -> AccessToken | None: scopes=[], expires_at=int(payload["exp"]), ) - except (ValueError, KeyError): + except (AuthError, KeyError): return None diff --git a/testgen/mcp/tools/common.py b/testgen/mcp/tools/common.py index c6f649ec..6c71c481 100644 --- a/testgen/mcp/tools/common.py +++ b/testgen/mcp/tools/common.py @@ -300,6 +300,22 @@ def parse_run_status_filter(value: str) -> list[JobStatus]: return statuses +class FailureGroupBy(StrEnum): + """User-facing values accepted for the ``group_by`` argument on ``get_failure_summary``.""" + + TEST_TYPE = "test_type" + TABLE = "table" + COLUMN = "column" + + +def parse_failure_group_by(value: str) -> FailureGroupBy: + try: + return FailureGroupBy(value) + except ValueError as err: + valid = ", ".join(g.value for g in FailureGroupBy) + raise MCPUserError(f"Invalid group_by `{value}`. Valid values: {valid}") from err + + def format_run_duration(started_at: datetime | None, completed_at: datetime | None) -> str | None: """Render an elapsed duration as ``Xs`` / ``Xm Ys`` / ``Xh Ym``. Returns ``None`` if either bound is missing.""" if not started_at or not completed_at: @@ -661,6 +677,52 @@ def resolve_notification(notification_id: str) -> NotificationSettings: return notif +def resolve_aggregate_scope( + project_code: str | None, + test_suite_id: str | None = None, + table_group_id: str | None = None, +) -> list[str]: + """Validate optional project / test-suite / table-group scope for a cross-run aggregation. + + Resolves any supplied suite or table group (existence + access via the + ``resolve_*`` helpers), and — when ``project_code`` is also given — requires the + resolved entity to belong to it, raising a clear error on a cross-project mismatch + so the query never silently returns empty. Returns the project codes to scope the + aggregation to. + """ + perms = get_project_permissions() + if project_code: + perms.verify_access(project_code, not_found=MCPResourceNotAccessible("Project", project_code)) + + scoped_projects: set[str] = set() + if test_suite_id: + suite = resolve_test_suite(test_suite_id) + if project_code and suite.project_code != project_code: + raise MCPUserError( + f"Test suite `{test_suite_id}` belongs to project `{suite.project_code}`, not `{project_code}`." + ) + scoped_projects.add(suite.project_code) + if table_group_id: + table_group = resolve_table_group(table_group_id) + if project_code and table_group.project_code != project_code: + raise MCPUserError( + f"Table group `{table_group_id}` belongs to project `{table_group.project_code}`, not `{project_code}`." + ) + scoped_projects.add(table_group.project_code) + + if len(scoped_projects) > 1: + # Suite and table group resolve to different projects (only reachable when no + # project_code pins the scope). The two filters would AND to an empty result — + # reject rather than silently return nothing. + raise MCPUserError("The test suite and table group belong to different projects — narrow to one scope.") + + if project_code: + return [project_code] + if scoped_projects: + return list(scoped_projects) + return perms.allowed_codes + + # Notification event-type labels. class NotificationEventLabel(StrEnum): diff --git a/testgen/mcp/tools/discovery.py b/testgen/mcp/tools/discovery.py index 05b7ab5b..978cc858 100644 --- a/testgen/mcp/tools/discovery.py +++ b/testgen/mcp/tools/discovery.py @@ -4,6 +4,7 @@ from testgen.common.models.project import Project from testgen.common.models.test_run import TestRun from testgen.common.models.test_suite import TestSuite +from testgen.mcp.exceptions import MCPResourceNotAccessible from testgen.mcp.permissions import get_project_permissions, mcp_permission from testgen.mcp.tools.common import DocGroup, resolve_table_group, validate_limit, validate_page from testgen.mcp.tools.markdown import MdDoc @@ -63,7 +64,7 @@ def list_test_suites(project_code: str) -> str: return "Missing required parameter `project_code`." perms = get_project_permissions() - perms.verify_access(project_code, not_found=f"No test suites found for project `{project_code}`.") + perms.verify_access(project_code, not_found=MCPResourceNotAccessible("Project", project_code)) summaries = TestSuite.select_summary(project_code) diff --git a/testgen/mcp/tools/test_results.py b/testgen/mcp/tools/test_results.py index d35d3f0a..b5b5b1c3 100644 --- a/testgen/mcp/tools/test_results.py +++ b/testgen/mcp/tools/test_results.py @@ -12,11 +12,14 @@ from testgen.mcp.permissions import get_project_permissions, mcp_permission from testgen.mcp.tools.common import ( DocGroup, + FailureGroupBy, format_page_footer, format_page_info, + parse_failure_group_by, parse_result_status, parse_since_arg, parse_uuid, + resolve_aggregate_scope, resolve_test_type, validate_limit, validate_page, @@ -27,6 +30,12 @@ _DEFAULT_SEARCH_STATUSES = [TestResultStatus.Failed, TestResultStatus.Warning] +_MODEL_GROUP_COLUMN = { + FailureGroupBy.TEST_TYPE: "test_type", + FailureGroupBy.TABLE: "table_name", + FailureGroupBy.COLUMN: "column_names", +} + @with_database_session @mcp_permission("view") @@ -167,20 +176,20 @@ def get_failure_summary( group_by: Group failures by 'test_type', 'table', or 'column' (default: 'test_type'). """ perms = get_project_permissions() + group = parse_failure_group_by(group_by) if not any((job_execution_id, test_suite_id, since)): raise MCPUserError( "Provide 'job_execution_id' for a single run, or 'test_suite_id' or 'project_code' " "to aggregate across runs. 'since' is required when 'test_suite_id' is not provided." ) - if group_by in ("table", "column") and not (job_execution_id or test_suite_id): + if group in (FailureGroupBy.TABLE, FailureGroupBy.COLUMN) and not (job_execution_id or test_suite_id): raise MCPUserError( - f"'{group_by}' grouping requires a single-suite scope. " + f"'{group}' grouping requires a single-suite scope. " "Provide 'job_execution_id' or 'test_suite_id'." ) - model_group_map = {"table": "table_name", "column": "column_names"} - model_group_by = model_group_map.get(group_by, group_by) + model_group_by = _MODEL_GROUP_COLUMN[group] scope_label: str test_run_id = None @@ -197,15 +206,7 @@ def get_failure_summary( scope_label = f"run `{job_execution_id}`" project_codes = perms.allowed_codes else: - if project_code: - perms.verify_access(project_code, not_found=MCPResourceNotAccessible("Project", project_code)) - project_codes = [project_code] - else: - project_codes = perms.allowed_codes - if test_suite_uuid is not None: - suite = TestSuite.get_regular(test_suite_uuid) - if suite is None or not perms.has_access(suite.project_code): - raise MCPResourceNotAccessible("Test suite", test_suite_id) + project_codes = resolve_aggregate_scope(project_code, test_suite_id=test_suite_id) scope_parts = [] if project_code: scope_parts.append(f"project `{project_code}`") @@ -227,14 +228,14 @@ def get_failure_summary( return f"No confirmed failures found for {scope_label}." total = sum(row[-1] for row in failures) - if group_by == "test_type": + if group is FailureGroupBy.TEST_TYPE: type_names = {tt.test_type: tt.test_name_short for tt in TestType.select_where(TestType.active == "Y")} doc = MdDoc() doc.heading(1, f"Failure Summary — {scope_label}") doc.text(f"**Total confirmed failures (Failed + Warning):** {total}") - if group_by == "test_type": + if group is FailureGroupBy.TEST_TYPE: headers = ["Test Type", "Severity", "Count"] rows = [] for row in failures: @@ -242,7 +243,7 @@ def get_failure_summary( name = type_names.get(code, code) severity = status.value if status else "Unknown" rows.append([name, severity, count]) - elif group_by == "column": + elif group is FailureGroupBy.COLUMN: headers = ["Column", "Count"] rows = [] for row in failures: @@ -253,9 +254,9 @@ def get_failure_summary( headers = ["Table Name", "Count"] rows = [[row[0], row[-1]] for row in failures] - doc.table(headers, rows, code=[0] if group_by == "table" else None) + doc.table(headers, rows, code=[0] if group is FailureGroupBy.TABLE else None) - if group_by == "test_type": + if group is FailureGroupBy.TEST_TYPE: doc.text( "Check `testgen://test-types` to understand what each test type checks " "and `get_test_type(test_type='...')` to fetch more details." @@ -450,12 +451,9 @@ def get_failure_trend( valid = ", ".join(v.value for v in BucketInterval) raise MCPUserError(f"Invalid `bucket`: `{bucket}`. Valid values: {valid}") from err - perms = get_project_permissions() - if project_code: - perms.verify_access(project_code, not_found=MCPResourceNotAccessible("Project", project_code)) - project_codes = [project_code] - else: - project_codes = perms.allowed_codes + project_codes = resolve_aggregate_scope( + project_code, test_suite_id=test_suite_id, table_group_id=table_group_id + ) anchor_today = datetime.now(UTC).date() if exclude_today: diff --git a/tests/unit/common/test_auth.py b/tests/unit/common/test_auth.py index 3d6bab34..6e926558 100644 --- a/tests/unit/common/test_auth.py +++ b/tests/unit/common/test_auth.py @@ -7,6 +7,7 @@ import pytest from testgen.common.auth import ( + AuthError, authorize_token, check_permission, create_jwt_token, @@ -50,14 +51,14 @@ def test_decode_jwt_token_decodes_valid_token(mock_settings): def test_decode_jwt_token_raises_for_expired_token(mock_settings): mock_settings.JWT_HASHING_KEY_B64 = JWT_KEY token = _make_token(exp_seconds=-3600) - with pytest.raises(ValueError, match="Invalid token"): + with pytest.raises(AuthError, match="Invalid token"): decode_jwt_token(token) @patch("testgen.common.auth.settings") def test_decode_jwt_token_raises_for_invalid_token(mock_settings): mock_settings.JWT_HASHING_KEY_B64 = JWT_KEY - with pytest.raises(ValueError, match="Invalid token"): + with pytest.raises(AuthError, match="Invalid token"): decode_jwt_token("not-a-valid-token") @@ -99,7 +100,7 @@ def test_authorize_token_rejects_revoked(): mock_token_record.access_token_revoked_at = 1700000000 _set_scalars_results(mock_session, mock_user, mock_token_record) - with pytest.raises(ValueError, match="Token has been revoked"): + with pytest.raises(AuthError, match="Token has been revoked"): authorize_token("revoked_token", "testuser", mock_session) @@ -118,7 +119,7 @@ def test_authorize_token_raises_when_user_not_found(): # User lookup returns None — token check is never reached _set_scalars_results(mock_session, None) - with pytest.raises(ValueError, match="User not found"): + with pytest.raises(AuthError, match="User not found"): authorize_token("some_token", "ghost", mock_session) diff --git a/tests/unit/mcp/test_auth.py b/tests/unit/mcp/test_auth.py index 29cafbdb..e8125948 100644 --- a/tests/unit/mcp/test_auth.py +++ b/tests/unit/mcp/test_auth.py @@ -7,6 +7,7 @@ import jwt import pytest +from testgen.common.auth import AuthError from testgen.mcp.auth import authenticate_user, validate_token from testgen.mcp.server import JWTTokenVerifier @@ -83,7 +84,7 @@ def test_validate_token_returns_user(mock_user_cls, mock_settings): def test_validate_token_raises_for_expired_token(mock_settings): mock_settings.JWT_HASHING_KEY_B64 = JWT_KEY - with pytest.raises(ValueError, match="Invalid token"): + with pytest.raises(AuthError, match="Invalid token"): validate_token(_make_token(exp_seconds=-3600)) @@ -91,7 +92,7 @@ def test_validate_token_raises_for_expired_token(mock_settings): def test_validate_token_raises_for_invalid_token(mock_settings): mock_settings.JWT_HASHING_KEY_B64 = JWT_KEY - with pytest.raises(ValueError, match="Invalid token"): + with pytest.raises(AuthError, match="Invalid token"): validate_token("not-a-valid-token") diff --git a/tests/unit/mcp/test_permissions.py b/tests/unit/mcp/test_permissions.py index 0f97cbe4..bad81379 100644 --- a/tests/unit/mcp/test_permissions.py +++ b/tests/unit/mcp/test_permissions.py @@ -3,7 +3,8 @@ import pytest -from testgen.mcp.exceptions import MCPPermissionDenied, MCPResourceNotAccessible +from testgen.common.auth import AuthError +from testgen.mcp.exceptions import MCPAuthenticationError, MCPPermissionDenied, MCPResourceNotAccessible from testgen.mcp.permissions import ( _NOT_SET, ProjectPermissions, @@ -38,12 +39,14 @@ def test_get_authorized_mcp_user_raises_when_no_username(): @patch("testgen.common.models.get_current_session") @patch("testgen.common.auth.authorize_token") -def test_get_authorized_mcp_user_raises_when_user_not_found(mock_authorize, mock_get_session): - mock_authorize.side_effect = ValueError("User not found") +def test_get_authorized_mcp_user_translates_unknown_user_to_auth_error(mock_authorize, mock_get_session): + """A valid-signature token whose user no longer exists surfaces a clean MCP auth error, + not the generic 'unexpected error' the boundary returns for raw exceptions.""" + mock_authorize.side_effect = AuthError("User not found") set_mcp_username("ghost") set_mcp_token("some_token") - with pytest.raises(ValueError, match="User not found"): + with pytest.raises(MCPAuthenticationError, match="no longer valid"): get_authorized_mcp_user() @@ -65,12 +68,12 @@ def test_get_authorized_mcp_user_returns_user(mock_authorize, mock_get_session): @patch("testgen.common.models.get_current_session") @patch("testgen.common.auth.authorize_token") -def test_get_authorized_mcp_user_rejects_revoked_token(mock_authorize, mock_get_session): - mock_authorize.side_effect = ValueError("Token has been revoked") +def test_get_authorized_mcp_user_translates_revoked_token_to_auth_error(mock_authorize, mock_get_session): + mock_authorize.side_effect = AuthError("Token has been revoked") set_mcp_username("admin") set_mcp_token("revoked_token") - with pytest.raises(ValueError, match="Token has been revoked"): + with pytest.raises(MCPAuthenticationError, match="no longer valid"): get_authorized_mcp_user() diff --git a/tests/unit/mcp/test_tools_discovery.py b/tests/unit/mcp/test_tools_discovery.py index 40146b21..409b20f2 100644 --- a/tests/unit/mcp/test_tools_discovery.py +++ b/tests/unit/mcp/test_tools_discovery.py @@ -181,7 +181,7 @@ def test_list_test_suites_raises_not_found_for_inaccessible_project( from testgen.mcp.tools.discovery import list_test_suites - with pytest.raises(MCPPermissionDenied, match="No test suites found for project `secret_project`"): + with pytest.raises(MCPResourceNotAccessible, match="Project `secret_project` not found or not accessible"): list_test_suites("secret_project") diff --git a/tests/unit/mcp/test_tools_test_results.py b/tests/unit/mcp/test_tools_test_results.py index 7b3d08e4..8ded1b5c 100644 --- a/tests/unit/mcp/test_tools_test_results.py +++ b/tests/unit/mcp/test_tools_test_results.py @@ -625,7 +625,7 @@ def test_get_failure_summary_cross_run_by_project(mock_compute, mock_result, db_ assert call_kwargs["since"] is not None -@patch("testgen.mcp.tools.test_results.TestSuite") +@patch("testgen.mcp.tools.common.TestSuite") @patch("testgen.mcp.tools.test_results.TestResult") @patch("testgen.mcp.permissions._compute_project_permissions") def test_get_failure_summary_cross_run_by_project_and_suite( @@ -636,7 +636,7 @@ def test_get_failure_summary_cross_run_by_project_and_suite( permission="view", username="test_user", ) - mock_suite_cls.get_regular.return_value = _mock_test_suite(project_code="proj_a") + mock_suite_cls.get.return_value = _mock_test_suite(project_code="proj_a") mock_result.select_failures.return_value = [] from testgen.mcp.tools.test_results import get_failure_summary @@ -664,12 +664,13 @@ def test_get_failure_summary_rejects_inaccessible_project(mock_compute, db_sessi get_failure_summary(project_code="proj_b", since="7 days") -@patch("testgen.mcp.tools.test_results.TestSuite") +@patch("testgen.mcp.tools.common.TestSuite") @patch("testgen.mcp.permissions._compute_project_permissions") def test_get_failure_summary_rejects_inaccessible_test_suite(mock_compute, mock_suite_cls, db_session_mock): - """test_suite_id branch validates suite access — same contract as list_test_results.""" + """test_suite_id branch validates suite access — inaccessible suites are filtered out by the + access-scoped query in resolve_test_suite, which returns None.""" mock_compute.return_value = ProjectPermissions(memberships={"proj_a": "role_a"}, permission="view", username="test_user") - mock_suite_cls.get_regular.return_value = _mock_test_suite(project_code="forbidden_project") + mock_suite_cls.get.return_value = None from testgen.mcp.tools.test_results import get_failure_summary @@ -677,10 +678,10 @@ def test_get_failure_summary_rejects_inaccessible_test_suite(mock_compute, mock_ get_failure_summary(test_suite_id=str(uuid4())) -@patch("testgen.mcp.tools.test_results.TestSuite") +@patch("testgen.mcp.tools.common.TestSuite") def test_get_failure_summary_rejects_unknown_or_monitor_test_suite(mock_suite_cls, db_session_mock): - # TestSuite.get_regular returns None for monitor suites and unknown ids alike. - mock_suite_cls.get_regular.return_value = None + # resolve_test_suite's access + is_monitor scoped query returns None for monitor suites and unknown ids alike. + mock_suite_cls.get.return_value = None from testgen.mcp.tools.test_results import get_failure_summary @@ -688,6 +689,29 @@ def test_get_failure_summary_rejects_unknown_or_monitor_test_suite(mock_suite_cl get_failure_summary(test_suite_id=str(uuid4())) +def test_get_failure_summary_rejects_invalid_group_by(db_session_mock): + from testgen.mcp.tools.test_results import get_failure_summary + + with pytest.raises(MCPUserError, match="Invalid group_by"): + get_failure_summary(job_execution_id=str(uuid4()), group_by="bogus") + + +@patch("testgen.mcp.tools.common.TestSuite") +@patch("testgen.mcp.permissions._compute_project_permissions") +def test_get_failure_summary_rejects_cross_project_suite(mock_compute, mock_suite_cls, db_session_mock): + """A suite the caller can access but in a different project than project_code is rejected, + not silently scoped away to an empty result.""" + mock_compute.return_value = ProjectPermissions( + memberships={"proj_a": "role_a", "proj_b": "role_a"}, permission="view", username="test_user", + ) + mock_suite_cls.get.return_value = _mock_test_suite(project_code="proj_b") + + from testgen.mcp.tools.test_results import get_failure_summary + + with pytest.raises(MCPUserError, match="belongs to project `proj_b`, not `proj_a`"): + get_failure_summary(project_code="proj_a", test_suite_id=str(uuid4())) + + # ---------------------------------------------------------------------- # search_test_results # ---------------------------------------------------------------------- @@ -866,6 +890,55 @@ def test_get_failure_trend_exclude_today_shifts_end_date(mock_compute, mock_fail assert mock_failure_trend.call_args.kwargs["end_date"] == real_today +@patch("testgen.mcp.tools.common.TestSuite") +@patch("testgen.mcp.permissions._compute_project_permissions") +def test_get_failure_trend_rejects_cross_project_suite(mock_compute, mock_suite_cls, db_session_mock): + mock_compute.return_value = ProjectPermissions( + memberships={"proj_a": "role_a", "proj_b": "role_a"}, permission="view", username="test_user", + ) + mock_suite_cls.get.return_value = _mock_test_suite(project_code="proj_b") + + from testgen.mcp.tools.test_results import get_failure_trend + + with pytest.raises(MCPUserError, match="belongs to project `proj_b`, not `proj_a`"): + get_failure_trend(project_code="proj_a", test_suite_id=str(uuid4())) + + +@patch("testgen.mcp.tools.common.TestSuite") +@patch("testgen.mcp.permissions._compute_project_permissions") +def test_get_failure_trend_rejects_inaccessible_suite(mock_compute, mock_suite_cls, db_session_mock): + """An inaccessible (or unknown/monitor) suite errors instead of silently returning an empty trend.""" + mock_compute.return_value = ProjectPermissions( + memberships={"proj_a": "role_a"}, permission="view", username="test_user", + ) + mock_suite_cls.get.return_value = None + + from testgen.mcp.tools.test_results import get_failure_trend + + with pytest.raises(MCPResourceNotAccessible, match="Test suite .* not found or not accessible"): + get_failure_trend(test_suite_id=str(uuid4())) + + +@patch("testgen.mcp.tools.common.TableGroup") +@patch("testgen.mcp.tools.common.TestSuite") +@patch("testgen.mcp.permissions._compute_project_permissions") +def test_get_failure_trend_rejects_suite_and_table_group_in_different_projects( + mock_compute, mock_suite_cls, mock_tg_cls, db_session_mock +): + """No project_code, but the suite and table group resolve to different accessible projects: + the two filters would AND to an empty result, so reject instead of silently returning empty.""" + mock_compute.return_value = ProjectPermissions( + memberships={"proj_a": "role_a", "proj_b": "role_a"}, permission="view", username="test_user", + ) + mock_suite_cls.get.return_value = _mock_test_suite(project_code="proj_a") + mock_tg_cls.get.return_value = MagicMock(project_code="proj_b") + + from testgen.mcp.tools.test_results import get_failure_trend + + with pytest.raises(MCPUserError, match="different projects"): + get_failure_trend(test_suite_id=str(uuid4()), table_group_id=str(uuid4())) + + # ---------------------------------------------------------------------- # compare_test_runs # ---------------------------------------------------------------------- From 3e4150643b2d3e22bf40e98bf92457b4d3155aca Mon Sep 17 00:00:00 2001 From: Ricardo Boni Date: Fri, 5 Jun 2026 16:43:10 -0400 Subject: [PATCH 11/78] feat(mcp): surface stable-pass count in compare_test_runs Add a stable_passes counter to RunDiff, incremented when a shared test definition is Passed in both the baseline and target runs, and render a Stable passes row in the compare_test_runs summary table. The count is the largest bucket in a healthy suite and the natural denominator for the failure counts. Co-Authored-By: Claude Opus 4.8 (1M context) --- testgen/common/models/test_result.py | 11 ++++-- testgen/mcp/tools/test_results.py | 1 + tests/unit/mcp/test_tools_test_results.py | 48 +++++++++++++++++++++++ 3 files changed, 57 insertions(+), 3 deletions(-) diff --git a/testgen/common/models/test_result.py b/testgen/common/models/test_result.py index c1d25e73..8bbd63ba 100644 --- a/testgen/common/models/test_result.py +++ b/testgen/common/models/test_result.py @@ -101,6 +101,7 @@ class RunDiff: total_baseline: int total_target: int + stable_passes: int = 0 regressions: list[DiffRow] = field(default_factory=list) improvements: list[DiffRow] = field(default_factory=list) persistent_failures: list[DiffRow] = field(default_factory=list) @@ -471,12 +472,16 @@ def _row(tid: UUID, baseline_info: dict | None, target_info: dict | None) -> Dif for tid in baseline_results.keys() & target_results.keys(): baseline_info, target_info = baseline_results[tid], target_results[tid] + baseline_status, target_status = baseline_info["status"], target_info["status"] + if baseline_status == TestResultStatus.Passed and target_status == TestResultStatus.Passed: + diff.stable_passes += 1 + continue row = _row(tid, baseline_info, target_info) - if baseline_info["status"] == TestResultStatus.Passed and target_info["status"] in failing: + if baseline_status == TestResultStatus.Passed and target_status in failing: diff.regressions.append(row) - elif baseline_info["status"] in failing and target_info["status"] == TestResultStatus.Passed: + elif baseline_status in failing and target_status == TestResultStatus.Passed: diff.improvements.append(row) - elif baseline_info["status"] in failing and target_info["status"] in failing: + elif baseline_status in failing and target_status in failing: diff.persistent_failures.append(row) for tid in target_results.keys() - baseline_results.keys(): diff --git a/testgen/mcp/tools/test_results.py b/testgen/mcp/tools/test_results.py index d35d3f0a..d6767eb8 100644 --- a/testgen/mcp/tools/test_results.py +++ b/testgen/mcp/tools/test_results.py @@ -595,6 +595,7 @@ def _require_completed(run: TestRun, label: str) -> None: doc.table( headers=["Category", "Count"], rows=[ + ["Stable passes (Baseline passed → Target passed)", diff.stable_passes], ["Regressions (Baseline passed → Target failed/warning)", len(diff.regressions)], ["Improvements (Baseline failed/warning → Target passed)", len(diff.improvements)], ["Persistent failures", len(diff.persistent_failures)], diff --git a/tests/unit/mcp/test_tools_test_results.py b/tests/unit/mcp/test_tools_test_results.py index 7b3d08e4..10607dfe 100644 --- a/tests/unit/mcp/test_tools_test_results.py +++ b/tests/unit/mcp/test_tools_test_results.py @@ -932,6 +932,7 @@ def test_compare_test_runs_happy_path( diff = MagicMock() diff.total_baseline = 100 diff.total_target = 100 + diff.stable_passes = 98 diff.regressions = [ _mock_diff_row( TestResultStatus.Passed, @@ -952,6 +953,8 @@ def test_compare_test_runs_happy_path( out = compare_test_runs(str(uuid4()), str(uuid4())) assert "Test Run Comparison" in out + assert "Stable passes (Baseline passed → Target passed)" in out + assert "| Stable passes (Baseline passed → Target passed) | 98 |" in out assert "Regressions" in out assert "Pattern Match" in out assert "Passed → Failed" in out @@ -961,6 +964,51 @@ def test_compare_test_runs_happy_path( mock_result.diff_with_details.assert_called_once_with(baseline_run.id, target_run.id) +def _fetch_row(test_definition_id, status): + row = MagicMock() + row.test_definition_id = test_definition_id + row.test_type = "Pattern_Match" + row.test_name_short = "Pattern Match" + row.table_name = "orders" + row.column_names = "customer_id" + row.status = status + row.result_measure = "0" + row.threshold_value = "0" + return row + + +def test_diff_with_details_counts_stable_passes(): + from testgen.common.models.test_result import TestResult + + stable_1, stable_2, regressed, improved = uuid4(), uuid4(), uuid4(), uuid4() + baseline_rows = [ + _fetch_row(stable_1, TestResultStatus.Passed), + _fetch_row(stable_2, TestResultStatus.Passed), + _fetch_row(regressed, TestResultStatus.Passed), + _fetch_row(improved, TestResultStatus.Failed), + ] + target_rows = [ + _fetch_row(stable_1, TestResultStatus.Passed), + _fetch_row(stable_2, TestResultStatus.Passed), + _fetch_row(regressed, TestResultStatus.Failed), + _fetch_row(improved, TestResultStatus.Passed), + ] + session = MagicMock() + session.execute.side_effect = [baseline_rows, target_rows] + + with patch("testgen.common.models.test_result.get_current_session", return_value=session): + diff = TestResult.diff_with_details(uuid4(), uuid4()) + + assert diff.stable_passes == 2 + assert len(diff.regressions) == 1 + assert len(diff.improvements) == 1 + assert len(diff.persistent_failures) == 0 + # Internal consistency: with no out-of-bucket statuses (Error/Log) in the fixture, the four + # named buckets account for every shared test definition. + shared = diff.total_target - len(diff.new_tests) + assert diff.stable_passes + len(diff.regressions) + len(diff.improvements) + len(diff.persistent_failures) == shared + + @patch("testgen.mcp.tools.test_results.TestSuite") @patch("testgen.mcp.tools.test_results.TestResult") @patch("testgen.mcp.tools.test_results.TestRun") From 89bb4ecf5060650bb318bd5c6baa09578b0cd3a9 Mon Sep 17 00:00:00 2001 From: Aarthy Adityan Date: Wed, 3 Jun 2026 22:10:01 -0400 Subject: [PATCH 12/78] fix(test-execution): render empty baseline/tolerance params as NULL MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit CAT measures substitute baseline params (counts, averages, standard deviations) and tolerances directly into SQL. When a test definition has an empty value, the substitution produced invalid SQL such as CAST( AS FLOAT) — erroring the test and poisoning the aggregated CAT batch (forcing a slow single rerun). Add a shared null_if_empty() helper and apply it to the numeric baseline params and tolerances in both the execution query builder and the source data lookup service. BASELINE_VALUE (used as a quoted literal, a number, or an IN-list) and Freshness_Trend's BASELINE_SUM (a quoted timestamp the template handles via NULLIF) are left as-is, since "NULL" is not valid SQL in those contexts. Co-Authored-By: Claude Opus 4.8 --- .../commands/queries/execute_tests_query.py | 29 +++++-- testgen/common/clean_sql.py | 13 +++ testgen/common/source_data_service.py | 14 ++-- .../queries/test_execute_tests_query.py | 84 ++++++++++++++++++- 4 files changed, 124 insertions(+), 16 deletions(-) diff --git a/testgen/commands/queries/execute_tests_query.py b/testgen/commands/queries/execute_tests_query.py index 751fd4d5..1316f6b6 100644 --- a/testgen/commands/queries/execute_tests_query.py +++ b/testgen/commands/queries/execute_tests_query.py @@ -7,7 +7,7 @@ import pandas as pd from testgen.common import read_template_sql_file -from testgen.common.clean_sql import concat_columns +from testgen.common.clean_sql import concat_columns, null_if_empty from testgen.common.database.database_service import ( fetch_dict_from_db, get_flavor_service, @@ -348,16 +348,27 @@ def _get_params(self, test_def: TestExecutionDef | None = None) -> dict: "CONCAT_COLUMNS": concat_columns(test_def.column_name, self.null_value) if test_def.column_name else "", "SKIP_ERRORS": test_def.skip_errors or 0, "CUSTOM_QUERY": test_def.custom_query, - "BASELINE_CT": test_def.baseline_ct, - "BASELINE_UNIQUE_CT": test_def.baseline_unique_ct, + # Numeric baseline params feed division/arithmetic in CAT measures. + # Render empty values as NULL so the measure evaluates to NULL instead of producing + # invalid SQL like CAST( AS FLOAT). Cannot default to 0 — several are denominators. + "BASELINE_CT": null_if_empty(test_def.baseline_ct), + "BASELINE_UNIQUE_CT": null_if_empty(test_def.baseline_unique_ct), + # BASELINE_VALUE is a quoted literal, an unquoted number, or an IN-list depending on + # test type — no single empty-substitution is valid SQL, so it is templated as-is. "BASELINE_VALUE": test_def.baseline_value, - "BASELINE_VALUE_CT": test_def.baseline_value_ct, + "BASELINE_VALUE_CT": null_if_empty(test_def.baseline_value_ct), "THRESHOLD_VALUE": test_def.threshold_value or 0, - "BASELINE_SUM": test_def.baseline_sum, - "BASELINE_AVG": test_def.baseline_avg, - "BASELINE_SD": test_def.baseline_sd, - "LOWER_TOLERANCE": "NULL" if test_def.lower_tolerance in (None, "") else test_def.lower_tolerance, - "UPPER_TOLERANCE": "NULL" if test_def.upper_tolerance in (None, "") else test_def.upper_tolerance, + # Freshness_Trend uses BASELINE_SUM as a quoted timestamp (the template's + # NULLIF('{...}', '') handles empty); Other tests use it as a numeric sum. + "BASELINE_SUM": ( + test_def.baseline_sum + if test_def.test_type == "Freshness_Trend" + else null_if_empty(test_def.baseline_sum) + ), + "BASELINE_AVG": null_if_empty(test_def.baseline_avg), + "BASELINE_SD": null_if_empty(test_def.baseline_sd), + "LOWER_TOLERANCE": null_if_empty(test_def.lower_tolerance), + "UPPER_TOLERANCE": null_if_empty(test_def.upper_tolerance), # SUBSET_CONDITION should be replaced after CUSTOM_QUERY # since the latter may contain the former "SUBSET_CONDITION": test_def.subset_condition or "1=1", diff --git a/testgen/common/clean_sql.py b/testgen/common/clean_sql.py index 8d1856e3..397fc9fc 100644 --- a/testgen/common/clean_sql.py +++ b/testgen/common/clean_sql.py @@ -50,6 +50,19 @@ def quote_identifiers(identifiers: str, flavor: str) -> str: return ", ".join(quoted_values) +def null_if_empty(value: object) -> object: + """Return the literal ``"NULL"`` when ``value`` is empty (``None`` or ``""``), else ``value``. + + For numeric test parameters substituted into SQL templates (baseline counts, + averages, standard deviations, tolerances). An empty substitution produces invalid SQL + such as ``CAST( AS FLOAT)``; ``"NULL"`` makes the surrounding expression evaluate to NULL + instead. A real ``0`` is preserved (``0 in (None, "")`` is ``False``). Not suitable for + params used as quoted literals or IN-lists (e.g. ``BASELINE_VALUE``, ``BASELINE_SUM`` for + Freshness_Trend), where ``"NULL"`` would not be valid SQL. + """ + return "NULL" if value in (None, "") else value + + def concat_columns(columns: str, null_value: str): # Prepares SQL expression to concatenate comma-separated column list expression = "" diff --git a/testgen/common/source_data_service.py b/testgen/common/source_data_service.py index 52cf2723..941679f1 100644 --- a/testgen/common/source_data_service.py +++ b/testgen/common/source_data_service.py @@ -11,7 +11,7 @@ import pandas as pd from sqlalchemy import text -from testgen.common.clean_sql import concat_columns +from testgen.common.clean_sql import concat_columns, null_if_empty from testgen.common.database.database_service import get_flavor_service, replace_params from testgen.common.date_service import parse_fuzzy_date from testgen.common.models import get_current_session @@ -207,12 +207,14 @@ def _build_query_standard(issue_data: dict, limit: int) -> str | None: if (parsed_test_date := parse_fuzzy_date(issue_data["test_date"])) else None, "CUSTOM_QUERY": test_definition.custom_query, + # BASELINE_VALUE varies (quoted literal / number / IN-list) by test type — templated as-is. + # The numeric baselines render empty as NULL to avoid invalid SQL like CAST( AS FLOAT). "BASELINE_VALUE": test_definition.baseline_value, - "BASELINE_CT": test_definition.baseline_ct, - "BASELINE_AVG": test_definition.baseline_avg, - "BASELINE_SD": test_definition.baseline_sd, - "LOWER_TOLERANCE": "NULL" if test_definition.lower_tolerance in (None, "") else test_definition.lower_tolerance, - "UPPER_TOLERANCE": "NULL" if test_definition.upper_tolerance in (None, "") else test_definition.upper_tolerance, + "BASELINE_CT": null_if_empty(test_definition.baseline_ct), + "BASELINE_AVG": null_if_empty(test_definition.baseline_avg), + "BASELINE_SD": null_if_empty(test_definition.baseline_sd), + "LOWER_TOLERANCE": null_if_empty(test_definition.lower_tolerance), + "UPPER_TOLERANCE": null_if_empty(test_definition.upper_tolerance), "THRESHOLD_VALUE": test_definition.threshold_value or 0, # SUBSET_CONDITION should be replaced after CUSTOM_QUERY # since the latter may contain the former diff --git a/tests/unit/commands/queries/test_execute_tests_query.py b/tests/unit/commands/queries/test_execute_tests_query.py index ca1c9555..0c1276a9 100644 --- a/tests/unit/commands/queries/test_execute_tests_query.py +++ b/tests/unit/commands/queries/test_execute_tests_query.py @@ -1,5 +1,5 @@ from datetime import UTC, datetime -from unittest.mock import patch +from unittest.mock import MagicMock, patch from uuid import uuid4 import pytest @@ -495,3 +495,85 @@ def test_aggregate_cat_tests_handles_null_max_query_chars(): assert grouped_defs == [[td]] +# --- TestExecutionSQL._get_params baseline guards --- + + +def _make_params_execution_sql() -> TestExecutionSQL: + """Build a minimal TestExecutionSQL for exercising _get_params without a database.""" + instance = TestExecutionSQL.__new__(TestExecutionSQL) + flavor_service = MagicMock() + flavor_service.quote_character = '"' + flavor_service.varchar_type = "VARCHAR" + instance.flavor_service = flavor_service + instance.flavor = "postgresql" + instance.table_group = MagicMock(id=uuid4()) + instance.test_run = MagicMock(test_suite_id=uuid4(), id=uuid4()) + instance.run_date = datetime(2026, 1, 1, tzinfo=UTC) + return instance + + +def test_get_params_empty_baseline_counts_become_null(): + """Empty baseline counts must render as NULL, not "", to avoid CAST( AS FLOAT) syntax errors.""" + instance = _make_params_execution_sql() + params = instance._get_params(_make_td(test_type="Missing_Pct", baseline_ct="", baseline_value_ct="")) + assert params["BASELINE_CT"] == "NULL" + assert params["BASELINE_VALUE_CT"] == "NULL" + + +def test_get_params_none_baseline_counts_become_null(): + instance = _make_params_execution_sql() + params = instance._get_params(_make_td(test_type="Missing_Pct", baseline_ct=None, baseline_value_ct=None)) + assert params["BASELINE_CT"] == "NULL" + assert params["BASELINE_VALUE_CT"] == "NULL" + + +def test_get_params_populated_baseline_counts_pass_through(): + instance = _make_params_execution_sql() + params = instance._get_params(_make_td(test_type="Missing_Pct", baseline_ct="1000", baseline_value_ct="950")) + assert params["BASELINE_CT"] == "1000" + assert params["BASELINE_VALUE_CT"] == "950" + + +def test_get_params_zero_baseline_count_is_not_nulled(): + """A real 0 is a meaningful value and must not be coerced to NULL.""" + instance = _make_params_execution_sql() + params = instance._get_params(_make_td(test_type="Row_Ct_Pct", baseline_ct=0)) + assert params["BASELINE_CT"] == 0 + + +def test_get_params_empty_numeric_baselines_become_null(): + """All numeric baseline params render NULL when empty.""" + instance = _make_params_execution_sql() + params = instance._get_params(_make_td( + test_type="Avg_Shift", + baseline_unique_ct="", baseline_avg="", baseline_sd="", baseline_sum="", + )) + assert params["BASELINE_UNIQUE_CT"] == "NULL" + assert params["BASELINE_AVG"] == "NULL" + assert params["BASELINE_SD"] == "NULL" + # Non-Freshness test types null-guard BASELINE_SUM (numeric use in Incr_Avg_Shift) + assert params["BASELINE_SUM"] == "NULL" + + +def test_get_params_freshness_baseline_sum_kept_raw_when_empty(): + """Freshness_Trend quotes BASELINE_SUM (NULLIF('', '') in template) — must stay empty, not 'NULL'.""" + instance = _make_params_execution_sql() + params = instance._get_params(_make_td(test_type="Freshness_Trend", baseline_sum="")) + assert params["BASELINE_SUM"] == "" + + +def test_get_params_baseline_value_left_unguarded(): + """BASELINE_VALUE has non-uniform usage (quoted/number/IN-list) — not coerced to NULL.""" + instance = _make_params_execution_sql() + params = instance._get_params(_make_td(test_type="Constant", baseline_value="")) + assert params["BASELINE_VALUE"] == "" + + +def test_get_params_empty_tolerances_become_null(): + """Tolerances use the same NULL guard.""" + instance = _make_params_execution_sql() + params = instance._get_params(_make_td(test_type="Volume_Trend", lower_tolerance="", upper_tolerance="")) + assert params["LOWER_TOLERANCE"] == "NULL" + assert params["UPPER_TOLERANCE"] == "NULL" + + From 4d2eacaaef465a6ed774a7f2efaf71d56e5cf136 Mon Sep 17 00:00:00 2001 From: Aarthy Adityan Date: Thu, 4 Jun 2026 14:05:17 -0400 Subject: [PATCH 13/78] fix(profiling): skip sampling on views where the sample clause rejects it MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit SQL Server and PostgreSQL reject TABLESAMPLE on views ("can only be used with local tables" / "tables and materialized views"), so profiling a table group containing a view with sampling enabled errored those columns. Add a samples_views flavor capability (False for mssql/postgresql) and an object_type signal from those flavors' DDF, and skip sampling for views on those flavors — they are profiled in full. Also compute the per-run sampling params once and share them between column profiling and frequency analysis. Frequency analysis previously never sampled (the secondary query's TABLESAMPLE branch was unreachable); it now samples the same tables, with the view-skip applied. Sample-scale frequency counts are left unscaled. Co-Authored-By: Claude Opus 4.8 --- .../queries/refresh_data_chars_query.py | 2 + testgen/commands/run_profiling.py | 58 ++++++++++++------- testgen/common/database/column_chars.py | 10 ++++ .../common/database/flavor/flavor_service.py | 3 +- .../database/flavor/mssql_flavor_service.py | 3 + .../flavor/postgresql_flavor_service.py | 3 + testgen/common/models/data_table.py | 3 + .../template/data_chars/data_chars_update.sql | 7 ++- .../030_initialize_new_schema_structure.sql | 2 + .../dbupgrade/0193_incremental_upgrade.sql | 7 +++ .../bigquery/data_chars/get_schema_ddf.sql | 7 +++ .../databricks/data_chars/get_schema_ddf.sql | 8 +++ .../mssql/data_chars/get_schema_ddf.sql | 2 + .../oracle/data_chars/get_schema_ddf.sql | 3 + .../postgresql/data_chars/get_schema_ddf.sql | 4 ++ .../redshift/data_chars/get_schema_ddf.sql | 4 ++ .../data_chars/get_schema_ddf.sql | 1 + .../sap_hana/data_chars/get_schema_ddf.sql | 16 ++++- .../snowflake/data_chars/get_schema_ddf.sql | 6 ++ tests/unit/common/test_flavor_sampling.py | 43 ++++++++++++++ 20 files changed, 168 insertions(+), 24 deletions(-) create mode 100644 testgen/template/dbupgrade/0193_incremental_upgrade.sql create mode 100644 tests/unit/common/test_flavor_sampling.py diff --git a/testgen/commands/queries/refresh_data_chars_query.py b/testgen/commands/queries/refresh_data_chars_query.py index c6d74a7c..68eb388f 100644 --- a/testgen/commands/queries/refresh_data_chars_query.py +++ b/testgen/commands/queries/refresh_data_chars_query.py @@ -30,6 +30,7 @@ class RefreshDataCharsSQL: "general_type", "column_type", "db_data_type", + "object_type", "approx_record_ct", "record_ct", ) @@ -157,6 +158,7 @@ def get_staging_data_chars(self, data_chars: list[ColumnChars], run_date: dateti column.general_type, column.column_type, column.db_data_type, + column.object_type, column.approx_record_ct, column.record_ct, ] diff --git a/testgen/commands/run_profiling.py b/testgen/commands/run_profiling.py index ffd1a58f..fd698eb5 100644 --- a/testgen/commands/run_profiling.py +++ b/testgen/commands/run_profiling.py @@ -19,7 +19,7 @@ write_to_app_db, ) from testgen.common.database.column_chars import ColumnChars -from testgen.common.database.database_service import ThreadedProgress +from testgen.common.database.database_service import ThreadedProgress, get_flavor_service from testgen.common.job_context import job_context from testgen.common.mixpanel_service import MixpanelService from testgen.common.models import get_current_session, with_database_session @@ -86,8 +86,9 @@ def run_profiling( if data_chars: sql_generator = ProfilingSQL(connection, table_group, profiling_run) - _run_column_profiling(sql_generator, data_chars) - _run_frequency_analysis(sql_generator) + sampling_params = _compute_sampling_params(sql_generator, data_chars) + _run_column_profiling(sql_generator, data_chars, sampling_params) + _run_frequency_analysis(sql_generator, sampling_params) _run_hygiene_issue_detection(sql_generator) # if table_group.profile_do_pair_rules == "Y": @@ -143,26 +144,40 @@ def _exclude_xde_columns(data_chars: list[ColumnChars], table_group_id: UUID) -> return filtered -def _run_column_profiling(sql_generator: ProfilingSQL, data_chars: list[ColumnChars]) -> None: +def _compute_sampling_params( + sql_generator: ProfilingSQL, data_chars: list[ColumnChars] +) -> dict[str, TableSampling]: + table_group = sql_generator.table_group + sampling_params: dict[str, TableSampling] = {} + if not table_group.profile_use_sampling: + return sampling_params + + sampleable_types = get_flavor_service(sql_generator.flavor).sampleable_object_types + for column in data_chars: + if sampling_params.get(column.table_name): + continue + if sampleable_types is not None and column.object_type not in sampleable_types: + continue + result = calculate_sampling_params( + table_name=column.table_name, + record_count=column.record_ct, + sample_percent_raw=table_group.profile_sample_percent, + min_sample=table_group.profile_sample_min_count, + ) + if result: + sampling_params[column.table_name] = result + return sampling_params + + +def _run_column_profiling( + sql_generator: ProfilingSQL, data_chars: list[ColumnChars], sampling_params: dict[str, TableSampling] +) -> None: profiling_run = sql_generator.profiling_run profiling_run.set_progress("col_profiling", "Running") profiling_run.save() get_current_session().commit() LOG.info(f"Running column profiling queries: {len(data_chars)}") - table_group = sql_generator.table_group - sampling_params: dict[str, TableSampling] = {} - if table_group.profile_use_sampling: - for column in data_chars: - if not sampling_params.get(column.table_name): - result = calculate_sampling_params( - table_name=column.table_name, - record_count=column.record_ct, - sample_percent_raw=table_group.profile_sample_percent, - min_sample=table_group.profile_sample_min_count, - ) - if result: - sampling_params[column.table_name] = result def update_column_progress(progress: ThreadedProgress) -> None: profiling_run.set_progress( @@ -218,7 +233,7 @@ def update_column_progress(progress: ThreadedProgress) -> None: ) -def _run_frequency_analysis(sql_generator: ProfilingSQL) -> None: +def _run_frequency_analysis(sql_generator: ProfilingSQL, sampling_params: dict[str, TableSampling]) -> None: profiling_run = sql_generator.profiling_run profiling_run.set_progress("freq_analysis", "Running") profiling_run.save() @@ -227,7 +242,7 @@ def _run_frequency_analysis(sql_generator: ProfilingSQL) -> None: error_data = None try: LOG.info("Selecting columns for frequency analysis") - frequency_columns = fetch_dict_from_db(*sql_generator.get_frequency_analysis_columns()) + frequency_columns = [ColumnChars(**column) for column in fetch_dict_from_db(*sql_generator.get_frequency_analysis_columns())] if frequency_columns: LOG.info(f"Running frequency analysis queries: {len(frequency_columns)}") @@ -240,7 +255,10 @@ def update_frequency_progress(progress: ThreadedProgress) -> None: get_current_session().commit() frequency_results, result_columns, error_data = fetch_from_db_threaded( - [sql_generator.run_frequency_analysis(ColumnChars(**column)) for column in frequency_columns], + [ + sql_generator.run_frequency_analysis(column, sampling_params.get(column.table_name)) + for column in frequency_columns + ], use_target_db=True, max_threads=sql_generator.connection.max_threads, progress_callback=update_frequency_progress, diff --git a/testgen/common/database/column_chars.py b/testgen/common/database/column_chars.py index 6faa08f7..014eb7f4 100644 --- a/testgen/common/database/column_chars.py +++ b/testgen/common/database/column_chars.py @@ -1,4 +1,13 @@ import dataclasses +from enum import StrEnum + + +class ObjectType(StrEnum): + TABLE = "TABLE" + VIEW = "VIEW" + MATERIALIZED_VIEW = "MATERIALIZED_VIEW" + EXTERNAL = "EXTERNAL" + OTHER = "OTHER" @dataclasses.dataclass @@ -11,6 +20,7 @@ class ColumnChars: column_type: str | None = None db_data_type: str | None = None is_decimal: bool = False + object_type: str | None = None # ObjectType values approx_record_ct: int | None = None # This should not default to 0 since we don't always retrieve actual row counts # UI relies on the null value to know that the approx_record_ct should be displayed instead diff --git a/testgen/common/database/flavor/flavor_service.py b/testgen/common/database/flavor/flavor_service.py index becb8b38..7e09406b 100644 --- a/testgen/common/database/flavor/flavor_service.py +++ b/testgen/common/database/flavor/flavor_service.py @@ -6,7 +6,7 @@ from sqlalchemy import create_engine as sqlalchemy_create_engine from sqlalchemy.engine.base import Engine -from testgen.common.database.column_chars import ColumnChars +from testgen.common.database.column_chars import ColumnChars, ObjectType from testgen.common.encrypt import DecryptText SQLFlavor = Literal["redshift", "redshift_spectrum", "snowflake", "mssql", "postgresql", "databricks", "bigquery", "oracle", "sap_hana", "salesforce_data360"] @@ -116,6 +116,7 @@ def row_limit_clauses(self, n: int) -> tuple[str, str]: qualifies_table_refs_with_schema = True metadata_via_api = False + sampleable_object_types: frozenset[ObjectType] | None = None def get_schema_columns(self, _params: ResolvedConnectionParams, _schema: str) -> list[ColumnChars] | None: """Return column metadata without querying information_schema. diff --git a/testgen/common/database/flavor/mssql_flavor_service.py b/testgen/common/database/flavor/mssql_flavor_service.py index 7f248e43..58928782 100644 --- a/testgen/common/database/flavor/mssql_flavor_service.py +++ b/testgen/common/database/flavor/mssql_flavor_service.py @@ -1,6 +1,7 @@ from sqlalchemy.engine import URL from testgen import settings +from testgen.common.database.column_chars import ObjectType from testgen.common.database.flavor.flavor_service import FlavorService, ResolvedConnectionParams @@ -10,6 +11,8 @@ class MssqlFlavorService(FlavorService): escaped_underscore = "[_]" row_limiting_clause = "top" url_scheme = "mssql+pyodbc" + # TABLESAMPLE applies only to tables and materialized views + sampleable_object_types = frozenset({ObjectType.TABLE, ObjectType.MATERIALIZED_VIEW}) def get_connection_string_from_fields(self, params: ResolvedConnectionParams) -> str: connection_url = URL.create( diff --git a/testgen/common/database/flavor/postgresql_flavor_service.py b/testgen/common/database/flavor/postgresql_flavor_service.py index 011ab05a..5f187ee0 100644 --- a/testgen/common/database/flavor/postgresql_flavor_service.py +++ b/testgen/common/database/flavor/postgresql_flavor_service.py @@ -1,5 +1,6 @@ from urllib.parse import quote_plus +from testgen.common.database.column_chars import ObjectType from testgen.common.database.flavor.flavor_service import ResolvedConnectionParams from testgen.common.database.flavor.redshift_flavor_service import RedshiftFlavorService @@ -8,6 +9,8 @@ class PostgresqlFlavorService(RedshiftFlavorService): escaped_underscore = "\\_" url_scheme = "postgresql" + # TABLESAMPLE applies only to tables and materialized views + sampleable_object_types = frozenset({ObjectType.TABLE, ObjectType.MATERIALIZED_VIEW}) def get_connection_string_from_fields(self, params: ResolvedConnectionParams) -> str: if params.host.startswith("/"): diff --git a/testgen/common/models/data_table.py b/testgen/common/models/data_table.py index 21ea22b3..3959d641 100644 --- a/testgen/common/models/data_table.py +++ b/testgen/common/models/data_table.py @@ -42,6 +42,7 @@ class TableProfilingOverview: table_groups_id: UUID schema_name: str | None table_name: str + object_type: str | None record_ct: int | None column_ct: int | None dq_score_profiling: float | None @@ -61,6 +62,7 @@ class DataTable(Entity): table_groups_id: UUID = Column(postgresql.UUID(as_uuid=True), ForeignKey("table_groups.id")) schema_name: str | None = Column(String) table_name: str = Column(String) + object_type: str | None = Column(String) column_ct: int | None = Column(BigInteger) record_ct: int | None = Column(BigInteger) approx_record_ct: int | None = Column(BigInteger) @@ -110,6 +112,7 @@ def get_profiling_overview( cls.table_groups_id, cls.schema_name, cls.table_name, + cls.object_type, cls.record_ct, cls.column_ct, cls.dq_score_profiling, diff --git a/testgen/template/data_chars/data_chars_update.sql b/testgen/template/data_chars/data_chars_update.sql index 38281dad..2084b587 100644 --- a/testgen/template/data_chars/data_chars_update.sql +++ b/testgen/template/data_chars/data_chars_update.sql @@ -8,6 +8,7 @@ WITH new_chars AS ( schema_name, table_name, run_date, + MAX(object_type) AS object_type, MAX(approx_record_ct) AS approx_record_ct, MAX(record_ct) AS record_ct, COUNT(*) AS column_ct @@ -20,7 +21,8 @@ WITH new_chars AS ( ), updated_records AS ( UPDATE data_table_chars - SET approx_record_ct = n.approx_record_ct, + SET object_type = n.object_type, + approx_record_ct = n.approx_record_ct, record_ct = n.record_ct, column_ct = n.column_ct, last_refresh_date = n.run_date, @@ -55,6 +57,7 @@ WITH new_chars AS ( schema_name, table_name, run_date, + MAX(object_type) AS object_type, MAX(approx_record_ct) AS approx_record_ct, MAX(record_ct) AS record_ct, COUNT(*) AS column_ct @@ -72,6 +75,7 @@ inserted_records AS ( table_name, add_date, last_refresh_date, + object_type, approx_record_ct, record_ct, column_ct @@ -81,6 +85,7 @@ inserted_records AS ( n.table_name, n.run_date, n.run_date, + n.object_type, n.approx_record_ct, n.record_ct, n.column_ct diff --git a/testgen/template/dbsetup/030_initialize_new_schema_structure.sql b/testgen/template/dbsetup/030_initialize_new_schema_structure.sql index 66e4db8b..2e381725 100644 --- a/testgen/template/dbsetup/030_initialize_new_schema_structure.sql +++ b/testgen/template/dbsetup/030_initialize_new_schema_structure.sql @@ -35,6 +35,7 @@ CREATE TABLE stg_data_chars_updates ( general_type VARCHAR(1), column_type VARCHAR(50), db_data_type VARCHAR(50), + object_type VARCHAR(20), approx_record_ct BIGINT, record_ct BIGINT ); @@ -419,6 +420,7 @@ CREATE TABLE data_table_chars ( table_groups_id UUID, schema_name VARCHAR(50), table_name VARCHAR(120), + object_type VARCHAR(20), functional_table_type VARCHAR(50), description VARCHAR(1000), critical_data_element BOOLEAN, diff --git a/testgen/template/dbupgrade/0193_incremental_upgrade.sql b/testgen/template/dbupgrade/0193_incremental_upgrade.sql new file mode 100644 index 00000000..29206eb2 --- /dev/null +++ b/testgen/template/dbupgrade/0193_incremental_upgrade.sql @@ -0,0 +1,7 @@ +SET SEARCH_PATH TO {SCHEMA_NAME}; + +ALTER TABLE data_table_chars + ADD COLUMN IF NOT EXISTS object_type VARCHAR(20); + +ALTER TABLE stg_data_chars_updates + ADD COLUMN IF NOT EXISTS object_type VARCHAR(20); diff --git a/testgen/template/flavors/bigquery/data_chars/get_schema_ddf.sql b/testgen/template/flavors/bigquery/data_chars/get_schema_ddf.sql index ee2165d3..f2a6e775 100644 --- a/testgen/template/flavors/bigquery/data_chars/get_schema_ddf.sql +++ b/testgen/template/flavors/bigquery/data_chars/get_schema_ddf.sql @@ -20,8 +20,15 @@ SELECT ELSE 'X' END AS general_type, REGEXP_CONTAINS(LOWER(c.data_type), r'(decimal|numeric|bignumeric)') AS is_decimal, + CASE it.table_type + WHEN 'BASE TABLE' THEN 'TABLE' + WHEN 'VIEW' THEN 'VIEW' + WHEN 'MATERIALIZED VIEW' THEN 'MATERIALIZED_VIEW' + WHEN 'EXTERNAL' THEN 'EXTERNAL' + ELSE 'OTHER' END AS object_type, t.row_count AS approx_record_ct FROM `{DATA_SCHEMA}.INFORMATION_SCHEMA.COLUMNS` c LEFT JOIN `{DATA_SCHEMA}.__TABLES__` t ON c.table_name = t.table_id + LEFT JOIN `{DATA_SCHEMA}.INFORMATION_SCHEMA.TABLES` it ON c.table_name = it.table_name WHERE c.table_schema = '{DATA_SCHEMA}' {TABLE_CRITERIA} ORDER BY c.table_schema, c.table_name, c.ordinal_position; diff --git a/testgen/template/flavors/databricks/data_chars/get_schema_ddf.sql b/testgen/template/flavors/databricks/data_chars/get_schema_ddf.sql index 6ae63a94..93a066ce 100644 --- a/testgen/template/flavors/databricks/data_chars/get_schema_ddf.sql +++ b/testgen/template/flavors/databricks/data_chars/get_schema_ddf.sql @@ -22,7 +22,15 @@ SELECT WHEN c.numeric_scale > 0 THEN 1 ELSE 0 END AS is_decimal, + CASE t.table_type + WHEN 'MANAGED' THEN 'TABLE' + WHEN 'VIEW' THEN 'VIEW' + WHEN 'MATERIALIZED_VIEW' THEN 'MATERIALIZED_VIEW' + WHEN 'EXTERNAL' THEN 'EXTERNAL' + WHEN 'FOREIGN' THEN 'EXTERNAL' + ELSE 'OTHER' END AS object_type, NULL AS approx_record_ct -- table statistics unavailable FROM information_schema.columns c + LEFT JOIN information_schema.tables t ON c.table_schema = t.table_schema AND c.table_name = t.table_name WHERE c.table_schema = '{DATA_SCHEMA}' {TABLE_CRITERIA} ORDER BY c.table_schema, c.table_name, c.ordinal_position; diff --git a/testgen/template/flavors/mssql/data_chars/get_schema_ddf.sql b/testgen/template/flavors/mssql/data_chars/get_schema_ddf.sql index 8a3ea74f..63423448 100644 --- a/testgen/template/flavors/mssql/data_chars/get_schema_ddf.sql +++ b/testgen/template/flavors/mssql/data_chars/get_schema_ddf.sql @@ -49,8 +49,10 @@ SELECT 'X' END AS general_type, CASE WHEN c.numeric_scale > 0 THEN 1 ELSE 0 END AS is_decimal, + CASE it.table_type WHEN 'BASE TABLE' THEN 'TABLE' WHEN 'VIEW' THEN 'VIEW' ELSE 'OTHER' END AS object_type, a.approx_record_ct AS approx_record_ct FROM information_schema.columns c LEFT JOIN approx_cts a ON c.table_schema = a.schema_name AND c.table_name = a.table_name + LEFT JOIN information_schema.tables it ON c.table_schema = it.table_schema AND c.table_name = it.table_name WHERE c.table_schema = '{DATA_SCHEMA}' {TABLE_CRITERIA} ORDER BY c.table_schema, c.table_name, c.ordinal_position; diff --git a/testgen/template/flavors/oracle/data_chars/get_schema_ddf.sql b/testgen/template/flavors/oracle/data_chars/get_schema_ddf.sql index d4f4f578..98acbc76 100644 --- a/testgen/template/flavors/oracle/data_chars/get_schema_ddf.sql +++ b/testgen/template/flavors/oracle/data_chars/get_schema_ddf.sql @@ -33,6 +33,9 @@ SELECT WHEN c.data_type = 'NUMBER' AND c.data_scale > 0 THEN 1 ELSE 0 END AS is_decimal, + CASE WHEN t.table_name IS NOT NULL THEN 'TABLE' + WHEN EXISTS (SELECT 1 FROM all_views v WHERE v.owner = c.owner AND v.view_name = c.table_name) THEN 'VIEW' + ELSE 'OTHER' END AS object_type, t.num_rows AS approx_record_ct FROM all_tab_columns c LEFT JOIN all_tables t ON c.owner = t.owner AND c.table_name = t.table_name diff --git a/testgen/template/flavors/postgresql/data_chars/get_schema_ddf.sql b/testgen/template/flavors/postgresql/data_chars/get_schema_ddf.sql index b5fcc322..14754855 100644 --- a/testgen/template/flavors/postgresql/data_chars/get_schema_ddf.sql +++ b/testgen/template/flavors/postgresql/data_chars/get_schema_ddf.sql @@ -44,6 +44,10 @@ SELECT WHEN c.data_type = 'numeric' THEN COALESCE(numeric_scale, 1) > 0 ELSE numeric_scale > 0 END as is_decimal, + CASE p.relkind + WHEN 'r' THEN 'TABLE' WHEN 'p' THEN 'TABLE' + WHEN 'v' THEN 'VIEW' WHEN 'm' THEN 'MATERIALIZED_VIEW' WHEN 'f' THEN 'EXTERNAL' + ELSE 'OTHER' END AS object_type, NULLIF(p.reltuples::BIGINT, -1) AS approx_record_ct FROM information_schema.columns c LEFT JOIN pg_namespace n ON c.table_schema = n.nspname diff --git a/testgen/template/flavors/redshift/data_chars/get_schema_ddf.sql b/testgen/template/flavors/redshift/data_chars/get_schema_ddf.sql index 6bda34ec..dfd4afea 100644 --- a/testgen/template/flavors/redshift/data_chars/get_schema_ddf.sql +++ b/testgen/template/flavors/redshift/data_chars/get_schema_ddf.sql @@ -39,6 +39,10 @@ SELECT WHEN c.data_type = 'numeric' THEN COALESCE(numeric_scale, 1) > 0 ELSE numeric_scale > 0 END AS is_decimal, + CASE p.relkind + WHEN 'r' THEN 'TABLE' WHEN 'p' THEN 'TABLE' + WHEN 'v' THEN 'VIEW' WHEN 'm' THEN 'MATERIALIZED_VIEW' + ELSE 'OTHER' END AS object_type, CASE WHEN reltuples > 0 AND reltuples < 1 THEN NULL ELSE reltuples::BIGINT diff --git a/testgen/template/flavors/redshift_spectrum/data_chars/get_schema_ddf.sql b/testgen/template/flavors/redshift_spectrum/data_chars/get_schema_ddf.sql index 3a6669f3..51460582 100644 --- a/testgen/template/flavors/redshift_spectrum/data_chars/get_schema_ddf.sql +++ b/testgen/template/flavors/redshift_spectrum/data_chars/get_schema_ddf.sql @@ -25,6 +25,7 @@ SELECT THEN 1 ELSE 0 END AS is_decimal, + 'EXTERNAL' AS object_type, NULL AS approx_record_ct -- Table statistics unavailable FROM svv_external_columns c WHERE c.schemaname = '{DATA_SCHEMA}' diff --git a/testgen/template/flavors/sap_hana/data_chars/get_schema_ddf.sql b/testgen/template/flavors/sap_hana/data_chars/get_schema_ddf.sql index 8a0838fd..ebb4e6f4 100644 --- a/testgen/template/flavors/sap_hana/data_chars/get_schema_ddf.sql +++ b/testgen/template/flavors/sap_hana/data_chars/get_schema_ddf.sql @@ -1,3 +1,14 @@ +WITH relations AS ( + SELECT SCHEMA_NAME, TABLE_NAME, COLUMN_NAME, DATA_TYPE_NAME, LENGTH, SCALE, POSITION, + 'TABLE' AS object_type + FROM SYS.TABLE_COLUMNS + WHERE SCHEMA_NAME = '{DATA_SCHEMA}' + UNION ALL + SELECT SCHEMA_NAME, VIEW_NAME AS TABLE_NAME, COLUMN_NAME, DATA_TYPE_NAME, LENGTH, SCALE, POSITION, + 'VIEW' AS object_type + FROM SYS.VIEW_COLUMNS + WHERE SCHEMA_NAME = '{DATA_SCHEMA}' +) SELECT c.SCHEMA_NAME AS schema_name, c.TABLE_NAME AS table_name, @@ -34,8 +45,9 @@ SELECT WHEN c.DATA_TYPE_NAME IN ('DOUBLE', 'REAL', 'SMALLDECIMAL') THEN 1 ELSE 0 END AS is_decimal, + c.object_type AS object_type, t.RECORD_COUNT AS approx_record_ct -FROM SYS.TABLE_COLUMNS c +FROM relations c LEFT JOIN SYS.M_TABLES t ON c.SCHEMA_NAME = t.SCHEMA_NAME AND c.TABLE_NAME = t.TABLE_NAME -WHERE c.SCHEMA_NAME = '{DATA_SCHEMA}' {TABLE_CRITERIA} +WHERE 1=1 {TABLE_CRITERIA} ORDER BY c.SCHEMA_NAME, c.TABLE_NAME, c.POSITION diff --git a/testgen/template/flavors/snowflake/data_chars/get_schema_ddf.sql b/testgen/template/flavors/snowflake/data_chars/get_schema_ddf.sql index 54940da8..3062875d 100644 --- a/testgen/template/flavors/snowflake/data_chars/get_schema_ddf.sql +++ b/testgen/template/flavors/snowflake/data_chars/get_schema_ddf.sql @@ -42,6 +42,12 @@ SELECT 'X' END AS general_type, numeric_scale > 0 AS is_decimal, + CASE t.table_type + WHEN 'BASE TABLE' THEN 'TABLE' + WHEN 'VIEW' THEN 'VIEW' + WHEN 'MATERIALIZED VIEW' THEN 'MATERIALIZED_VIEW' + WHEN 'EXTERNAL TABLE' THEN 'EXTERNAL' + ELSE 'OTHER' END AS object_type, t.row_count AS approx_record_ct FROM information_schema.columns c LEFT JOIN information_schema.tables t ON c.table_schema = t.table_schema AND c.table_name = t.table_name diff --git a/tests/unit/common/test_flavor_sampling.py b/tests/unit/common/test_flavor_sampling.py new file mode 100644 index 00000000..f1bb1443 --- /dev/null +++ b/tests/unit/common/test_flavor_sampling.py @@ -0,0 +1,43 @@ +"""Per-flavor sampleable object types. + +Profiling skips sampling for object types a flavor's sample clause cannot handle +(see run_profiling._compute_sampling_params). These tests lock in the per-flavor +sampleable_object_types and the object-type signal the skip depends on. +""" +import pytest + +from testgen.common.database.column_chars import ColumnChars, ObjectType +from testgen.common.database.database_service import get_flavor_service + +pytestmark = pytest.mark.unit + + +@pytest.mark.parametrize("flavor", ["mssql", "postgresql"]) +def test_sampleable_restricted_to_physical_where_clause_rejects_views(flavor): + # Verified live: TABLESAMPLE errors on a view for these flavors, so only physical + # relations are sampleable; views/external/other are profiled in full. + assert get_flavor_service(flavor).sampleable_object_types == { + ObjectType.TABLE, + ObjectType.MATERIALIZED_VIEW, + } + + +@pytest.mark.parametrize("flavor", ["redshift", "snowflake", "databricks", "oracle", "bigquery"]) +def test_sampleable_unrestricted_where_sampling_works_on_views(flavor): + # None = no restriction. Row-based samplers, or engines that accept the sample clause on + # views (verified live for snowflake/databricks/oracle). + assert get_flavor_service(flavor).sampleable_object_types is None + + +def test_object_type_enum_values_match_ddf_strings(): + # The DDF emits these literal strings; the skip compares column.object_type against them. + assert ObjectType.TABLE == "TABLE" + assert ObjectType.VIEW == "VIEW" + assert ObjectType.MATERIALIZED_VIEW == "MATERIALIZED_VIEW" + assert ObjectType.EXTERNAL == "EXTERNAL" + assert ObjectType.OTHER == "OTHER" + + +def test_column_chars_object_type_defaults_none(): + # Flavors whose DDF does not supply object_type leave it None. + assert ColumnChars(schema_name="s", table_name="t", column_name="c").object_type is None From a13cce0e4ed89d448f1c1aed6d1bff2bc2f8e9d5 Mon Sep 17 00:00:00 2001 From: Aarthy Adityan Date: Fri, 5 Jun 2026 00:24:55 -0400 Subject: [PATCH 14/78] feat(catalog): surface object type and warn when views skip sampling - Data Catalog: show 'Object Type' (title-cased) as the first attribute on the table Characteristics card, mirroring 'Data Type' for columns. Sourced from data_table_chars.object_type via get_tables_by_condition. - Table group form: under Sampling Parameters, show a yellow warning icon + note that views are profiled in full, only for flavors whose sample clause skips views (SQL Server, PostgreSQL) and only when sampling is enabled. Co-Authored-By: Claude Opus 4.8 --- .../js/data_profiling/data_characteristics.js | 7 ++++++ .../js/data_profiling/data_profiling_utils.js | 1 + testgen/ui/queries/profiling_queries.py | 1 + .../static/js/components/table_group_form.js | 22 +++++++++++++++++-- 4 files changed, 29 insertions(+), 2 deletions(-) diff --git a/testgen/ui/components/frontend/js/data_profiling/data_characteristics.js b/testgen/ui/components/frontend/js/data_profiling/data_characteristics.js index 8689fefc..6aa6a46e 100644 --- a/testgen/ui/components/frontend/js/data_profiling/data_characteristics.js +++ b/testgen/ui/components/frontend/js/data_profiling/data_characteristics.js @@ -36,6 +36,9 @@ const DataCharacteristicsCard = (/** @type Properties */ props, /** @type Column ); } } else { + if (item.object_type) { + attributes.push({ key: 'object_type', label: 'Object Type' }); + } attributes.push( { key: 'functional_table_type', label: `Semantic Table Type ${item.is_latest_profile ? '*' : ''}` }, ); @@ -77,6 +80,10 @@ const DataCharacteristicsCard = (/** @type Properties */ props, /** @type Column ); } else if (key === 'datatype_suggestion') { value = (value || '').toLowerCase(); + } else if (key === 'object_type') { + value = (value || '').split('_') + .map(word => word ? (word[0].toUpperCase() + word.substring(1).toLowerCase()) : '') + .join(' '); } else if (key === 'functional_table_type') { value = (value || '').split('-') .map(word => word ? (word[0].toUpperCase() + word.substring(1)) : '') diff --git a/testgen/ui/components/frontend/js/data_profiling/data_profiling_utils.js b/testgen/ui/components/frontend/js/data_profiling/data_profiling_utils.js index 39342edc..5089bea5 100644 --- a/testgen/ui/components/frontend/js/data_profiling/data_profiling_utils.js +++ b/testgen/ui/components/frontend/js/data_profiling/data_profiling_utils.js @@ -150,6 +150,7 @@ * @property {string} connection_id * @property {string} project_code * * Characteristics + * @property {string} object_type * @property {string} functional_table_type * @property {number} approx_record_ct * @property {number} record_ct diff --git a/testgen/ui/queries/profiling_queries.py b/testgen/ui/queries/profiling_queries.py index 70e5f681..ce7ff242 100644 --- a/testgen/ui/queries/profiling_queries.py +++ b/testgen/ui/queries/profiling_queries.py @@ -255,6 +255,7 @@ def get_tables_by_condition( table_chars.schema_name, table_chars.table_groups_id::VARCHAR AS table_group_id, -- Characteristics + table_chars.object_type, functional_table_type, approx_record_ct, table_chars.record_ct, diff --git a/testgen/ui/static/js/components/table_group_form.js b/testgen/ui/static/js/components/table_group_form.js index 93becaa9..85f53b66 100644 --- a/testgen/ui/static/js/components/table_group_form.js +++ b/testgen/ui/static/js/components/table_group_form.js @@ -55,8 +55,12 @@ import { required } from '../form_validators.js'; import { Select } from './select.js'; import { Caption } from './caption.js'; import { Textarea } from './textarea.js'; +import { Icon } from './icon.js'; -const { div } = van.tags; +// Flavors whose profiling sample clause cannot sample views — views are profiled in full. +const SAMPLE_SKIPS_VIEWS_FLAVORS = ['mssql', 'postgresql']; + +const { div, span } = van.tags; const normalizeTableSet = (value) => { return value?.split(/[,\n]/) @@ -119,6 +123,13 @@ const TableGroupForm = (props) => { return flavor === 'salesforce_data360'; }); + const sampleSkipsViews = van.derive(() => { + const connections = getValue(props.connections) ?? []; + const selected = connections.find(c => c.connection_id === tableGroupConnectionId.val); + const flavor = selected?.sql_flavor ?? getValue(props.sqlFlavor); + return SAMPLE_SKIPS_VIEWS_FLAVORS.includes(flavor); + }); + const updatedTableGroup = van.derive(() => { return { id: tableGroup.id, @@ -206,7 +217,7 @@ const TableGroupForm = (props) => { addScorecardDefinition, ), SamplingForm( - { setValidity: setFieldValidity }, + { setValidity: setFieldValidity, sampleSkipsViews }, profileUseSampling, profileSamplePercent, profileSampleMinCount, @@ -447,6 +458,13 @@ const SamplingForm = ( }, }), ), + () => (getValue(options.sampleSkipsViews) && profileUseSampling.val) + ? div( + { class: 'flex-row fx-align-center fx-gap-1' }, + Icon({ style: 'color: var(--orange); font-size: 18px;' }, 'warning'), + span('Views are profiled in full on this database — sampling applies only to tables and materialized views.'), + ) + : '', ), ); }; From 10f8a595c7c011fe0f6ba0425a93f7a3a1914900 Mon Sep 17 00:00:00 2001 From: Aarthy Adityan Date: Mon, 8 Jun 2026 11:16:10 -0400 Subject: [PATCH 15/78] refactor(catalog): tighten object_type typing and mssql sampleable set MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Address review feedback on MR !547: - Narrow mssql sampleable_object_types to {TABLE} — SQL Server has no materialized views, so the MATERIALIZED_VIEW entry never matched. - Annotate object_type as ObjectType (not str) on ColumnChars and the DataTable model/overview, per the StrEnum-everywhere convention. Co-Authored-By: Claude Opus 4.8 --- testgen/common/database/column_chars.py | 2 +- .../common/database/flavor/mssql_flavor_service.py | 4 ++-- testgen/common/models/data_table.py | 5 +++-- tests/unit/common/test_flavor_sampling.py | 14 +++++++++----- 4 files changed, 15 insertions(+), 10 deletions(-) diff --git a/testgen/common/database/column_chars.py b/testgen/common/database/column_chars.py index 014eb7f4..b9c79d08 100644 --- a/testgen/common/database/column_chars.py +++ b/testgen/common/database/column_chars.py @@ -20,7 +20,7 @@ class ColumnChars: column_type: str | None = None db_data_type: str | None = None is_decimal: bool = False - object_type: str | None = None # ObjectType values + object_type: ObjectType | None = None approx_record_ct: int | None = None # This should not default to 0 since we don't always retrieve actual row counts # UI relies on the null value to know that the approx_record_ct should be displayed instead diff --git a/testgen/common/database/flavor/mssql_flavor_service.py b/testgen/common/database/flavor/mssql_flavor_service.py index 58928782..356cba09 100644 --- a/testgen/common/database/flavor/mssql_flavor_service.py +++ b/testgen/common/database/flavor/mssql_flavor_service.py @@ -11,8 +11,8 @@ class MssqlFlavorService(FlavorService): escaped_underscore = "[_]" row_limiting_clause = "top" url_scheme = "mssql+pyodbc" - # TABLESAMPLE applies only to tables and materialized views - sampleable_object_types = frozenset({ObjectType.TABLE, ObjectType.MATERIALIZED_VIEW}) + # TABLESAMPLE is rejected on views; SQL Server has no materialized views, so only base tables. + sampleable_object_types = frozenset({ObjectType.TABLE}) def get_connection_string_from_fields(self, params: ResolvedConnectionParams) -> str: connection_url = URL.create( diff --git a/testgen/common/models/data_table.py b/testgen/common/models/data_table.py index 3959d641..ce149b94 100644 --- a/testgen/common/models/data_table.py +++ b/testgen/common/models/data_table.py @@ -17,6 +17,7 @@ ) from sqlalchemy.dialects import postgresql +from testgen.common.database.column_chars import ObjectType from testgen.common.models import get_current_session from testgen.common.models.data_column import DataColumnChars from testgen.common.models.entity import Entity @@ -42,7 +43,7 @@ class TableProfilingOverview: table_groups_id: UUID schema_name: str | None table_name: str - object_type: str | None + object_type: ObjectType | None record_ct: int | None column_ct: int | None dq_score_profiling: float | None @@ -62,7 +63,7 @@ class DataTable(Entity): table_groups_id: UUID = Column(postgresql.UUID(as_uuid=True), ForeignKey("table_groups.id")) schema_name: str | None = Column(String) table_name: str = Column(String) - object_type: str | None = Column(String) + object_type: ObjectType | None = Column(String) column_ct: int | None = Column(BigInteger) record_ct: int | None = Column(BigInteger) approx_record_ct: int | None = Column(BigInteger) diff --git a/tests/unit/common/test_flavor_sampling.py b/tests/unit/common/test_flavor_sampling.py index f1bb1443..06177d4f 100644 --- a/tests/unit/common/test_flavor_sampling.py +++ b/tests/unit/common/test_flavor_sampling.py @@ -12,11 +12,15 @@ pytestmark = pytest.mark.unit -@pytest.mark.parametrize("flavor", ["mssql", "postgresql"]) -def test_sampleable_restricted_to_physical_where_clause_rejects_views(flavor): - # Verified live: TABLESAMPLE errors on a view for these flavors, so only physical - # relations are sampleable; views/external/other are profiled in full. - assert get_flavor_service(flavor).sampleable_object_types == { +# Verified live: TABLESAMPLE errors on a view for these flavors, so only physical relations +# are sampleable; views/external/other are profiled in full. +def test_mssql_samples_only_base_tables(): + # SQL Server has no materialized views, so the sampleable set is just base tables. + assert get_flavor_service("mssql").sampleable_object_types == {ObjectType.TABLE} + + +def test_postgresql_samples_tables_and_materialized_views(): + assert get_flavor_service("postgresql").sampleable_object_types == { ObjectType.TABLE, ObjectType.MATERIALIZED_VIEW, } From 3133a1fd13f9b9265a5e941a11a550d5d8b83892 Mon Sep 17 00:00:00 2001 From: Luis Date: Mon, 1 Jun 2026 18:34:05 -0400 Subject: [PATCH 16/78] refactor: extract reusable test-result disposition service MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Move test-result disposition logic — the disposition write, the parent test-definition active/lock coupling, and the Passed-row exclusion — out of the test_results UI view into a Streamlit-free service under common/ so it can be reused outside the UI. The view now delegates to it; behavior is unchanged. Co-Authored-By: Claude Opus 4.8 (1M context) --- .../common/test_result_disposition_service.py | 101 ++++++++++++++++++ testgen/ui/views/test_results.py | 47 ++------ .../test_test_result_disposition_service.py | 83 ++++++++++++++ 3 files changed, 190 insertions(+), 41 deletions(-) create mode 100644 testgen/common/test_result_disposition_service.py create mode 100644 tests/unit/common/test_test_result_disposition_service.py diff --git a/testgen/common/test_result_disposition_service.py b/testgen/common/test_result_disposition_service.py new file mode 100644 index 00000000..4e1c2033 --- /dev/null +++ b/testgen/common/test_result_disposition_service.py @@ -0,0 +1,101 @@ +"""Shared test-result disposition service. + +Sets the disposition on test results and keeps the parent test definition's +active/lock state coupled: a "Muted" (``Disposition.INACTIVE``) disposition +deactivates the test definition and locks it against auto-regeneration; any +other value — or clearing the disposition — reactivates and unlocks it. Passed +test results are never dispositioned. Used by both the Streamlit UI +(test_results page) and the MCP tools. Must not import Streamlit. +""" +from collections.abc import Sequence +from dataclasses import dataclass +from datetime import UTC, datetime +from uuid import UUID + +from sqlalchemy import func, select, update + +from testgen.common.enums import Disposition +from testgen.common.models import get_current_session +from testgen.common.models.test_definition import TestDefinition +from testgen.common.models.test_result import TestResult, TestResultStatus + + +@dataclass(frozen=True) +class DispositionUpdate: + """Outcome of a disposition write. + + ``matched`` is the number of non-Passed results updated; ``passed_skipped`` is + the number of Passed results that matched but were left unchanged. + """ + + matched: int + passed_skipped: int + + +def coupled_test_definition_state(disposition: Disposition | None) -> tuple[bool, bool]: + """Return the parent test definition's ``(test_active, lock_refresh)`` for a disposition. + + Muted (``INACTIVE``) → ``(False, True)``: deactivate and lock against + auto-regeneration. Any other value, or cleared (``None``) → ``(True, False)``. + """ + deactivate = disposition == Disposition.INACTIVE + return (not deactivate, deactivate) + + +def coerce_ui_disposition(value: str | None) -> Disposition | None: + """Map a UI disposition string to the stored value. + + The UI passes ``Confirmed`` / ``Dismissed`` / ``Inactive``, or ``No Decision`` / + empty / ``None`` to clear. Unknown values raise ``ValueError`` (caller's bug). + """ + if value in (None, "", "No Decision"): + return None + return Disposition(value) + + +def set_test_results_disposition( + test_result_ids: Sequence[str | UUID], + disposition: Disposition | None, +) -> DispositionUpdate: + """Set ``disposition`` on the given results and couple their parent test definitions. + + Passed results are excluded from the write. ``disposition=None`` clears it (NULL). + Returns the matched (non-Passed, updated) and passed-skipped counts. + """ + ids = [UUID(str(rid)) for rid in test_result_ids] + if not ids: + return DispositionUpdate(matched=0, passed_skipped=0) + + session = get_current_session() + + passed_skipped = session.scalar( + select(func.count()) + .select_from(TestResult) + .where(TestResult.id.in_(ids), TestResult.status == TestResultStatus.Passed) + ) or 0 + + # NULL result_status rows (e.g. training-mode results) are excluded by `!= Passed`, + # matching the prior UI behavior — disposition only applies to evaluated results. + non_passed = (TestResult.id.in_(ids), TestResult.status != TestResultStatus.Passed) + + tr_stmt = ( + update(TestResult) + .where(*non_passed) + .values(disposition=disposition.value if disposition is not None else None) + ) + matched = session.execute(tr_stmt).rowcount + + test_active, lock_refresh = coupled_test_definition_state(disposition) + affected_td_ids = select(TestResult.test_definition_id).where(*non_passed) + td_stmt = ( + update(TestDefinition) + .where(TestDefinition.id.in_(affected_td_ids)) + .values( + test_active=test_active, + lock_refresh=lock_refresh, + last_manual_update=datetime.now(UTC).replace(tzinfo=None), + ) + ) + session.execute(td_stmt) + + return DispositionUpdate(matched=matched, passed_skipped=passed_skipped) diff --git a/testgen/ui/views/test_results.py b/testgen/ui/views/test_results.py index fd4d7b2c..b62473aa 100644 --- a/testgen/ui/views/test_results.py +++ b/testgen/ui/views/test_results.py @@ -13,6 +13,10 @@ from testgen.common.models.test_definition import TestDefinition, TestDefinitionNote, TestDefinitionSummary from testgen.common.models.test_suite import TestSuiteMinimal from testgen.common.pii_masking import get_pii_columns, mask_profiling_pii +from testgen.common.test_result_disposition_service import ( + coerce_ui_disposition, + set_test_results_disposition, +) from testgen.ui.components import widgets as testgen from testgen.ui.components.widgets.download_dialog import ( FILE_DATA_TYPE, @@ -31,7 +35,7 @@ get_test_issue_source_query, get_test_issue_source_query_custom, ) -from testgen.ui.services.database_service import execute_db_query, fetch_df_from_db, fetch_one_from_db +from testgen.ui.services.database_service import fetch_df_from_db, fetch_one_from_db from testgen.ui.services.query_cache import ( get_table_group_minimal, get_test_definition, @@ -911,43 +915,4 @@ def update_result_disposition( test_result_ids: list[str], disposition: str, ) -> None: - execute_db_query( - """ - WITH selects - AS (SELECT UNNEST(ARRAY [:test_result_ids]) AS selected_id) - UPDATE test_results - SET disposition = NULLIF(:disposition, 'No Decision') - FROM test_results r - INNER JOIN selects s - ON (r.id = s.selected_id::UUID) - WHERE r.id = test_results.id - AND r.result_status != 'Passed'; - """, - { - "test_result_ids": test_result_ids, - "disposition": disposition, - }, - ) - - execute_db_query( - """ - WITH selects - AS (SELECT UNNEST(ARRAY [:test_result_ids]) AS selected_id) - UPDATE test_definitions - SET test_active = :test_active, - last_manual_update = CURRENT_TIMESTAMP AT TIME ZONE 'UTC', - lock_refresh = :lock_refresh - FROM test_definitions d - INNER JOIN test_results r - ON (d.id = r.test_definition_id) - INNER JOIN selects s - ON (r.id = s.selected_id::UUID) - WHERE d.id = test_definitions.id - AND r.result_status != 'Passed'; - """, - { - "test_result_ids": test_result_ids, - "test_active": "N" if disposition == "Inactive" else "Y", - "lock_refresh": "Y" if disposition == "Inactive" else "N", - }, - ) + set_test_results_disposition(test_result_ids, coerce_ui_disposition(disposition)) diff --git a/tests/unit/common/test_test_result_disposition_service.py b/tests/unit/common/test_test_result_disposition_service.py new file mode 100644 index 00000000..83d3de62 --- /dev/null +++ b/tests/unit/common/test_test_result_disposition_service.py @@ -0,0 +1,83 @@ +from unittest.mock import MagicMock, patch +from uuid import uuid4 + +import pytest + +from testgen.common.enums import Disposition +from testgen.common.test_result_disposition_service import ( + DispositionUpdate, + coerce_ui_disposition, + coupled_test_definition_state, + set_test_results_disposition, +) + +pytestmark = pytest.mark.unit + +MODULE = "testgen.common.test_result_disposition_service" + + +class Test_coupled_test_definition_state: + def test_muted_deactivates_and_locks(self): + # (test_active, lock_refresh) + assert coupled_test_definition_state(Disposition.INACTIVE) == (False, True) + + @pytest.mark.parametrize("disposition", [Disposition.CONFIRMED, Disposition.DISMISSED, None]) + def test_other_values_reactivate_and_unlock(self, disposition): + assert coupled_test_definition_state(disposition) == (True, False) + + +class Test_coerce_ui_disposition: + @pytest.mark.parametrize( + "value,expected", + [ + ("Confirmed", Disposition.CONFIRMED), + ("Dismissed", Disposition.DISMISSED), + ("Inactive", Disposition.INACTIVE), + ("No Decision", None), + (None, None), + ("", None), + ], + ) + def test_maps_ui_string_to_stored_value(self, value, expected): + assert coerce_ui_disposition(value) is expected + + +class Test_set_test_results_disposition: + def test_empty_ids_is_noop(self): + with patch(f"{MODULE}.get_current_session") as get_session: + result = set_test_results_disposition([], Disposition.CONFIRMED) + assert result == DispositionUpdate(matched=0, passed_skipped=0) + get_session.assert_not_called() + + def test_returns_matched_and_passed_counts(self): + session = MagicMock() + # First call: COUNT of passed rows -> 2. Then the TR update returns rowcount 3. + session.scalar.return_value = 2 + tr_update_result = MagicMock(rowcount=3) + session.execute.return_value = tr_update_result + with patch(f"{MODULE}.get_current_session", return_value=session): + result = set_test_results_disposition([uuid4(), uuid4()], Disposition.DISMISSED) + assert result == DispositionUpdate(matched=3, passed_skipped=2) + # One TR update + one TD update. + assert session.execute.call_count == 2 + + def test_muted_couples_td_to_inactive(self): + session = MagicMock() + session.scalar.return_value = 0 + session.execute.return_value = MagicMock(rowcount=1) + with patch(f"{MODULE}.get_current_session", return_value=session): + set_test_results_disposition([uuid4()], Disposition.INACTIVE) + td_stmt = session.execute.call_args_list[1].args[0] + params = td_stmt.compile().params + # YNString stores 'N'/'Y'; bound values are the Python bools before type processing. + assert params["test_active"] is False + assert params["lock_refresh"] is True + + def test_clear_passes_null_disposition(self): + session = MagicMock() + session.scalar.return_value = 0 + session.execute.return_value = MagicMock(rowcount=1) + with patch(f"{MODULE}.get_current_session", return_value=session): + set_test_results_disposition([uuid4()], None) + tr_stmt = session.execute.call_args_list[0].args[0] + assert tr_stmt.compile().params["disposition"] is None From 01124a6c763e4c29ceb5e35d7512e47d04fbdc8e Mon Sep 17 00:00:00 2001 From: Luis Date: Mon, 1 Jun 2026 18:34:05 -0400 Subject: [PATCH 17/78] feat(mcp): add test-result disposition tools Add update_test_result and bulk_update_test_results to disposition test results (confirm, dismiss, mute, or clear), gated on the disposition permission and reusing the shared disposition service. Adds the parse_test_result_disposition and resolve_test_result helpers, surfaces test_result_id in list_test_results, and registers both tools. Co-Authored-By: Claude Opus 4.8 (1M context) --- testgen/mcp/server.py | 4 + testgen/mcp/tools/common.py | 43 +++- testgen/mcp/tools/test_results.py | 123 +++++++++++ tests/unit/mcp/test_tools_common.py | 62 ++++++ tests/unit/mcp/test_tools_test_results.py | 251 +++++++++++++++++++++- 5 files changed, 481 insertions(+), 2 deletions(-) diff --git a/testgen/mcp/server.py b/testgen/mcp/server.py index f7f474e3..ad8cc276 100644 --- a/testgen/mcp/server.py +++ b/testgen/mcp/server.py @@ -220,12 +220,14 @@ def build_mcp_server( validate_custom_test, ) from testgen.mcp.tools.test_results import ( + bulk_update_test_results, compare_test_runs, get_failure_summary, get_failure_trend, list_test_result_history, list_test_results, search_test_results, + update_test_result, ) from testgen.mcp.tools.test_runs import get_test_run, list_test_runs @@ -266,6 +268,8 @@ def safe_prompt(fn): safe_tool(search_test_results) safe_tool(get_failure_trend) safe_tool(compare_test_runs) + safe_tool(update_test_result) + safe_tool(bulk_update_test_results) safe_tool(get_test_type) safe_tool(get_source_data) safe_tool(get_source_data_query) diff --git a/testgen/mcp/tools/common.py b/testgen/mcp/tools/common.py index 6c71c481..d05511ce 100644 --- a/testgen/mcp/tools/common.py +++ b/testgen/mcp/tools/common.py @@ -30,7 +30,7 @@ from testgen.common.models.scores import ScoreCategory, ScoreDefinition from testgen.common.models.table_group import TableGroup from testgen.common.models.test_definition import TestDefinition, TestDefinitionNote, TestType -from testgen.common.models.test_result import TestResultStatus +from testgen.common.models.test_result import TestResult, TestResultStatus from testgen.common.models.test_suite import TestSuite from testgen.mcp.exceptions import MCPResourceNotAccessible, MCPUserError from testgen.mcp.permissions import get_project_permissions @@ -353,6 +353,25 @@ def parse_disposition(value: str) -> Disposition: return db_value +_NO_DECISION = "No Decision" + + +def parse_test_result_disposition(value: str) -> Disposition | None: + """Validate a user-facing test-result disposition and return the stored value. + + Accepts ``Confirmed``, ``Dismissed``, ``Muted``, and ``No Decision``. ``Muted`` + maps to ``Disposition.INACTIVE``; ``No Decision`` clears the disposition (returns + ``None`` → NULL). + """ + if value == _NO_DECISION: + return None + db_value = _DISPOSITION_USER_TO_DB.get(value) + if db_value is None: + valid = ", ".join([*_DISPOSITION_USER_TO_DB, _NO_DECISION]) + raise MCPUserError(f"Invalid disposition `{value}`. Valid values: {valid}") + return db_value + + def format_disposition(value: Disposition | str) -> str: """Map a stored disposition to its user-facing label (``INACTIVE`` → "Muted").""" try: @@ -623,6 +642,28 @@ def resolve_test_definition(test_definition_id: str) -> TestDefinition: return td +def resolve_test_result(test_result_id: str) -> TestResult: + """Resolve a test result ID to the live ORM model, collapsing missing-or-inaccessible. + + Filters monitor suites and project access via the result's parent test suite. + """ + result_uuid = parse_uuid(test_result_id, "test_result_id") + perms = get_project_permissions() + query = ( + select(TestResult) + .join(TestSuite, TestResult.test_suite_id == TestSuite.id) + .where( + TestResult.id == result_uuid, + TestSuite.is_monitor.isnot(True), + TestSuite.project_code.in_(perms.allowed_codes), + ) + ) + result = get_current_session().scalars(query).first() + if result is None: + raise MCPResourceNotAccessible("Test result", test_result_id) + return result + + def resolve_test_note(test_note_id: str) -> TestDefinitionNote: """Resolve a test note ID to the live ORM model, collapsing missing-or-inaccessible. diff --git a/testgen/mcp/tools/test_results.py b/testgen/mcp/tools/test_results.py index 411f8413..d4d215cf 100644 --- a/testgen/mcp/tools/test_results.py +++ b/testgen/mcp/tools/test_results.py @@ -1,6 +1,8 @@ from datetime import UTC, datetime, timedelta from uuid import UUID +from sqlalchemy import select + from testgen.common.enums import JobStatus from testgen.common.models import get_current_session, with_database_session from testgen.common.models.job_execution import JobExecution @@ -8,6 +10,10 @@ from testgen.common.models.test_result import BucketInterval, TestResult, TestResultStatus from testgen.common.models.test_run import TestRun, TestRunSummary from testgen.common.models.test_suite import TestSuite +from testgen.common.test_result_disposition_service import ( + DispositionUpdate, + set_test_results_disposition, +) from testgen.mcp.exceptions import MCPResourceNotAccessible, MCPUserError from testgen.mcp.permissions import get_project_permissions, mcp_permission from testgen.mcp.tools.common import ( @@ -18,8 +24,11 @@ parse_failure_group_by, parse_result_status, parse_since_arg, + parse_test_result_disposition, parse_uuid, resolve_aggregate_scope, + resolve_test_result, + resolve_test_suite, resolve_test_type, validate_limit, validate_page, @@ -134,6 +143,7 @@ def list_test_results( doc.heading(2, f"[{status_str}] {test_name} on `{r.column_names}` in `{r.table_name}`") else: doc.heading(2, f"[{status_str}] {test_name} on `{r.table_name}`") + doc.field("Test result", r.id, code=True) doc.field("Test definition", r.test_definition_id, code=True) if r.column_names: doc.field("Column", r.column_names, code=True) @@ -636,3 +646,116 @@ def _section(title: str, rows: list) -> None: _section("Removed Tests", diff.removed_tests) return doc.render() + + +@with_database_session +@mcp_permission("disposition") +def update_test_result(test_result_id: str, disposition: str) -> str: + """Set the disposition on a single test result (confirm, dismiss, mute, or clear). + + Args: + test_result_id: UUID of the test result, e.g. from ``list_test_results``. + disposition: New disposition. One of 'Confirmed', 'Dismissed', 'Muted', + 'No Decision' (clears it). 'Muted' deactivates the parent test and locks + it against auto-regeneration; any other value reactivates and unlocks it. + """ + result = resolve_test_result(test_result_id) + db_disposition = parse_test_result_disposition(disposition) + + update: DispositionUpdate = set_test_results_disposition([result.id], db_disposition) + + doc = MdDoc() + if update.matched == 0: + doc.text( + f"Test result {MdDoc.code(test_result_id)} was not dispositioned — disposition does " + f"not apply to passed results. No change made." + ) + return doc.render() + + doc.text(f"Updated test result {MdDoc.code(test_result_id)} disposition to **{disposition}**.") + return doc.render() + + +@with_database_session +@mcp_permission("disposition") +def bulk_update_test_results( + test_suite_id: str, + disposition: str, + job_execution_id: str | None = None, + table_name: str | None = None, + test_type: str | None = None, + status: str | None = None, + test_definition_id: str | None = None, +) -> str: + """Set the disposition on every matching test result in a suite (confirm, dismiss, mute, clear). + + Args: + test_suite_id: UUID of the test suite, e.g. from ``list_test_suites``. + disposition: New disposition. One of 'Confirmed', 'Dismissed', 'Muted', + 'No Decision' (clears it). 'Muted' deactivates the parent tests and locks + them against auto-regeneration; any other value reactivates and unlocks. + job_execution_id: UUID of a test run within the suite. Defaults to the suite's + latest completed run when omitted. + table_name: Optional table-name filter. Case-sensitive. + test_type: Optional test type name (e.g. 'Alpha Truncation'). + status: Optional result-status filter (Passed, Failed, Warning, Error, Log). + test_definition_id: Optional single test-definition filter. + """ + suite = resolve_test_suite(test_suite_id) + db_disposition = parse_test_result_disposition(disposition) + + if job_execution_id: + run = TestRun.get_by_id_or_job(parse_uuid(job_execution_id, "job_execution_id")) + if run is None or run.test_suite_id != suite.id: + raise MCPResourceNotAccessible("Test run", job_execution_id) + else: + run = ( + TestRun.get_by_id_or_job(suite.last_complete_test_run_id) + if suite.last_complete_test_run_id + else None + ) + if run is None: + raise MCPUserError(f"No completed test runs found for test suite `{test_suite_id}`.") + + clauses = [TestResult.test_suite_id == suite.id, TestResult.test_run_id == run.id] + if status: + clauses.append(TestResult.status == parse_result_status(status)) + if table_name: + clauses.append(TestResult.table_name == table_name) + if test_type: + clauses.append(TestResult.test_type == resolve_test_type(test_type)) + if test_definition_id: + clauses.append(TestResult.test_definition_id == parse_uuid(test_definition_id, "test_definition_id")) + + result_ids = list(get_current_session().scalars(select(TestResult.id).where(*clauses)).all()) + update = set_test_results_disposition(result_ids, db_disposition) + + filters = [] + if table_name: + filters.append(f"table_name=`{table_name}`") + if test_type: + filters.append(f"test_type=`{test_type}`") + if status: + filters.append(f"status=`{status}`") + if test_definition_id: + filters.append(f"test_definition_id=`{test_definition_id}`") + filter_str = ", ".join(filters) if filters else "no filter" + + doc = MdDoc() + if update.matched == 0 and update.passed_skipped == 0: + doc.heading(1, "No test results matched") + doc.text( + f"No test results in suite `{suite.test_suite}` matched the filter ({filter_str}). Nothing changed." + ) + return doc.render() + + doc.heading(1, f"Updated {update.matched} test results in suite `{suite.test_suite}`") + doc.field("Disposition", disposition) + doc.field("Run", run.job_execution_id, code=True) + doc.field("Filter", filter_str) + if update.passed_skipped: + doc.text( + f"Left {update.passed_skipped} passed results unchanged — disposition is not " + f"applied to passed results." + ) + return doc.render() diff --git a/tests/unit/mcp/test_tools_common.py b/tests/unit/mcp/test_tools_common.py index 79c83f2e..068d1096 100644 --- a/tests/unit/mcp/test_tools_common.py +++ b/tests/unit/mcp/test_tools_common.py @@ -772,3 +772,65 @@ def test_parse_sql_flavor_invalid_lists_display_values(): msg = str(exc.value) for member in SqlFlavorLabel: assert member.value in msg + + +# --- parse_test_result_disposition --- + + +@pytest.mark.parametrize( + "user_label,expected", + [ + ("Confirmed", Disposition.CONFIRMED), + ("Dismissed", Disposition.DISMISSED), + ("Muted", Disposition.INACTIVE), + ("No Decision", None), + ], +) +def test_parse_test_result_disposition_user_labels(user_label, expected): + from testgen.mcp.tools.common import parse_test_result_disposition + + assert parse_test_result_disposition(user_label) is expected + + +def test_parse_test_result_disposition_rejects_unknown_and_lists_accepted(): + from testgen.mcp.tools.common import parse_test_result_disposition + + with pytest.raises(MCPUserError) as exc: + parse_test_result_disposition("Inactive") # DB value, not user-facing + msg = str(exc.value) + for label in ("Confirmed", "Dismissed", "Muted", "No Decision"): + assert label in msg + + +# --- resolve_test_result --- + + +@patch("testgen.mcp.tools.common.get_project_permissions") +@patch("testgen.mcp.tools.common.get_current_session") +def test_resolve_test_result_happy_path(mock_session, mock_perms, db_session_mock): + from testgen.mcp.tools.common import resolve_test_result + + result = MagicMock() + mock_session.return_value.scalars.return_value.first.return_value = result + mock_perms.return_value = _mock_perms(allowed_projects=("demo",)) + + assert resolve_test_result(str(uuid4())) is result + + +@patch("testgen.mcp.tools.common.get_project_permissions") +@patch("testgen.mcp.tools.common.get_current_session") +def test_resolve_test_result_missing_or_inaccessible(mock_session, mock_perms, db_session_mock): + from testgen.mcp.tools.common import resolve_test_result + + mock_session.return_value.scalars.return_value.first.return_value = None + mock_perms.return_value = _mock_perms(allowed_projects=("demo",)) + + with pytest.raises(MCPResourceNotAccessible, match=r"Test result .* not found or not accessible"): + resolve_test_result(str(uuid4())) + + +def test_resolve_test_result_invalid_uuid(): + from testgen.mcp.tools.common import resolve_test_result + + with pytest.raises(MCPUserError, match="Invalid test_result_id"): + resolve_test_result("not-a-uuid") diff --git a/tests/unit/mcp/test_tools_test_results.py b/tests/unit/mcp/test_tools_test_results.py index 037e5424..dc4f9084 100644 --- a/tests/unit/mcp/test_tools_test_results.py +++ b/tests/unit/mcp/test_tools_test_results.py @@ -4,8 +4,9 @@ import pytest -from testgen.common.enums import JobStatus +from testgen.common.enums import Disposition, JobStatus from testgen.common.models.test_result import TestResultStatus +from testgen.common.test_result_disposition_service import DispositionUpdate from testgen.mcp.exceptions import MCPResourceNotAccessible, MCPUserError from testgen.mcp.permissions import ProjectPermissions @@ -53,6 +54,40 @@ def test_list_test_results_basic(mock_result, mock_tt_cls, mock_test_run_cls, mo assert "Truncation detected" in result +@patch("testgen.mcp.tools.test_results.TestSuite") +@patch("testgen.mcp.tools.test_results.TestRun") +@patch("testgen.mcp.tools.test_results.TestType") +@patch("testgen.mcp.tools.test_results.TestResult") +def test_list_test_results_emits_test_result_id(mock_result, mock_tt_cls, mock_test_run_cls, mock_suite_cls, db_session_mock): + mock_test_run_cls.get_by_id_or_job.return_value = _mock_test_run() + mock_suite_cls.get_regular.return_value = _mock_test_suite() + + result_id = uuid4() + r1 = MagicMock() + r1.id = result_id + r1.status = TestResultStatus.Failed + r1.test_type = "Alpha_Trunc" + r1.test_definition_id = uuid4() + r1.table_name = "orders" + r1.column_names = "customer_name" + r1.result_measure = "15.3" + r1.threshold_value = "10.0" + r1.message = "Truncation detected" + mock_result.select_results.return_value = [r1] + + tt = MagicMock() + tt.test_type = "Alpha_Trunc" + tt.test_name_short = "Alpha Truncation" + mock_tt_cls.select_where.return_value = [tt] + + from testgen.mcp.tools.test_results import list_test_results + + result = list_test_results(str(uuid4())) + + assert "Test result" in result + assert str(result_id) in result + + @patch("testgen.mcp.tools.test_results.TestSuite") @patch("testgen.mcp.tools.test_results.TestRun") @patch("testgen.mcp.tools.test_results.TestType") @@ -1317,3 +1352,217 @@ def test_compare_test_runs_rejects_baseline_not_completed( with _patch_test_results_session([_je(), _je(status=JobStatus.ERROR)]), \ pytest.raises(MCPUserError, match=r"Baseline run is in `Error` state"): compare_test_runs(str(uuid4()), str(uuid4())) + + +# --------------------------------------------------------------------------- +# update_test_result +# --------------------------------------------------------------------------- + + +@pytest.fixture +def disposition_perms(): + """Grant 'disposition' permission on demo (the conftest matrix omits it).""" + perms = MagicMock(spec=ProjectPermissions) + perms.memberships = {"demo": "role_a"} + perms.permission = "disposition" + perms.username = "test_user" + perms.allowed_codes = ["demo"] + perms.codes_allowed_to.return_value = ["demo"] + perms.has_access.side_effect = lambda code: code in ["demo"] + with patch("testgen.mcp.permissions._compute_project_permissions", return_value=perms): + yield perms + + +def test_update_test_result_invalid_uuid(db_session_mock, disposition_perms): + from testgen.mcp.tools.test_results import update_test_result + + with pytest.raises(MCPUserError, match="not a valid UUID"): + update_test_result(test_result_id="bogus", disposition="Confirmed") + + +@patch("testgen.mcp.tools.test_results.set_test_results_disposition") +@patch("testgen.mcp.tools.test_results.resolve_test_result") +def test_update_test_result_muted_maps_to_inactive(mock_resolve, mock_set, db_session_mock, disposition_perms): + from testgen.mcp.tools.test_results import update_test_result + + mock_resolve.return_value = MagicMock(id=uuid4()) + mock_set.return_value = DispositionUpdate(matched=1, passed_skipped=0) + update_test_result(test_result_id=str(uuid4()), disposition="Muted") + + assert mock_set.call_args.args[1] == Disposition.INACTIVE + + +@patch("testgen.mcp.tools.test_results.set_test_results_disposition") +@patch("testgen.mcp.tools.test_results.resolve_test_result") +def test_update_test_result_no_decision_clears(mock_resolve, mock_set, db_session_mock, disposition_perms): + from testgen.mcp.tools.test_results import update_test_result + + mock_resolve.return_value = MagicMock(id=uuid4()) + mock_set.return_value = DispositionUpdate(matched=1, passed_skipped=0) + update_test_result(test_result_id=str(uuid4()), disposition="No Decision") + + assert mock_set.call_args.args[1] is None + + +@patch("testgen.mcp.tools.test_results.set_test_results_disposition") +@patch("testgen.mcp.tools.test_results.resolve_test_result") +def test_update_test_result_success_message(mock_resolve, mock_set, db_session_mock, disposition_perms): + from testgen.mcp.tools.test_results import update_test_result + + rid = str(uuid4()) + mock_resolve.return_value = MagicMock(id=uuid4()) + mock_set.return_value = DispositionUpdate(matched=1, passed_skipped=0) + out = update_test_result(test_result_id=rid, disposition="Dismissed") + + assert rid in out and "Dismissed" in out + + +@patch("testgen.mcp.tools.test_results.set_test_results_disposition") +@patch("testgen.mcp.tools.test_results.resolve_test_result") +def test_update_test_result_passed_row_is_noop(mock_resolve, mock_set, db_session_mock, disposition_perms): + from testgen.mcp.tools.test_results import update_test_result + + mock_resolve.return_value = MagicMock(id=uuid4()) + mock_set.return_value = DispositionUpdate(matched=0, passed_skipped=1) + out = update_test_result(test_result_id=str(uuid4()), disposition="Confirmed") + + assert "passed" in out and "No change" in out + + +def test_update_test_result_uses_disposition_permission(): + import testgen.mcp.tools.test_results as mod + + closure = {c.cell_contents for c in mod.update_test_result.__wrapped__.__closure__} + assert "disposition" in closure + + +# --------------------------------------------------------------------------- +# bulk_update_test_results +# --------------------------------------------------------------------------- + + +@patch("testgen.mcp.tools.test_results.set_test_results_disposition") +@patch("testgen.mcp.tools.test_results.get_current_session") +@patch("testgen.mcp.tools.test_results.resolve_test_suite") +def test_bulk_update_uses_latest_run_when_run_omitted( + mock_resolve_suite, mock_session, mock_set, db_session_mock, disposition_perms +): + from testgen.mcp.tools.test_results import bulk_update_test_results + + run_id = uuid4() + suite = MagicMock(id=uuid4(), test_suite="suite_a", last_complete_test_run_id=run_id) + mock_resolve_suite.return_value = suite + matched_ids = [uuid4(), uuid4()] + mock_session.return_value.scalars.return_value.all.return_value = matched_ids + mock_set.return_value = DispositionUpdate(matched=2, passed_skipped=0) + + with patch("testgen.mcp.tools.test_results.TestRun.get_by_id_or_job", + return_value=MagicMock(id=run_id, job_execution_id=run_id, test_suite_id=suite.id)): + out = bulk_update_test_results(test_suite_id=str(uuid4()), disposition="Dismissed") + + assert mock_set.call_args.args[0] == matched_ids + assert mock_set.call_args.args[1] == Disposition.DISMISSED + assert "2" in out + + +@patch("testgen.mcp.tools.test_results.resolve_test_suite") +def test_bulk_update_no_completed_run_errors(mock_resolve_suite, db_session_mock, disposition_perms): + from testgen.mcp.tools.test_results import bulk_update_test_results + + mock_resolve_suite.return_value = MagicMock(last_complete_test_run_id=None, test_suite="suite_a") + with pytest.raises(MCPUserError, match="No completed test runs"): + bulk_update_test_results(test_suite_id=str(uuid4()), disposition="Confirmed") + + +@patch("testgen.mcp.tools.test_results.set_test_results_disposition") +@patch("testgen.mcp.tools.test_results.get_current_session") +@patch("testgen.mcp.tools.test_results.resolve_test_suite") +def test_bulk_update_reports_passed_exclusions( + mock_resolve_suite, mock_session, mock_set, db_session_mock, disposition_perms +): + from testgen.mcp.tools.test_results import bulk_update_test_results + + run_id = uuid4() + suite = MagicMock(id=uuid4(), test_suite="suite_a", last_complete_test_run_id=run_id) + mock_resolve_suite.return_value = suite + mock_session.return_value.scalars.return_value.all.return_value = [uuid4(), uuid4(), uuid4()] + mock_set.return_value = DispositionUpdate(matched=2, passed_skipped=1) + + with patch("testgen.mcp.tools.test_results.TestRun.get_by_id_or_job", + return_value=MagicMock(id=run_id, job_execution_id=run_id, test_suite_id=suite.id)): + out = bulk_update_test_results(test_suite_id=str(uuid4()), disposition="Muted") + + assert "2" in out # matched + assert "1" in out and "passed" in out # exclusions surfaced + + +@patch("testgen.mcp.tools.test_results.set_test_results_disposition") +@patch("testgen.mcp.tools.test_results.get_current_session") +@patch("testgen.mcp.tools.test_results.resolve_test_suite") +def test_bulk_update_no_matches(mock_resolve_suite, mock_session, mock_set, db_session_mock, disposition_perms): + from testgen.mcp.tools.test_results import bulk_update_test_results + + run_id = uuid4() + suite = MagicMock(id=uuid4(), test_suite="suite_a", last_complete_test_run_id=run_id) + mock_resolve_suite.return_value = suite + mock_session.return_value.scalars.return_value.all.return_value = [] + mock_set.return_value = DispositionUpdate(matched=0, passed_skipped=0) + + with patch("testgen.mcp.tools.test_results.TestRun.get_by_id_or_job", + return_value=MagicMock(id=run_id, job_execution_id=run_id, test_suite_id=suite.id)): + out = bulk_update_test_results(test_suite_id=str(uuid4()), disposition="Confirmed") + + assert "No test results matched" in out + + +def test_bulk_update_uses_disposition_permission(): + import testgen.mcp.tools.test_results as mod + + closure = {c.cell_contents for c in mod.bulk_update_test_results.__wrapped__.__closure__} + assert "disposition" in closure + + +@patch("testgen.mcp.tools.test_results.set_test_results_disposition") +@patch("testgen.mcp.tools.test_results.get_current_session") +@patch("testgen.mcp.tools.test_results.resolve_test_suite") +def test_bulk_update_explicit_run_in_suite( + mock_resolve_suite, mock_session, mock_set, db_session_mock, disposition_perms +): + from testgen.mcp.tools.test_results import bulk_update_test_results + + suite = MagicMock(id=uuid4(), test_suite="suite_a") + mock_resolve_suite.return_value = suite + run_id = uuid4() + matched_ids = [uuid4()] + mock_session.return_value.scalars.return_value.all.return_value = matched_ids + mock_set.return_value = DispositionUpdate(matched=1, passed_skipped=0) + + with patch( + "testgen.mcp.tools.test_results.TestRun.get_by_id_or_job", + return_value=MagicMock(id=uuid4(), job_execution_id=run_id, test_suite_id=suite.id), + ): + out = bulk_update_test_results( + test_suite_id=str(uuid4()), disposition="Confirmed", job_execution_id=str(run_id) + ) + + assert mock_set.call_args.args[0] == matched_ids + assert "1" in out + + +@patch("testgen.mcp.tools.test_results.resolve_test_suite") +def test_bulk_update_explicit_run_from_other_suite_rejected( + mock_resolve_suite, db_session_mock, disposition_perms +): + from testgen.mcp.tools.test_results import bulk_update_test_results + + suite = MagicMock(id=uuid4(), test_suite="suite_a") + mock_resolve_suite.return_value = suite + + with patch( + "testgen.mcp.tools.test_results.TestRun.get_by_id_or_job", + return_value=MagicMock(test_suite_id=uuid4()), # different suite + ): + with pytest.raises(MCPResourceNotAccessible, match=r"Test run .* not found or not accessible"): + bulk_update_test_results( + test_suite_id=str(uuid4()), disposition="Confirmed", job_execution_id=str(uuid4()) + ) From ed82df7876d9e05ccfe73d72e5555859d516b002 Mon Sep 17 00:00:00 2001 From: Aarthy Adityan Date: Tue, 9 Jun 2026 01:28:35 -0400 Subject: [PATCH 18/78] fix(mcp): support modern TLS ciphers and port-less origins for hosted gateways Hosted MCP gateways (e.g. Databricks Genie and AI Playground) could not connect to the MCP server: - uvicorn served only legacy TLS 1.2 ciphers with no AEAD, so the gateway's hardened TLS client failed the handshake. Serve a modern AEAD cipher set when API TLS is enabled. - The DNS-rebinding allowlist rejected the gateway's port-less Origin header. Allow the bare origin for each TG_MCP_EXTRA_ALLOWED_HOSTS entry. Co-Authored-By: Claude Opus 4.8 (1M context) --- testgen/mcp/server.py | 3 +++ testgen/server/__init__.py | 11 +++++++++++ testgen/settings.py | 3 ++- tests/unit/mcp/test_transport_security.py | 12 +++++++++--- 4 files changed, 25 insertions(+), 4 deletions(-) diff --git a/testgen/mcp/server.py b/testgen/mcp/server.py index ad8cc276..8f94fdac 100644 --- a/testgen/mcp/server.py +++ b/testgen/mcp/server.py @@ -110,6 +110,9 @@ def _build_transport_security() -> TransportSecuritySettings: host_pattern = host if ":" in host else f"{host}:*" allowed_hosts.add(host_pattern) allowed_origins.update({f"http://{host_pattern}", f"https://{host_pattern}"}) + bare = host.split(":", 1)[0] + allowed_hosts.add(bare) + allowed_origins.update({f"http://{bare}", f"https://{bare}"}) return TransportSecuritySettings( enable_dns_rebinding_protection=True, diff --git a/testgen/server/__init__.py b/testgen/server/__init__.py index a1c5afef..9d8dd1d8 100644 --- a/testgen/server/__init__.py +++ b/testgen/server/__init__.py @@ -180,6 +180,17 @@ def run_server() -> None: if settings.API_TLS_ENABLED: ssl_kwargs["ssl_certfile"] = settings.SSL_CERT_FILE ssl_kwargs["ssl_keyfile"] = settings.SSL_KEY_FILE + # uvicorn defaults ssl_ciphers to "TLSv1", which restricts TLS 1.2 to legacy + # CBC/SHA-1 suites with no AEAD. Hardened TLS clients (e.g. the Databricks + # Unity Catalog HTTP-connection gateway) negotiate TLS 1.2 with an AEAD-only + # policy and find no shared cipher, so the handshake fails before any request + # is sent. Offer a modern AEAD cipher set (Mozilla intermediate). + ssl_kwargs["ssl_ciphers"] = ( + "ECDHE-ECDSA-AES128-GCM-SHA256:ECDHE-RSA-AES128-GCM-SHA256:" + "ECDHE-ECDSA-AES256-GCM-SHA384:ECDHE-RSA-AES256-GCM-SHA384:" + "ECDHE-ECDSA-CHACHA20-POLY1305:ECDHE-RSA-CHACHA20-POLY1305:" + "DHE-RSA-AES128-GCM-SHA256:DHE-RSA-AES256-GCM-SHA384" + ) LOG.info( "Starting server on %s:%s (TLS: %s, MCP: %s)", diff --git a/testgen/settings.py b/testgen/settings.py index 69e922dc..ad3ed162 100644 --- a/testgen/settings.py +++ b/testgen/settings.py @@ -634,7 +634,8 @@ def _default_ui_base_url() -> str: Extra Host header values accepted by MCP DNS rebinding protection (comma-separated). BASE_URL's hostname and loopback are always allowed; this adds more for multi-domain deployments or reverse proxies that rewrite Host. Entries without a port (`tg.example.com`) -get an automatic `:*` wildcard; entries with a port are matched literally +are allowed both with an automatic `:*` port wildcard and bare, so gateways that send an +Origin with no port are accepted; entries with a port are matched literally (`tg.example.com:8080`) or with explicit wildcard (`tg.example.com:*`). Only affects MCP routes — the parent FastAPI app does not validate Host headers. diff --git a/tests/unit/mcp/test_transport_security.py b/tests/unit/mcp/test_transport_security.py index 22b101ec..94b250f8 100644 --- a/tests/unit/mcp/test_transport_security.py +++ b/tests/unit/mcp/test_transport_security.py @@ -30,14 +30,20 @@ def test_loopback_and_base_url_always_present(): assert "https://localhost:*" in settings.allowed_origins -def test_extra_host_without_port_gets_wildcard(): - """An extras entry without `:` gets `:*` automatically appended.""" +def test_extra_host_without_port_gets_wildcard_and_bare(): + """An extras entry without `:` is allowed both with a `:*` port wildcard and bare. + + The bare (port-less) form is required because some MCP gateways (e.g. Databricks) + send an Origin with no port, which the `:*` wildcard does not match. + """ settings = _build_with("http://localhost:8530", extras=["tg.example.com"]) assert "tg.example.com:*" in settings.allowed_hosts - assert "tg.example.com" not in settings.allowed_hosts # bare entry should NOT be present + assert "tg.example.com" in settings.allowed_hosts assert "http://tg.example.com:*" in settings.allowed_origins assert "https://tg.example.com:*" in settings.allowed_origins + assert "http://tg.example.com" in settings.allowed_origins + assert "https://tg.example.com" in settings.allowed_origins def test_extra_host_with_explicit_port_preserved_literally(): From fad9a27f13055e8522b2e26b36ad819e384e5db3 Mon Sep 17 00:00:00 2001 From: Luis Date: Mon, 8 Jun 2026 09:57:49 -0400 Subject: [PATCH 19/78] feat(mcp): add data catalog convenience tools - generate_create_table_script - get_table_sample Co-Authored-By: Claude Opus 4.8 (1M context) --- testgen/common/data_catalog_service.py | 119 +++++++++++ testgen/common/models/data_column.py | 53 +++++ testgen/mcp/server.py | 5 +- testgen/mcp/tools/profiling.py | 19 ++ testgen/mcp/tools/source_data.py | 54 ++++- testgen/ui/views/data_catalog.py | 47 ++--- .../dialogs/table_create_script_dialog.py | 26 --- .../unit/common/test_data_catalog_service.py | 194 ++++++++++++++++++ tests/unit/mcp/test_tools_profiling.py | 30 +++ tests/unit/mcp/test_tools_source_data.py | 106 ++++++++++ 10 files changed, 594 insertions(+), 59 deletions(-) create mode 100644 testgen/common/data_catalog_service.py delete mode 100644 testgen/ui/views/dialogs/table_create_script_dialog.py create mode 100644 tests/unit/common/test_data_catalog_service.py diff --git a/testgen/common/data_catalog_service.py b/testgen/common/data_catalog_service.py new file mode 100644 index 00000000..936e7544 --- /dev/null +++ b/testgen/common/data_catalog_service.py @@ -0,0 +1,119 @@ +"""Shared data catalog convenience service. + +Generates CREATE TABLE scripts from profiled column metadata and fetches +sample rows from source tables. Used by both the Streamlit UI and MCP tools. +""" +import logging +from dataclasses import dataclass +from typing import Literal +from uuid import UUID + +import pandas as pd + +from testgen.common.database.database_service import get_flavor_service +from testgen.common.database.flavor.flavor_service import FlavorService +from testgen.common.models.connection import Connection +from testgen.common.models.data_column import CreateScriptColumn, DataColumnChars +from testgen.common.pii_masking import get_pii_columns, mask_source_data_pii +from testgen.ui.services.database_service import fetch_from_target_db +from testgen.utils import to_dataframe + +LOG = logging.getLogger("testgen") + + +@dataclass +class TableSampleResult: + status: Literal["OK", "ND", "ERR"] + message: str | None = None + df: pd.DataFrame | None = None + pii_redacted: bool = False + + +def render_create_table_script( + schema_name: str, + table_name: str, + columns: list[CreateScriptColumn], + flavor_service: FlavorService, + *, + annotate_changes: bool = False, +) -> str: + """Render CREATE TABLE DDL from profiled columns, quoting identifiers for the flavor. + + Column types use the profiling-derived suggestion, falling back to the original + database type. ``annotate_changes`` appends ``-- WAS `` comments where the + suggestion differs from the original type. + """ + quote = flavor_service.quote_character + table_ref = flavor_service.get_table_ref(schema_name, table_name) + quoted_names = [f"{quote}{col.column_name}{quote}" for col in columns] + + name_width = max(len(name) for name in quoted_names) + type_width = max(len(col.datatype_suggestion or col.db_data_type or "") for col in columns) + + col_defs = [] + for index, (col, name) in enumerate(zip(columns, quoted_names, strict=True)): + col_type = col.datatype_suggestion or col.db_data_type or "" + separator = "" if index == len(columns) - 1 else "," + line = f"{name:<{name_width}} {col_type:<{type_width}}{separator}" + if ( + annotate_changes + and col.db_data_type + and col.datatype_suggestion + and col.db_data_type.lower() != col.datatype_suggestion.lower() + ): + line = f"{line} -- WAS {col.db_data_type}" + col_defs.append(line.rstrip()) + + body = "\n ".join(col_defs) + return f"CREATE TABLE {table_ref} (\n {body}\n);" + + +def build_create_table_script( + table_group_id: UUID, table_name: str, *, annotate_changes: bool = False, +) -> str | None: + """Build a CREATE TABLE script for a profiled table, or ``None`` if it is not in the catalog.""" + schema_name, columns = DataColumnChars.list_for_create_script(table_group_id, table_name) + if not columns or schema_name is None: + return None + connection = Connection.get_by_table_group(table_group_id) + if connection is None: + return None + flavor_service = get_flavor_service(connection.sql_flavor) + return render_create_table_script( + schema_name, table_name, columns, flavor_service, annotate_changes=annotate_changes, + ) + + +def fetch_table_sample( + connection: Connection, + table_group_id: UUID, + schema_name: str, + table_name: str, + *, + limit: int, + mask_pii: bool, + column_name: str | None = None, +) -> TableSampleResult: + """Fetch distinct sample rows from a source table, masking PII columns when requested.""" + flavor_service = get_flavor_service(connection.sql_flavor) + prefix, suffix = flavor_service.row_limit_clauses(limit) + quote = flavor_service.quote_character + table_ref = flavor_service.get_table_ref(schema_name, table_name) + columns_expr = f"{quote}{column_name}{quote}" if column_name else "*" + query = f"SELECT DISTINCT {prefix} {columns_expr} FROM {table_ref} {suffix}".strip() + + try: + results = fetch_from_target_db(connection, query) + except Exception: + LOG.exception("Table sample fetch encountered an error.") + return TableSampleResult("ERR", message="The sample data could not be loaded.") + + if not results: + return TableSampleResult("ND") + + df = to_dataframe(results) + pii_redacted = False + if mask_pii: + pii_columns = get_pii_columns(str(table_group_id), schema_name, table_name) + pii_redacted = mask_source_data_pii(df, pii_columns) + return TableSampleResult("OK", df=df, pii_redacted=pii_redacted) diff --git a/testgen/common/models/data_column.py b/testgen/common/models/data_column.py index cee7d088..ed21e312 100644 --- a/testgen/common/models/data_column.py +++ b/testgen/common/models/data_column.py @@ -236,6 +236,13 @@ class ColumnSearchHit(EntityMinimal): column_name: str +@dataclass +class CreateScriptColumn: + column_name: str + db_data_type: str | None + datatype_suggestion: str | None + + class DataColumnChars(Entity): __tablename__ = "data_column_chars" @@ -268,6 +275,52 @@ class DataColumnChars(Entity): # warnings_7_days_prior, warnings_30_days_prior, valid_profile_issue_ct, # valid_test_issue_ct + @classmethod + def list_for_create_script( + cls, table_groups_id: UUID, table_name: str, + ) -> tuple[str | None, list[CreateScriptColumn]]: + """Return ``(schema_name, columns)`` for a table's CREATE TABLE script. + + Columns are ordered by ordinal position and carry the profiling-derived type + suggestion from their latest complete profiling run. Returns ``(None, [])`` when + the table is not in the table group's profiled catalog. + """ + query = ( + select( + cls.schema_name, + cls.column_name, + cls.db_data_type, + ProfileResult.datatype_suggestion, + ) + .outerjoin( + ProfileResult, + and_( + ProfileResult.profile_run_id == cls.last_complete_profile_run_id, + ProfileResult.schema_name == cls.schema_name, + ProfileResult.table_name == cls.table_name, + ProfileResult.column_name == cls.column_name, + ), + ) + .where( + cls.table_groups_id == table_groups_id, + cls.table_name == table_name, + cls.drop_date.is_(None), + ) + .order_by(asc(cls.ordinal_position), asc(cls.column_name)) + ) + rows = get_current_session().execute(query).mappings().all() + if not rows: + return None, [] + columns = [ + CreateScriptColumn( + column_name=row["column_name"], + db_data_type=row["db_data_type"], + datatype_suggestion=row["datatype_suggestion"], + ) + for row in rows + ] + return rows[0]["schema_name"], columns + @classmethod def list_for_table_group( cls, diff --git a/testgen/mcp/server.py b/testgen/mcp/server.py index 8f94fdac..cd673604 100644 --- a/testgen/mcp/server.py +++ b/testgen/mcp/server.py @@ -168,6 +168,7 @@ def build_mcp_server( get_schema_history, ) from testgen.mcp.tools.profiling import ( + generate_create_table_script, get_column_frequent_values, get_column_patterns, get_column_profile_detail, @@ -203,7 +204,7 @@ def build_mcp_server( list_schedules, update_schedule, ) - from testgen.mcp.tools.source_data import get_source_data, get_source_data_query + from testgen.mcp.tools.source_data import get_source_data, get_source_data_query, get_table_sample from testgen.mcp.tools.table_groups import ( create_table_group, preview_table_group, @@ -281,6 +282,8 @@ def safe_prompt(fn): safe_tool(list_test_notes) safe_tool(list_test_types) safe_tool(get_table) + safe_tool(generate_create_table_script) + safe_tool(get_table_sample) safe_tool(list_column_profiles) safe_tool(list_profiling_summaries) safe_tool(list_profiling_runs) diff --git a/testgen/mcp/tools/profiling.py b/testgen/mcp/tools/profiling.py index c861e802..874cd5c3 100644 --- a/testgen/mcp/tools/profiling.py +++ b/testgen/mcp/tools/profiling.py @@ -3,6 +3,7 @@ from sqlalchemy import func, or_ +from testgen.common.data_catalog_service import build_create_table_script from testgen.common.models import with_database_session from testgen.common.models.data_column import ( SUGGESTED_DATA_TYPE_TO_PREFIX, @@ -1136,3 +1137,21 @@ def search_columns( if footer: doc.text(footer) return doc.render() + + +@with_database_session +@mcp_permission("catalog") +def generate_create_table_script(table_group_id: str, table_name: str) -> str: + """Generate a CREATE TABLE script for a profiled table from its columns and suggested data types. + + Args: + table_group_id: UUID of the table group, e.g. from `get_data_inventory`. + table_name: Table name exactly as stored in TestGen (case-sensitive). + """ + tg = resolve_table_group(table_group_id) + + script = build_create_table_script(tg.id, table_name) + if script is None: + raise MCPResourceNotAccessible("Table", table_name) + + return MdDoc().code_block(script, language="sql").render() diff --git a/testgen/mcp/tools/source_data.py b/testgen/mcp/tools/source_data.py index bcab0191..91321860 100644 --- a/testgen/mcp/tools/source_data.py +++ b/testgen/mcp/tools/source_data.py @@ -1,6 +1,9 @@ from datetime import datetime +from testgen.common.data_catalog_service import fetch_table_sample from testgen.common.models import with_database_session +from testgen.common.models.connection import Connection +from testgen.common.models.data_column import DataColumnChars from testgen.common.models.profiling_run import ProfilingRun from testgen.common.models.test_definition import TestDefinition from testgen.common.source_data_service import ( @@ -12,7 +15,13 @@ ) 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, resolve_hygiene_issue, validate_limit +from testgen.mcp.tools.common import ( + DocGroup, + parse_uuid, + resolve_hygiene_issue, + resolve_table_group, + validate_limit, +) from testgen.mcp.tools.markdown import MdDoc _DOC_GROUP = DocGroup.INVESTIGATE @@ -206,3 +215,46 @@ def get_source_data( doc.code_block(result.query, language="sql") return doc.render() + + +@with_database_session +@mcp_permission("catalog") +def get_table_sample(table_group_id: str, table_name: str, limit: int = 100) -> str: + """Fetch sample rows from a source table for inspection. + + Args: + table_group_id: UUID of the table group, e.g. from `get_data_inventory`. + table_name: Table name exactly as stored in TestGen (case-sensitive). + limit: Maximum rows to return (default 100, max 500). + """ + validate_limit(limit, 500) + tg = resolve_table_group(table_group_id) + + schema_name, _ = DataColumnChars.list_for_create_script(tg.id, table_name) + if schema_name is None: + raise MCPResourceNotAccessible("Table", table_name) + + connection = Connection.get_by_table_group(tg.id) + if connection is None: + raise MCPResourceNotAccessible("Table", table_name) + + mask_pii = not get_project_permissions().has_permission("view_pii", tg.project_code) + result = fetch_table_sample( + connection, tg.id, schema_name, table_name, limit=limit, mask_pii=mask_pii, + ) + + if result.status == "ERR": + raise MCPUserError( + f"Could not read from the source database for connection `{connection.connection_name}`." + ) + + doc = MdDoc() + if result.status == "ND": + return doc.text("Table has no rows.").render() + + row_count = len(result.df) if result.df is not None else 0 + doc.field("Rows returned", row_count) + if result.pii_redacted: + doc.text("_PII columns have been redacted._") + doc.table_from_dataframe(result.df) + return doc.render() diff --git a/testgen/ui/views/data_catalog.py b/testgen/ui/views/data_catalog.py index d316f237..f5db0893 100644 --- a/testgen/ui/views/data_catalog.py +++ b/testgen/ui/views/data_catalog.py @@ -10,7 +10,7 @@ from sqlalchemy.sql.expression import func as sa_func from streamlit.delta_generator import DeltaGenerator -from testgen.common.database.database_service import get_flavor_service +from testgen.common.data_catalog_service import build_create_table_script, fetch_table_sample from testgen.common.enums import JobSource from testgen.common.models import database_session, with_database_session from testgen.common.models.connection import Connection @@ -22,7 +22,6 @@ get_pii_columns, mask_hygiene_detail, mask_profiling_pii, - mask_source_data_pii, ) from testgen.common.profile_top_values import parse_top_freq_values, parse_top_patterns from testgen.ui.components import widgets as testgen @@ -46,7 +45,7 @@ get_tables_by_id, get_tables_by_table_group, ) -from testgen.ui.services.database_service import execute_db_query, fetch_all_from_db, fetch_from_target_db +from testgen.ui.services.database_service import execute_db_query, fetch_all_from_db from testgen.ui.services.query_cache import ( get_profiling_run_summaries, get_project_summary, @@ -61,7 +60,6 @@ build_import_preview_props, parse_import_csv, ) -from testgen.ui.views.dialogs.table_create_script_dialog import generate_create_script from testgen.utils import friendly_score, is_uuid4, make_json_safe, score LOG = logging.getLogger("testgen") @@ -255,12 +253,6 @@ def on_data_preview_clicked(item) -> None: item["table_name"], item.get("column_name"), ) - if preview_data.get("rows") and not session.auth.user_has_permission("view_pii"): - pii_columns = get_pii_columns(item["table_group_id"], item["schema_name"], item["table_name"]) - if pii_columns: - df = pd.DataFrame(preview_data["rows"], columns=preview_data["columns"]) - mask_source_data_pii(df, pii_columns) - preview_data["rows"] = make_json_safe(df.values.tolist()) st.session_state[DC_DATA_PREVIEW_DIALOG_KEY] = preview_data def on_data_preview_dialog_closed(*_) -> None: @@ -302,7 +294,7 @@ def on_history_dialog_closed(*_) -> None: create_script_dialog_data = None if create_script_item := st.session_state.get(DC_CREATE_SCRIPT_DIALOG_KEY): - script = generate_create_script(create_script_item["table_name"], columns) + script = build_create_table_script(table_group_id, create_script_item["table_name"], annotate_changes=True) create_script_dialog_data = { "title": f"Table CREATE Script: {create_script_item['table_name']}", "table_name": create_script_item["table_name"], @@ -906,32 +898,25 @@ def get_preview_data( if not connection: return {"title": title, "status": "ERR", "message": "Connection not found."} - flavor_service = get_flavor_service(connection.sql_flavor) - prefix, suffix = flavor_service.row_limit_clauses(100) - quote = flavor_service.quote_character - table_ref = flavor_service.get_table_ref(schema_name, table_name) - query = f""" - SELECT DISTINCT - {prefix} - {f"{quote}{column_name}{quote}" if column_name else "*"} - FROM {table_ref} - {suffix} - """ + result = fetch_table_sample( + connection, + table_group_id, + schema_name, + table_name, + limit=100, + mask_pii=not session.auth.user_has_permission("view_pii"), + column_name=column_name, + ) - try: - results = fetch_from_target_db(connection, query) - except Exception: + if result.status == "ERR": return {"title": title, "status": "ERR", "message": "The preview data could not be loaded."} - - if not results: + if result.status == "ND" or result.df is None: return {"title": title, "status": "ND", "message": "No data found."} - columns_list = list(results[0].keys()) - rows = [list(row.values()) for row in results] return { "title": title, - "columns": columns_list, - "rows": make_json_safe(rows), + "columns": list(result.df.columns), + "rows": make_json_safe(result.df.values.tolist()), } diff --git a/testgen/ui/views/dialogs/table_create_script_dialog.py b/testgen/ui/views/dialogs/table_create_script_dialog.py deleted file mode 100644 index b121e4e4..00000000 --- a/testgen/ui/views/dialogs/table_create_script_dialog.py +++ /dev/null @@ -1,26 +0,0 @@ -def generate_create_script(table_name: str, data: list[dict]) -> str | None: - table_data = [col for col in data if col["table_name"] == table_name] - if not table_data: - return None - - max_name = max(len(col["column_name"]) for col in table_data) + 3 - max_type = max(len(col["datatype_suggestion"] or "") for col in table_data) + 3 - - col_defs = [] - for index, col in enumerate(table_data): - comment = ( - f"-- WAS {col['db_data_type']}" - if isinstance(col["db_data_type"], str) - and isinstance(col["datatype_suggestion"], str) - and col["db_data_type"].lower() != col["datatype_suggestion"].lower() - else "" - ) - col_type = col["datatype_suggestion"] or col["db_data_type"] or "" - separator = " " if index == len(table_data) - 1 else "," - col_defs.append(f"{col['column_name']:<{max_name}} {(col_type):<{max_type}}{separator} {comment}") - - col_defs_joined = "\n ".join(col_defs) - return f""" -CREATE TABLE {table_data[0]['schema_name']}.{table_data[0]['table_name']} ( - {col_defs_joined} -);""" diff --git a/tests/unit/common/test_data_catalog_service.py b/tests/unit/common/test_data_catalog_service.py new file mode 100644 index 00000000..335ae6ce --- /dev/null +++ b/tests/unit/common/test_data_catalog_service.py @@ -0,0 +1,194 @@ +from unittest.mock import MagicMock, patch + +import pytest + +from testgen.common.data_catalog_service import ( + TableSampleResult, + build_create_table_script, + fetch_table_sample, + render_create_table_script, +) +from testgen.common.models.data_column import CreateScriptColumn + +pytestmark = pytest.mark.unit + +MODULE = "testgen.common.data_catalog_service" + + +def _flavor_service(table_ref='"demo"."orders"', limit_clauses=("", "LIMIT 100")): + fs = MagicMock() + fs.quote_character = '"' + fs.get_table_ref.return_value = table_ref + fs.row_limit_clauses.return_value = limit_clauses + return fs + + +def _column(name, db_type="varchar(50)", suggestion="VARCHAR(20)"): + return CreateScriptColumn(column_name=name, db_data_type=db_type, datatype_suggestion=suggestion) + + +# --------------------------------------------------------------------------- +# render_create_table_script +# --------------------------------------------------------------------------- + +def test_render_uses_suggested_type_and_quotes_identifiers(): + fs = _flavor_service() + columns = [_column("id", "int", "INTEGER"), _column("name", "text", "VARCHAR(20)")] + + script = render_create_table_script("demo", "orders", columns, fs) + + assert script.startswith('CREATE TABLE "demo"."orders" (') + assert '"id"' in script + assert '"name"' in script + assert "INTEGER" in script + assert "VARCHAR(20)" in script + + +def test_render_omits_was_annotations_by_default(): + fs = _flavor_service() + columns = [_column("name", db_type="text", suggestion="VARCHAR(20)")] + + script = render_create_table_script("demo", "orders", columns, fs) + + assert "-- WAS" not in script + + +def test_render_includes_was_annotations_when_enabled(): + fs = _flavor_service() + columns = [_column("name", db_type="text", suggestion="VARCHAR(20)")] + + script = render_create_table_script("demo", "orders", columns, fs, annotate_changes=True) + + assert "-- WAS text" in script + + +def test_render_no_was_annotation_when_type_unchanged(): + fs = _flavor_service() + columns = [_column("name", db_type="VARCHAR(20)", suggestion="VARCHAR(20)")] + + script = render_create_table_script("demo", "orders", columns, fs, annotate_changes=True) + + assert "-- WAS" not in script + + +def test_render_falls_back_to_db_type_when_no_suggestion(): + fs = _flavor_service() + columns = [_column("name", db_type="text", suggestion=None)] + + script = render_create_table_script("demo", "orders", columns, fs) + + assert "text" in script + + +def test_render_only_last_column_has_no_trailing_comma(): + fs = _flavor_service() + columns = [_column("a", "int", "INTEGER"), _column("b", "int", "INTEGER")] + + script = render_create_table_script("demo", "orders", columns, fs) + + body_lines = [line for line in script.splitlines() if '"a"' in line or '"b"' in line] + assert body_lines[0].rstrip().endswith(",") + assert not body_lines[1].rstrip().endswith(",") + + +# --------------------------------------------------------------------------- +# build_create_table_script +# --------------------------------------------------------------------------- + +@patch(f"{MODULE}.DataColumnChars") +def test_build_returns_none_when_no_columns(mock_columns): + mock_columns.list_for_create_script.return_value = (None, []) + + assert build_create_table_script("tg-id", "orders") is None + + +@patch(f"{MODULE}.get_flavor_service") +@patch(f"{MODULE}.Connection") +@patch(f"{MODULE}.DataColumnChars") +def test_build_renders_script_for_existing_table(mock_columns, mock_conn, mock_flavor): + mock_columns.list_for_create_script.return_value = ("demo", [_column("id", "int", "INTEGER")]) + mock_conn.get_by_table_group.return_value = MagicMock(sql_flavor="postgresql") + mock_flavor.return_value = _flavor_service() + + script = build_create_table_script("tg-id", "orders") + + assert script is not None + assert "CREATE TABLE" in script + assert "INTEGER" in script + + +# --------------------------------------------------------------------------- +# fetch_table_sample +# --------------------------------------------------------------------------- + +@patch(f"{MODULE}.fetch_from_target_db") +@patch(f"{MODULE}.get_flavor_service") +def test_fetch_sample_ok(mock_flavor, mock_fetch): + mock_flavor.return_value = _flavor_service() + mock_fetch.return_value = [{"id": 1, "name": "a"}, {"id": 2, "name": "b"}] + connection = MagicMock(sql_flavor="postgresql") + + result = fetch_table_sample(connection, "tg-id", "demo", "orders", limit=100, mask_pii=False) + + assert result.status == "OK" + assert len(result.df) == 2 + assert result.pii_redacted is False + + +@patch(f"{MODULE}.fetch_from_target_db") +@patch(f"{MODULE}.get_flavor_service") +def test_fetch_sample_builds_distinct_limited_query(mock_flavor, mock_fetch): + mock_flavor.return_value = _flavor_service(limit_clauses=("", "LIMIT 50")) + mock_fetch.return_value = [{"id": 1}] + connection = MagicMock(sql_flavor="postgresql") + + fetch_table_sample(connection, "tg-id", "demo", "orders", limit=50, mask_pii=False) + + query = mock_fetch.call_args.args[1] + assert "SELECT DISTINCT" in query + assert "LIMIT 50" in query + assert '"demo"."orders"' in query + + +@patch(f"{MODULE}.fetch_from_target_db") +@patch(f"{MODULE}.get_flavor_service") +def test_fetch_sample_empty_returns_nd(mock_flavor, mock_fetch): + mock_flavor.return_value = _flavor_service() + mock_fetch.return_value = [] + connection = MagicMock(sql_flavor="postgresql") + + result = fetch_table_sample(connection, "tg-id", "demo", "orders", limit=100, mask_pii=False) + + assert result.status == "ND" + + +@patch(f"{MODULE}.fetch_from_target_db", side_effect=Exception("connection refused")) +@patch(f"{MODULE}.get_flavor_service") +def test_fetch_sample_error_returns_err(mock_flavor, _mock_fetch): + mock_flavor.return_value = _flavor_service() + connection = MagicMock(sql_flavor="postgresql") + + result = fetch_table_sample(connection, "tg-id", "demo", "orders", limit=100, mask_pii=False) + + assert result.status == "ERR" + + +@patch(f"{MODULE}.get_pii_columns", return_value={"name"}) +@patch(f"{MODULE}.fetch_from_target_db") +@patch(f"{MODULE}.get_flavor_service") +def test_fetch_sample_masks_pii_columns(mock_flavor, mock_fetch, _mock_pii): + mock_flavor.return_value = _flavor_service() + mock_fetch.return_value = [{"id": 1, "name": "secret"}] + connection = MagicMock(sql_flavor="postgresql") + + result = fetch_table_sample(connection, "tg-id", "demo", "orders", limit=100, mask_pii=True) + + assert result.status == "OK" + assert result.pii_redacted is True + assert "secret" not in result.df["name"].tolist() + + +def test_table_sample_result_defaults(): + result = TableSampleResult("ND") + assert result.df is None + assert result.pii_redacted is False diff --git a/tests/unit/mcp/test_tools_profiling.py b/tests/unit/mcp/test_tools_profiling.py index 588bd42d..733902f8 100644 --- a/tests/unit/mcp/test_tools_profiling.py +++ b/tests/unit/mcp/test_tools_profiling.py @@ -1898,3 +1898,33 @@ def test_search_columns_table_group_scope_skips_per_project_summary( assert "Matches by project" not in result mock_dcc_cls.summarize_matches_by_project.assert_not_called() + + +# ---------------------------------------------------------------------- +# generate_create_table_script +# ---------------------------------------------------------------------- + +@patch("testgen.mcp.tools.profiling.build_create_table_script") +@patch("testgen.mcp.tools.common.TableGroup") +def test_generate_create_table_script_happy_path(mock_tg_cls, mock_build, db_session_mock): + mock_tg_cls.get.return_value = _mock_table_group() + mock_build.return_value = 'CREATE TABLE "demo"."orders" (\n "id" INTEGER\n);' + + from testgen.mcp.tools.profiling import generate_create_table_script + result = generate_create_table_script(str(uuid4()), "orders") + + assert "CREATE TABLE" in result + assert '"id"' in result + assert "-- WAS" not in result + # MCP tool requests clean DDL (no change annotations) + assert mock_build.call_args.kwargs.get("annotate_changes", False) is False + + +@patch("testgen.mcp.tools.profiling.build_create_table_script", return_value=None) +@patch("testgen.mcp.tools.common.TableGroup") +def test_generate_create_table_script_unknown_table(mock_tg_cls, _mock_build, db_session_mock): + mock_tg_cls.get.return_value = _mock_table_group() + + from testgen.mcp.tools.profiling import generate_create_table_script + with pytest.raises(MCPResourceNotAccessible): + generate_create_table_script(str(uuid4()), "missing") diff --git a/tests/unit/mcp/test_tools_source_data.py b/tests/unit/mcp/test_tools_source_data.py index 1f1f26b1..a7370982 100644 --- a/tests/unit/mcp/test_tools_source_data.py +++ b/tests/unit/mcp/test_tools_source_data.py @@ -4,6 +4,7 @@ import pandas as pd import pytest +from testgen.common.data_catalog_service import TableSampleResult from testgen.common.source_data_service import SourceDataResult from testgen.mcp.exceptions import MCPResourceNotAccessible, MCPUserError from testgen.mcp.permissions import ProjectPermissions @@ -493,3 +494,108 @@ def test_get_source_data_hygiene_invalid_uuid(db_session_mock): with pytest.raises(MCPUserError, match="not a valid UUID"): get_source_data(issue_id="bad-uuid") + + +# ---------------------------------------------------------------------- +# get_table_sample +# ---------------------------------------------------------------------- + +def _mock_table_group(tg_id=None, project_code="demo"): + tg = MagicMock() + tg.id = tg_id or uuid4() + tg.project_code = project_code + return tg + + +@patch("testgen.mcp.tools.source_data.fetch_table_sample") +@patch("testgen.mcp.tools.source_data.Connection") +@patch("testgen.mcp.tools.source_data.DataColumnChars") +@patch("testgen.mcp.tools.common.TableGroup") +def test_get_table_sample_happy_path(mock_tg_cls, mock_dc, mock_conn, mock_fetch, db_session_mock): + mock_tg_cls.get.return_value = _mock_table_group() + mock_dc.list_for_create_script.return_value = ("demo", []) + mock_conn.get_by_table_group.return_value = MagicMock(connection_name="main") + mock_fetch.return_value = TableSampleResult( + "OK", df=pd.DataFrame([{"id": 1, "name": "a"}]), pii_redacted=False, + ) + + from testgen.mcp.tools.source_data import get_table_sample + result = get_table_sample(str(uuid4()), "orders") + + assert "id" in result + assert "name" in result + assert "PII" not in result + + +@patch("testgen.mcp.tools.source_data.fetch_table_sample") +@patch("testgen.mcp.tools.source_data.Connection") +@patch("testgen.mcp.tools.source_data.DataColumnChars") +@patch("testgen.mcp.tools.common.TableGroup") +def test_get_table_sample_empty_table(mock_tg_cls, mock_dc, mock_conn, mock_fetch, db_session_mock): + mock_tg_cls.get.return_value = _mock_table_group() + mock_dc.list_for_create_script.return_value = ("demo", []) + mock_conn.get_by_table_group.return_value = MagicMock(connection_name="main") + mock_fetch.return_value = TableSampleResult("ND") + + from testgen.mcp.tools.source_data import get_table_sample + result = get_table_sample(str(uuid4()), "orders") + + assert result.strip() == "Table has no rows." + + +@patch("testgen.mcp.tools.source_data.fetch_table_sample") +@patch("testgen.mcp.tools.source_data.Connection") +@patch("testgen.mcp.tools.source_data.DataColumnChars") +@patch("testgen.mcp.tools.common.TableGroup") +def test_get_table_sample_redaction_note(mock_tg_cls, mock_dc, mock_conn, mock_fetch, db_session_mock): + mock_tg_cls.get.return_value = _mock_table_group() + mock_dc.list_for_create_script.return_value = ("demo", []) + mock_conn.get_by_table_group.return_value = MagicMock(connection_name="main") + mock_fetch.return_value = TableSampleResult( + "OK", df=pd.DataFrame([{"id": 1}]), pii_redacted=True, + ) + + from testgen.mcp.tools.source_data import get_table_sample + result = get_table_sample(str(uuid4()), "orders") + + assert "redacted" in result.lower() + + +@patch("testgen.mcp.tools.source_data.fetch_table_sample") +@patch("testgen.mcp.tools.source_data.Connection") +@patch("testgen.mcp.tools.source_data.DataColumnChars") +@patch("testgen.mcp.tools.common.TableGroup") +def test_get_table_sample_connection_failure_names_connection( + mock_tg_cls, mock_dc, mock_conn, mock_fetch, db_session_mock, +): + mock_tg_cls.get.return_value = _mock_table_group() + mock_dc.list_for_create_script.return_value = ("demo", []) + mock_conn.get_by_table_group.return_value = MagicMock(connection_name="prod-warehouse") + mock_fetch.return_value = TableSampleResult("ERR", message="boom") + + from testgen.mcp.tools.source_data import get_table_sample + with pytest.raises(MCPUserError) as exc: + get_table_sample(str(uuid4()), "orders") + + assert "prod-warehouse" in str(exc.value) + + +@patch("testgen.mcp.tools.source_data.DataColumnChars") +@patch("testgen.mcp.tools.common.TableGroup") +def test_get_table_sample_unknown_table(mock_tg_cls, mock_dc, db_session_mock): + mock_tg_cls.get.return_value = _mock_table_group() + mock_dc.list_for_create_script.return_value = (None, []) + + from testgen.mcp.tools.source_data import get_table_sample + with pytest.raises(MCPResourceNotAccessible): + get_table_sample(str(uuid4()), "missing") + + +@pytest.mark.parametrize("bad_limit", [0, 501]) +@patch("testgen.mcp.tools.common.TableGroup") +def test_get_table_sample_rejects_out_of_range_limit(mock_tg_cls, bad_limit, db_session_mock): + mock_tg_cls.get.return_value = _mock_table_group() + + from testgen.mcp.tools.source_data import get_table_sample + with pytest.raises(MCPUserError): + get_table_sample(str(uuid4()), "orders", limit=bad_limit) From 832f7b531807940db966836f6fda214536678244 Mon Sep 17 00:00:00 2001 From: Aarthy Adityan Date: Tue, 9 Jun 2026 17:22:08 -0400 Subject: [PATCH 20/78] fix: match SQL Server schema names case-sensitively in profiling DDF On SQL Server's default (case-insensitive) collation, a table group whose configured schema case differs from the database's actual case (e.g. 'common' vs 'Common') silently broke Freshness monitor generation and Data Catalog stats. Profiling stamps profile_results.schema_name with the entered literal while the DDF stamps data_table_chars/data_column_chars with the DB's actual case, so the case-sensitive PostgreSQL joins between them returned nothing; Volume_Trend (data_table_chars only) still worked, masking it. Add COLLATE Latin1_General_BIN to the schema comparison in the mssql get_schema_ddf so matching is case-sensitive, consistent with TestGen's case-sensitive joins and the other flavors. Migration 0194 realigns existing table_group_schema values to the DB's actual case (from data_table_chars), scoped to mssql connections, single unambiguous schema, case-only differences. Co-Authored-By: Claude Opus 4.8 (1M context) --- .../dbupgrade/0194_incremental_upgrade.sql | 34 +++++++++++++++++++ .../mssql/data_chars/get_schema_ddf.sql | 2 +- 2 files changed, 35 insertions(+), 1 deletion(-) create mode 100644 testgen/template/dbupgrade/0194_incremental_upgrade.sql diff --git a/testgen/template/dbupgrade/0194_incremental_upgrade.sql b/testgen/template/dbupgrade/0194_incremental_upgrade.sql new file mode 100644 index 00000000..683afc8c --- /dev/null +++ b/testgen/template/dbupgrade/0194_incremental_upgrade.sql @@ -0,0 +1,34 @@ +SET SEARCH_PATH TO {SCHEMA_NAME}; + +-- Schema-name matching in the SQL Server DDF query is now case-sensitive, +-- consistent with TestGen's case-sensitive joins and the other flavors. +-- A case-insensitive source previously accepted a table +-- group whose configured schema name differed in case from the database; +-- such a group would now match no tables on its next profiling run. +-- +-- data_table_chars.schema_name holds the case the source database actually +-- reported (it is populated from the DDF's c.table_schema, not the entered +-- value), so realign table_group_schema to that case. Guards: +-- * only mssql connections -- the DDF change was made only for that flavor, +-- * only when the catalog reports a single unambiguous schema for the group, +-- * only a case-only difference (LOWER() match), never a different schema, +-- * '<>' is case-sensitive in PostgreSQL, so correct rows are left untouched. +-- After this, the group's next profiling run stamps profile_results with the +-- corrected case, realigning the case-sensitive downstream joins. +UPDATE table_groups tg +SET table_group_schema = actual.schema_name +FROM ( + SELECT table_groups_id, MIN(schema_name) AS schema_name + FROM data_table_chars + WHERE drop_date IS NULL + GROUP BY table_groups_id + HAVING COUNT(DISTINCT schema_name) = 1 +) actual +WHERE actual.table_groups_id = tg.id + AND tg.table_group_schema <> actual.schema_name + AND LOWER(tg.table_group_schema) = LOWER(actual.schema_name) + AND EXISTS ( + SELECT 1 FROM connections cn + WHERE cn.connection_id = tg.connection_id + AND cn.sql_flavor = 'mssql' + ); diff --git a/testgen/template/flavors/mssql/data_chars/get_schema_ddf.sql b/testgen/template/flavors/mssql/data_chars/get_schema_ddf.sql index 63423448..8fbda543 100644 --- a/testgen/template/flavors/mssql/data_chars/get_schema_ddf.sql +++ b/testgen/template/flavors/mssql/data_chars/get_schema_ddf.sql @@ -54,5 +54,5 @@ SELECT FROM information_schema.columns c LEFT JOIN approx_cts a ON c.table_schema = a.schema_name AND c.table_name = a.table_name LEFT JOIN information_schema.tables it ON c.table_schema = it.table_schema AND c.table_name = it.table_name -WHERE c.table_schema = '{DATA_SCHEMA}' {TABLE_CRITERIA} +WHERE c.table_schema = '{DATA_SCHEMA}' COLLATE Latin1_General_BIN {TABLE_CRITERIA} ORDER BY c.table_schema, c.table_name, c.ordinal_position; From 1334876b5d0d1db82985da14e93e331b630be795 Mon Sep 17 00:00:00 2001 From: Ricardo Boni Date: Sat, 6 Jun 2026 07:42:18 -0400 Subject: [PATCH 21/78] refactor(runs): deduplicate run tables against job_executions Make the run primary key the job execution id: promote job_execution_id to the primary key (also a foreign key to job_executions with cascade on delete) and drop the duplicated status, end-time, log-message, and process-id columns. Lifecycle status and timestamps are read through the job_execution relationship; the run tables keep only their result and progress data. - Stop writing the duplicated columns in the profiling and test handlers - Collapse dual-id lookups to a single id lookup (get_by_id_or_job -> get, drop get_job_execution_ids); the run id is the job execution id - Collapse the hand-rolled cascade delete to a single job_executions delete now that the foreign-key chain cascades - Migration 0193: rewrite external references to the job execution id, promote the primary key, add the cascade foreign keys, drop the duplicated columns - Repoint score-rollup and entity-list SQL templates to job_executions for run status Co-Authored-By: Claude Opus 4.8 --- testgen/api/runs.py | 4 +- testgen/commands/job_registry.py | 6 +- testgen/commands/job_runner.py | 4 +- testgen/commands/run_data_cleanup.py | 10 +- testgen/commands/run_profiling.py | 20 +-- testgen/commands/run_score_update.py | 4 +- testgen/commands/run_test_execution.py | 24 +-- testgen/common/models/data_column.py | 10 +- testgen/common/models/data_table.py | 2 +- testgen/common/models/hygiene_issue.py | 10 +- testgen/common/models/profiling_run.py | 97 ++++------- testgen/common/models/table_group.py | 9 +- testgen/common/models/test_result.py | 2 +- testgen/common/models/test_run.py | 81 +++------- testgen/common/models/test_suite.py | 7 +- testgen/common/notifications/profiling_run.py | 4 +- testgen/common/notifications/test_run.py | 4 +- testgen/mcp/tools/common.py | 2 +- testgen/mcp/tools/discovery.py | 8 +- testgen/mcp/tools/hygiene_issues.py | 4 +- testgen/mcp/tools/profile_history.py | 8 +- testgen/mcp/tools/profiling.py | 12 +- testgen/mcp/tools/test_results.py | 22 +-- .../030_initialize_new_schema_structure.sql | 35 ++-- .../dbupgrade/0194_incremental_upgrade.sql | 150 ++++++++++++++++++ .../get_errored_autogen_monitors.sql | 7 +- .../get_entities/get_profile_list.sql | 6 +- .../get_entities/get_test_run_list.sql | 8 +- .../rollup_scores_profile_table_group.sql | 3 +- .../rollup_scores_test_table_group.sql | 3 +- .../template/score_cards/add_latest_runs.sql | 11 +- testgen/ui/queries/profiling_queries.py | 2 +- testgen/ui/queries/scoring_queries.py | 4 +- testgen/ui/queries/test_result_queries.py | 2 +- testgen/ui/views/profiling_runs.py | 5 +- testgen/ui/views/test_runs.py | 5 +- tests/unit/api/test_runs.py | 8 +- tests/unit/commands/test_run_data_cleanup.py | 17 +- .../test_profiling_run_notifications.py | 1 - .../test_test_run_notifications.py | 1 - tests/unit/mcp/test_model_hygiene_issue.py | 6 +- tests/unit/mcp/test_model_profiling_run.py | 6 +- tests/unit/mcp/test_tools_common.py | 6 +- tests/unit/mcp/test_tools_discovery.py | 8 +- tests/unit/mcp/test_tools_hygiene_issues.py | 9 +- tests/unit/mcp/test_tools_profiling.py | 28 ++-- tests/unit/mcp/test_tools_test_results.py | 86 +++++----- 47 files changed, 411 insertions(+), 360 deletions(-) create mode 100644 testgen/template/dbupgrade/0194_incremental_upgrade.sql diff --git a/testgen/api/runs.py b/testgen/api/runs.py index be2f62f1..4e6b7fb5 100644 --- a/testgen/api/runs.py +++ b/testgen/api/runs.py @@ -35,7 +35,7 @@ ) def get_test_run(job: JobExecution = resolve_job("view", JobExecution.job_key == JobKey.run_tests)): # noqa: B008 """Get a test run by the job execution ID that created it.""" - test_run = TestRun.get_by_id_or_job(job.id) + test_run = TestRun.get(job.id) result = None if test_run: @@ -76,7 +76,7 @@ def get_test_run(job: JobExecution = resolve_job("view", JobExecution.job_key == ) def get_profiling_run(job: JobExecution = resolve_job("view", JobExecution.job_key == JobKey.run_profile)): # noqa: B008 """Get a profiling run by the job execution ID that created it.""" - profiling_run = ProfilingRun.get_by_id_or_job(job.id) + profiling_run = ProfilingRun.get(job.id) result = None if profiling_run: diff --git a/testgen/commands/job_registry.py b/testgen/commands/job_registry.py index 3fa3dceb..f41ee4e1 100644 --- a/testgen/commands/job_registry.py +++ b/testgen/commands/job_registry.py @@ -80,7 +80,7 @@ def run_final_callbacks(job_exec: JobExecution) -> None: def _notify_profiling_run(job_exec: JobExecution) -> None: with database_session() as session: profiling_run = session.scalars( - select(ProfilingRun).where(ProfilingRun.job_execution_id == job_exec.id) + select(ProfilingRun).where(ProfilingRun.id == job_exec.id) ).first() if not profiling_run: LOG.warning("No profiling_run found for job %s; skipping notification", job_exec.id) @@ -90,7 +90,7 @@ def _notify_profiling_run(job_exec: JobExecution) -> None: def _notify_test_run(job_exec: JobExecution) -> None: with database_session() as session: - test_run = session.scalars(select(TestRun).where(TestRun.job_execution_id == job_exec.id)).first() + test_run = session.scalars(select(TestRun).where(TestRun.id == job_exec.id)).first() if not test_run: LOG.warning("No test_run found for job %s; skipping notification", job_exec.id) return @@ -99,7 +99,7 @@ def _notify_test_run(job_exec: JobExecution) -> None: def _notify_monitor_run(job_exec: JobExecution) -> None: with database_session() as session: - test_run = session.scalars(select(TestRun).where(TestRun.job_execution_id == job_exec.id)).first() + test_run = session.scalars(select(TestRun).where(TestRun.id == job_exec.id)).first() if not test_run: LOG.warning("No test_run found for job %s; skipping monitor notification", job_exec.id) return diff --git a/testgen/commands/job_runner.py b/testgen/commands/job_runner.py index d1b7b068..5643999e 100644 --- a/testgen/commands/job_runner.py +++ b/testgen/commands/job_runner.py @@ -71,12 +71,12 @@ def _print_run_summary(job_id: UUID, job_key: str) -> None: session = get_current_session() match job_key: case "run-profile": - run = session.scalars(select(ProfilingRun).where(ProfilingRun.job_execution_id == job_id)).first() + run = session.scalars(select(ProfilingRun).where(ProfilingRun.id == job_id)).first() if run: status_msg = "Profiling encountered an error. Check log for details." if run.job_execution.status == JobStatus.ERROR else "Profiling completed." click.echo(f"\n {status_msg}\n Run ID: {run.id}\n ") case "run-tests" | "run-monitors": - run = session.scalars(select(TestRun).where(TestRun.job_execution_id == job_id)).first() + run = session.scalars(select(TestRun).where(TestRun.id == job_id)).first() if run: status_msg = "Test execution encountered an error. Check log for details." if run.job_execution.status == JobStatus.ERROR else "Test execution completed." click.echo(f"\n {status_msg}\n Run ID: {run.id}\n ") diff --git a/testgen/commands/run_data_cleanup.py b/testgen/commands/run_data_cleanup.py index 08e6341c..0dbb0799 100644 --- a/testgen/commands/run_data_cleanup.py +++ b/testgen/commands/run_data_cleanup.py @@ -42,13 +42,9 @@ def run_data_cleanup(project_code: str, retention_days: int) -> None: with database_session(): protected_profiling_ids = ProfilingRun.find_latest_per_table_group(project_code) protected_test_run_ids = TestRun.find_latest_per_test_suite(project_code) - # Translate protected run ids → their job_execution_ids so the JE sweep - # can carve them out. Nulls (older runs without a JE) are filtered here. - je_map = { - **ProfilingRun.get_job_execution_ids(list(protected_profiling_ids)), - **TestRun.get_job_execution_ids(list(protected_test_run_ids)), - } - protected_job_execution_ids = {je for je in je_map.values() if je is not None} + # The run id is the job execution id, so the protected run ids are + # already the job executions the sweep must carve out. + protected_job_execution_ids = protected_profiling_ids | protected_test_run_ids LOG.info( "Protected latest runs: profiling=%d test=%d job_executions=%d", diff --git a/testgen/commands/run_profiling.py b/testgen/commands/run_profiling.py index fd698eb5..0748f57b 100644 --- a/testgen/commands/run_profiling.py +++ b/testgen/commands/run_profiling.py @@ -1,5 +1,4 @@ import logging -import os from datetime import UTC, datetime, timedelta from uuid import UUID @@ -52,14 +51,12 @@ def run_profiling( LOG.info("Creating profiling run record") profiling_run = ProfilingRun( + id=job_context.get().job_id, project_code=table_group.project_code, connection_id=connection.connection_id, table_groups_id=table_group.id, profiling_starttime=datetime.now(UTC) + time_delta, - process_id=os.getpid(), ) - if job_id := job_context.get().job_id: - profiling_run.job_execution_id = job_id # This runs in a subprocess — commit after every save so progress is visible # to the UI (separate session) and to execute_db_queries (independent connection). @@ -96,19 +93,12 @@ def run_profiling( # run_pairwise_contingency_check(profiling_run.id, table_group.profile_pair_rule_pct) else: LOG.info("No columns were selected to profile.") - except Exception as e: + except Exception: LOG.exception("Profiling encountered an error.") - LOG.info("Setting profiling run status to Error") - profiling_run.log_message = get_exception_message(e) - profiling_run.profiling_endtime = datetime.now(UTC) + time_delta - profiling_run.status = "Error" - profiling_run.save() - session.commit() + end_time = datetime.now(UTC) + time_delta raise else: - LOG.info("Setting profiling run status to Completed") - profiling_run.profiling_endtime = datetime.now(UTC) + time_delta - profiling_run.status = "Complete" + end_time = datetime.now(UTC) + time_delta profiling_run.save() session.commit() @@ -122,7 +112,7 @@ def run_profiling( sampling=table_group.profile_use_sampling, table_count=profiling_run.table_ct or 0, column_count=profiling_run.column_ct or 0, - run_duration=(profiling_run.profiling_endtime - profiling_run.profiling_starttime).total_seconds(), + run_duration=(end_time - profiling_run.profiling_starttime.replace(tzinfo=UTC)).total_seconds(), ) return profiling_run.id diff --git a/testgen/commands/run_score_update.py b/testgen/commands/run_score_update.py index 230a4a9d..caa9125e 100644 --- a/testgen/commands/run_score_update.py +++ b/testgen/commands/run_score_update.py @@ -36,7 +36,7 @@ def run_score_update(parent_job_id: str, parent_job_key: str) -> None: def _rollup_profiling(parent_je_id: UUID) -> None: with database_session() as session: profiling_run = session.scalars( - select(ProfilingRun).where(ProfilingRun.job_execution_id == parent_je_id) + select(ProfilingRun).where(ProfilingRun.id == parent_je_id) ).first() if not profiling_run: LOG.error("No profiling_run found for job execution %s; skipping score rollup", parent_je_id) @@ -60,7 +60,7 @@ def _rollup_test(parent_je_id: UUID) -> None: row = session.execute( select(TestRun.id, TestRun.test_starttime, TestSuite.table_groups_id, TestSuite.project_code) .join(TestSuite, TestRun.test_suite_id == TestSuite.id) - .where(TestRun.job_execution_id == parent_je_id) + .where(TestRun.id == parent_je_id) ).first() if not row: LOG.error("No test_run found for job execution %s; skipping score rollup", parent_je_id) diff --git a/testgen/commands/run_test_execution.py b/testgen/commands/run_test_execution.py index 65d759df..381dc2f0 100644 --- a/testgen/commands/run_test_execution.py +++ b/testgen/commands/run_test_execution.py @@ -1,5 +1,4 @@ import logging -import os from collections import defaultdict from datetime import UTC, datetime, timedelta from functools import partial @@ -24,7 +23,6 @@ from testgen.common.models.table_group import TableGroup from testgen.common.models.test_run import TestRun from testgen.common.models.test_suite import TestSuite -from testgen.utils import get_exception_message from .run_refresh_data_chars import run_data_chars_refresh from .run_test_validation import run_test_validation @@ -52,12 +50,10 @@ def run_test_execution( LOG.info("Creating test run record") test_run = TestRun( + id=job_context.get().job_id, test_suite_id=test_suite_id, test_starttime=datetime.now(UTC) + time_delta, - process_id=os.getpid(), ) - if job_id := job_context.get().job_id: - test_run.job_execution_id = job_id # This runs in a subprocess — commit after every save so progress is visible # to the UI (separate session) and to execute_db_queries (independent connection). @@ -134,22 +130,12 @@ def run_test_execution( # Refresh needed because previous query updates the test run too test_run.refresh() - except Exception as e: + except Exception: LOG.exception("Test execution encountered an error.") - LOG.info("Setting test run status to Error") - test_run.log_message = get_exception_message(e) - test_run.test_endtime = datetime.now(UTC) + time_delta - test_run.status = "Error" - test_run.save() - session.commit() + end_time = datetime.now(UTC) + time_delta raise else: - LOG.info("Setting test run status to Completed") - test_run.test_endtime = datetime.now(UTC) + time_delta - test_run.status = "Complete" - test_run.save() - session.commit() - + end_time = datetime.now(UTC) + time_delta LOG.info("Updating latest run for test suite") test_suite.last_complete_test_run_id = test_run.id test_suite.save() @@ -167,7 +153,7 @@ def run_test_execution( username=username, sql_flavor=connection.sql_flavor_code, test_count=test_run.test_ct, - run_duration=(test_run.test_endtime - test_run.test_starttime.replace(tzinfo=UTC)).total_seconds(), + run_duration=(end_time - test_run.test_starttime.replace(tzinfo=UTC)).total_seconds(), prediction_duration=(datetime.now(UTC) + time_delta - prediction_start).total_seconds(), ) diff --git a/testgen/common/models/data_column.py b/testgen/common/models/data_column.py index ed21e312..29aa626a 100644 --- a/testgen/common/models/data_column.py +++ b/testgen/common/models/data_column.py @@ -22,6 +22,7 @@ from testgen.common.models import get_current_session from testgen.common.models.entity import Entity, EntityMinimal from testgen.common.models.hygiene_issue import HygieneIssue +from testgen.common.models.job_execution import JobExecution from testgen.common.models.profile_result import ProfileResult from testgen.common.models.profiling_run import ProfilingRun @@ -552,11 +553,11 @@ def get_column_detail( cls.dq_score_testing, func.coalesce(hygiene_subq.c.hygiene_issue_count, 0).label("hygiene_issue_count"), ProfilingRun.id.label("profile_run_id"), - ProfilingRun.job_execution_id.label("profile_run_je_id"), - ProfilingRun.status.label("profile_run_status"), + ProfilingRun.id.label("profile_run_je_id"), + JobExecution.status.label("profile_run_status"), ProfilingRun.profiling_starttime.label("profile_run_started_at"), - ProfilingRun.profiling_endtime.label("profile_run_ended_at"), - ProfilingRun.log_message.label("profile_run_log_message"), + JobExecution.completed_at.label("profile_run_ended_at"), + JobExecution.error_message.label("profile_run_log_message"), ) .outerjoin(DataTable, DataTable.id == cls.table_id) .outerjoin( @@ -578,6 +579,7 @@ def get_column_detail( ), ) .outerjoin(ProfilingRun, ProfilingRun.id == ProfileResult.profile_run_id) + .outerjoin(JobExecution, JobExecution.id == ProfilingRun.id) .where( cls.table_groups_id == table_groups_id, cls.table_name == table_name, diff --git a/testgen/common/models/data_table.py b/testgen/common/models/data_table.py index ce149b94..95141f29 100644 --- a/testgen/common/models/data_table.py +++ b/testgen/common/models/data_table.py @@ -123,7 +123,7 @@ def get_profiling_overview( JobExecution.id.label("latest_profile_job_execution_id"), ) .outerjoin(ProfilingRun, ProfilingRun.id == cls.last_complete_profile_run_id) - .outerjoin(JobExecution, JobExecution.id == ProfilingRun.job_execution_id) + .outerjoin(JobExecution, JobExecution.id == ProfilingRun.id) .where( cls.table_groups_id == table_groups_id, cls.table_name == table_name, diff --git a/testgen/common/models/hygiene_issue.py b/testgen/common/models/hygiene_issue.py index c7683479..c77a2606 100644 --- a/testgen/common/models/hygiene_issue.py +++ b/testgen/common/models/hygiene_issue.py @@ -275,7 +275,7 @@ def list_for_run( ProfileResult.column_name == cls.column_name, ), ) - .where(ProfilingRun.job_execution_id == job_execution_id, *clauses) + .where(ProfilingRun.id == job_execution_id, *clauses) .order_by(cls._priority_order(), cls.table_name, cls.column_name, cls.id) ) return cls._paginate(query, page=page, limit=limit, data_class=HygieneIssueListRow) @@ -299,7 +299,7 @@ def search( cls.project_code.label("project_code"), HygieneIssueType.name.label("issue_type_name"), TableGroup.table_groups_name.label("table_groups_name"), - ProfilingRun.job_execution_id.label("job_execution_id"), + ProfilingRun.id.label("job_execution_id"), JobExecution.started_at.label("started_at"), cls.schema_name.label("schema_name"), cls.table_name.label("table_name"), @@ -314,7 +314,7 @@ def search( ) .join(HygieneIssueType, HygieneIssueType.id == cls.type_id) .join(ProfilingRun, ProfilingRun.id == cls.profile_run_id) - .outerjoin(JobExecution, JobExecution.id == ProfilingRun.job_execution_id) + .outerjoin(JobExecution, JobExecution.id == ProfilingRun.id) .join(TableGroup, TableGroup.id == cls.table_groups_id) .outerjoin( ProfileResult, @@ -358,7 +358,7 @@ def get_with_context(cls, issue_id: UUID, *clauses) -> HygieneIssueDetail | None cls.detail.label("detail"), HygieneIssueType.detail_redactable.label("detail_redactable"), ProfileResult.pii_flag.label("pii_flag"), - ProfilingRun.job_execution_id.label("job_execution_id"), + ProfilingRun.id.label("job_execution_id"), JobExecution.started_at.label("started_at"), ProfileResult.general_type.label("column_general_type"), ProfileResult.db_data_type.label("column_db_data_type"), @@ -368,7 +368,7 @@ def get_with_context(cls, issue_id: UUID, *clauses) -> HygieneIssueDetail | None ) .join(HygieneIssueType, HygieneIssueType.id == cls.type_id) .join(ProfilingRun, ProfilingRun.id == cls.profile_run_id) - .outerjoin(JobExecution, JobExecution.id == ProfilingRun.job_execution_id) + .outerjoin(JobExecution, JobExecution.id == ProfilingRun.id) .outerjoin( ProfileResult, and_( diff --git a/testgen/common/models/profiling_run.py b/testgen/common/models/profiling_run.py index 0c898172..1e2b38d3 100644 --- a/testgen/common/models/profiling_run.py +++ b/testgen/common/models/profiling_run.py @@ -1,10 +1,10 @@ from collections.abc import Iterable from dataclasses import dataclass -from datetime import UTC, datetime +from datetime import datetime from typing import ClassVar, Literal, NamedTuple, Self, TypedDict -from uuid import UUID, uuid4 +from uuid import UUID -from sqlalchemy import BigInteger, Column, Float, Integer, String, desc, func, select, text, update +from sqlalchemy import BigInteger, Column, Float, ForeignKey, String, delete, desc, func, select, text from sqlalchemy.dialects import postgresql from sqlalchemy.orm import InstrumentedAttribute, foreign, relationship from sqlalchemy.orm.attributes import flag_modified @@ -20,7 +20,6 @@ from testgen.common.models.table_group import TableGroup from testgen.utils import is_uuid4 -ProfilingRunStatus = Literal["Running", "Complete", "Error", "Cancelled"] ProgressKey = Literal["data_chars", "col_profiling", "freq_analysis", "hygiene_issues"] ProgressStatus = Literal["Pending", "Running", "Completed", "Warning"] @@ -57,8 +56,6 @@ class ProfilingRunSummary(EntityMinimal): progress: list[ProgressStep] table_groups_name: str | None table_group_schema: str | None - process_id: int | None - log_message: str | None table_ct: int | None column_ct: int | None record_ct: int | None @@ -95,15 +92,16 @@ class LatestProfilingRun(NamedTuple): class ProfilingRun(Entity): __tablename__ = "profiling_runs" - id: UUID = Column(postgresql.UUID(as_uuid=True), primary_key=True, default=uuid4) + id: UUID = Column( + postgresql.UUID(as_uuid=True), + ForeignKey("job_executions.id", ondelete="CASCADE"), + primary_key=True, + ) project_code: str = Column(String, nullable=False) connection_id: str = Column(BigInteger, nullable=False) table_groups_id: UUID = Column(postgresql.UUID(as_uuid=True), nullable=False) profiling_starttime: datetime = Column(postgresql.TIMESTAMP) - profiling_endtime: datetime = Column(postgresql.TIMESTAMP) - status: ProfilingRunStatus = Column(String, default="Running") progress: list[ProgressStep] = Column(postgresql.JSONB, default=[]) - log_message: str = Column(String) table_ct: int = Column(BigInteger) column_ct: int = Column(BigInteger) record_ct: int = Column(BigInteger) @@ -114,12 +112,10 @@ class ProfilingRun(Entity): dq_affected_data_points: int = Column(BigInteger) dq_total_data_points: int = Column(BigInteger) dq_score_profiling: float = Column(Float) - process_id: int = Column(Integer) - job_execution_id: UUID | None = Column(postgresql.UUID(as_uuid=True), nullable=True) job_execution = relationship( JobExecution, - primaryjoin=foreign(job_execution_id) == JobExecution.id, + primaryjoin=foreign(id) == JobExecution.id, uselist=False, viewonly=True, ) @@ -139,12 +135,6 @@ class ProfilingRun(Entity): ).label("is_latest_run"), ) - @classmethod - def get_by_id_or_job(cls, identifier: UUID) -> Self | None: - """Look up a profiling run by its own ID or by job_execution_id.""" - query = select(cls).where((cls.id == identifier) | (cls.job_execution_id == identifier)) - return get_current_session().scalars(query).first() - @classmethod def get_minimal(cls, run_id: str | UUID) -> ProfilingRunMinimal | None: if not is_uuid4(run_id): @@ -153,7 +143,7 @@ def get_minimal(cls, run_id: str | UUID) -> ProfilingRunMinimal | None: query = ( select(*cls._minimal_columns) .join(TableGroup, cls.table_groups_id == TableGroup.id) - .where((cls.id == run_id) | (cls.job_execution_id == run_id)) + .where(cls.id == run_id) ) result = get_current_session().execute(query).mappings().first() return ProfilingRunMinimal(**result) if result else None @@ -162,7 +152,7 @@ def get_minimal(cls, run_id: str | UUID) -> ProfilingRunMinimal | None: def get_latest_run(cls, project_code: str) -> LatestProfilingRun | None: query = ( select(ProfilingRun.id, JobExecution.started_at.label("run_time")) - .join(JobExecution, ProfilingRun.job_execution_id == JobExecution.id) + .join(JobExecution, ProfilingRun.id == JobExecution.id) .where(ProfilingRun.project_code == project_code, JobExecution.status == JobStatus.COMPLETED) .order_by(desc(JobExecution.started_at)) .limit(1) @@ -174,15 +164,14 @@ def get_latest_run(cls, project_code: str) -> LatestProfilingRun | None: @classmethod def get_latest_complete_je_id_for_table_group(cls, table_groups_id: UUID) -> UUID | None: - """Return the ``job_execution_id`` of the latest completed profiling run for a table group. + """Return the id of the latest completed profiling run for a table group. - Computed live from ``profiling_runs`` joined to ``job_executions`` — does not read the - legacy ``table_groups.last_complete_profile_run_id`` cache, which points at the internal - run PK rather than the JE id. + Computed live from ``profiling_runs`` joined to ``job_executions`` rather than + reading the ``table_groups.last_complete_profile_run_id`` cache. """ query = ( - select(cls.job_execution_id) - .join(JobExecution, cls.job_execution_id == JobExecution.id) + select(cls.id) + .join(JobExecution, cls.id == JobExecution.id) .where(cls.table_groups_id == table_groups_id, JobExecution.status == JobStatus.COMPLETED) .order_by(desc(JobExecution.started_at)) .limit(1) @@ -250,8 +239,6 @@ def select_summary( COALESCE(pr.progress, '[]'::jsonb) AS progress, tg.table_groups_name, tg.table_group_schema, - pr.process_id, - pr.log_message, pr.table_ct, pr.column_ct, pr.record_ct, @@ -264,7 +251,7 @@ def select_summary( pr.dq_score_profiling, COUNT(*) OVER() AS total_count FROM job_executions je - LEFT JOIN profiling_runs pr ON pr.job_execution_id = je.id + LEFT JOIN profiling_runs pr ON pr.id = je.id LEFT JOIN table_groups tg ON tg.id = pr.table_groups_id LEFT JOIN profile_anomalies pa ON pa.profile_run_id = pr.id WHERE je.job_key = 'run-profile' @@ -345,7 +332,7 @@ def has_active_job_for(cls, entity_cls: type[Entity], *entity_ids: str | int | U """Check whether any active profiling job exists for the given entity or entities.""" query = ( select(func.count(cls.id)) - .join(JobExecution, cls.job_execution_id == JobExecution.id) + .join(JobExecution, cls.id == JobExecution.id) .where(JobExecution.status.in_(cls._ACTIVE_JOB_STATUSES)) ) if entity_cls is cls: @@ -360,33 +347,15 @@ def has_active_job_for(cls, entity_cls: type[Entity], *entity_ids: str | int | U raise ValueError(f"Unsupported entity: {entity_cls.__name__}") return get_current_session().execute(query).scalar() > 0 - @classmethod - def cancel_run(cls, run_id: str | UUID) -> None: - query = update(cls).where(cls.id == run_id).values(status="Cancelled", profiling_endtime=datetime.now(UTC)) - db_session = get_current_session() - db_session.execute(query) - @classmethod def cascade_delete(cls, ids: list[str]) -> None: - query = """ - DELETE FROM profile_pair_rules - WHERE profile_run_id IN :profiling_run_ids; - - DELETE FROM profile_anomaly_results - WHERE profile_run_id IN :profiling_run_ids; + """Delete runs and their results by removing the parent job executions. - DELETE FROM profile_results - WHERE profile_run_id IN :profiling_run_ids; - - DELETE FROM job_executions - WHERE id IN ( - SELECT job_execution_id FROM profiling_runs - WHERE id IN :profiling_run_ids AND job_execution_id IS NOT NULL - ); + The run id is the job execution id, and the foreign-key chain + (profile_results / profile_anomaly_results -> profiling_runs -> + job_executions) cascades on delete. """ - db_session = get_current_session() - db_session.execute(text(query), {"profiling_run_ids": tuple(ids)}) - cls.delete_where(cls.id.in_(ids)) + get_current_session().execute(delete(JobExecution).where(JobExecution.id.in_(ids))) @classmethod def find_latest_per_table_group(cls, project_code: str) -> set[UUID]: @@ -403,7 +372,7 @@ def find_latest_per_table_group(cls, project_code: str) -> set[UUID]: """ rows = get_current_session().scalars( select(cls.id) - .join(JobExecution, cls.job_execution_id == JobExecution.id) + .join(JobExecution, cls.id == JobExecution.id) .where( cls.project_code == project_code, JobExecution.status == JobStatus.COMPLETED, @@ -442,7 +411,7 @@ def delete_older_than( if protected_ids: where_clauses.append(cls.id.notin_(protected_ids)) - base_select = select(cls.id).join(JobExecution, cls.job_execution_id == JobExecution.id) + base_select = select(cls.id).join(JobExecution, cls.id == JobExecution.id) if dry_run: return get_current_session().scalar( @@ -459,18 +428,6 @@ def delete_older_than( total += len(ids) return total - @classmethod - def get_job_execution_ids(cls, profiling_run_ids: list[UUID]) -> dict[UUID, UUID | None]: - """Map profiling_run PKs to their job_execution_ids (batch lookup). - - Mirrors TestRun.get_job_execution_ids. - """ - if not profiling_run_ids: - return {} - query = select(cls.id, cls.job_execution_id).where(cls.id.in_(profiling_run_ids)) - rows = get_current_session().execute(query).all() - return {row.id: row.job_execution_id for row in rows} - def init_progress(self) -> None: self._progress = { "data_chars": {"label": "Refreshing data catalog"}, @@ -494,7 +451,7 @@ def set_progress(self, key: ProgressKey, status: ProgressStatus, detail: str | N def get_previous(self) -> Self | None: query = ( select(ProfilingRun) - .join(JobExecution, ProfilingRun.job_execution_id == JobExecution.id) + .join(JobExecution, ProfilingRun.id == JobExecution.id) .where( ProfilingRun.table_groups_id == self.table_groups_id, JobExecution.status == JobStatus.COMPLETED, @@ -510,7 +467,7 @@ def list_recent_complete(cls, table_groups_id: UUID, limit: int) -> list[Self]: """Return the most recent completed profiling runs for a table group, newest first.""" query = ( select(cls) - .join(JobExecution, cls.job_execution_id == JobExecution.id) + .join(JobExecution, cls.id == JobExecution.id) .where( cls.table_groups_id == table_groups_id, JobExecution.status == JobStatus.COMPLETED, diff --git a/testgen/common/models/table_group.py b/testgen/common/models/table_group.py index 251dcc46..98e2c858 100644 --- a/testgen/common/models/table_group.py +++ b/testgen/common/models/table_group.py @@ -224,7 +224,6 @@ def select_summary( latest_profile AS ( SELECT latest_run.table_groups_id, latest_run.id, - latest_run.job_execution_id, MAX(latest_je.started_at) AS started_at, latest_run.anomaly_ct, SUM( @@ -259,7 +258,7 @@ def select_summary( groups.last_complete_profile_run_id = latest_run.id ) LEFT JOIN job_executions latest_je ON ( - latest_run.job_execution_id = latest_je.id + latest_run.id = latest_je.id ) LEFT JOIN profile_anomaly_results latest_anomalies ON ( latest_run.id = latest_anomalies.profile_run_id @@ -334,7 +333,7 @@ def select_summary( groups.dq_score_profiling, groups.dq_score_testing, latest_profile.id AS latest_profile_id, - latest_profile.job_execution_id AS latest_profile_job_execution_id, + latest_profile.id AS latest_profile_job_execution_id, latest_profile.started_at AS latest_profile_start, latest_profile.anomaly_ct AS latest_hygiene_issues_ct, latest_profile.definite_ct AS latest_hygiene_issues_definite_ct, @@ -417,8 +416,8 @@ def cascade_delete(cls, ids: list[str]) -> None: DELETE FROM job_executions WHERE id IN ( - SELECT pr.job_execution_id FROM profiling_runs pr - WHERE pr.table_groups_id IN :table_group_ids AND pr.job_execution_id IS NOT NULL + SELECT pr.id FROM profiling_runs pr + WHERE pr.table_groups_id IN :table_group_ids ); DELETE FROM profiling_runs pr diff --git a/testgen/common/models/test_result.py b/testgen/common/models/test_result.py index 8bbd63ba..751e896f 100644 --- a/testgen/common/models/test_result.py +++ b/testgen/common/models/test_result.py @@ -309,7 +309,7 @@ def search_results( select( cls.test_definition_id.label("test_definition_id"), cls.test_run_id.label("test_run_id"), - TestRun.job_execution_id.label("job_execution_id"), + TestRun.id.label("job_execution_id"), cls.test_time.label("test_time"), TestSuite.id.label("test_suite_id"), TestSuite.test_suite.label("test_suite_name"), diff --git a/testgen/common/models/test_run.py b/testgen/common/models/test_run.py index 1887d886..501e990e 100644 --- a/testgen/common/models/test_run.py +++ b/testgen/common/models/test_run.py @@ -1,9 +1,9 @@ from dataclasses import dataclass -from datetime import UTC, datetime +from datetime import datetime from typing import ClassVar, Literal, NamedTuple, Self, TypedDict -from uuid import UUID, uuid4 +from uuid import UUID -from sqlalchemy import BigInteger, Column, Float, ForeignKey, Integer, String, Text, desc, func, select, text, update +from sqlalchemy import BigInteger, Column, Float, ForeignKey, Integer, delete, desc, func, select, text from sqlalchemy.dialects import postgresql from sqlalchemy.orm import foreign, relationship from sqlalchemy.orm.attributes import flag_modified @@ -20,7 +20,6 @@ from testgen.common.models.test_suite import TestSuite from testgen.utils import is_uuid4 -TestRunStatus = Literal["Running", "Complete", "Error", "Cancelled"] ProgressKey = Literal["data_chars", "validation", "QUERY", "CAT", "METADATA"] ProgressStatus = Literal["Pending", "Running", "Completed", "Warning"] @@ -59,8 +58,6 @@ class TestRunSummary(EntityMinimal): test_suite: str | None project_code: str project_name: str - process_id: int | None - log_message: str | None test_ct: int | None passed_ct: int | None warning_ct: int | None @@ -99,13 +96,14 @@ class LatestTestRun(NamedTuple): class TestRun(Entity): __tablename__ = "test_runs" - id: UUID = Column(postgresql.UUID(as_uuid=True), primary_key=True, nullable=False, default=uuid4) + id: UUID = Column( + postgresql.UUID(as_uuid=True), + ForeignKey("job_executions.id", ondelete="CASCADE"), + primary_key=True, + ) test_suite_id: UUID = Column(postgresql.UUID(as_uuid=True), ForeignKey("test_suites.id"), nullable=False) test_starttime: datetime = Column(postgresql.TIMESTAMP) - test_endtime: datetime = Column(postgresql.TIMESTAMP) - status: TestRunStatus = Column(String, default="Running") progress: list[ProgressStep] = Column(postgresql.JSONB, default=[]) - log_message: str = Column(Text) test_ct: int = Column(Integer) passed_ct: int = Column(Integer) failed_ct: int = Column(Integer) @@ -119,12 +117,10 @@ class TestRun(Entity): dq_affected_data_points: int = Column(BigInteger) dq_total_data_points: int = Column(BigInteger) dq_score_test_run: float = Column(Float) - process_id: int = Column(Integer) - job_execution_id: UUID | None = Column(postgresql.UUID(as_uuid=True), nullable=True) job_execution = relationship( JobExecution, - primaryjoin=foreign(job_execution_id) == JobExecution.id, + primaryjoin=foreign(id) == JobExecution.id, uselist=False, viewonly=True, ) @@ -144,21 +140,6 @@ class TestRun(Entity): ).label("is_latest_run"), ) - @classmethod - def get_by_id_or_job(cls, identifier: UUID) -> Self | None: - """Look up a test run by its own ID or by job_execution_id.""" - query = select(cls).where((cls.id == identifier) | (cls.job_execution_id == identifier)) - return get_current_session().scalars(query).first() - - @classmethod - def get_job_execution_ids(cls, test_run_ids: list[UUID]) -> dict[UUID, UUID | None]: - """Map test_run PKs to their job_execution_ids (batch lookup).""" - if not test_run_ids: - return {} - query = select(cls.id, cls.job_execution_id).where(cls.id.in_(test_run_ids)) - rows = get_current_session().execute(query).all() - return {row.id: row.job_execution_id for row in rows} - @classmethod def get_minimal(cls, run_id: str | UUID) -> TestRunMinimal | None: if not is_uuid4(run_id): @@ -167,7 +148,7 @@ def get_minimal(cls, run_id: str | UUID) -> TestRunMinimal | None: query = ( select(*cls._minimal_columns) .join(TestSuite) - .where((cls.id == run_id) | (cls.job_execution_id == run_id)) + .where(cls.id == run_id) ) result = get_current_session().execute(query).mappings().first() return TestRunMinimal(**result) if result else None @@ -176,7 +157,7 @@ def get_minimal(cls, run_id: str | UUID) -> TestRunMinimal | None: def get_latest_run(cls, project_code: str) -> LatestTestRun | None: query = ( select(TestRun.id, JobExecution.started_at.label("run_time")) - .join(JobExecution, TestRun.job_execution_id == JobExecution.id) + .join(JobExecution, TestRun.id == JobExecution.id) .join(TestSuite) .where(TestSuite.project_code == project_code, JobExecution.status == JobStatus.COMPLETED) .order_by(desc(JobExecution.started_at)) @@ -190,7 +171,7 @@ def get_latest_run(cls, project_code: str) -> LatestTestRun | None: def get_previous(self) -> Self | None: query = ( select(TestRun) - .join(JobExecution, TestRun.job_execution_id == JobExecution.id) + .join(JobExecution, TestRun.id == JobExecution.id) .where( TestRun.test_suite_id == self.test_suite_id, JobExecution.status == JobStatus.COMPLETED, @@ -266,8 +247,6 @@ def select_summary( ts.test_suite, je.project_code, p.project_name, - tr.process_id, - tr.log_message, tr.test_ct, rr.passed_ct, rr.warning_ct, @@ -278,7 +257,7 @@ def select_summary( tr.dq_score_test_run AS dq_score_testing, COUNT(*) OVER() AS total_count FROM job_executions je - LEFT JOIN test_runs tr ON tr.job_execution_id = je.id + LEFT JOIN test_runs tr ON tr.id = je.id LEFT JOIN test_suites ts ON ts.id = tr.test_suite_id LEFT JOIN table_groups tg ON tg.id = ts.table_groups_id LEFT JOIN projects p ON p.project_code = je.project_code @@ -325,7 +304,7 @@ def get_monitoring_summary(self, table_name: str | None = None) -> TestRunMonito )) projection = [ TestRun.id.label("test_run_id"), - TestRun.test_endtime, + JobExecution.completed_at.label("test_endtime"), TableGroup.id.label("table_group_id"), TableGroup.table_groups_name, Project.project_name, @@ -335,7 +314,7 @@ def get_monitoring_summary(self, table_name: str | None = None) -> TestRunMonito ] group_by = [ TestRun.id, - TestRun.test_endtime, + JobExecution.completed_at, TableGroup.id, TableGroup.table_groups_name, Project.project_name, @@ -349,6 +328,7 @@ def get_monitoring_summary(self, table_name: str | None = None) -> TestRunMonito .join(TableGroup, TableGroup.monitor_test_suite_id == TestRun.test_suite_id) .join(Project, Project.project_code == TableGroup.project_code) .join(TestResult, TestResult.test_run_id == TestRun.id) + .join(JobExecution, JobExecution.id == TestRun.id) .where( TestRun.id == self.id, (TestResult.table_name == table_name) if table_name else True, @@ -365,7 +345,7 @@ def has_active_job_for(cls, entity_cls: type[Entity], *entity_ids: str | int | U """Check whether any active test run job exists for the given entity or entities.""" query = ( select(func.count(cls.id)) - .join(JobExecution, cls.job_execution_id == JobExecution.id) + .join(JobExecution, cls.id == JobExecution.id) .where(JobExecution.status.in_(cls._ACTIVE_JOB_STATUSES)) ) if entity_cls is cls: @@ -382,27 +362,14 @@ def has_active_job_for(cls, entity_cls: type[Entity], *entity_ids: str | int | U raise ValueError(f"Unsupported entity: {entity_cls.__name__}") return get_current_session().execute(query).scalar() > 0 - @classmethod - def cancel_run(cls, run_id: str | UUID) -> None: - query = update(cls).where(cls.id == run_id).values(status="Cancelled", test_endtime=datetime.now(UTC)) - db_session = get_current_session() - db_session.execute(query) - @classmethod def cascade_delete(cls, ids: list[str]) -> None: - query = """ - DELETE FROM test_results - WHERE test_run_id IN :test_run_ids; - - DELETE FROM job_executions - WHERE id IN ( - SELECT job_execution_id FROM test_runs - WHERE id IN :test_run_ids AND job_execution_id IS NOT NULL - ); + """Delete runs and their results by removing the parent job executions. + + The run id is the job execution id, and the foreign-key chain + (test_results -> test_runs -> job_executions) cascades on delete. """ - db_session = get_current_session() - db_session.execute(text(query), {"test_run_ids": tuple(ids)}) - cls.delete_where(cls.id.in_(ids)) + get_current_session().execute(delete(JobExecution).where(JobExecution.id.in_(ids))) @classmethod def find_latest_per_test_suite(cls, project_code: str) -> set[UUID]: @@ -417,7 +384,7 @@ def find_latest_per_test_suite(cls, project_code: str) -> set[UUID]: rows = get_current_session().scalars( select(cls.id) .join(TestSuite, cls.test_suite_id == TestSuite.id) - .join(JobExecution, cls.job_execution_id == JobExecution.id) + .join(JobExecution, cls.id == JobExecution.id) .where( TestSuite.project_code == project_code, JobExecution.status == JobStatus.COMPLETED, @@ -459,7 +426,7 @@ def delete_older_than( base_select = ( select(cls.id) .join(TestSuite, cls.test_suite_id == TestSuite.id) - .join(JobExecution, cls.job_execution_id == JobExecution.id) + .join(JobExecution, cls.id == JobExecution.id) ) if dry_run: diff --git a/testgen/common/models/test_suite.py b/testgen/common/models/test_suite.py index ac29ddce..fad51a0e 100644 --- a/testgen/common/models/test_suite.py +++ b/testgen/common/models/test_suite.py @@ -113,7 +113,6 @@ def select_summary(cls, project_code: str, table_group_id: str | UUID | None = N WITH last_run AS ( SELECT test_runs.test_suite_id, test_runs.id, - test_runs.job_execution_id, test_runs.test_starttime, test_runs.test_ct, SUM( @@ -184,7 +183,7 @@ def select_summary(cls, project_code: str, table_group_id: str | UUID | None = N test_defs.count AS test_ct, last_complete_profile_run_id, last_run.id AS latest_run_id, - last_run.job_execution_id AS latest_run_job_execution_id, + last_run.id AS latest_run_job_execution_id, last_run.test_starttime AS latest_run_start, last_run.test_ct AS last_run_test_ct, last_run.passed_ct AS last_run_passed_ct, @@ -229,8 +228,8 @@ def cascade_delete(cls, ids: list[str]) -> None: query = """ DELETE FROM job_executions WHERE id IN ( - SELECT job_execution_id FROM test_runs - WHERE test_suite_id IN :test_suite_ids AND job_execution_id IS NOT NULL + SELECT id FROM test_runs + WHERE test_suite_id IN :test_suite_ids ); DELETE FROM test_runs diff --git a/testgen/common/notifications/profiling_run.py b/testgen/common/notifications/profiling_run.py index 4d7b2197..1ceec969 100644 --- a/testgen/common/notifications/profiling_run.py +++ b/testgen/common/notifications/profiling_run.py @@ -266,7 +266,7 @@ def send_profiling_run_notifications(profiling_run: ProfilingRun, result_list_ct settings.UI_BASE_URL, "/profiling-runs:hygiene?project_code=", str(profiling_run.project_code), - "&run_id=", str(profiling_run.job_execution_id), + "&run_id=", str(profiling_run.id), "&source=email" ) ) @@ -319,7 +319,7 @@ def send_profiling_run_notifications(profiling_run: ProfilingRun, result_list_ct "/profiling-runs:results?project_code=", str(profiling_run.project_code), "&run_id=", - str(profiling_run.job_execution_id), + str(profiling_run.id), "&source=email" ) ), diff --git a/testgen/common/notifications/test_run.py b/testgen/common/notifications/test_run.py index 9ad2e08e..72ad5140 100644 --- a/testgen/common/notifications/test_run.py +++ b/testgen/common/notifications/test_run.py @@ -137,7 +137,7 @@ def get_main_content_template(self): {{/if}} {{#if (eq test_run.status 'error')}}
{{test_run.log_message}}
{{test_run.error_message}}