improve DataValidationError output: use identifying columns if present

Author: djudjuu

dlt/common/schema/exceptions.py

@@ -178,6 +178,25 @@ def __init__(
             + f" . Contract on `{schema_entity}` with `{contract_mode=:}` is violated. "
             + (extended_info or "")
         )
+        data_item_str = ""
+        if (
+            data_item
+            and hasattr(data_item, "get")
+            and table_schema
+            and hasattr(table_schema, "get")
+        ):
+            identifier_columns = [
+                x.get("name")
+                for x in table_schema.get("columns", {}).values()
+                if x.get("primary_key") or x.get("merge_key") or x.get("unique")
+            ]
+            if identifier_columns:
+                data_item_str += "Offending data item: "
+                data_item_keys = [
+                    f"{column}: {data_item.get(column)}" for column in identifier_columns
+                ]
+                data_item_str += ", ".join(data_item_keys)
+                msg += f" {data_item_str}"
         super().__init__(schema_name, msg)
         self.table_name = table_name
         self.column_name = column_name

tests/common/schema/test_schema_contract.py

@@ -354,3 +354,80 @@ def test_check_adding_new_variant() -> None:
             cast(TSchemaContractDict, {**DEFAULT_SCHEMA_CONTRACT_MODE, **{"data_type": "freeze"}}),
             copy.deepcopy(table_update),
         )
+
+
+def test_data_validation_error_message_with_primary_key() -> None:
+    """Test that DataValidationError includes primary key information in the message"""
+    schema = get_schema()
+
+    table_update: TTableSchema = {
+        "name": "tables",
+        "columns": {"column_1": {"name": "column_1", "data_type": "text", "primary_key": True}},
+    }
+
+    partial_table, _ = schema.apply_schema_contract(DEFAULT_SCHEMA_CONTRACT_MODE, table_update)
+    schema.update_table(partial_table)
+
+    # Create a table update that tries to add a variant column
+    variant_table_update: TTableSchema = {
+        "name": "tables",
+        "columns": {
+            "column_2_variant": {
+                "name": "column_2_variant",
+                "data_type": "boolean",
+                "variant": True,
+            }
+        },
+    }
+
+    # apply update with data_type freeze mode, providing data item as evidence
+    with pytest.raises(DataValidationError) as val_ex:
+        schema.apply_schema_contract(
+            {**DEFAULT_SCHEMA_CONTRACT_MODE, **{"data_type": "freeze"}},
+            variant_table_update,
+            data_item={"column_1": "test", "column_2": 123},
+        )
+
+    # data item should be included by primary key
+    assert "Offending data item: column_1: test" in str(val_ex.value)
+
+
+def test_data_validation_error_message_with_multiple_identifiers() -> None:
+    """Test that DataValidationError includes multiple identifier columns in the message"""
+    schema = get_schema()
+
+    table_update: TTableSchema = {
+        "name": "tables",
+        "columns": {
+            "column_1": {"name": "column_1", "data_type": "text", "primary_key": True},
+            "column_2": {"name": "column_2", "data_type": "bigint", "merge_key": True},
+            "column_3": {"name": "column_3", "data_type": "text", "unique": True},
+        },
+    }
+
+    partial_table, _ = schema.apply_schema_contract(DEFAULT_SCHEMA_CONTRACT_MODE, table_update)
+    schema.update_table(partial_table)
+
+    # Now create a table update that tries to add a variant column (data type evolution)
+    variant_table_update: TTableSchema = {
+        "name": "tables",
+        "columns": {
+            "column_4_variant": {"name": "column_4_variant", "data_type": "text", "variant": True}
+        },
+    }
+
+    with pytest.raises(DataValidationError) as val_ex:
+        schema.apply_schema_contract(
+            {**DEFAULT_SCHEMA_CONTRACT_MODE, **{"data_type": "freeze"}},
+            variant_table_update,
+            data_item={
+                "column_1": "test",
+                "column_2": 123,
+                "column_3": "unique_value",
+                "column_4": "some_text",
+            },
+        )
+
+    # Check that the message includes all identifier columns
+    error_msg = str(val_ex.value)
+    assert "Offending data item: column_1: test, column_2: 123, column_3: unique_value" in error_msg