Fork test: PR #1259 - Migrate Configs and Rules to Pydantic

.build-constraints.txt

@@ -1,9 +1,9 @@
 hatch-fancy-pypi-readme==25.1.0 \
     --hash=sha256:9c58ed3dff90d51f43414ce37009ad1d5b0f08ffc9fc216998a06380f01c0045 \
     --hash=sha256:ce0134c40d63d874ac48f48ccc678b8f3b62b8e50e9318520d2bffc752eedaf3
-hatchling==1.29.0 \
-    --hash=sha256:50af9343281f34785fab12da82e445ed987a6efb34fd8c2fc0f6e6630dbcc1b0 \
-    --hash=sha256:793c31816d952cee405b83488ce001c719f325d9cda69f1fc4cd750527640ea6
+hatchling==1.30.1 \
+    --hash=sha256:161eacafb3c6f91526e92116d21426369f2c36e98c36a864f11a96345ad4ee31 \
+    --hash=sha256:eee4fd45357f72ebb3d7a42e5d72cfb5e29ed426d79e8836288926c4258d5f2e
     # via hatch-fancy-pypi-readme
 packaging==26.2 \
     --hash=sha256:5fc45236b9446107ff2415ce77c807cee2862cb6fac22b8a73826d0693b0980e \
@@ -17,7 +17,7 @@ pluggy==1.6.0 \
     --hash=sha256:7dcc130b76258d33b90f61b658791dede3486c3e6bfb003ee5c9bfb396dd22f3 \
     --hash=sha256:e920276dd6813095e9377c0bc5566d94c932c33b27a3e3945d8389c374dd4746
     # via hatchling
-trove-classifiers==2026.1.14.14 \
-    --hash=sha256:00492545a1402b09d4858605ba190ea33243d361e2b01c9c296ce06b5c3325f3 \
-    --hash=sha256:1f9553927f18d0513d8e5ff80ab8980b8202ce37ecae0e3274ed2ef11880e74d
+trove-classifiers==2026.6.1.19 \
+    --hash=sha256:ab4c4ec93cc4a4e7815fa759906e05e6bb3f2fbd92ea0f897288c6a43efd15b3 \
+    --hash=sha256:c5132b4b61a829d11cfbd2d72e97f20a45ed6edb95e45c5efdeb5e00836b2745
     # via hatchling

Makefile

@@ -279,7 +279,11 @@ lock-dependencies: ## Regenerate top-level uv.lock and .build-constraints.txt
 	$(UV_RUN) --group yq tomlq -r '.["build-system"].requires[]' pyproject.toml | \
 	  uv pip compile --generate-hashes --universal --no-header - > build-constraints-new.txt
 	mv build-constraints-new.txt .build-constraints.txt
-	perl -pi -e 's|registry = "https://[^"]*"|registry = "https://pypi.org/simple"|g' uv.lock
+	# Normalize the lock so contributors inside Databricks (private proxy) and outside (public PyPI)
+	# produce an identical file. A proxy mirrors PyPI with identical paths, so rewrite the registry
+	# index and every per-package "/packages/..." download URL to the public hosts. Also drop the
+	# "size" field: the private proxy never reports it, so it is the only form both can reproduce.
+	perl -pi -e 's|registry = "https://[^"]*"|registry = "https://pypi.org/simple"|g; s|url = "https://[^/"]+/packages/|url = "https://files.pythonhosted.org/packages/|g; s|, size = \d+||g' uv.lock
 
 ##@ Misc
 

pyproject.toml

@@ -40,6 +40,7 @@ dependencies = [
     "databricks-sdk~=0.73",
     "sqlalchemy>=2.0,<3.0",
     "PyYAML~=6.0.3",
+    "pydantic>=2.8.2,<3",
 ]
 
 [project.optional-dependencies]
