cognite-neat 0.98.0__py3-none-any.whl → 0.99.0__py3-none-any.whl

cognite/neat/_utils/cdf/loaders/_data_modeling.py

@@ -1,339 +0,0 @@
-from collections.abc import Callable, Sequence
-from graphlib import TopologicalSorter
-from typing import Any, Literal, cast
-
-from cognite.client import CogniteClient
-from cognite.client.data_classes import filters
-from cognite.client.data_classes._base import (
-    T_CogniteResourceList,
-    T_WritableCogniteResource,
-    T_WriteClass,
-)
-from cognite.client.data_classes.data_modeling import (
-    Container,
-    ContainerApply,
-    ContainerApplyList,
-    ContainerList,
-    DataModel,
-    DataModelApply,
-    DataModelApplyList,
-    DataModelList,
-    RequiresConstraint,
-    Space,
-    SpaceApply,
-    SpaceApplyList,
-    SpaceList,
-    View,
-    ViewApply,
-    ViewApplyList,
-    ViewList,
-)
-from cognite.client.data_classes.data_modeling.ids import (
-    ContainerId,
-    DataModelId,
-    NodeId,
-    ViewId,
-)
-from cognite.client.exceptions import CogniteAPIError
-from cognite.client.utils.useful_types import SequenceNotStr
-
-from ._base import T_ID, ResourceLoader, T_WritableCogniteResourceList
-
-
-class DataModelingLoader(
-    ResourceLoader[T_ID, T_WriteClass, T_WritableCogniteResource, T_CogniteResourceList, T_WritableCogniteResourceList]
-):
-    @classmethod
-    def in_space(cls, item: T_WriteClass | T_WritableCogniteResource | T_ID, space: set[str]) -> bool:
-        if hasattr(item, "space"):
-            return item.space in space
-        raise ValueError(f"Item {item} does not have a space attribute")
-
-    def sort_by_dependencies(self, items: list[T_WriteClass]) -> list[T_WriteClass]:
-        return items
-
-    def _create_force(
-        self,
-        items: Sequence[T_WriteClass],
-        tried_force_deploy: set[T_ID],
-        create_method: Callable[[Sequence[T_WriteClass]], T_WritableCogniteResourceList],
-    ) -> T_WritableCogniteResourceList:
-        try:
-            return create_method(items)
-        except CogniteAPIError as e:
-            failed_items = {failed.as_id() for failed in e.failed if hasattr(failed, "as_id")}
-            to_redeploy = [
-                item
-                for item in items
-                if item.as_id() in failed_items and item.as_id() not in tried_force_deploy  # type: ignore[attr-defined]
-            ]
-            if not to_redeploy:
-                # Avoid infinite loop
-                raise e
-            ids = [item.as_id() for item in to_redeploy]  # type: ignore[attr-defined]
-            tried_force_deploy.update(ids)
-            self.delete(ids)
-            return self._create_force(to_redeploy, tried_force_deploy, create_method)
-
-
-class SpaceLoader(DataModelingLoader[str, SpaceApply, Space, SpaceApplyList, SpaceList]):
-    resource_name = "spaces"
-
-    @classmethod
-    def get_id(cls, item: Space | SpaceApply | str | dict) -> str:
-        if isinstance(item, Space | SpaceApply):
-            return item.space
-        if isinstance(item, dict):
-            return item["space"]
-        return item
-
-    def create(self, items: Sequence[SpaceApply]) -> SpaceList:
-        return self.client.data_modeling.spaces.apply(items)
-
-    def retrieve(self, ids: SequenceNotStr[str]) -> SpaceList:
-        return self.client.data_modeling.spaces.retrieve(ids)
-
-    def update(self, items: Sequence[SpaceApply]) -> SpaceList:
-        return self.create(items)
-
-    def delete(self, ids: SequenceNotStr[str] | Sequence[Space | SpaceApply]) -> list[str]:
-        if all(isinstance(item, Space) for item in ids) or all(isinstance(item, SpaceApply) for item in ids):
-            ids = [cast(Space | SpaceApply, item).space for item in ids]
-        return self.client.data_modeling.spaces.delete(cast(SequenceNotStr[str], ids))
-
-    def clean(self, space: str) -> None:
-        """Deletes all data in a space.
-
-        This means all nodes, edges, views, containers, and data models located in the given space.
-
-        Args:
-            client: Connected CogniteClient
-            space: The space to delete.
-
-        """
-        edges = self.client.data_modeling.instances.list(
-            "edge", limit=-1, filter=filters.Equals(["edge", "space"], space)
-        )
-        if edges:
-            instances = self.client.data_modeling.instances.delete(edges=edges.as_ids())
-            print(f"Deleted {len(instances.edges)} edges")
-        nodes = self.client.data_modeling.instances.list(
-            "node", limit=-1, filter=filters.Equals(["node", "space"], space)
-        )
-        node_types = {NodeId(node.type.space, node.type.external_id) for node in nodes if node.type}
-        node_data = set(nodes.as_ids()) - node_types
-        if node_data:
-            instances = self.client.data_modeling.instances.delete(nodes=list(node_data))
-            print(f"Deleted {len(instances.nodes)} nodes")
-        if node_types:
-            instances = self.client.data_modeling.instances.delete(nodes=list(node_types))
-            print(f"Deleted {len(instances.nodes)} node types")
-        views = self.client.data_modeling.views.list(limit=-1, space=space)
-        if views:
-            deleted_views = self.client.data_modeling.views.delete(views.as_ids())
-            print(f"Deleted {len(deleted_views)} views")
-        containers = self.client.data_modeling.containers.list(limit=-1, space=space)
-        if containers:
-            deleted_containers = self.client.data_modeling.containers.delete(containers.as_ids())
-            print(f"Deleted {len(deleted_containers)} containers")
-        if data_models := self.client.data_modeling.data_models.list(limit=-1, space=space):
-            deleted_data_models = self.client.data_modeling.data_models.delete(data_models.as_ids())
-            print(f"Deleted {len(deleted_data_models)} data models")
-        deleted_space = self.client.data_modeling.spaces.delete(space)
-        print(f"Deleted space {deleted_space}")
-
-
-class ViewLoader(DataModelingLoader[ViewId, ViewApply, View, ViewApplyList, ViewList]):
-    resource_name = "views"
-
-    def __init__(self, client: CogniteClient, existing_handling: Literal["fail", "skip", "update", "force"] = "fail"):
-        super().__init__(client)
-        self.existing_handling = existing_handling
-        self._cache_view_by_id: dict[ViewId, View] = {}
-        self._tried_force_deploy: set[ViewId] = set()
-
-    @classmethod
-    def get_id(cls, item: View | ViewApply | ViewId | dict) -> ViewId:
-        if isinstance(item, View | ViewApply):
-            return item.as_id()
-        if isinstance(item, dict):
-            return ViewId.load(item)
-        return item
-
-    def create(self, items: Sequence[ViewApply]) -> ViewList:
-        if self.existing_handling == "force":
-            return self._create_force(items, self._tried_force_deploy, self.client.data_modeling.views.apply)
-        else:
-            return self.client.data_modeling.views.apply(items)
-
-    def retrieve(self, ids: SequenceNotStr[ViewId]) -> ViewList:
-        return self.client.data_modeling.views.retrieve(cast(Sequence, ids))
-
-    def update(self, items: Sequence[ViewApply]) -> ViewList:
-        return self.create(items)
-
-    def delete(self, ids: SequenceNotStr[ViewId]) -> list[ViewId]:
-        return self.client.data_modeling.views.delete(cast(Sequence, ids))
-
-    def _as_write_raw(self, view: View) -> dict[str, Any]:
-        dumped = view.as_write().dump()
-        if view.properties:
-            # All read version of views have all the properties of their parent views.
-            # We need to remove these properties to compare with the local view.
-            parents = self._retrieve_view_ancestors(view.implements or [], self._cache_view_by_id)
-            for parent in parents:
-                for prop_name in parent.properties.keys():
-                    dumped["properties"].pop(prop_name, None)
-
-        if "properties" in dumped and not dumped["properties"]:
-            # All properties were removed, so we remove the properties key.
-            dumped.pop("properties", None)
-        return dumped
-
-    def are_equal(self, local: ViewApply, remote: View) -> bool:
-        local_dumped = local.dump()
-        if not remote.implements:
-            return local_dumped == remote.as_write().dump()
-
-        cdf_resource_dumped = self._as_write_raw(remote)
-
-        if "properties" in local_dumped and not local_dumped["properties"]:
-            # In case the local properties are set to an empty dict.
-            local_dumped.pop("properties", None)
-
-        return local_dumped == cdf_resource_dumped
-
-    def as_write(self, view: View) -> ViewApply:
-        return ViewApply.load(self._as_write_raw(view))
-
-    def retrieve_all_parents(self, views: list[ViewId]) -> list[View]:
-        return self._retrieve_view_ancestors(views, self._cache_view_by_id)
-
-    def _retrieve_view_ancestors(self, parents: list[ViewId], cache: dict[ViewId, View]) -> list[View]:
-        """Retrieves all ancestors of a view.
-
-        This will mutate the cache passed in, and return a list of views that are the ancestors
-        of the views in the parents list.
-
-        Args:
-            parents: The parents of the view to retrieve all ancestors for
-            cache: The cache to store the views in
-        """
-        parent_ids = parents.copy()
-        found: list[View] = []
-        found_ids: set[ViewId] = set()
-        while parent_ids:
-            to_lookup: set[ViewId] = set()
-            grand_parent_ids = []
-            for parent in parent_ids:
-                if parent in found_ids:
-                    continue
-                elif parent in cache:
-                    found.append(cache[parent])
-                    grand_parent_ids.extend(cache[parent].implements or [])
-                else:
-                    to_lookup.add(parent)
-
-            if to_lookup:
-                looked_up = self.client.data_modeling.views.retrieve(list(to_lookup))
-                cache.update({view.as_id(): view for view in looked_up})
-                found.extend(looked_up)
-                found_ids.update({view.as_id() for view in looked_up})
-                for view in looked_up:
-                    grand_parent_ids.extend(view.implements or [])
-
-            parent_ids = grand_parent_ids
-        return found
-
-
-class ContainerLoader(DataModelingLoader[ContainerId, ContainerApply, Container, ContainerApplyList, ContainerList]):
-    resource_name = "containers"
-
-    def __init__(self, client: CogniteClient, existing_handling: Literal["fail", "skip", "update", "force"] = "fail"):
-        super().__init__(client)
-        self.existing_handling = existing_handling
-        self._tried_force_deploy: set[ContainerId] = set()
-
-    @classmethod
-    def get_id(cls, item: Container | ContainerApply | ContainerId | dict) -> ContainerId:
-        if isinstance(item, Container | ContainerApply):
-            return item.as_id()
-        if isinstance(item, dict):
-            return ContainerId.load(item)
-        return item
-
-    def sort_by_dependencies(self, items: Sequence[ContainerApply]) -> list[ContainerApply]:
-        container_by_id = {container.as_id(): container for container in items}
-        container_dependencies = {
-            container.as_id(): {
-                const.require
-                for const in container.constraints.values()
-                if isinstance(const, RequiresConstraint) and const.require in container_by_id
-            }
-            for container in items
-        }
-        return [
-            container_by_id[container_id] for container_id in TopologicalSorter(container_dependencies).static_order()
-        ]
-
-    def create(self, items: Sequence[ContainerApply]) -> ContainerList:
-        if self.existing_handling == "force":
-            return self._create_force(items, self._tried_force_deploy, self.client.data_modeling.containers.apply)
-        else:
-            return self.client.data_modeling.containers.apply(items)
-
-    def retrieve(self, ids: SequenceNotStr[ContainerId]) -> ContainerList:
-        return self.client.data_modeling.containers.retrieve(cast(Sequence, ids))
-
-    def update(self, items: Sequence[ContainerApply]) -> ContainerList:
-        return self.create(items)
-
-    def delete(self, ids: SequenceNotStr[ContainerId]) -> list[ContainerId]:
-        return self.client.data_modeling.containers.delete(cast(Sequence, ids))
-
-    def are_equal(self, local: ContainerApply, remote: Container) -> bool:
-        local_dumped = local.dump(camel_case=True)
-        if "usedFor" not in local_dumped:
-            # Setting used_for to "node" as it is the default value in the CDF.
-            local_dumped["usedFor"] = "node"
-
-        return local_dumped == remote.as_write().dump(camel_case=True)
-
-
-class DataModelLoader(DataModelingLoader[DataModelId, DataModelApply, DataModel, DataModelApplyList, DataModelList]):
-    resource_name = "data_models"
-
-    @classmethod
-    def get_id(cls, item: DataModel | DataModelApply | DataModelId | dict) -> DataModelId:
-        if isinstance(item, DataModel | DataModelApply):
-            return item.as_id()
-        if isinstance(item, dict):
-            return DataModelId.load(item)
-        return item
-
-    def create(self, items: Sequence[DataModelApply]) -> DataModelList:
-        return self.client.data_modeling.data_models.apply(items)
-
-    def retrieve(self, ids: SequenceNotStr[DataModelId]) -> DataModelList:
-        return self.client.data_modeling.data_models.retrieve(cast(Sequence, ids))
-
-    def update(self, items: Sequence[DataModelApply]) -> DataModelList:
-        return self.create(items)
-
-    def delete(self, ids: SequenceNotStr[DataModelId]) -> list[DataModelId]:
-        return self.client.data_modeling.data_models.delete(cast(Sequence, ids))
-
-    def are_equal(self, local: DataModelApply, remote: DataModel) -> bool:
-        local_dumped = local.dump()
-        cdf_resource_dumped = remote.as_write().dump()
-
-        # Data models that have the same views, but in different order, are considered equal.
-        # We also account for whether views are given as IDs or View objects.
-        local_dumped["views"] = sorted(
-            (v if isinstance(v, ViewId) else v.as_id()).as_tuple() for v in local.views or []
-        )
-        cdf_resource_dumped["views"] = sorted(
-            (v if isinstance(v, ViewId) else v.as_id()).as_tuple() for v in remote.views or []
-        )
-
-        return local_dumped == cdf_resource_dumped

