feat(Data Modeling): support usedFor=record on containers (#2621)

Author: evertoncolling

cognite/client/_api/data_modeling/containers.py

@@ -10,6 +10,7 @@
     Container,
     ContainerApply,
     ContainerList,
+    ContainerUsedFor,
     _ContainerFilter,
 )
 from cognite.client.data_classes.data_modeling.ids import (
@@ -48,6 +49,7 @@ def __call__(
         chunk_size: None = None,
         space: str | None = None,
         include_global: bool = False,
+        used_for: ContainerUsedFor | Sequence[ContainerUsedFor] | None = None,
         limit: int | None = None,
     ) -> AsyncIterator[Container]: ...
 
@@ -57,6 +59,7 @@ def __call__(
         chunk_size: int,
         space: str | None = None,
         include_global: bool = False,
+        used_for: ContainerUsedFor | Sequence[ContainerUsedFor] | None = None,
         limit: int | None = None,
     ) -> AsyncIterator[ContainerList]: ...
 
@@ -65,6 +68,7 @@ async def __call__(
         chunk_size: int | None = None,
         space: str | None = None,
         include_global: bool = False,
+        used_for: ContainerUsedFor | Sequence[ContainerUsedFor] | None = None,
         limit: int | None = None,
     ) -> AsyncIterator[Container] | AsyncIterator[ContainerList]:
         """Iterate over containers
@@ -75,12 +79,13 @@ async def __call__(
             chunk_size (int | None): Number of containers to return in each chunk. Defaults to yielding one container a time.
             space (str | None): The space to query.
             include_global (bool): Whether the global containers should be returned.
+            used_for (ContainerUsedFor | Sequence[ContainerUsedFor] | None): Only include containers marked for these purposes. If omitted, containers of every kind (nodes, edges, and records) are returned.
             limit (int | None): Maximum number of containers to return. Defaults to returning all items.
 
         Yields:
             Container | ContainerList: yields Container one by one if chunk_size is not specified, else ContainerList objects.
         """  # noqa: DOC404
-        flt = _ContainerFilter(space, include_global)
+        flt = _ContainerFilter(space, include_global, used_for)
         async for item in self._list_generator(
             list_cls=ContainerList,
             resource_cls=Container,
@@ -226,13 +231,15 @@ async def list(
         space: str | None = None,
         limit: int | None = DATA_MODELING_DEFAULT_LIMIT_READ,
         include_global: bool = False,
+        used_for: ContainerUsedFor | Sequence[ContainerUsedFor] | None = None,
     ) -> ContainerList:
         """`List containers <https://api-docs.cognite.com/20230101/tag/Containers/operation/listContainers>`_.
 
         Args:
             space (str | None): The space to query
             limit (int | None): Maximum number of containers to return. Defaults to 10. Set to -1, float("inf") or None to return all items.
             include_global (bool): Whether the global containers should be returned.
+            used_for (ContainerUsedFor | Sequence[ContainerUsedFor] | None): Only include containers marked for these purposes. If omitted, containers of every kind (nodes, edges, and records) are returned.
 
         Returns:
             ContainerList: List of requested containers
@@ -246,6 +253,16 @@ async def list(
                 >>> # async_client = AsyncCogniteClient()  # another option
                 >>> container_list = client.data_modeling.containers.list(limit=5)
 
+            Filter containers by `used_for`. Note that "all" refers to containers that stores *both*
+            nodes and edges:
+
+                >>> # High-volume data containers (records)
+                >>> record_containers = client.data_modeling.containers.list(used_for="record")
+                >>> # Containers that store ONLY nodes or ONLY edges (excludes "all"):
+                >>> containers = client.data_modeling.containers.list(used_for=["node", "edge"])
+                >>> # All containers that can store nodes:
+                >>> containers = client.data_modeling.containers.list(used_for=["node", "all"])
+
             Iterate over containers, one-by-one:
 
                 >>> for container in client.data_modeling.containers():
@@ -256,7 +273,7 @@ async def list(
                 >>> for container_list in client.data_modeling.containers(chunk_size=10):
                 ...     container_list  # do something with the containers
         """
-        flt = _ContainerFilter(space, include_global)
+        flt = _ContainerFilter(space, include_global, used_for)
         return await self._list(
             list_cls=ContainerList,
             resource_cls=Container,

cognite/client/_sync_api/data_modeling/containers.py

@@ -1,6 +1,6 @@
 """
 ===============================================================================
-465575f7283dac5a1831298bc22b3bda
+fa523a777c9b728e0712f75d938fbb0c
 This file is auto-generated from the Async API modules, - do not edit manually!
 ===============================================================================
 """
@@ -17,6 +17,7 @@
     Container,
     ContainerApply,
     ContainerList,
+    ContainerUsedFor,
 )
 from cognite.client.data_classes.data_modeling.ids import (
     ConstraintIdentifier,
@@ -38,19 +39,30 @@ def __init__(self, async_client: AsyncCogniteClient) -> None:
 
     @overload
     def __call__(
-        self, chunk_size: None = None, space: str | None = None, include_global: bool = False, limit: int | None = None
+        self,
+        chunk_size: None = None,
+        space: str | None = None,
+        include_global: bool = False,
+        used_for: ContainerUsedFor | Sequence[ContainerUsedFor] | None = None,
+        limit: int | None = None,
     ) -> Iterator[Container]: ...
 
     @overload
     def __call__(
-        self, chunk_size: int, space: str | None = None, include_global: bool = False, limit: int | None = None
+        self,
+        chunk_size: int,
+        space: str | None = None,
+        include_global: bool = False,
+        used_for: ContainerUsedFor | Sequence[ContainerUsedFor] | None = None,
+        limit: int | None = None,
     ) -> Iterator[ContainerList]: ...
 
     def __call__(
         self,
         chunk_size: int | None = None,
         space: str | None = None,
         include_global: bool = False,
+        used_for: ContainerUsedFor | Sequence[ContainerUsedFor] | None = None,
         limit: int | None = None,
     ) -> Iterator[Container] | Iterator[ContainerList]:
         """
@@ -62,14 +74,15 @@ def __call__(
             chunk_size (int | None): Number of containers to return in each chunk. Defaults to yielding one container a time.
             space (str | None): The space to query.
             include_global (bool): Whether the global containers should be returned.
+            used_for (ContainerUsedFor | Sequence[ContainerUsedFor] | None): Only include containers marked for these purposes. If omitted, containers of every kind (nodes, edges, and records) are returned.
             limit (int | None): Maximum number of containers to return. Defaults to returning all items.
 
         Yields:
             Container | ContainerList: yields Container one by one if chunk_size is not specified, else ContainerList objects.
         """  # noqa: DOC404
         yield from SyncIterator(
             self.__async_client.data_modeling.containers(
-                chunk_size=chunk_size, space=space, include_global=include_global, limit=limit
+                chunk_size=chunk_size, space=space, include_global=include_global, used_for=used_for, limit=limit
             )
         )  # type: ignore [misc]
 
@@ -171,6 +184,7 @@ def list(
         space: str | None = None,
         limit: int | None = DATA_MODELING_DEFAULT_LIMIT_READ,
         include_global: bool = False,
+        used_for: ContainerUsedFor | Sequence[ContainerUsedFor] | None = None,
     ) -> ContainerList:
         """
         `List containers <https://api-docs.cognite.com/20230101/tag/Containers/operation/listContainers>`_.
@@ -179,6 +193,7 @@ def list(
             space (str | None): The space to query
             limit (int | None): Maximum number of containers to return. Defaults to 10. Set to -1, float("inf") or None to return all items.
             include_global (bool): Whether the global containers should be returned.
+            used_for (ContainerUsedFor | Sequence[ContainerUsedFor] | None): Only include containers marked for these purposes. If omitted, containers of every kind (nodes, edges, and records) are returned.
 
         Returns:
             ContainerList: List of requested containers
@@ -192,6 +207,16 @@ def list(
                 >>> # async_client = AsyncCogniteClient()  # another option
                 >>> container_list = client.data_modeling.containers.list(limit=5)
 
+            Filter containers by `used_for`. Note that "all" refers to containers that stores *both*
+            nodes and edges:
+
+                >>> # High-volume data containers (records)
+                >>> record_containers = client.data_modeling.containers.list(used_for="record")
+                >>> # Containers that store ONLY nodes or ONLY edges (excludes "all"):
+                >>> containers = client.data_modeling.containers.list(used_for=["node", "edge"])
+                >>> # All containers that can store nodes:
+                >>> containers = client.data_modeling.containers.list(used_for=["node", "all"])
+
             Iterate over containers, one-by-one:
 
                 >>> for container in client.data_modeling.containers():
@@ -203,7 +228,9 @@ def list(
                 ...     container_list  # do something with the containers
         """
         return run_sync(
-            self.__async_client.data_modeling.containers.list(space=space, limit=limit, include_global=include_global)
+            self.__async_client.data_modeling.containers.list(
+                space=space, limit=limit, include_global=include_global, used_for=used_for
+            )
         )
 
     @overload

cognite/client/data_classes/data_modeling/__init__.py

@@ -14,6 +14,7 @@
     ContainerList,
     ContainerProperty,
     ContainerPropertyApply,
+    ContainerUsedFor,
     Index,
     IndexApply,
     InvertedIndex,
@@ -168,6 +169,7 @@
     "ContainerList",
     "ContainerProperty",
     "ContainerPropertyApply",
+    "ContainerUsedFor",
     "DataModel",
     "DataModelApply",
     "DataModelApplyList",

cognite/client/data_classes/data_modeling/containers.py

@@ -2,7 +2,7 @@
 
 from abc import ABC, abstractmethod
 from builtins import type as type_
-from collections.abc import Mapping
+from collections.abc import Mapping, Sequence
 from dataclasses import dataclass, field
 from typing import Any, Literal, TypeVar, cast
 
@@ -30,6 +30,9 @@
 _T_Constraint = TypeVar("_T_Constraint", bound="ConstraintCore")
 _T_Index = TypeVar("_T_Index", bound="IndexCore")
 
+ContainerUsedFor = Literal["node", "edge", "record", "all"]
+_ALL_CONTAINER_USED_FOR: tuple[ContainerUsedFor, ...] = ("node", "edge", "record", "all")
+
 
 @dataclass
 class ContainerCore(DataModelingSchemaResource["ContainerApply"], ABC):
@@ -107,13 +110,13 @@ class ContainerApply(ContainerCore):
         name (str | None): Human readable name for the container.
         constraints (Mapping[str, ConstraintApply]): Set of constraints to apply to the container
         indexes (Mapping[str, IndexApply]): Set of indexes to apply to the container.
-        used_for (Literal['node', 'edge', 'all'] | None): Should this operation apply to nodes, edges or both.
+        used_for (ContainerUsedFor | None): Whether the container is for nodes, edges, records, or both nodes and edges (``all``).
     """
 
     properties: Mapping[str, ContainerPropertyApply]
     constraints: Mapping[str, ConstraintApply] = field(default_factory=dict)
     indexes: Mapping[str, IndexApply] = field(default_factory=dict)
-    used_for: Literal["node", "edge", "all"] | None = None
+    used_for: ContainerUsedFor | None = None
 
     def __post_init__(self) -> None:
         validate_data_modeling_identifier(self.space, self.external_id)
@@ -149,14 +152,14 @@ class Container(ContainerCore):
         is_global (bool): Whether this is a global container, i.e., one of the out-of-the-box models.
         last_updated_time (int): The number of milliseconds since 00:00:00 Thursday, 1 January 1970, Coordinated Universal Time (UTC), minus leap seconds.
         created_time (int): The number of milliseconds since 00:00:00 Thursday, 1 January 1970, Coordinated Universal Time (UTC), minus leap seconds.
-        used_for (Literal['node', 'edge', 'all']): Should this operation apply to nodes, edges or both.
+        used_for (ContainerUsedFor): Whether the container is for nodes, edges, records, or both nodes and edges (``all``).
     """
 
     properties: Mapping[str, ContainerProperty]
     is_global: bool
     last_updated_time: int
     created_time: int
-    used_for: Literal["node", "edge", "all"]
+    used_for: ContainerUsedFor
     constraints: Mapping[str, Constraint] = field(default_factory=dict)
     indexes: Mapping[str, Index] = field(default_factory=dict)
 
@@ -226,9 +229,23 @@ def as_write(self) -> ContainerApplyList:
 
 
 class _ContainerFilter(CogniteFilter):
-    def __init__(self, space: str | None = None, include_global: bool = False) -> None:
+    def __init__(
+        self,
+        space: str | None = None,
+        include_global: bool = False,
+        used_for: ContainerUsedFor | Sequence[ContainerUsedFor] | None = None,
+    ) -> None:
         self.space = space
         self.include_global = include_global
+        self.used_for: Sequence[ContainerUsedFor]
+        if used_for is None:
+            self.used_for = list(_ALL_CONTAINER_USED_FOR)
+        elif isinstance(used_for, str):
+            self.used_for = [used_for]
+        elif isinstance(used_for, Sequence):
+            self.used_for = cast("Sequence[ContainerUsedFor]", used_for)
+        else:
+            raise TypeError(f"Invalid value for 'used_for': {used_for!r}")
 
 
 @dataclass(frozen=True)

tests/tests_unit/test_api/test_data_modeling/test_containers.py

@@ -2,13 +2,18 @@
 
 import re
 from typing import TYPE_CHECKING, Any, cast
+from urllib.parse import parse_qs, urlparse
 
 import pytest
 from pytest_httpx import HTTPXMock
 
 from cognite.client import CogniteClient
 from cognite.client.data_classes.data_modeling import ContainerApply, ContainerId, ContainerPropertyApply, Text
-from cognite.client.data_classes.data_modeling.containers import BTreeIndexApply, RequiresConstraintApply
+from cognite.client.data_classes.data_modeling.containers import (
+    BTreeIndexApply,
+    RequiresConstraintApply,
+    _ContainerFilter,
+)
 from tests.utils import get_url
 
 if TYPE_CHECKING:
@@ -116,3 +121,50 @@ def test_apply_retrieve_and_delete_index(
 
         deleted_indexes = cognite_client.data_modeling.containers.delete_indexes([(new_container.as_id(), "index1")])
         assert deleted_indexes == [(new_container.as_id(), "index1")]
+
+    def test_list_request_includes_used_for(
+        self, httpx_mock: Any, cognite_client: CogniteClient, async_client: AsyncCogniteClient
+    ) -> None:
+        base = get_url(async_client.data_modeling.containers) + "/models/containers"
+        httpx_mock.add_response(url=re.compile("^" + re.escape(base)), json={"items": []})
+
+        cognite_client.data_modeling.containers.list(used_for=["node", "record"], limit=5)
+
+        req = httpx_mock.get_requests()[0]
+        raw_query = urlparse(str(req.url)).query
+        # Spec defines `usedFor` as a query array with no explicit style/explode, so OpenAPI's
+        # default (style=form, explode=true) applies: each value gets its own `usedFor=` pair.
+        assert "usedFor=node" in raw_query and "usedFor=record" in raw_query
+        assert "usedFor=node%2Crecord" not in raw_query and "usedFor=node,record" not in raw_query
+        qs = parse_qs(raw_query)
+        assert qs.get("usedFor") == ["node", "record"]
+
+    def test_list_request_used_for_single_value(
+        self, httpx_mock: Any, cognite_client: CogniteClient, async_client: AsyncCogniteClient
+    ) -> None:
+        base = get_url(async_client.data_modeling.containers) + "/models/containers"
+        httpx_mock.add_response(url=re.compile("^" + re.escape(base)), json={"items": []})
+
+        cognite_client.data_modeling.containers.list(used_for="record")
+
+        req = httpx_mock.get_requests()[0]
+        qs = parse_qs(urlparse(str(req.url)).query)
+        assert qs.get("usedFor") == ["record"]
+
+    def test_list_request_default_returns_all_kinds(
+        self, httpx_mock: Any, cognite_client: CogniteClient, async_client: AsyncCogniteClient
+    ) -> None:
+        # When the caller does not specify `used_for`, the SDK should request all container kinds,
+        # including records, rather than rely on the server's default of `all` which excludes records.
+        base = get_url(async_client.data_modeling.containers) + "/models/containers"
+        httpx_mock.add_response(url=re.compile("^" + re.escape(base)), json={"items": []})
+
+        cognite_client.data_modeling.containers.list()
+
+        req = httpx_mock.get_requests()[0]
+        qs = parse_qs(urlparse(str(req.url)).query)
+        assert qs.get("usedFor") == ["node", "edge", "record", "all"]
+
+    def test_container_filter_rejects_invalid_used_for(self) -> None:
+        with pytest.raises(TypeError, match="Invalid value for 'used_for'"):
+            _ContainerFilter(used_for=123)  # type: ignore[arg-type]