feat(query): add _iter_sorts traversal and _prepare_sorts to Query(Sync) class hierarchy

Author: haakonvt

cognite/client/data_classes/data_modeling/query.py

@@ -2,7 +2,7 @@
 
 from abc import ABC
 from collections import UserDict
-from collections.abc import Mapping, MutableMapping, Sequence
+from collections.abc import Iterator, Mapping, MutableMapping, Sequence
 from dataclasses import dataclass, field
 from typing import TYPE_CHECKING, Any, Generic, Literal, TypeVar
 
@@ -85,6 +85,9 @@ def _load_list(
 class SelectBase(CogniteResource, ABC):
     sources: list[SourceSelector] = field(default_factory=list)
 
+    def _iter_sorts(self) -> Iterator[InstanceSort]:
+        yield from ()
+
     def dump(self, camel_case: bool = True) -> dict[str, Any]:
         output: dict[str, Any] = {}
         if self.sources:
@@ -133,6 +136,9 @@ class Select(SelectBase):
     sort: list[InstanceSort] = field(default_factory=list)
     limit: int | None = None
 
+    def _iter_sorts(self) -> Iterator[InstanceSort]:
+        yield from self.sort
+
     def dump(self, camel_case: bool = True) -> dict[str, Any]:
         output = super().dump(camel_case)
         if self.sort:
@@ -179,6 +185,16 @@ def instance_type_by_result_expression(self) -> dict[str, type[NodeListWithCurso
             for k, v in self.with_.items()
         }
 
+    def _iter_sorts(self) -> Iterator[InstanceSort]:
+        for expr in self.with_.values():
+            yield from expr._iter_sorts()
+        for sel in self.select.values():
+            yield from sel._iter_sorts()
+
+    def _prepare_sorts(self) -> None:
+        for sort in self._iter_sorts():
+            sort._apply_postgres_defaults_or_maybe_warn()
+
     def dump(self, camel_case: bool = True) -> dict[str, Any]:
         output: dict[str, Any] = {
             "with": {k: v.dump(camel_case) for k, v in self.with_.items()},
@@ -281,6 +297,9 @@ class ResultSetExpressionBase(CogniteResource, ABC):
     def _load_sort(resource: dict[str, Any], name: str) -> list[InstanceSort]:
         return [InstanceSort.load(sort) for sort in resource.get(name, [])]
 
+    def _iter_sorts(self) -> Iterator[InstanceSort]:
+        yield from ()
+
     @staticmethod
     def _init_through(through: list[str] | tuple[str, str, str] | PropertyId | None) -> PropertyId | None:
         def error() -> Never:
@@ -336,6 +355,9 @@ def __eq__(self, other: object) -> bool:
             return NotImplemented
         return type(self) is type(other) and self.dump() == other.dump()
 
+    def _iter_sorts(self) -> Iterator[InstanceSort]:
+        yield from self.sort
+
 
 @dataclass(eq=False)  # Prevents @dataclass from generating its own __eq__, so the parent's is used
 class NodeResultSetExpression(NodeOrEdgeResultSetExpression):
@@ -429,6 +451,10 @@ class EdgeResultSetExpression(NodeOrEdgeResultSetExpression):
     limit_each: int | None = None
     post_sort: list[InstanceSort] = field(default_factory=list)
 
+    def _iter_sorts(self) -> Iterator[InstanceSort]:
+        yield from self.sort
+        yield from self.post_sort
+
     @classmethod
     def _load(cls, resource: dict[str, Any]) -> Self:
         query_edge = resource["edges"]
@@ -494,6 +520,9 @@ def __eq__(self, other: object) -> bool:
             return NotImplemented
         return type(self) is type(other) and self.dump() == other.dump()
 
+    def _iter_sorts(self) -> Iterator[InstanceSort]:
+        yield from self.backfill_sort
+
     @classmethod
     def _load(cls, resource: dict[str, Any]) -> ResultSetExpressionSync:
         if "nodes" in resource: