diff --git a/cognite/client/_api/data_modeling/containers.py b/cognite/client/_api/data_modeling/containers.py index 7acf419d9a..c616f428a8 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. 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,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. If omitted, containers of every kind (nodes, edges, and records) are returned. 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..1a53fe5804 100644 --- a/cognite/client/_sync_api/data_modeling/containers.py +++ b/cognite/client/_sync_api/data_modeling/containers.py @@ -1,6 +1,6 @@ """ =============================================================================== -465575f7283dac5a1831298bc22b3bda +945df424d0b546e88f29acb8e105cdc0 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. 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: @@ -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. If omitted, containers of every kind (nodes, edges, and records) are returned. 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..18a8570839 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,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) 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..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 @@ -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]