Feature/pydantic serialization (#1302)

## Changes

Follow-up to #1259, completing the Pydantic migration for the
declarative checks path.
Also addresses the review comment on #1259 about covering Pydantic
`ValidationError`
behaviour with tests: DQX callers never see raw
`pydantic_core.ValidationError` —
every entry point raises DQX error types with the pre-migration message
format.

**Validation & deserialization through Pydantic:**
- New `CheckSpec` / `CheckBlock` Pydantic schemas in
`checks_validator.py` own the
structural validation of check metadata dicts (required `check` key,
criticality enum,
`for_each_column` non-empty list, argument dict shape). Signature-based
argument
validation (function exists, required params, types) is unchanged and
runs after.
- Pydantic errors are translated back into the pre-existing
human-readable messages
(`_translate_pydantic_error`), so users and existing tests see identical
error strings.
Unmatched errors name the offending field path (e.g. `user_metadata ->
confidence`)
instead of a generic message. Translation matches on stable Pydantic
error *types*,
  not version-dependent message text.
- `ChecksDeserializer.deserialize` reads typed attributes off the
validated `CheckSpec`
instead of raw-dict `.get()` calls, so schema, defaults, and validation
live in one place.
Note: it re-parses via `model_validate` after `validate_checks` —
structural parsing is
cheap, and having the validator return parsed specs would be a larger
API change.
- New `project_to_check_schema` strips storage-only columns
(`run_config_name`,
`created_at`, `rule_fingerprint`, `rule_set_fingerprint`) when loading
checks from
Lakebase, so stored checks round-trip cleanly into
`apply_checks_by_metadata` and match
the Delta path shape. The allow-list is derived from
`CheckSpec.model_fields`.

**`extra` policy on the migrated models (deliberate split):**
- `extra="forbid"` on
`DQRule`/`DQForEachColRule`/`BaseChecksStorageConfig`: a misspelled
kwarg raises `InvalidCheckError`/`InvalidConfigError` at construction.
This *restores
released behaviour* — v0.15.0 (last release) predates the Pydantic
migration and its
dataclasses raised `TypeError` on unknown kwargs; the lenient `ignore`
default was never
  shipped.
- `extra="ignore"` on `CheckSpec`/`CheckBlock`: check *metadata*
legitimately carries
extra keys (storage columns, user files), and rejecting them would break
the
  load -> apply round-trip.

**Cleanup:**
- Removed dead `checks_formats.py` (+ its tests):
`FILE_SERIALIZERS`/`FILE_DESERIALIZERS`
had no consumers anywhere in `src/` (already unused on `main`);
`SerializerFactory` is
  the single source of truth for file formats.

### Linked issues

Resolves #1285

### Tests

- [x] manually tested
- [ ] added unit tests
- [ ] added integration tests
- [ ] added end-to-end tests
- [ ] added performance tests

New `tests/unit/test_pydantic_serialization.py` covers:
- **Fingerprint byte-parity corpus**: SHA-256 values captured *before*
the migration,
asserted byte-identical after — stored fingerprints in existing tables
stay valid.
- Every translated Pydantic error path (missing `check`, bad criticality
value *and* type,
non-list/empty `for_each_column`, wrong-typed fields reported with their
location).
- YAML/JSON round-trip incl. the `__decimal__` wire format; deserializer
dispatch to
`DQRowRule`/`DQDatasetRule`/`for_each_column`; storage-column tolerance
+ projection.
- Regression tests for `extra="forbid"` (unknown kwargs rejected,
wrapped in DQX errors).

---------

Co-authored-by: Marcin Wojtyczka <marcin.wojtyczka@databricks.com>

Author: fedeflowers

AGENTS.md

@@ -64,7 +64,7 @@ src/databricks/labs/dqx/
   ├── manager.py           # DQRuleManager — build/manage rule collections
   ├── config.py            # WorkspaceConfig, RunConfig, AnomalyParams, LLMModelConfig, ExtraParams
   ├── checks_storage.py    # WorkspaceFileChecksStorageHandler, VolumeFileChecksStorageHandler
-  ├── checks_serializer.py / checks_resolver.py / checks_validator.py / checks_formats.py
+  ├── checks_serializer.py / checks_resolver.py / checks_validator.py
   ├── config_serializer.py # ConfigSerializer — use instead of dataclasses.asdict()
   ├── cli.py               # Databricks Labs CLI commands (@dqx.command)
   ├── errors.py            # For example: MissingParameterError, InvalidParameterError, UnsafeSqlQueryError — use instead of built-in exceptions

src/databricks/labs/dqx/checks_formats.py

@@ -8,7 +8,7 @@
 import yaml
 
 from databricks.labs.dqx.checks_resolver import resolve_check_function
-from databricks.labs.dqx.checks_validator import ChecksValidator
+from databricks.labs.dqx.checks_validator import ChecksValidator, CheckSpec
 from databricks.labs.dqx.rule import (
     DQRule,
     DQRowRule,
@@ -22,6 +22,31 @@
 logger = logging.getLogger(__name__)
 
 
+def project_to_check_schema(check: dict) -> dict:
+    """Return a copy of *check* containing only the logical check-metadata keys.
+
+    Storage backends persist columns alongside the check that are not part of the check metadata
+    accepted by *apply_checks_by_metadata* (e.g. *run_config_name*, *created_at*, *rule_fingerprint*,
+    *rule_set_fingerprint*). Loading a check therefore yields those extra keys. This helper drops
+    them so a loaded check round-trips cleanly.
+
+    Both load paths funnel their assembled check dict through this single projection so they agree
+    on the loaded shape: the Lakebase path (*LakebaseChecksStorageHandler._load_checks_from_lakebase*)
+    and the Delta path (*DataFrameConverter.from_dataframe*). The allowed keys are derived from
+    *CheckSpec.model_fields*, so a new logical field added to the schema is retained by both paths
+    automatically — provided the backing store also carries it (adding a *CheckSpec* field still
+    requires a matching Delta table column and Lakebase column for the value to survive a round-trip).
+
+    Args:
+        check: A check dict that may carry storage-only keys.
+
+    Returns:
+        A new dict with only the keys defined by *CheckSpec*.
+    """
+    allowed_keys = set(CheckSpec.model_fields)
+    return {key: value for key, value in check.items() if key in allowed_keys}
+
+
 class ChecksNormalizer:
     """
     Handles normalization and denormalization of check dictionaries.
@@ -226,7 +251,7 @@ def __init__(self, custom_checks: dict[str, Callable] | None = None):
 
     def deserialize(self, checks: list[dict]) -> list[DQRule]:
         """
-        Converts a list of quality checks defined as Python dictionaries to a list of `DQRule` objects.
+        Converts a list of quality checks defined as Python dictionaries to a list of *DQRule* objects.
 
         Args:
             checks: list of dictionaries describing checks. Each check is a dictionary
@@ -245,29 +270,30 @@ def deserialize(self, checks: list[dict]) -> list[DQRule]:
         Raises:
             InvalidCheckError: If any dictionary is invalid or unsupported.
         """
-        status = ChecksValidator.validate_checks(checks, self.custom_checks)
+        status, specs = ChecksValidator.validate_and_parse(checks, self.custom_checks)
         if status.has_errors:
             raise InvalidCheckError(str(status))
 
         dq_rule_checks: list[DQRule] = []
-        for check_def in checks:
-            logger.debug(f"Processing check definition: {check_def}")
+        for spec in specs:
+            # No errors means every check parsed, so specs holds a CheckSpec for each entry.
+            assert spec is not None  # validated above; narrows the type for the reads below
+            logger.debug(f"Processing check definition: {spec}")
 
-            check = check_def.get("check", {})
-            name = check_def.get("name") or ""
-            func_name = check.get("function")
+            func_name = spec.check.function
             func = resolve_check_function(func_name, self.custom_checks, fail_on_missing=True)
             assert func  # should already be validated
 
-            func_args = check.get("arguments", {})
-            for_each_column = check.get("for_each_column")
+            func_args: dict[str, Any] = spec.check.arguments
+            for_each_column = spec.check.for_each_column
             column = func_args.get("column")  # should be defined for single-column checks only
             columns = func_args.get("columns")  # should be defined for multi-column checks only
             assert not (column and columns)  # should already be validated
-            criticality = check_def.get("criticality", "error")
-            filter_str = check_def.get("filter")
-            user_metadata = check_def.get("user_metadata")
-            message_expr = check_def.get("message_expr")
+            name = spec.name or ""
+            criticality = spec.criticality
+            filter_str = spec.filter
+            user_metadata = spec.user_metadata
+            message_expr = spec.message_expr
 
             # Exclude `column` and `columns` from check_func_kwargs
             # as these are always included in the check function call

src/databricks/labs/dqx/checks_serializer.py

@@ -60,6 +60,7 @@
     SerializerFactory,
     ChecksNormalizer,
     deserialize_checks,
+    project_to_check_schema,
 )
 from databricks.labs.dqx.checks_validator import ChecksValidator
 from databricks.labs.dqx.rule_fingerprint import compute_rule_set_fingerprint_by_metadata
@@ -178,8 +179,10 @@ def from_dataframe(
                 check_dict["filter"] = row.filter
             if row.user_metadata is not None:
                 check_dict["user_metadata"] = row.user_metadata
-            # Denormalize special markers back to objects
-            checks.append(ChecksNormalizer.denormalize_value(check_dict))
+            # Route through the same projection as the Lakebase load path so both loaders agree on
+            # the loaded shape (the allow-list is derived from CheckSpec.model_fields). Then
+            # denormalize special markers back to objects.
+            checks.append(ChecksNormalizer.denormalize_value(project_to_check_schema(check_dict)))
         return checks
 
     @staticmethod
@@ -535,7 +538,7 @@ def get_engine(self, config: LakebaseChecksStorageConfig) -> Engine:
         engine = create_engine(
             engine_url,
             pool_recycle=45 * 60,  # recycle connections every 45 minutes
-            connect_args={'sslmode': 'require'},
+            connect_args={"sslmode": "require"},
             pool_size=4,
         )
         event.listen(engine, "do_connect", self._prepare_before_connect(config))
@@ -779,7 +782,7 @@ def _load_checks_from_lakebase(self, config: LakebaseChecksStorageConfig, engine
                     f"for run_config_name='{config.run_config_name}'. "
                     f"Make sure the profiler has run successfully and saved checks to this location."
                 )
-            checks_dict = [dict(check) for check in checks]
+            checks_dict = [project_to_check_schema(dict(check)) for check in checks]
 
             # Denormalize special markers back to objects
             return ChecksNormalizer.denormalize(checks_dict)
@@ -802,7 +805,7 @@ def _rule_set_columns_exists(engine: Engine, schema: str, table: str) -> bool:
         if not inspector.has_table(table, schema=schema):
             return False
         cols = inspector.get_columns(table, schema=schema)
-        existing_column_names = {col['name'] for col in cols}
+        existing_column_names = {col["name"] for col in cols}
         return all(x in existing_column_names for x in _VERSIONING_COLUMNS)
 
     def _check_for_undefined_table_error(self, e: ProgrammingError, config: LakebaseChecksStorageConfig) -> NoReturn:
@@ -819,8 +822,8 @@ def _check_for_undefined_table_error(self, e: ProgrammingError, config: Lakebase
             NotFound: If the table does not exist in the Lakebase instance (pgcode 42P01).
             ProgrammingError: Re-raises the original error if it's not an undefined table error.
         """
-        pgcode = getattr(getattr(e, 'orig', None), 'pgcode', None)
-        postgres_undefined_table_error = '42P01'
+        pgcode = getattr(getattr(e, "orig", None), "pgcode", None)
+        postgres_undefined_table_error = "42P01"
         if pgcode == postgres_undefined_table_error:
             raise NotFound(f"Table '{config.location}' does not exist in the Lakebase instance") from e
         raise e