add /models/statistics

Author: haakonvt

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)

cognite/client/_api/data_modeling/statistics.py

@@ -0,0 +1,43 @@
+from __future__ import annotations
+
+import itertools
+from collections.abc import Sequence
+from typing import TYPE_CHECKING
+
+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, ProjectStatsAndLimits
+
+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:
+        """Usage data and limits for a project's data modelling usage, including data model schemas and graph instances"""
+        return ProjectStatsAndLimits._load(
+            self._get(self._RESOURCE_PATH).json(), project=self._cognite_client._config.project
+        )
+
+    def per_space(self, space: str | Sequence[str] | None = None, return_all: bool = False) -> InstanceStatsList:
+        """Statistics and limits for all spaces in the project, or for a select subset"""
+        if return_all:
+            return InstanceStatsList._load(self._get(self._RESOURCE_PATH + "/spaces").json()["items"])
+
+        elif space is not None:
+            # Statistics and limits for spaces by their space identifiers
+            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)
+                )
+            )
+        raise ValueError("Either 'space' or 'return_all' must be specified")

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)