diff --git a/cognite/client/_api/data_modeling/__init__.py b/cognite/client/_api/data_modeling/__init__.py index 4ace03d4f..2c0ed3782 100644 --- a/cognite/client/_api/data_modeling/__init__.py +++ b/cognite/client/_api/data_modeling/__init__.py @@ -7,6 +7,7 @@ from cognite.client._api.data_modeling.graphql import DataModelingGraphQLAPI from cognite.client._api.data_modeling.instances import InstancesAPI from cognite.client._api.data_modeling.spaces import SpacesAPI +from cognite.client._api.data_modeling.statistics import StatisticsAPI from cognite.client._api.data_modeling.views import ViewsAPI from cognite.client._api_client import APIClient @@ -24,3 +25,4 @@ def __init__(self, config: ClientConfig, api_version: str | None, cognite_client self.views = ViewsAPI(config, api_version, cognite_client) self.instances = InstancesAPI(config, api_version, cognite_client) self.graphql = DataModelingGraphQLAPI(config, api_version, cognite_client) + self.statistics = StatisticsAPI(config, api_version, cognite_client) diff --git a/cognite/client/_api/data_modeling/statistics.py b/cognite/client/_api/data_modeling/statistics.py new file mode 100644 index 000000000..45d74728d --- /dev/null +++ b/cognite/client/_api/data_modeling/statistics.py @@ -0,0 +1,101 @@ +from __future__ import annotations + +import itertools +from typing import TYPE_CHECKING, Any, Literal, overload + +from cognite.client._api_client import APIClient +from cognite.client.data_classes.data_modeling.ids import _load_space_identifier +from cognite.client.data_classes.data_modeling.statistics import ( + InstanceStatsList, + InstanceStatsPerSpace, + ProjectStatsAndLimits, +) +from cognite.client.utils.useful_types import SequenceNotStr + +if TYPE_CHECKING: + from cognite.client._cognite_client import CogniteClient + from cognite.client.config import ClientConfig + + +class StatisticsAPI(APIClient): + _RESOURCE_PATH = "/models/statistics" + + def __init__(self, config: ClientConfig, api_version: str | None, cognite_client: CogniteClient) -> None: + super().__init__(config, api_version, cognite_client) + self._RETRIEVE_LIMIT = 100 # may need to be renamed, but fine for now + + def project(self) -> ProjectStatsAndLimits: + """`Retrieve project-wide usage data and limits `_ + + Returns the usage data and limits for a project's data modelling usage, including data model schemas and graph instances + + Returns: + ProjectStatsAndLimits: The requested statistics and limits + + Examples: + + Fetch project statistics (and limits) and check the current number of data models vs. + and how many more can be created: + + >>> from cognite.client import CogniteClient + >>> client = CogniteClient() + >>> stats = client.data_modeling.statistics.project() + >>> num_dm = stats.data_models.current + >>> num_dm_left = stats.data_models.limit - num_dm + """ + return ProjectStatsAndLimits._load( + self._get(self._RESOURCE_PATH).json(), project=self._cognite_client._config.project + ) + + @overload + def per_space(self, space: str, return_all: Literal[False]) -> InstanceStatsPerSpace: ... + + @overload + def per_space(self, space: Any, return_all: Literal[True]) -> InstanceStatsList: ... + + @overload + def per_space(self, space: SequenceNotStr[str], return_all: bool) -> InstanceStatsList: ... + + def per_space( + self, space: str | SequenceNotStr[str] | None = None, return_all: bool = False + ) -> InstanceStatsPerSpace | InstanceStatsList: + """`Retrieve usage data and limits per space `_ + + See also: `Retrieve statistics and limits for all spaces `_ + + Args: + space (str | SequenceNotStr[str] | None): The space or spaces to retrieve statistics for. + return_all (bool): If True, fetch statistics for all spaces. If False, fetch statistics for the specified space(s). + + Returns: + InstanceStatsPerSpace | InstanceStatsList: InstanceStatsPerSpace if a single space is given, else InstanceStatsList (which is a list of InstanceStatsPerSpace) + + Examples: + + Fetch statistics for a single space: + + >>> from cognite.client import CogniteClient + >>> client = CogniteClient() + >>> res = client.data_modeling.statistics.per_space("my-space") + + Fetch statistics for multiple spaces: + >>> res = client.data_modeling.statistics.per_space( + ... ["my-space1", "my-space2"] + ... ) + + Fetch statistics for all spaces (ignores the 'space' argument): + >>> res = client.data_modeling.statistics.per_space(return_all=True) + """ + if return_all: + return InstanceStatsList._load(self._get(self._RESOURCE_PATH + "/spaces").json()["items"]) + + elif space is None: + raise ValueError("Either 'space' or 'return_all' must be specified") + + ids = _load_space_identifier(space) + return InstanceStatsList._load( + itertools.chain.from_iterable( + self._post(self._RESOURCE_PATH + "/spaces/byids", json={"items": chunk.as_dicts()}).json()["items"] + for chunk in ids.chunked(self._RETRIEVE_LIMIT) + ) + ) diff --git a/cognite/client/data_classes/data_modeling/ids.py b/cognite/client/data_classes/data_modeling/ids.py index 5eda8c625..0d8250fa9 100644 --- a/cognite/client/data_classes/data_modeling/ids.py +++ b/cognite/client/data_classes/data_modeling/ids.py @@ -190,7 +190,7 @@ def version(self) -> str | None: ... Id = tuple[str, str] | tuple[str, str, str] | IdLike | VersionedIdLike -def _load_space_identifier(ids: str | SequenceNotStr[str]) -> DataModelingIdentifierSequence: +def _load_space_identifier(ids: str | Sequence[str] | SequenceNotStr[str]) -> DataModelingIdentifierSequence: is_sequence = isinstance(ids, Sequence) and not isinstance(ids, str) spaces = [ids] if isinstance(ids, str) else ids return DataModelingIdentifierSequence( diff --git a/cognite/client/data_classes/data_modeling/statistics.py b/cognite/client/data_classes/data_modeling/statistics.py new file mode 100644 index 000000000..efc066693 --- /dev/null +++ b/cognite/client/data_classes/data_modeling/statistics.py @@ -0,0 +1,138 @@ +from __future__ import annotations + +from collections import UserList +from collections.abc import Iterable +from dataclasses import asdict, dataclass +from typing import TYPE_CHECKING, Any + +from typing_extensions import Self + +from cognite.client.utils._importing import local_import +from cognite.client.utils._pandas_helpers import notebook_display_with_fallback + +if TYPE_CHECKING: + import pandas as pd + + +@dataclass +class InstanceStatsPerSpace: + space: str + nodes: int + edges: int + soft_deleted_nodes: int + soft_deleted_edges: int + + @classmethod + def _load(cls, data: dict[str, Any]) -> Self: + return cls( + space=data["space"], + nodes=data["nodes"], + edges=data["edges"], + soft_deleted_nodes=data["softDeletedNodes"], + soft_deleted_edges=data["softDeletedEdges"], + ) + + def _repr_html_(self) -> str: + return notebook_display_with_fallback(self) + + def to_pandas(self) -> pd.DataFrame: + pd = local_import("pandas") + space = (dumped := asdict(self)).pop("space") + return pd.Series(dumped).to_frame(name=space) + + +class InstanceStatsList(UserList): + def __init__(self, items: list[InstanceStatsPerSpace]): + super().__init__(items) + + @classmethod + def _load(cls, data: Iterable[dict[str, Any]]) -> Self: + return cls([InstanceStatsPerSpace._load(item) for item in data]) + + def _repr_html_(self) -> str: + return notebook_display_with_fallback(self) + + def to_pandas(self) -> pd.DataFrame: + pd = local_import("pandas") + df = pd.DataFrame([asdict(item) for item in self]).set_index("space") + order_by_total = (df["nodes"] + df["edges"]).sort_values(ascending=False).index + return df.loc[order_by_total] + + +@dataclass +class CountLimit: + count: int + limit: int + + @classmethod + def _load(cls, data: dict[str, Any]) -> Self: + return cls(count=data["count"], limit=data["limit"]) + + +@dataclass +class InstanceStatsAndLimits: + nodes: int + edges: int + instances: int + instances_limit: int + soft_deleted_nodes: int + soft_deleted_edges: int + soft_deleted_instances: int + soft_deleted_instances_limit: int + + @classmethod + def _load(cls, data: dict[str, Any]) -> Self: + return cls( + nodes=data["nodes"], + edges=data["edges"], + instances=data["instances"], + instances_limit=data["instancesLimit"], + soft_deleted_nodes=data["softDeletedNodes"], + soft_deleted_edges=data["softDeletedEdges"], + soft_deleted_instances=data["softDeletedInstances"], + soft_deleted_instances_limit=data["softDeletedInstancesLimit"], + ) + + def _repr_html_(self) -> str: + return notebook_display_with_fallback(self) + + def to_pandas(self) -> pd.DataFrame: + pd = local_import("pandas") + return pd.Series(asdict(self)).to_frame() + + +@dataclass +class ProjectStatsAndLimits: + project: str + spaces: CountLimit + containers: CountLimit + views: CountLimit + data_models: CountLimit + container_properties: CountLimit + instances: InstanceStatsAndLimits + concurrent_read_limit: int + concurrent_write_limit: int + concurrent_delete_limit: int + + @classmethod + def _load(cls, data: dict[str, Any], project: str) -> Self: + return cls( + project=project, + spaces=CountLimit._load(data["spaces"]), + containers=CountLimit._load(data["containers"]), + views=CountLimit._load(data["views"]), + data_models=CountLimit._load(data["dataModels"]), + container_properties=CountLimit._load(data["containerProperties"]), + instances=InstanceStatsAndLimits._load(data["instances"]), + concurrent_read_limit=data["concurrentReadLimit"], + concurrent_write_limit=data["concurrentWriteLimit"], + concurrent_delete_limit=data["concurrentDeleteLimit"], + ) + + def _repr_html_(self) -> str: + return notebook_display_with_fallback(self) + + def to_pandas(self) -> pd.DataFrame: + pd = local_import("pandas") + project = (dumped := asdict(self)).pop("project") + return pd.Series(dumped).to_frame(name=project) diff --git a/cognite/client/testing.py b/cognite/client/testing.py index 2a56ea0e5..3ac930664 100644 --- a/cognite/client/testing.py +++ b/cognite/client/testing.py @@ -17,6 +17,7 @@ from cognite.client._api.data_modeling.graphql import DataModelingGraphQLAPI from cognite.client._api.data_modeling.instances import InstancesAPI from cognite.client._api.data_modeling.spaces import SpacesAPI +from cognite.client._api.data_modeling.statistics import StatisticsAPI from cognite.client._api.data_modeling.views import ViewsAPI from cognite.client._api.data_sets import DataSetsAPI from cognite.client._api.datapoints import DatapointsAPI @@ -114,6 +115,7 @@ def __init__(self, *args: Any, **kwargs: Any) -> None: self.data_modeling.views = MagicMock(spec_set=ViewsAPI) self.data_modeling.instances = MagicMock(spec_set=InstancesAPI) self.data_modeling.graphql = MagicMock(spec_set=DataModelingGraphQLAPI) + self.data_modeling.statistics = MagicMock(spec_set=StatisticsAPI) self.data_sets = MagicMock(spec_set=DataSetsAPI) diff --git a/cognite/client/utils/_pandas_helpers.py b/cognite/client/utils/_pandas_helpers.py index c0cba2e3c..ccad5bff3 100644 --- a/cognite/client/utils/_pandas_helpers.py +++ b/cognite/client/utils/_pandas_helpers.py @@ -7,7 +7,7 @@ from inspect import signature from itertools import chain from numbers import Integral -from typing import TYPE_CHECKING, Any, Literal +from typing import TYPE_CHECKING, Any, Literal, Protocol from cognite.client.exceptions import CogniteImportError from cognite.client.utils._importing import local_import @@ -18,7 +18,6 @@ import pandas as pd from cognite.client.data_classes import Datapoints, DatapointsArray, DatapointsArrayList, DatapointsList - from cognite.client.data_classes._base import T_CogniteResource, T_CogniteResourceList NULLABLE_INT_COLS = { @@ -88,7 +87,11 @@ def concat_dps_dataframe_list( return concat_dataframes_with_nullable_int_cols(dfs) -def notebook_display_with_fallback(inst: T_CogniteResource | T_CogniteResourceList, **kwargs: Any) -> str: +class PandasConvertible(Protocol): + def to_pandas(self) -> pd.DataFrame: ... + + +def notebook_display_with_fallback(inst: PandasConvertible, **kwargs: Any) -> str: params = signature(inst.to_pandas).parameters # Default of False enforced (when accepted by method): if "camel_case" in params: diff --git a/tests/tests_unit/test_docstring_examples.py b/tests/tests_unit/test_docstring_examples.py index 3890aa526..04aaf7c81 100644 --- a/tests/tests_unit/test_docstring_examples.py +++ b/tests/tests_unit/test_docstring_examples.py @@ -25,7 +25,7 @@ units, workflows, ) -from cognite.client._api.data_modeling import containers, data_models, graphql, instances, spaces, views +from cognite.client._api.data_modeling import containers, data_models, graphql, instances, spaces, statistics, views from cognite.client._api.hosted_extractors import destinations, jobs, mappings, sources from cognite.client._api.postgres_gateway import tables as postgres_gateway_tables from cognite.client._api.postgres_gateway import users as postgres_gateway_users @@ -114,6 +114,7 @@ def test_data_modeling(self): run_docstring_tests(data_models) run_docstring_tests(spaces) run_docstring_tests(graphql) + run_docstring_tests(statistics) def test_datapoint_subscriptions(self): run_docstring_tests(datapoints_subscriptions)