warn on non-index-aligned InstanceSort at PostgreSQL-backed API call sites

Author: haakonvt

cognite/client/_api/data_modeling/instances.py

@@ -894,6 +894,7 @@ async def subscribe(
                 >>> subscription_context.cancel()
 
         """
+        query.warn_non_index_aligned_sorts()
         subscription_context = SubscriptionContext()
 
         async def _poll_loop() -> None:
@@ -968,10 +969,11 @@ def _create_other_params(
                     f"Received in `sources` argument for views: {with_properties}."
                 )
         if sort:
-            if isinstance(sort, (InstanceSort, dict)):
-                other_params["sort"] = [cls._dump_instance_sort(sort)]
-            else:
-                other_params["sort"] = [cls._dump_instance_sort(s) for s in sort]
+            sorts_seq = [sort] if isinstance(sort, (InstanceSort, dict)) else list(sort)
+            for s in sorts_seq:
+                if isinstance(s, InstanceSort):
+                    s.warn_if_not_index_aligned()
+            other_params["sort"] = [cls._dump_instance_sort(s) for s in sorts_seq]
         if instance_type:
             other_params["instanceType"] = instance_type
         if debug:
@@ -1647,6 +1649,7 @@ async def query(
                 >>> res = client.data_modeling.instances.query(query, debug=debug_params)
                 >>> print(res.debug)
         """
+        query.warn_non_index_aligned_sorts()
         return await self._query_or_sync(query, "query", include_typing=include_typing, debug=debug)
 
     async def sync(
@@ -1749,6 +1752,7 @@ async def sync(
                 >>> res = client.data_modeling.instances.sync(query, debug=debug_params)
                 >>> print(res.debug)
         """
+        query.warn_non_index_aligned_sorts()
         return await self._query_or_sync(query, "sync", include_typing=include_typing, debug=debug)
 
     async def _query_or_sync(

cognite/client/_sync_api/data_modeling/instances.py

@@ -1,6 +1,6 @@
 """
 ===============================================================================
-c0af4f83cd8ffa0aaed6fa641bde9a98
+28cd64f8d9cdd590b9601c562f1b69a1
 This file is auto-generated from the Async API modules, - do not edit manually!
 ===============================================================================
 """

tests/tests_unit/test_data_classes/test_data_models/test_instances.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import warnings
 from datetime import date, datetime
 from typing import Any, cast
 
@@ -12,6 +13,7 @@
     EdgeApply,
     EdgeList,
     Float64,
+    InstanceSort,
     Node,
     NodeApply,
     NodeId,
@@ -638,3 +640,156 @@ def test_to_pandas(self) -> None:
         df = info.to_pandas()
 
         pd.testing.assert_frame_equal(df, expected)
+
+
+class TestInstanceSort:
+    @pytest.mark.parametrize(
+        "direction, expected_nulls_first",
+        [
+            ("ascending", False),
+            ("descending", True),
+        ],
+    )
+    def test_nulls_first_auto_derived_from_direction(self, direction: str, expected_nulls_first: bool) -> None:
+        sort = InstanceSort(["node", "externalId"], direction=direction)  # type: ignore[arg-type]
+        assert sort.nulls_first is expected_nulls_first
+
+    def test_no_warning_at_construction_for_any_combination(self) -> None:
+        with warnings.catch_warnings(record=True) as w:
+            warnings.simplefilter("always")
+            InstanceSort(["node", "externalId"], direction="ascending")
+            InstanceSort(["node", "externalId"], direction="descending")
+            InstanceSort(["node", "externalId"], direction="ascending", nulls_first=False)
+            InstanceSort(["node", "externalId"], direction="ascending", nulls_first=True)
+            InstanceSort(["node", "externalId"], direction="descending", nulls_first=True)
+            InstanceSort(["node", "externalId"], direction="descending", nulls_first=False)
+        assert len(w) == 0
+
+    @pytest.mark.parametrize(
+        "direction, nulls_first",
+        [
+            ("ascending", True),
+            ("descending", False),
+        ],
+    )
+    def test_non_index_aligned_nulls_first_is_stored_as_given(self, direction: str, nulls_first: bool) -> None:
+        sort = InstanceSort(["node", "externalId"], direction=direction, nulls_first=nulls_first)  # type: ignore[arg-type]
+        assert sort.nulls_first is nulls_first
+
+    def test_load_does_not_override_nulls_first(self) -> None:
+        raw = {"property": ["node", "externalId"], "direction": "ascending", "nullsFirst": True}
+        sort = InstanceSort._load(raw)
+        assert sort.direction == "ascending"
+        assert sort.nulls_first is True
+
+    def test_dump_round_trip(self) -> None:
+        sort = InstanceSort(["space", "externalId"], direction="descending")
+        dumped = sort.dump(camel_case=True)
+        assert dumped["direction"] == "descending"
+        assert dumped["nullsFirst"] is True
+
+    @pytest.mark.parametrize(
+        "direction, nulls_first, expected",
+        [
+            ("ascending", False, True),
+            ("descending", True, True),
+            ("ascending", True, False),
+            ("descending", False, False),
+        ],
+    )
+    def test_is_index_aligned(self, direction: str, nulls_first: bool, expected: bool) -> None:
+        sort = InstanceSort(["node", "externalId"], direction=direction, nulls_first=nulls_first)  # type: ignore[arg-type]
+        assert sort.is_index_aligned is expected
+
+    def test_warn_if_not_index_aligned_fires(self) -> None:
+        sort = InstanceSort(["node", "externalId"], direction="ascending", nulls_first=True)
+        with warnings.catch_warnings(record=True) as w:
+            warnings.simplefilter("always")
+            sort.warn_if_not_index_aligned()
+        assert len(w) == 1
+        assert issubclass(w[0].category, UserWarning)
+        assert "not index-aligned" in str(w[0].message)
+
+    def test_warn_if_not_index_aligned_silent_when_aligned(self) -> None:
+        sort = InstanceSort(["node", "externalId"], direction="ascending")
+        with warnings.catch_warnings(record=True) as w:
+            warnings.simplefilter("always")
+            sort.warn_if_not_index_aligned()
+        assert len(w) == 0
+
+    @pytest.mark.parametrize("bad_direction", ["asc", "desc", "random"])
+    def test_invalid_direction_raises(self, bad_direction: str) -> None:
+        with pytest.raises(ValueError, match="direction must be"):
+            InstanceSort(["node", "externalId"], direction=bad_direction)  # type: ignore[arg-type]
+
+    def test_invalid_direction_error_shows_original_value(self) -> None:
+        with pytest.raises(ValueError, match="'Asc'"):
+            InstanceSort(["node", "externalId"], direction="Asc")  # type: ignore[arg-type]
+
+
+class TestInstanceSortAPIWarning:
+    @pytest.mark.parametrize(
+        "direction, nulls_first",
+        [
+            ("ascending", True),
+            ("descending", False),
+        ],
+    )
+    def test_warn_fires_for_non_index_aligned_sorts(self, direction: str, nulls_first: bool) -> None:
+        sort = InstanceSort(["node", "externalId"], direction=direction, nulls_first=nulls_first)  # type: ignore[arg-type]
+        with warnings.catch_warnings(record=True) as w:
+            warnings.simplefilter("always")
+            sort.warn_if_not_index_aligned()
+        assert len(w) == 1
+        assert issubclass(w[0].category, UserWarning)
+        assert "not index-aligned" in str(w[0].message)
+        assert "PostgreSQL" in str(w[0].message)
+
+    @pytest.mark.parametrize(
+        "direction, nulls_first",
+        [
+            ("ascending", False),
+            ("descending", True),
+            ("ascending", None),
+            ("descending", None),
+        ],
+    )
+    def test_no_warn_for_aligned_or_auto_derived_sorts(self, direction: str, nulls_first: bool | None) -> None:
+        sort = InstanceSort(["node", "externalId"], direction=direction, nulls_first=nulls_first)  # type: ignore[arg-type]
+        with warnings.catch_warnings(record=True) as w:
+            warnings.simplefilter("always")
+            sort.warn_if_not_index_aligned()
+        assert len(w) == 0
+
+    def test_query__iter_sorts_collects_all_sorts(self) -> None:
+        from cognite.client.data_classes.data_modeling import ViewId
+        from cognite.client.data_classes.data_modeling.query import (
+            NodeResultSetExpression,
+            Query,
+            Select,
+            SourceSelector,
+        )
+
+        view = ViewId("s", "v", "1")
+        sort_a = InstanceSort(view.as_property_ref("a"))
+        sort_b = InstanceSort(view.as_property_ref("b"))
+        query = Query(
+            with_={"nodes": NodeResultSetExpression(sort=[sort_a])},
+            select={"nodes": Select([SourceSelector(view, ["a"])], sort=[sort_b])},
+        )
+        result = list(query._iter_sorts())
+        assert sort_a in result
+        assert sort_b in result
+        assert len(result) == 2
+
+    def test_query_sync__iter_sorts_collects_backfill_sort(self) -> None:
+        from cognite.client.data_classes.data_modeling.query import (
+            NodeResultSetExpressionSync,
+            QuerySync,
+        )
+
+        sort = InstanceSort(["node", "externalId"])
+        query = QuerySync(with_={"nodes": NodeResultSetExpressionSync(backfill_sort=[sort])}, select={})
+        result = list(query._iter_sorts())
+        assert sort in result
+        assert len(result) == 1