@@ -884,6 +885,7 @@ redefining-builtins-modules = ["six.moves", "past.builtins", "future.builtins",
 "tests/integration/test_app_backend.py" = "import-outside-toplevel"
 "src/databricks/labs/dqx/anomaly/anomaly_workflow.py" = "import-outside-toplevel"
 "src/databricks/labs/dqx/check_funcs.py" = "protected-access"
+"tests/unit/test_engine.py" = "protected-access"
 "tests/unit/test_anomaly_llm_explainer.py" = "protected-access"
 "tests/unit/test_anomaly_check_funcs_validation.py" = "protected-access"
 "tests/unit/test_anomaly_scoring_run.py" = "protected-access"

src/databricks/labs/dqx/anomaly/core.py

@@ -37,7 +37,6 @@
 logger = logging.getLogger(__name__)
 
 DEFAULT_SAMPLE_FRACTION = 0.3
-DEFAULT_MAX_ROWS = 1_000_000
 DEFAULT_TRAIN_RATIO = 0.8
 SCORE_QUANTILE_PROBS = [0.0, 0.01, 0.05, 0.1, 0.25, 0.5, 0.75, 0.9, 0.95, 0.99, 1.0]
 SCORE_QUANTILE_KEYS = ["p00", "p01", "p05", "p10", "p25", "p50", "p75", "p90", "p95", "p99", "p100"]

src/databricks/labs/dqx/checks_serializer.py

@@ -254,7 +254,7 @@ def deserialize(self, checks: list[dict]) -> list[DQRule]:
             logger.debug(f"Processing check definition: {check_def}")
 
             check = check_def.get("check", {})
-            name = check_def.get("name", None)
+            name = check_def.get("name") or ""
             func_name = check.get("function")
             func = resolve_check_function(func_name, self.custom_checks, fail_on_missing=True)
             assert func  # should already be validated

src/databricks/labs/dqx/checks_validator.py

@@ -2,44 +2,48 @@
 import functools as ft
 import inspect
 from collections.abc import Callable
-from dataclasses import dataclass, field
 from types import UnionType
 from typing import Any, get_origin, get_args
 
+from pydantic import BaseModel, ConfigDict
+
 from databricks.labs.dqx.checks_resolver import resolve_check_function
 from databricks.labs.dqx.rule import Criticality
 
 logger = logging.getLogger(__name__)
 
 
-@dataclass(frozen=True)
-class ChecksValidationStatus:
-    """Class to represent the validation status."""
+class ChecksValidationStatus(BaseModel):
+    """Class to represent the validation status.
+
+    This model is used as a mutable accumulator: *add_error* and *add_errors* append to the
+    *errors* list in place.  Pydantic instantiates a fresh copy of the ``[]`` default for each
+    instance, so the list is never shared between instances created via the constructor.  The
+    only sharing risk is a shallow ``model_copy()`` (without ``deep=True``); this model is never
+    shallow-copied, but use ``model_copy(deep=True)`` if that ever changes.
+    """
+
+    model_config = ConfigDict(extra="forbid")
 
-    _errors: list[str] = field(default_factory=list)
+    errors: list[str] = []
 
     def add_error(self, error: str):
         """Add an error to the validation status."""
-        self._errors.append(error)
+        self.errors.append(error)
 
     def add_errors(self, errors: list[str]):
-        """Add an error to the validation status."""
-        self._errors.extend(errors)
+        """Add errors to the validation status."""
+        self.errors.extend(errors)
 
     @property
     def has_errors(self) -> bool:
         """Check if there are any errors in the validation status."""
-        return bool(self._errors)
-
-    @property
-    def errors(self) -> list[str]:
-        """Get the list of errors in the validation status."""
-        return self._errors
+        return bool(self.errors)
 
     def to_string(self) -> str:
         """Convert the validation status to a string."""
         if self.has_errors:
-            return "\n".join(self._errors)
+            return "\n".join(self.errors)
         return "No errors found"
 
     def __str__(self) -> str:

src/databricks/labs/dqx/config.py

@@ -1,6 +1,8 @@
-import abc
+from abc import ABC
 from dataclasses import asdict, dataclass, field
-from functools import cached_property
+from typing import Any
+
+from pydantic import BaseModel, ValidationError, model_validator
 
 from databricks.labs.dqx.checks_serializer import SerializerFactory
 from databricks.labs.dqx.errors import InvalidConfigError, InvalidParameterError
@@ -326,8 +328,7 @@ def get_run_config(self, run_config_name: str | None = "default") -> RunConfig:
         raise InvalidConfigError("No run configurations available")
 
 
-@dataclass
-class BaseChecksStorageConfig(abc.ABC):
+class BaseChecksStorageConfig(BaseModel, ABC):
     """Marker base class for storage configuration.
 
     Args:
@@ -336,8 +337,32 @@ class BaseChecksStorageConfig(abc.ABC):
 
     location: str
 
+    def __init__(self, **data: Any) -> None:
+        try:
+            super().__init__(**data)
+        except ValidationError as exc:
+            raise InvalidConfigError(str(exc)) from exc
+
+    def replace(self, **changes: Any) -> "BaseChecksStorageConfig":
+        """Return a new config instance with the given field overrides, fully re-validated.
+
+        Unlike *model_copy(update=...)*, which shallow-copies the instance and skips all
+        validators, this rebuilds the config through the constructor so the *model_validator*
+        checks (e.g. *mode* and *location* format) re-run against the updated fields. Use this
+        instead of *model_copy* when overriding fields, so an invalid override is rejected at
+        the point of the change rather than failing later during a save/load operation.
+
+        Args:
+            **changes: Field values to override on the new instance.
+
+        Returns:
+            A new, fully validated config of the same concrete type.
+        """
+        fields = {name: getattr(self, name) for name in type(self).model_fields}
+        fields.update(changes)
+        return type(self)(**fields)
+
 
-@dataclass
 class FileChecksStorageConfig(BaseChecksStorageConfig):
     """
     Configuration class for storing checks in a file.
@@ -346,14 +371,20 @@ class FileChecksStorageConfig(BaseChecksStorageConfig):
         location: The file path where the checks are stored.
     """
 
-    location: str
-
-    def __post_init__(self):
-        if not self.location:
+    @model_validator(mode='before')
+    @classmethod
+    def validate_location(cls, data: Any) -> Any:
+        # `cls is FileChecksStorageConfig` intentionally excludes subclasses.
+        # InstallationChecksStorageConfig uses diamond inheritance from this class and
+        # WorkspaceFileChecksStorageConfig; both validators are inherited and would fire
+        # for any subclass without the exact-class guard, giving confusing duplicate errors.
+        # InstallationChecksStorageConfig has its own validate_installation_location that
+        # covers the empty-location case for itself.
+        if cls is FileChecksStorageConfig and isinstance(data, dict) and not data.get("location"):
             raise InvalidConfigError("The file path ('location' field) must not be empty or None.")
+        return data
 
 
-@dataclass
 class WorkspaceFileChecksStorageConfig(BaseChecksStorageConfig):
     """
     Configuration class for storing checks in a workspace file.
@@ -362,14 +393,15 @@ class WorkspaceFileChecksStorageConfig(BaseChecksStorageConfig):
         location: The workspace file path where the checks are stored.
     """
 
-    location: str
-
-    def __post_init__(self):
-        if not self.location:
+    @model_validator(mode='before')
+    @classmethod
+    def validate_location(cls, data: Any) -> Any:
+        # See FileChecksStorageConfig.validate_location for why the exact-class guard is needed.
+        if cls is WorkspaceFileChecksStorageConfig and isinstance(data, dict) and not data.get("location"):
             raise InvalidConfigError("The workspace file path ('location' field) must not be empty or None.")
+        return data
 
 
-@dataclass
 class TableChecksStorageConfig(BaseChecksStorageConfig):
     """
     Configuration class for storing checks in a table.
@@ -387,17 +419,19 @@ class TableChecksStorageConfig(BaseChecksStorageConfig):
             When None (default), loads the latest batch.
     """
 
-    location: str
     run_config_name: str = "default"  # to filter checks by run config
     mode: str = "append"
     rule_set_fingerprint: str | None = None  # to filter checks by rule set fingerprint
 
-    def __post_init__(self):
-        if not self.location:
+    @model_validator(mode='before')
+    @classmethod
+    def validate_location(cls, data: Any) -> Any:
+        # See FileChecksStorageConfig.validate_location for why the exact-class guard is needed.
+        if cls is TableChecksStorageConfig and isinstance(data, dict) and not data.get("location"):
             raise InvalidConfigError("The table name ('location' field) must not be empty or None.")
+        return data
 
 
-@dataclass
 class LakebaseChecksStorageConfig(BaseChecksStorageConfig):
     """
     Configuration class for storing checks in a Lakebase table.
@@ -418,19 +452,26 @@ class LakebaseChecksStorageConfig(BaseChecksStorageConfig):
             When None (default), loads the latest batch.
     """
 
-    location: str
     instance_name: str | None = None
     client_id: str | None = None
     port: str = "5432"
     run_config_name: str = "default"
     mode: str = "append"
     rule_set_fingerprint: str | None = None
 
-    def __post_init__(self):
-        if not self.location or self.location == "":
+    @model_validator(mode='before')
+    @classmethod
+    def validate_location_present(cls, data: Any) -> Any:
+        if cls is LakebaseChecksStorageConfig and isinstance(data, dict) and not data.get("location"):
             raise InvalidParameterError("Location must not be empty or None.")
+        return data
+
+    @model_validator(mode='after')
+    def validate_lakebase_config(self) -> 'LakebaseChecksStorageConfig':
+        if self.__class__ is not LakebaseChecksStorageConfig:
+            return self
 
-        if not self.instance_name or self.instance_name == "":
+        if not self.instance_name:
             raise InvalidParameterError("Instance name must not be empty or None.")
 
         if len(self.location.split(".")) != 3:
@@ -441,26 +482,27 @@ def __post_init__(self):
         if self.mode not in ("append", "overwrite"):
             raise InvalidConfigError(f"Invalid mode '{self.mode}'. Must be 'append' or 'overwrite'.")
 
+        return self
+
     def _split_location(self) -> tuple[str, ...]:
         """Splits 'database.schema.table' into components."""
         if not self.location:
             raise InvalidConfigError("location must be set before accessing database components.")
         return tuple(self.location.split("."))
 
-    @cached_property
+    @property
     def database_name(self) -> str:
         return self._split_location()[0]
 
-    @cached_property
+    @property
     def schema_name(self) -> str:
         return self._split_location()[1]
 
-    @cached_property
+    @property
     def table_name(self) -> str:
         return self._split_location()[2]
 
 
-@dataclass
 class VolumeFileChecksStorageConfig(BaseChecksStorageConfig):
     """
     Configuration class for storing checks in a Unity Catalog volume file.
@@ -469,13 +511,19 @@ class VolumeFileChecksStorageConfig(BaseChecksStorageConfig):
         location: The Unity Catalog volume file path where the checks are stored.
     """
 
-    location: str
-
-    def __post_init__(self):
-        if not self.location:
+    @model_validator(mode='before')
+    @classmethod
+    def validate_location_present(cls, data: Any) -> Any:
+        if cls is VolumeFileChecksStorageConfig and isinstance(data, dict) and not data.get("location"):
             raise InvalidParameterError(
                 "The Unity Catalog volume file path ('location' field) must not be empty or None."
             )
+        return data
+
+    @model_validator(mode='after')
+    def validate_location(self) -> 'VolumeFileChecksStorageConfig':
+        if self.__class__ is not VolumeFileChecksStorageConfig:
+            return self
 
         # Expected format: /Volumes/{catalog}/{schema}/{volume}/{path/to/file}
         if not self.location.startswith("/Volumes/"):
@@ -492,8 +540,9 @@ def __post_init__(self):
         if len(parts) < 6 or not parts[-1].lower().endswith(SerializerFactory.get_supported_extensions()):
             raise InvalidParameterError("Invalid path: Path must include a file name after the volume")
 
+        return self
+
 
-@dataclass
 class InstallationChecksStorageConfig(
     WorkspaceFileChecksStorageConfig,
     TableChecksStorageConfig,
@@ -523,3 +572,11 @@ class InstallationChecksStorageConfig(
     assume_user: bool = True
     install_folder: str | None = None
     overwrite_location: bool = False
+
+    @model_validator(mode='before')
+    @classmethod
+    def validate_installation_location(cls, data: Any) -> Any:
+        # Only guard against an explicitly empty/None location; an absent key keeps the default above.
+        if isinstance(data, dict) and "location" in data and not data["location"]:
+            raise InvalidConfigError("The workspace file path ('location' field) must not be empty or None.")
+        return data

src/databricks/labs/dqx/engine.py

@@ -5,7 +5,6 @@
 import re
 from concurrent import futures
 from collections.abc import Callable
-from dataclasses import replace
 from datetime import datetime
 from functools import cached_property
 from typing import Any
@@ -484,7 +483,7 @@ def _preselect_original_columns(self, df: DataFrame, check: DQRule) -> DQRule:
         # preselect original columns
         rule_kwargs = check.check_func_kwargs.copy()
         rule_kwargs["columns"] = [col for col in df.columns if col not in set(self._result_column_names.values())]
-        return replace(check, check_func_kwargs=rule_kwargs)
+        return check.replace(check_func_kwargs=rule_kwargs)
 
     def _append_empty_checks(self, df: DataFrame) -> DataFrame:
         """Append empty checks at the end of DataFrame.