Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add /models/statistics #2095

Open
wants to merge 4 commits into
base: master
Choose a base branch
from
Open
Show file tree
Hide file tree
Changes from 2 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 2 additions & 0 deletions cognite/client/_api/data_modeling/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand All @@ -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)
43 changes: 43 additions & 0 deletions cognite/client/_api/data_modeling/statistics.py
Original file line number Diff line number Diff line change
@@ -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")
2 changes: 1 addition & 1 deletion cognite/client/data_classes/data_modeling/ids.py
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand Down
138 changes: 138 additions & 0 deletions cognite/client/data_classes/data_modeling/statistics.py
Original file line number Diff line number Diff line change
@@ -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)
9 changes: 6 additions & 3 deletions cognite/client/utils/_pandas_helpers.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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 = {
Expand Down Expand Up @@ -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:
Expand Down
Loading