Merge remote-tracking branch 'origin/enterprise' into fix/mssql-schema-case-sensitive

Author: testgen-ci-bot

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

testgen/common/models/data_column.py

@@ -260,16 +260,23 @@ 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)
     dq_score_testing: float | None = Column(Float)
 
     _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,

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

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)

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()