From 1c934d0cba7a7b9d4e931ad4fe38c9c9b22a36a4 Mon Sep 17 00:00:00 2001 From: Everton Colling Date: Tue, 12 May 2026 15:03:39 +0200 Subject: [PATCH 1/3] add support for record containers --- .../client/_api/data_modeling/containers.py | 16 ++++++++-- .../_sync_api/data_modeling/containers.py | 32 ++++++++++++++++--- .../data_classes/data_modeling/__init__.py | 2 ++ .../data_classes/data_modeling/containers.py | 25 +++++++++++---- .../test_data_modeling/test_containers.py | 30 +++++++++++++++++ 5 files changed, 92 insertions(+), 13 deletions(-) diff --git a/cognite/client/_api/data_modeling/containers.py b/cognite/client/_api/data_modeling/containers.py index 7acf419d9a..c5bca352f7 100644 --- a/cognite/client/_api/data_modeling/containers.py +++ b/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 (e.g. ``"record"`` for record containers). If omitted, the API default applies (typically node, edge, and all — not records). 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,6 +231,7 @@ 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 `_. @@ -233,6 +239,7 @@ async 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 (e.g. ``"record"`` for record containers). If omitted, the API default applies (typically node, edge, and all — not records). Returns: ContainerList: List of requested containers @@ -246,6 +253,11 @@ async def list( >>> # async_client = AsyncCogniteClient() # another option >>> container_list = client.data_modeling.containers.list(limit=5) + Filter containers by `used_for`: + + >>> record_containers = client.data_modeling.containers.list(used_for="record") + >>> all_containers = client.data_modeling.containers.list(used_for=["all", "record"]) + Iterate over containers, one-by-one: >>> for container in client.data_modeling.containers(): @@ -256,7 +268,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, diff --git a/cognite/client/_sync_api/data_modeling/containers.py b/cognite/client/_sync_api/data_modeling/containers.py index ed7817391e..694ee62a8d 100644 --- a/cognite/client/_sync_api/data_modeling/containers.py +++ b/cognite/client/_sync_api/data_modeling/containers.py @@ -1,6 +1,6 @@ """ =============================================================================== -465575f7283dac5a1831298bc22b3bda +523c9372ce1b89556aebab31a1969e3e 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,12 +39,22 @@ 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__( @@ -51,6 +62,7 @@ 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, ) -> Iterator[Container] | Iterator[ContainerList]: """ @@ -62,6 +74,7 @@ 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 (e.g. ``"record"`` for record containers). If omitted, the API default applies (typically node, edge, and all — not records). limit (int | None): Maximum number of containers to return. Defaults to returning all items. Yields: @@ -69,7 +82,7 @@ def __call__( """ # 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 `_. @@ -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 (e.g. ``"record"`` for record containers). If omitted, the API default applies (typically node, edge, and all — not records). Returns: ContainerList: List of requested containers @@ -192,6 +207,11 @@ def list( >>> # async_client = AsyncCogniteClient() # another option >>> container_list = client.data_modeling.containers.list(limit=5) + Filter containers by `used_for`: + + >>> record_containers = client.data_modeling.containers.list(used_for="record") + >>> all_containers = client.data_modeling.containers.list(used_for=["all", "record"]) + Iterate over containers, one-by-one: >>> for container in client.data_modeling.containers(): @@ -203,7 +223,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 diff --git a/cognite/client/data_classes/data_modeling/__init__.py b/cognite/client/data_classes/data_modeling/__init__.py index 0fab81400f..c773aaa63e 100644 --- a/cognite/client/data_classes/data_modeling/__init__.py +++ b/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", diff --git a/cognite/client/data_classes/data_modeling/containers.py b/cognite/client/data_classes/data_modeling/containers.py index 4fe916e538..cbefd4feb1 100644 --- a/cognite/client/data_classes/data_modeling/containers.py +++ b/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,8 @@ _T_Constraint = TypeVar("_T_Constraint", bound="ConstraintCore") _T_Index = TypeVar("_T_Index", bound="IndexCore") +ContainerUsedFor = Literal["node", "edge", "record", "all"] + @dataclass class ContainerCore(DataModelingSchemaResource["ContainerApply"], ABC): @@ -107,13 +109,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 +151,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 +228,20 @@ 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 + if used_for is None: + self.used_for: list[ContainerUsedFor] | None = None + elif isinstance(used_for, str): + self.used_for = [used_for] + else: + self.used_for = list(used_for) @dataclass(frozen=True) diff --git a/tests/tests_unit/test_api/test_data_modeling/test_containers.py b/tests/tests_unit/test_api/test_data_modeling/test_containers.py index dead7d0e1a..658e5229eb 100644 --- a/tests/tests_unit/test_api/test_data_modeling/test_containers.py +++ b/tests/tests_unit/test_api/test_data_modeling/test_containers.py @@ -2,6 +2,7 @@ import re from typing import TYPE_CHECKING, Any, cast +from urllib.parse import parse_qs, urlparse import pytest from pytest_httpx import HTTPXMock @@ -116,3 +117,32 @@ 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"] From 0ef2e15a3424e244e094c354b52ef4e6a4058fa1 Mon Sep 17 00:00:00 2001 From: Everton Colling Date: Tue, 12 May 2026 16:58:52 +0200 Subject: [PATCH 2/3] changes based on review --- .../client/_api/data_modeling/containers.py | 4 ++-- .../_sync_api/data_modeling/containers.py | 6 ++--- .../data_classes/data_modeling/containers.py | 8 +++++-- .../test_data_modeling/test_containers.py | 24 ++++++++++++++++++- 4 files changed, 34 insertions(+), 8 deletions(-) diff --git a/cognite/client/_api/data_modeling/containers.py b/cognite/client/_api/data_modeling/containers.py index c5bca352f7..c616f428a8 100644 --- a/cognite/client/_api/data_modeling/containers.py +++ b/cognite/client/_api/data_modeling/containers.py @@ -79,7 +79,7 @@ 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 (e.g. ``"record"`` for record containers). If omitted, the API default applies (typically node, edge, and all — not records). + 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: @@ -239,7 +239,7 @@ async 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 (e.g. ``"record"`` for record containers). If omitted, the API default applies (typically node, edge, and all — not records). + 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 diff --git a/cognite/client/_sync_api/data_modeling/containers.py b/cognite/client/_sync_api/data_modeling/containers.py index 694ee62a8d..1a53fe5804 100644 --- a/cognite/client/_sync_api/data_modeling/containers.py +++ b/cognite/client/_sync_api/data_modeling/containers.py @@ -1,6 +1,6 @@ """ =============================================================================== -523c9372ce1b89556aebab31a1969e3e +945df424d0b546e88f29acb8e105cdc0 This file is auto-generated from the Async API modules, - do not edit manually! =============================================================================== """ @@ -74,7 +74,7 @@ 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 (e.g. ``"record"`` for record containers). If omitted, the API default applies (typically node, edge, and all — not records). + 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: @@ -193,7 +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 (e.g. ``"record"`` for record containers). If omitted, the API default applies (typically node, edge, and all — not records). + 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 diff --git a/cognite/client/data_classes/data_modeling/containers.py b/cognite/client/data_classes/data_modeling/containers.py index cbefd4feb1..984fe0f7e8 100644 --- a/cognite/client/data_classes/data_modeling/containers.py +++ b/cognite/client/data_classes/data_modeling/containers.py @@ -31,6 +31,7 @@ _T_Index = TypeVar("_T_Index", bound="IndexCore") ContainerUsedFor = Literal["node", "edge", "record", "all"] +_ALL_CONTAINER_USED_FOR: tuple[ContainerUsedFor, ...] = ("all", "record") @dataclass @@ -236,12 +237,15 @@ def __init__( ) -> None: self.space = space self.include_global = include_global + self.used_for: Sequence[ContainerUsedFor] if used_for is None: - self.used_for: list[ContainerUsedFor] | None = 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: - self.used_for = list(used_for) + raise TypeError(f"Invalid value for 'used_for': {used_for!r}") @dataclass(frozen=True) diff --git a/tests/tests_unit/test_api/test_data_modeling/test_containers.py b/tests/tests_unit/test_api/test_data_modeling/test_containers.py index 658e5229eb..bd35a78e13 100644 --- a/tests/tests_unit/test_api/test_data_modeling/test_containers.py +++ b/tests/tests_unit/test_api/test_data_modeling/test_containers.py @@ -9,7 +9,11 @@ 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: @@ -146,3 +150,21 @@ def test_list_request_used_for_single_value( 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") == ["all", "record"] + + 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] From bfe0ca8db0e7e1e85282940745ed9a3abf57f3eb Mon Sep 17 00:00:00 2001 From: Everton Colling Date: Tue, 12 May 2026 21:01:48 +0200 Subject: [PATCH 3/3] return all variants --- cognite/client/data_classes/data_modeling/containers.py | 2 +- tests/tests_unit/test_api/test_data_modeling/test_containers.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/cognite/client/data_classes/data_modeling/containers.py b/cognite/client/data_classes/data_modeling/containers.py index 984fe0f7e8..18a8570839 100644 --- a/cognite/client/data_classes/data_modeling/containers.py +++ b/cognite/client/data_classes/data_modeling/containers.py @@ -31,7 +31,7 @@ _T_Index = TypeVar("_T_Index", bound="IndexCore") ContainerUsedFor = Literal["node", "edge", "record", "all"] -_ALL_CONTAINER_USED_FOR: tuple[ContainerUsedFor, ...] = ("all", "record") +_ALL_CONTAINER_USED_FOR: tuple[ContainerUsedFor, ...] = ("node", "edge", "record", "all") @dataclass diff --git a/tests/tests_unit/test_api/test_data_modeling/test_containers.py b/tests/tests_unit/test_api/test_data_modeling/test_containers.py index bd35a78e13..ec805602aa 100644 --- a/tests/tests_unit/test_api/test_data_modeling/test_containers.py +++ b/tests/tests_unit/test_api/test_data_modeling/test_containers.py @@ -163,7 +163,7 @@ def test_list_request_default_returns_all_kinds( req = httpx_mock.get_requests()[0] qs = parse_qs(urlparse(str(req.url)).query) - assert qs.get("usedFor") == ["all", "record"] + 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'"):