cognite-neat 0.87.6__py3-none-any.whl → 0.88.1__py3-none-any.whl

cognite/neat/legacy/graph/loaders/rdf_to_dms.py

@@ -1,309 +0,0 @@
-import logging
-from collections.abc import Iterable
-from itertools import islice
-from typing import Literal, cast, overload
-
-from cognite.client import CogniteClient
-from cognite.client.data_classes.data_modeling import EdgeApply, InstanceApply, NodeApply
-from pydantic_core import ErrorDetails
-
-from cognite.neat.exceptions import NeatException
-from cognite.neat.legacy.graph.stores import NeatGraphStoreBase
-from cognite.neat.legacy.graph.transformations.query_generator.sparql import triples2dictionary
-from cognite.neat.legacy.rules.exporters._rules2dms import DMSSchemaComponents
-from cognite.neat.legacy.rules.exporters._rules2pydantic_models import add_class_prefix_to_xid, rules_to_pydantic_models
-from cognite.neat.legacy.rules.models.rules import Rules
-from cognite.neat.utils.auxiliary import retry_decorator
-from cognite.neat.utils.collection_ import chunker
-from cognite.neat.utils.time_ import datetime_utc_now
-
-from ._base import CogniteLoader
-
-
-class DMSLoader(CogniteLoader[InstanceApply]):
-    """Loads a Neat Graph into CDF as nodes and edges.
-
-    Args:
-        rules: Rules object
-        graph_store: Graph store
-        add_class_prefix: Add class prefix to external_id. Defaults to False.
-
-    """
-
-    def __init__(self, rules: Rules, graph_store: NeatGraphStoreBase, add_class_prefix: bool = False):
-        super().__init__(rules, graph_store)
-        self.add_class_prefix = add_class_prefix
-
-    @overload
-    def load(self, stop_on_exception: Literal[True]) -> Iterable[InstanceApply]: ...
-
-    @overload
-    def load(self, stop_on_exception: Literal[False] = False) -> Iterable[InstanceApply | ErrorDetails]: ...
-
-    def load(self, stop_on_exception: bool = False) -> Iterable[InstanceApply | ErrorDetails]:
-        """Load the graph with data."""
-        if self.rules.metadata.namespace is None:
-            raise ValueError("Namespace is not defined in transformation rules metadata")
-
-        data_model = DMSSchemaComponents.from_rules(self.rules)
-        pydantic_models = rules_to_pydantic_models(self.rules)
-
-        exclude = {
-            class_name
-            for class_name in self.rules.classes
-            if f"{self.rules.space}:{class_name}" not in data_model.containers
-        }
-
-        for class_name, triples in self._iterate_class_triples(exclude_classes=exclude):
-            logging.info(f"<DMSLoader> Processing class : {class_name}")
-            counter = 0
-            start_time = datetime_utc_now()
-            for instance_dict in triples2dictionary(triples).values():
-                counter += 1
-                try:
-                    instance = pydantic_models[class_name].from_dict(instance_dict)  # type: ignore[attr-defined]
-                    if self.add_class_prefix:
-                        instance.external_id = add_class_prefix_to_xid(
-                            class_name=type(instance).__name__, external_id=instance.external_id
-                        )
-                    new_node = instance.to_node(data_model, self.add_class_prefix)  # type: ignore[attr-defined]
-                    is_valid, reason = is_node_valid(new_node)
-                    if is_valid:
-                        yield new_node
-                    else:
-                        yield ErrorDetails(
-                            input=instance_dict["external_id"],
-                            loc=tuple(["Nodes"]),
-                            msg=f"Not valid node {new_node.external_id}. Reason: {reason}",
-                            type="Node validation error",
-                        )
-                        continue
-
-                    new_edges = instance.to_edge(data_model, self.add_class_prefix)
-                    for new_edge in new_edges:
-                        is_valid, reason = is_edge_valid(new_edge)
-                        if is_valid:
-                            yield new_edge
-                        else:
-                            yield ErrorDetails(
-                                input=instance_dict["external_id"],
-                                loc=tuple(["Edges"]),
-                                msg=f"Not valid edge {new_edge.external_id}. Reason: {reason}",
-                                type="Edge validation error",
-                            )
-                            continue
-
-                    delta_time = datetime_utc_now() - start_time
-                    delta_time = (delta_time.seconds * 1000000 + delta_time.microseconds) / 1000
-
-                except Exception as e:
-                    logging.error(
-                        f"Instance {instance_dict['external_id']} of {class_name}"
-                        f" cannot be resolved to nodes and edges. Reason: {e}"
-                    )
-                    if stop_on_exception:
-                        raise e
-
-                    if isinstance(e, NeatException):
-                        yield e.to_error_dict()
-                    else:
-                        yield ErrorDetails(
-                            input=instance_dict["external_id"],
-                            loc=tuple(["rdf2nodes_and_edges"]),
-                            msg=str(e),
-                            type=f"Exception of type {type(e).__name__} occurred  \
-                                        when processing instance of {class_name}",
-                        )
-
-    def load_to_cdf(
-        self, client: CogniteClient, batch_size: int | None = 1000, max_retries: int = 1, retry_delay: int = 3
-    ) -> None:
-        """Uploads nodes to CDF
-
-        Args:
-            client: Instance of CogniteClient
-            batch_size: Size of batch. Default to 1000.
-            max_retries: Maximum times to retry the upload. Default to 1.
-            retry_delay: Time delay before retrying the upload. Default to 3.
-
-        !!! note "batch_size"
-            If batch size is set to 1 or None, all nodes will be pushed to CDF in one go.
-        """
-        if batch_size is None:
-            logging.info("Batch size not set, pushing all nodes and edges to CDF in one go!")
-            nodes, edges, errors = self.as_nodes_and_edges(stop_on_exception=False)
-
-            @retry_decorator(max_retries=max_retries, retry_delay=retry_delay, component_name="create-instances")
-            def create_instances():
-                client.data_modeling.instances.apply(
-                    nodes=nodes, edges=edges, auto_create_start_nodes=True, auto_create_end_nodes=True
-                )
-
-            create_instances()
-            return
-        logging.info(f"Uploading nodes in batches of {batch_size}")
-        for instances in _batched(self.load(stop_on_exception=False), batch_size):
-            nodes = [instance for instance in instances if isinstance(instance, NodeApply)]
-            edges = [instance for instance in instances if isinstance(instance, EdgeApply)]
-            # Todo make _micro_batch_push handle both nodes and edges simultaneously
-            _micro_batch_push(
-                client, nodes, batch_size, message="Upload", max_retries=max_retries, retry_delay=retry_delay
-            )
-            _micro_batch_push(
-                client, edges, batch_size, message="Upload", max_retries=max_retries, retry_delay=retry_delay
-            )
-
-    def as_nodes_and_edges(
-        self, stop_on_exception: bool = False
-    ) -> tuple[list[NodeApply], list[EdgeApply], list[ErrorDetails]]:
-        nodes = []
-        edges = []
-        exceptions: list[ErrorDetails] = []
-        for instance in self.load(stop_on_exception):  # type: ignore[call-overload]
-            if isinstance(instance, NodeApply):
-                nodes.append(instance)
-            elif isinstance(instance, EdgeApply):
-                edges.append(instance)
-            elif isinstance(instance, dict):
-                exceptions.append(cast(ErrorDetails, instance))
-            else:
-                raise ValueError(f"Unknown instance type: {type(instance)}")
-        return nodes, edges, exceptions
-
-
-def _batched(iterable: Iterable, size: int):
-    "Batch data into lists of length n. The last batch may be shorter."
-    # batched('ABCDEFG', 3) --> ABC DEF G
-    it = iter(iterable)
-    while True:
-        batch = list(islice(it, size))
-        if not batch:
-            return
-        yield batch
-
-
-def is_node_valid(node: NodeApply) -> tuple[bool, str]:
-    return is_valid_external_id(node.external_id)
-
-
-def is_edge_valid(edge: EdgeApply) -> tuple[bool, str]:
-    for external_id in [edge.external_id, edge.start_node.external_id, edge.end_node.external_id]:
-        is_valid, reason = is_valid_external_id(external_id)
-        if not is_valid:
-            return False, reason
-    return True, ""
-
-
-def is_valid_external_id(external_id: str) -> tuple[bool, str]:
-    if external_id is None or external_id == "" or len(external_id) >= 255:
-        return False, f"external_id {external_id} is empty of too long"
-    return True, ""
-
-
-def upload_nodes(
-    client: CogniteClient, nodes: list[NodeApply], batch_size: int = 1000, max_retries: int = 1, retry_delay: int = 3
-):
-    """Uploads nodes to CDF
-
-    Args:
-        client: Instance of CogniteClient
-        nodes: List of nodes to upload to CDF
-        batch_size: Size of batch. Defaults to 1000.
-        max_retries: Maximum times to retry the upload. Defaults to 1.
-        retry_delay: Time delay before retrying the upload. Defaults to 3.
-
-    !!! note "batch_size"
-        If batch size is set to 1 or None, all nodes will be pushed to CDF in one go.
-    """
-    if batch_size:
-        logging.info(f"Uploading nodes in batches of {batch_size}")
-        _micro_batch_push(client, nodes, batch_size, message="Upload", max_retries=max_retries, retry_delay=retry_delay)
-
-    else:
-        logging.info("Batch size not set, pushing all nodes to CDF in one go!")
-
-        @retry_decorator(max_retries=max_retries, retry_delay=retry_delay, component_name="create-nodes")
-        def create_nodes():
-            client.data_modeling.instances.apply(nodes=nodes)
-
-        create_nodes()
-
-
-def upload_edges(
-    client: CogniteClient, edges: list[EdgeApply], batch_size: int = 5000, max_retries: int = 1, retry_delay: int = 3
-):
-    """Uploads edges to CDF
-
-    Args:
-        client: Instance of CogniteClient
-        edges: List of edges to upload to CDF
-        batch_size: Size of batch. Defaults to 5000.
-        max_retries: Maximum times to retry the upload. Defaults to 1.
-        retry_delay: Time delay before retrying the upload. Defaults to 3.
-
-    !!! note "batch_size"
-        If batch size is set to 1 or None, all edges will be pushed to CDF in one go.
-
-    """
-    if batch_size:
-        logging.info(f"Uploading edges in batches of {batch_size}")
-        _micro_batch_push(client, edges, batch_size, message="Upload", max_retries=max_retries, retry_delay=retry_delay)
-
-    else:
-        logging.info("Batch size not set, pushing all edges to CDF in one go!")
-
-        @retry_decorator(max_retries=max_retries, retry_delay=retry_delay, component_name="create-edges")
-        def create_nodes():
-            client.data_modeling.instances.apply(edges=edges, auto_create_start_nodes=True, auto_create_end_nodes=True)
-
-        create_nodes()
-
-
-def _micro_batch_push(
-    client: CogniteClient,
-    nodes_or_edges: list[NodeApply] | list[EdgeApply],
-    batch_size: int = 1000,
-    message: str = "Upload",
-    max_retries: int = 1,
-    retry_delay: int = 3,
-):
-    """Uploads nodes or edges in batches
-
-    Args:
-        client: Instance of CogniteClient
-        nodes_or_edges: List of nodes or edges
-        batch_size: Size of batch. Defaults to 1000.
-        message: Message to logged. Defaults to "Upload".
-        max_retries: Maximum times to retry the upload. Defaults to 1.
-        retry_delay: Time delay before retrying the upload. Defaults to 3.
-    """
-    total = len(nodes_or_edges)
-    counter = 0
-
-    if nodes_or_edges and isinstance(nodes_or_edges[0], NodeApply):
-        push_type = "nodes"
-    elif nodes_or_edges and isinstance(nodes_or_edges[0], EdgeApply):
-        push_type = "edges"
-    else:
-        raise ValueError("nodes_or_edges must be a list of NodeApply or EdgeApply objects")
-
-    for batch in chunker(nodes_or_edges, batch_size):
-        counter += len(batch)
-        start_time = datetime_utc_now()
-
-        @retry_decorator(max_retries=max_retries, retry_delay=retry_delay, component_name=f"microbatch-{push_type}")
-        def upsert_nodes_or_edges(upload_batch):
-            if push_type == "nodes":
-                client.data_modeling.instances.apply(nodes=upload_batch)
-            elif push_type == "edges":
-                client.data_modeling.instances.apply(
-                    edges=upload_batch, auto_create_start_nodes=True, auto_create_end_nodes=True
-                )
-
-        upsert_nodes_or_edges(batch)
-
-        delta_time = (datetime_utc_now() - start_time).seconds
-
-        msg = f"{message} {counter} of {total} {push_type}, batch processing time: {delta_time:.2f} "
-        msg += f"seconds ETC: {delta_time * (total - counter) / (60*batch_size) :.2f} minutes"
-        logging.info(msg)

