Merge branch 'feat/TG-1085-data-classification' into 'enterprise'

feat(TG-1085): add data_classification tag field

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

Author: ci bot

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.

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)

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)

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)

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)

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",
 }
 
 

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:

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.
     """

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:

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