cognite/neat/_utils/cdf/loaders/_ingestion.py

@@ -1,167 +0,0 @@
-from collections.abc import Iterable, Sequence
-from itertools import groupby
-from typing import cast, overload
-
-from cognite.client.data_classes import (
-    Database,
-    DatabaseList,
-    DatabaseWrite,
-    DatabaseWriteList,
-    Transformation,
-    TransformationList,
-    TransformationWrite,
-    TransformationWriteList,
-)
-from cognite.client.exceptions import CogniteAPIError
-from cognite.client.utils.useful_types import SequenceNotStr
-
-from cognite.neat._utils.cdf.data_classes import RawTable, RawTableID, RawTableList, RawTableWrite, RawTableWriteList
-
-from ._base import ResourceLoader
-
-
-class TransformationLoader(
-    ResourceLoader[str, TransformationWrite, Transformation, TransformationWriteList, TransformationList]
-):
-    resource_name = "transformations"
-
-    @classmethod
-    def get_id(cls, item: Transformation | TransformationWrite | str | dict) -> str:
-        if isinstance(item, Transformation | TransformationWrite):
-            if item.external_id is None:
-                raise ValueError(f"Transformation {item} does not have an external_id")
-            return item.external_id
-        if isinstance(item, dict):
-            if item.get("externalId") is None:
-                raise ValueError(f"Transformation {item} does not have an external_id")
-            return item["externalId"]
-        return item
-
-    def create(self, items: Sequence[TransformationWrite]) -> TransformationList:
-        return self.client.transformations.create(items)
-
-    def retrieve(self, ids: SequenceNotStr[str]) -> TransformationList:
-        return self.client.transformations.retrieve_multiple(external_ids=ids, ignore_unknown_ids=True)
-
-    def update(self, items: Sequence[TransformationWrite]) -> TransformationList:
-        return self.client.transformations.update(items)
-
-    def delete(self, ids: SequenceNotStr[str]) -> list[str]:
-        existing = self.retrieve(ids)
-        self.client.transformations.delete(external_id=ids, ignore_unknown_ids=True)
-        return existing.as_external_ids()
-
-
-class RawDatabaseLoader(ResourceLoader[str, DatabaseWrite, Database, DatabaseWriteList, DatabaseList]):
-    resource_name = "databases"
-
-    @classmethod
-    def get_id(cls, item: Database | DatabaseWrite | str | dict) -> str:
-        if isinstance(item, Database | DatabaseWrite):
-            if item.name is None:
-                raise ValueError(f"Database {item} does not have a name")
-            return item.name
-        if isinstance(item, dict):
-            if item.get("name") is None:
-                raise ValueError(f"Database {item} does not have a name")
-            return item["name"]
-        return item
-
-    def create(self, items: Sequence[DatabaseWrite]) -> DatabaseList:
-        return self.client.raw.databases.create([item.name for item in items if item.name is not None])
-
-    def retrieve(self, ids: SequenceNotStr[str]) -> DatabaseList:
-        all_databases = self.client.raw.databases.list(limit=-1)
-        return DatabaseList([db for db in all_databases if db.name in ids])
-
-    def update(self, items: Sequence[DatabaseWrite]) -> DatabaseList:
-        if not items:
-            return DatabaseList([])
-        raise NotImplementedError("The CDF API does not support updating a RAW database.")
-
-    def delete(self, ids: SequenceNotStr[str]) -> list[str]:
-        existing_databases = self.retrieve(ids)
-        existing_names = {item.name for item in existing_databases}
-        self.client.raw.databases.delete([name for name in ids if name in existing_names])
-        return existing_databases.as_names()
-
-
-class RawTableLoader(ResourceLoader[RawTableID, RawTableWrite, RawTable, RawTableWriteList, RawTableList]):
-    resource_name = "tables"
-
-    @classmethod
-    def get_id(cls, item: RawTable | RawTableWrite | RawTableID | dict) -> RawTableID:
-        if isinstance(item, RawTable | RawTableWrite):
-            return item.as_id()
-        if isinstance(item, dict):
-            return RawTableID(database=item["database"], table=item["name"])
-        return item
-
-    @overload
-    def _groupby_database(self, items: Sequence[RawTableWrite]) -> Iterable[tuple[str, Iterable[RawTableWrite]]]: ...
-
-    @overload
-    def _groupby_database(self, items: SequenceNotStr[RawTableID]) -> Iterable[tuple[str, Iterable[RawTableID]]]: ...
-
-    def _groupby_database(
-        self, items: Sequence[RawTableWrite] | SequenceNotStr[RawTableID]
-    ) -> Iterable[tuple[str, Iterable[RawTableWrite] | Iterable[RawTableID]]]:
-        return cast(
-            Iterable[tuple[str, Iterable[RawTableID] | Iterable[RawTableWrite]]],
-            groupby(sorted(items, key=lambda x: x.database or ""), lambda x: x.database or ""),
-        )
-
-    def create(self, items: Sequence[RawTableWrite]) -> RawTableList:
-        existing = set(self.retrieve([table.as_id() for table in items]).as_ids())
-        output = RawTableList([])
-        for db_name, tables in self._groupby_database(items):
-            to_create = [table.name for table in tables if table.name if table.as_id() not in existing]
-            if not to_create:
-                continue
-            created = self.client.raw.tables.create(db_name=db_name, name=to_create)
-            for table in created:
-                output.append(
-                    RawTable(
-                        name=table.name, database=db_name, created_time=table.created_time, cognite_client=self.client
-                    )
-                )
-        return output
-
-    def retrieve(self, ids: SequenceNotStr[RawTableID]) -> RawTableList:
-        output = RawTableList([])
-        for db_name, id_group in self._groupby_database(ids):
-            try:
-                all_tables = self.client.raw.tables.list(db_name, limit=-1)
-            except CogniteAPIError as e:
-                if e.code == 404 and e.message.startswith("Following databases not found"):
-                    continue
-            looking_for = {table_id.table for table_id in id_group if table_id.table is not None}
-            output.extend(
-                [
-                    RawTable(
-                        name=table.name, database=db_name, created_time=table.created_time, cognite_client=self.client
-                    )
-                    for table in all_tables
-                    if table.name in looking_for
-                ]
-            )
-        return output
-
-    def update(self, items: Sequence[RawTableWrite]) -> RawTableList:
-        if not items:
-            return RawTableList([])
-        raise NotImplementedError("The CDF API does not support updating a RAW table.")
-
-    def delete(self, ids: SequenceNotStr[RawTableID]) -> list[RawTableID]:
-        existing_tables = self.retrieve(ids)
-        existing_names = {item.name for item in existing_tables}
-        for db_name, id_group in self._groupby_database(ids):
-            self.client.raw.tables.delete(
-                db_name=db_name,
-                name=[
-                    table_id.table
-                    for table_id in id_group
-                    if table_id.table is not None and table_id.table in existing_names
-                ],
-            )
-        return existing_tables.as_ids()