cognite/neat/legacy/graph/loaders/validator.py

@@ -1,87 +0,0 @@
-"""Should contain methods to validate Graph Transformation Rules sheet,
-as well App Data Model (RDF)
-"""
-
-import logging
-from typing import Any
-
-
-def _find_circular_reference_path(
-    asset: dict[str, Any], assets: dict[str, dict[str, Any]], max_hierarchy_depth: int = 10000
-) -> list:
-    original_external_id = asset.get("external_id", "")
-    circle: list[str] = [original_external_id]
-    parent_external_id = asset.get("parent_external_id")
-    if isinstance(parent_external_id, str):
-        ref = assets.get(parent_external_id)
-    else:
-        ref = None
-
-    hop = 0
-    while ref is not None and hop < max_hierarchy_depth:
-        hop += 1
-        if external_id := ref.get("external_id"):
-            circle.append(external_id)
-        if len(circle) != len(set(circle)):
-            msg = (
-                f"Found circular reference in asset hierarchy which starts with "
-                f"{original_external_id} and enters loop at {circle[-1]}. "
-            )
-            logging.error(msg)
-            return circle
-        if parent_external_id := ref.get("parent_external_id"):
-            ref = assets.get(parent_external_id)
-        else:
-            ref = None
-
-    if hop >= max_hierarchy_depth:
-        msg = (
-            f"Your asset hierarchy is too deep. Max depth is {max_hierarchy_depth}. "
-            "You probably have a circular reference."
-        )
-        logging.error(msg)
-        return circle
-    else:
-        return []
-
-
-def validate_asset_hierarchy(
-    assets: dict[str, dict[str, Any]],
-) -> tuple[list[str], list[list[str]], dict[str, list[str]]]:
-    """Validates asset hierarchy and reports on orphan assets and circular dependency
-
-    Args:
-        assets : A dictionary of assets with external_id as key
-
-    Returns:
-        List of orphan assets external ids and list of circular path of external ids.
-        If both lists are empty, the hierarchy is healthy.
-    """
-    orphan_assets: list[str] = []
-    circular_reference_paths: list[list[str]] = []
-    parent_children_map: dict[str, list[str]] = {}
-
-    for asset in assets.values():
-        parent_external_id = asset.get("parent_external_id")
-        asset_extarnal_id = asset.get("external_id")
-        if asset_extarnal_id and parent_external_id:
-            if parent_external_id in parent_children_map:
-                parent_children_map[parent_external_id].append(asset_extarnal_id)
-            else:
-                parent_children_map[parent_external_id] = [asset_extarnal_id]
-        if parent_external_id is not None and parent_external_id not in assets:
-            msg = (
-                f"Found orphan asset {asset.get('external_id')} with parent {parent_external_id} which does not exist."
-            )
-            logging.error(msg)
-            if external_id := asset.get("external_id"):
-                orphan_assets.append(external_id)
-        circular_reference_path = _find_circular_reference_path(asset, assets)
-        if not len(circular_reference_path):
-            continue
-
-        # Save the circle only once, not once for every asset
-        if set(circular_reference_path) in [set(path) for path in circular_reference_paths]:
-            continue
-        circular_reference_paths.append(circular_reference_path)
-    return orphan_assets, circular_reference_paths, parent_children_map

cognite/neat/legacy/graph/models.py

@@ -1,6 +0,0 @@
-from typing import TypeAlias
-
-from rdflib import Literal
-from rdflib.term import URIRef
-
-Triple: TypeAlias = tuple[URIRef, URIRef, Literal | URIRef]

cognite/neat/legacy/graph/stores/__init__.py

@@ -1,13 +0,0 @@
-from ._base import NeatGraphStoreBase
-from ._graphdb_store import GraphDBStore
-from ._memory_store import MemoryStore
-from ._oxigraph_store import OxiGraphStore
-
-STORE_BY_TYPE: dict[str, type[NeatGraphStoreBase]] = {}
-for store in NeatGraphStoreBase.__subclasses__():
-    STORE_BY_TYPE[store.rdf_store_type] = store  # type: ignore[type-abstract]
-
-del store  # Cleanup namespace
-AVAILABLE_STORES = set(STORE_BY_TYPE.keys())
-
-__all__ = ["NeatGraphStoreBase", "MemoryStore", "OxiGraphStore", "GraphDBStore", "STORE_BY_TYPE", "AVAILABLE_STORES"]