From e5e8b2d7fca61495a969e628b82f59925296b1aa Mon Sep 17 00:00:00 2001
From: Tim Jenness <tjenness@lsst.org>
Date: Thu, 3 Aug 2023 17:04:35 -0700
Subject: [PATCH] Use v2 model validator for DatasetRef

Now use a single model validator that checks everything because
we are checking multiple keys.
---
 python/lsst/daf/butler/core/datasets/ref.py | 56 +++++++++++++--------
 1 file changed, 36 insertions(+), 20 deletions(-)

diff --git a/python/lsst/daf/butler/core/datasets/ref.py b/python/lsst/daf/butler/core/datasets/ref.py
index 49af583d5b..d2321d4b8e 100644
--- a/python/lsst/daf/butler/core/datasets/ref.py
+++ b/python/lsst/daf/butler/core/datasets/ref.py
@@ -35,9 +35,9 @@
 from collections.abc import Iterable
 from typing import TYPE_CHECKING, Any, ClassVar, Protocol, TypeAlias, runtime_checkable
 
-from lsst.daf.butler._compat import _BaseModelCompat
+from lsst.daf.butler._compat import PYDANTIC_V2, _BaseModelCompat
 from lsst.utils.classes import immutable
-from pydantic import StrictStr, validator
+from pydantic import StrictStr
 
 from ..configSupport import LookupKey
 from ..dimensions import DataCoordinate, DimensionGraph, DimensionUniverse, SerializedDataCoordinate
@@ -179,24 +179,40 @@ class SerializedDatasetRef(_BaseModelCompat):
     run: StrictStr | None = None
     component: StrictStr | None = None
 
-    @validator("dataId")
-    def _check_dataId(cls, v: Any, values: dict[str, Any]) -> Any:  # noqa: N805
-        if (d := "datasetType") in values and values[d] is None:
-            raise ValueError("Can not specify 'dataId' without specifying 'datasetType'")
-        return v
-
-    @validator("run")
-    def _check_run(cls, v: Any, values: dict[str, Any]) -> Any:  # noqa: N805
-        if v and (i := "id") in values and values[i] is None:
-            raise ValueError("'run' cannot be provided unless 'id' is.")
-        return v
-
-    @validator("component")
-    def _check_component(cls, v: Any, values: dict[str, Any]) -> Any:  # noqa: N805
-        # Component should not be given if datasetType is given
-        if v and (d := "datasetType") in values and values[d] is not None:
-            raise ValueError(f"datasetType ({values[d]}) can not be set if component is given ({v}).")
-        return v
+    if PYDANTIC_V2:
+        from pydantic import model_validator
+
+        # Can not use "after" validator since in some cases the validator
+        # seems to trigger with the datasetType field not yet set.
+        @model_validator(mode="before")
+        @classmethod
+        def check_consistent_parameters(cls, data: Any) -> Any:
+            has_datasetType = data.get("datasetType") is not None
+            has_dataId = data.get("dataId") is not None
+            if has_datasetType is not has_dataId:
+                raise ValueError("If specifying datasetType or dataId, must specify both.")
+
+            if data.get("component") is not None and has_datasetType:
+                raise ValueError("datasetType can not be set if component is given.")
+            return data
+
+    else:
+        from pydantic import validator
+
+        @validator("dataId")
+        def _check_dataId(cls, v: Any, values: dict[str, Any]) -> Any:  # noqa: N805
+            print("---\nGetting dataId ", v)
+            print("Dataset Type: ", values["datasetType"])
+            if v and (d := "datasetType") in values and values[d] is None:
+                raise ValueError("Can not specify 'dataId' without specifying 'datasetType'")
+            return v
+
+        @validator("component")
+        def _check_component(cls, v: Any, values: dict[str, Any]) -> Any:  # noqa: N805
+            # Component should not be given if datasetType is given
+            if v and (d := "datasetType") in values and values[d] is not None:
+                raise ValueError(f"datasetType ({values[d]}) can not be set if component is given ({v}).")
+            return v
 
     @classmethod
     def direct(