cognite-neat 0.107.0__tar.gz → 0.109.0__tar.gz

{cognite_neat-0.107.0 → cognite_neat-0.109.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: cognite-neat
-Version: 0.107.0
+Version: 0.109.0
 Summary: Knowledge graph transformation
 License: Apache-2.0
 Author: Nikola Vasiljevic

{cognite_neat-0.107.0 → cognite_neat-0.109.0}/cognite/neat/_constants.py

@@ -5,6 +5,7 @@ from typing import TYPE_CHECKING
 from cognite.client import data_modeling as dm
 from cognite.client.data_classes.data_modeling.ids import DataModelId
 from rdflib import DC, DCTERMS, FOAF, OWL, RDF, RDFS, SH, SKOS, XSD, Namespace, URIRef
+from rdflib.namespace import DefinedNamespace
 
 from cognite import neat
 
@@ -73,10 +74,22 @@ DEFAULT_NAMESPACE = Namespace("http://purl.org/cognite/neat/")
 CDF_NAMESPACE = Namespace("https://cognitedata.com/")
 DEFAULT_BASE_URI = URIRef(DEFAULT_NAMESPACE)
 CLASSIC_CDF_NAMESPACE = Namespace("http://purl.org/cognite/cdf-classic#")
-UNKNOWN_TYPE = DEFAULT_NAMESPACE.UnknownType
 XML_SCHEMA_NAMESPACE = Namespace("http://www.w3.org/2001/XMLSchema#")
 
 
+class NEAT(DefinedNamespace):
+    """
+    NEAT internal data model used for internal purposes of the NEAT library
+
+    """
+
+    _fail = True
+    _NS = Namespace("http://thisisneat.io/internal/")
+
+    UnknownType: URIRef  # Unknown type used to express that the type of a subject is unknown
+    EmptyType: URIRef  # Empty type used to express that the type of a subject is empty
+
+
 def get_default_prefixes_and_namespaces() -> dict[str, Namespace]:
     return {
         "owl": OWL._NS,
@@ -154,3 +167,24 @@ READONLY_PROPERTIES_BY_CONTAINER: Mapping[dm.ContainerId, frozenset[str]] = {
 
 def is_readonly_property(container: dm.ContainerId, property_: str) -> bool:
     return container in READONLY_PROPERTIES_BY_CONTAINER and property_ in READONLY_PROPERTIES_BY_CONTAINER[container]
+
+
+DMS_RESERVED_PROPERTIES = frozenset(
+    {
+        "createdTime",
+        "deletedTime",
+        "edge_id",
+        "extensions",
+        "externalId",
+        "lastUpdatedTime",
+        "node_id",
+        "project_id",
+        "property_group",
+        "seq",
+        "space",
+        "version",
+        "tg_table_name",
+        "startNode",
+        "endNode",
+    }
+)

{cognite_neat-0.107.0 → cognite_neat-0.109.0}/cognite/neat/_graph/_shared.py

@@ -7,6 +7,10 @@ MIMETypes: TypeAlias = Literal[
 RDFTypes: TypeAlias = Literal["xml", "rdf", "owl", "n3", "ttl", "turtle", "nt", "nq", "nquads", "trig"]
 
 
+def quad_formats() -> list[str]:
+    return ["trig", "nquads", "nq", "nt"]
+
+
 def rdflib_to_oxi_type(rdflib_format: str) -> str | None:
     """Convert an RDFlib format to a MIME type.
 

{cognite_neat-0.107.0 → cognite_neat-0.109.0}/cognite/neat/_graph/extractors/_classic_cdf/_base.py

@@ -1,6 +1,8 @@
 import json
 import re
 import sys
+import typing
+import urllib.parse
 import warnings
 from abc import ABC, abstractmethod
 from collections.abc import Callable, Iterable, Sequence, Set
@@ -16,7 +18,8 @@ from rdflib import RDF, XSD, Literal, Namespace, URIRef
 
 from cognite.neat._constants import DEFAULT_NAMESPACE
 from cognite.neat._graph.extractors._base import BaseExtractor
-from cognite.neat._issues.warnings import CDFAuthWarning
+from cognite.neat._issues.errors import NeatValueError
+from cognite.neat._issues.warnings import CDFAuthWarning, NeatValueWarning
 from cognite.neat._shared import Triple
 from cognite.neat._utils.auxiliary import string_to_ideal_type
 from cognite.neat._utils.collection_ import iterate_progress_bar_if_above_config_threshold
@@ -72,6 +75,8 @@ class ClassicCDFBaseExtractor(BaseExtractor, ABC, Generic[T_CogniteResource]):
         camel_case (bool, optional): Whether to use camelCase instead of snake_case for property names.
             Defaults to True.
         as_write (bool, optional): Whether to use the write/request format of the items. Defaults to False.
+        prefix (str, optional): A prefix to add to the rdf type. Defaults to None.
+        identifier (Literal["id", "externalId"], optional): The identifier to use. Defaults to "id".
     """
 
     _default_rdf_type: str
@@ -90,6 +95,7 @@ class ClassicCDFBaseExtractor(BaseExtractor, ABC, Generic[T_CogniteResource]):
         camel_case: bool = True,
         as_write: bool = False,
         prefix: str | None = None,
+        identifier: typing.Literal["id", "externalId"] = "id",
     ):
         self.namespace = namespace or DEFAULT_NAMESPACE
         self.items = items
@@ -101,9 +107,18 @@ class ClassicCDFBaseExtractor(BaseExtractor, ABC, Generic[T_CogniteResource]):
         self.camel_case = camel_case
         self.as_write = as_write
         self.prefix = prefix
+        self.identifier = identifier
+        # If identifier=externalId, we need to keep track of the external ids
+        # and use them in linking of Files, Sequences, TimeSeries, and Events.
+        self.asset_external_ids_by_id: dict[int, str] = {}
+        self.lookup_dataset_external_id: Callable[[int], str] | None = None
+        # Used by the ClassicGraphExtractor to log URIRefs
+        self._log_urirefs = False
+        self._uriref_by_external_id: dict[str, URIRef] = {}
 
     def extract(self) -> Iterable[Triple]:
         """Extracts an asset with the given asset_id."""
+        from ._assets import AssetsExtractor
 
         if self.total is not None and self.total > 0:
             to_iterate = iterate_progress_bar_if_above_config_threshold(
@@ -111,21 +126,40 @@ class ClassicCDFBaseExtractor(BaseExtractor, ABC, Generic[T_CogniteResource]):
             )
         else:
             to_iterate = self.items
+        if self.identifier == "externalId" and isinstance(self, AssetsExtractor):
+            to_iterate = self._store_asset_external_ids(to_iterate)  # type: ignore[attr-defined]
+
         for no, asset in enumerate(to_iterate):
             yield from self._item2triples(asset)
             if self.limit and no >= self.limit:
                 break
 
+    def _store_asset_external_ids(self, items: Iterable[T_CogniteResource]) -> Iterable[T_CogniteResource]:
+        for item in items:
+            if hasattr(item, "id") and hasattr(item, "external_id"):
+                self.asset_external_ids_by_id[item.id] = item.external_id
+            yield item
+
     def _item2triples(self, item: T_CogniteResource) -> list[Triple]:
-        id_value: str | None
-        if hasattr(item, "id"):
-            id_value = str(item.id)
+        if self.identifier == "id":
+            id_value: str | None
+            if hasattr(item, "id"):
+                id_value = str(item.id)
+            else:
+                id_value = self._fallback_id(item)
+            if id_value is None:
+                return []
+            id_suffix = id_value
+        elif self.identifier == "externalId":
+            if not hasattr(item, "external_id"):
+                return []
+            id_suffix = self._external_id_as_uri_suffix(item.external_id)
         else:
-            id_value = self._fallback_id(item)
-        if id_value is None:
-            return []
+            raise NeatValueError(f"Unknown identifier {self.identifier}")
 
-        id_ = self.namespace[f"{self._instance_id_prefix}{id_value}"]
+        id_ = self.namespace[f"{self._instance_id_prefix}{id_suffix}"]
+        if self._log_urirefs and hasattr(item, "external_id"):
+            self._uriref_by_external_id[item.external_id] = id_
 
         type_ = self._get_rdf_type(item)
 
@@ -154,10 +188,25 @@ class ClassicCDFBaseExtractor(BaseExtractor, ABC, Generic[T_CogniteResource]):
         """This can be overridden to handle special cases for the item."""
         return []
 
+    @classmethod
+    def _external_id_as_uri_suffix(cls, external_id: str | None) -> str:
+        if external_id == "":
+            warnings.warn(NeatValueWarning(f"Empty external id in {cls._default_rdf_type}"), stacklevel=2)
+            return "empty"
+        elif external_id == "\x00":
+            warnings.warn(NeatValueWarning(f"Null external id in {cls._default_rdf_type}"), stacklevel=2)
+            return "null"
+        elif external_id is None:
+            warnings.warn(NeatValueWarning(f"None external id in {cls._default_rdf_type}"), stacklevel=2)
+            return "None"
+        # The external ID needs to pass the ^[^\\x00]{1,256}$ regex for the DMS API.
+        # In addition, neat internals requires the external ID to be a valid URI.
+        return urllib.parse.quote(external_id)
+
     def _fallback_id(self, item: T_CogniteResource) -> str | None:
         raise AttributeError(
             f"Item of type {type(item)} does not have an id attribute. "
-            f"Please implement the _fallback_id method in the extractor."
+            "Please implement the _fallback_id method in the extractor."
         )
 
     def _metadata_to_triples(self, id_: URIRef, metadata: dict[str, str]) -> Iterable[Triple]:
@@ -181,10 +230,29 @@ class ClassicCDFBaseExtractor(BaseExtractor, ABC, Generic[T_CogniteResource]):
         return self._SPACE_PATTERN.sub("_", type_)
 
     def _as_object(self, raw: Any, key: str) -> Literal | URIRef:
+        """Return properly formatted object part of s-p-o triple"""
         if key in {"data_set_id", "dataSetId"}:
-            return self.namespace[f"{InstanceIdPrefix.data_set}{raw}"]
+            if self.identifier == "externalId" and self.lookup_dataset_external_id:
+                try:
+                    data_set_external_id = self.lookup_dataset_external_id(raw)
+                except KeyError:
+                    return Literal("Unknown data set")
+                else:
+                    return self.namespace[
+                        f"{InstanceIdPrefix.data_set}{self._external_id_as_uri_suffix(data_set_external_id)}"
+                    ]
+            else:
+                return self.namespace[f"{InstanceIdPrefix.data_set}{raw}"]
         elif key in {"assetId", "asset_id", "assetIds", "asset_ids", "parentId", "rootId", "parent_id", "root_id"}:
-            return self.namespace[f"{InstanceIdPrefix.asset}{raw}"]
+            if self.identifier == "id":
+                return self.namespace[f"{InstanceIdPrefix.asset}{raw}"]
+            else:
+                try:
+                    asset_external_id = self._external_id_as_uri_suffix(self.asset_external_ids_by_id[raw])
+                except KeyError:
+                    return Literal("Unknown asset", datatype=XSD.string)
+                else:
+                    return self.namespace[f"{InstanceIdPrefix.asset}{asset_external_id}"]
         elif key in {
             "startTime",
             "endTime",
@@ -223,10 +291,21 @@ class ClassicCDFBaseExtractor(BaseExtractor, ABC, Generic[T_CogniteResource]):
         camel_case: bool = True,
         as_write: bool = False,
         prefix: str | None = None,
+        identifier: typing.Literal["id", "externalId"] = "id",
     ):
         total, items = cls._handle_no_access(lambda: cls._from_dataset(client, data_set_external_id))
         return cls(
-            items, namespace, to_type, total, limit, unpack_metadata, skip_metadata_values, camel_case, as_write, prefix
+            items,
+            namespace,
+            to_type,
+            total,
+            limit,
+            unpack_metadata,
+            skip_metadata_values,
+            camel_case,
+            as_write,
+            prefix,
+            identifier,
         )
 
     @classmethod
@@ -249,10 +328,21 @@ class ClassicCDFBaseExtractor(BaseExtractor, ABC, Generic[T_CogniteResource]):
         camel_case: bool = True,
         as_write: bool = False,
         prefix: str | None = None,
+        identifier: typing.Literal["id", "externalId"] = "id",
     ):
         total, items = cls._handle_no_access(lambda: cls._from_hierarchy(client, root_asset_external_id))
         return cls(
-            items, namespace, to_type, total, limit, unpack_metadata, skip_metadata_values, camel_case, as_write, prefix
+            items,
+            namespace,
+            to_type,
+            total,
+            limit,
+            unpack_metadata,
+            skip_metadata_values,
+            camel_case,
+            as_write,
+            prefix,
+            identifier,
         )
 
     @classmethod
@@ -274,10 +364,21 @@ class ClassicCDFBaseExtractor(BaseExtractor, ABC, Generic[T_CogniteResource]):
         camel_case: bool = True,
         as_write: bool = False,
         prefix: str | None = None,
+        identifier: typing.Literal["id", "externalId"] = "id",
     ):
         total, items = cls._from_file(file_path)
         return cls(
-            items, namespace, to_type, total, limit, unpack_metadata, skip_metadata_values, camel_case, as_write, prefix
+            items,
+            namespace,
+            to_type,
+            total,
+            limit,
+            unpack_metadata,
+            skip_metadata_values,
+            camel_case,
+            as_write,
+            prefix,
+            identifier,
         )
 
     @classmethod

{cognite_neat-0.107.0 → cognite_neat-0.109.0}/cognite/neat/_graph/extractors/_classic_cdf/_classic.py

@@ -1,3 +1,5 @@
+import typing
+import urllib.parse
 import warnings
 from collections import defaultdict
 from collections.abc import Iterable, Sequence
@@ -9,8 +11,8 @@ from rdflib import Namespace, URIRef
 
 from cognite.neat._constants import CLASSIC_CDF_NAMESPACE, DEFAULT_NAMESPACE, get_default_prefixes_and_namespaces
 from cognite.neat._graph.extractors._base import KnowledgeGraphExtractor
-from cognite.neat._issues.errors import NeatValueError
-from cognite.neat._issues.warnings import CDFAuthWarning
+from cognite.neat._issues.errors import NeatValueError, ResourceNotFoundError
+from cognite.neat._issues.warnings import CDFAuthWarning, NeatValueWarning
 from cognite.neat._rules._shared import ReadRules
 from cognite.neat._rules.catalog import classic_model
 from cognite.neat._rules.models import InformationInputRules, InformationRules
@@ -101,6 +103,7 @@ class ClassicGraphExtractor(KnowledgeGraphExtractor):
         namespace: Namespace | None = None,
         limit_per_type: int | None = None,
         prefix: str | None = None,
+        identifier: typing.Literal["id", "externalId"] = "id",
     ):
         self._client = client
         if sum([bool(data_set_external_id), bool(root_asset_external_id)]) != 1:
@@ -115,16 +118,23 @@ class ClassicGraphExtractor(KnowledgeGraphExtractor):
             camel_case=True,
             limit=limit_per_type,
             prefix=prefix,
+            identifier=identifier,
         )
+        self._identifier = identifier
         self._prefix = prefix
         self._limit_per_type = limit_per_type
 
+        self._uris_by_external_id_by_type: dict[InstanceIdPrefix, dict[str, URIRef]] = defaultdict(dict)
         self._source_external_ids_by_type: dict[InstanceIdPrefix, set[str]] = defaultdict(set)
         self._target_external_ids_by_type: dict[InstanceIdPrefix, set[str]] = defaultdict(set)
+        self._relationship_subject_predicate_type_external_id: list[tuple[URIRef, URIRef, str, str]] = []
         self._labels: set[str] = set()
         self._data_set_ids: set[int] = set()
+        self._data_set_external_ids: set[str] = set()
         self._extracted_labels = False
         self._extracted_data_sets = False
+        self._asset_external_ids_by_id: dict[int, str] = {}
+        self._dataset_external_ids_by_id: dict[int, str] = {}
 
     def _get_activity_names(self) -> list[str]:
         activities = [data_access_object.extractor_cls.__name__ for data_access_object in self._classic_node_types] + [
@@ -138,12 +148,17 @@ class ClassicGraphExtractor(KnowledgeGraphExtractor):
 
     def extract(self) -> Iterable[Triple]:
         """Extracts all classic CDF Resources."""
+        self._validate_exists()
+
         yield from self._extract_core_start_nodes()
 
         yield from self._extract_start_node_relationships()
 
         yield from self._extract_core_end_nodes()
 
+        if self._identifier == "id":
+            yield from self._extract_relationship_target_triples()
+
         try:
             yield from self._extract_labels()
         except CogniteAPIError as e:
@@ -209,7 +224,17 @@ class ClassicGraphExtractor(KnowledgeGraphExtractor):
         else:
             resource = "unknown"
             external_id = "unknown"
-        return DEFAULT_NAMESPACE[f"{self._client.config.project}/{resource}/{external_id}"]
+        return DEFAULT_NAMESPACE[f"{self._client.config.project}/{resource}/{urllib.parse.quote(external_id)}"]
+
+    def _validate_exists(self) -> None:
+        if self._data_set_external_id:
+            if self._client.data_sets.retrieve(external_id=self._data_set_external_id) is None:
+                raise ResourceNotFoundError(self._data_set_external_id, "data set")
+        elif self._root_asset_external_id:
+            if self._client.assets.retrieve(external_id=self._root_asset_external_id) is None:
+                raise ResourceNotFoundError(self._root_asset_external_id, "root asset")
+        else:
+            raise ValueError("Exactly one of data_set_external_id or root_asset_external_id must be set.")
 
     def _extract_core_start_nodes(self):
         for core_node in self._classic_node_types:
@@ -224,8 +249,20 @@ class ClassicGraphExtractor(KnowledgeGraphExtractor):
             else:
                 raise ValueError("Exactly one of data_set_external_id or root_asset_external_id must be set.")
 
+            if self._identifier == "externalId":
+                if isinstance(extractor, AssetsExtractor):
+                    self._asset_external_ids_by_id = extractor.asset_external_ids_by_id
+                else:
+                    extractor.asset_external_ids_by_id = self._asset_external_ids_by_id
+                extractor.lookup_dataset_external_id = self._lookup_dataset
+            elif self._identifier == "id":
+                extractor._log_urirefs = True
+
             yield from self._extract_with_logging_label_dataset(extractor, core_node.resource_type)
 
+            if self._identifier == "id":
+                self._uris_by_external_id_by_type[core_node.resource_type].update(extractor._uriref_by_external_id)
+
     def _extract_start_node_relationships(self):
         for start_resource_type, source_external_ids in self._source_external_ids_by_type.items():
             start_type = start_resource_type.removesuffix("_")
@@ -236,6 +273,8 @@ class ClassicGraphExtractor(KnowledgeGraphExtractor):
                 extractor = RelationshipsExtractor(relationship_iterator, **self._extractor_args)
                 # This is a private attribute, but we need to set it to log the target nodes.
                 extractor._log_target_nodes = True
+                if self._identifier == "id":
+                    extractor._uri_by_external_id_by_by_type = self._uris_by_external_id_by_type
 
                 yield from extractor.extract()
 
@@ -254,6 +293,11 @@ class ClassicGraphExtractor(KnowledgeGraphExtractor):
                         ):
                             self._target_external_ids_by_type[end_type].add(external_id)
 
+                if self._identifier == "id":
+                    # We need to store all future target triples which we will lookup after fetching
+                    # the target nodes.
+                    self._relationship_subject_predicate_type_external_id.extend(extractor._target_triples)
+
     def _extract_core_end_nodes(self):
         for core_node in self._classic_node_types:
             target_external_ids = self._target_external_ids_by_type[core_node.resource_type]
@@ -264,8 +308,26 @@ class ClassicGraphExtractor(KnowledgeGraphExtractor):
             ):
                 resource_iterator = api.retrieve_multiple(external_ids=list(chunk), ignore_unknown_ids=True)
                 extractor = core_node.extractor_cls(resource_iterator, **self._extractor_args)
+
+                extractor.asset_external_ids_by_id = self._asset_external_ids_by_id
+                extractor.lookup_dataset_external_id = self._lookup_dataset
+                if self._identifier == "id":
+                    extractor._log_urirefs = True
+
                 yield from self._extract_with_logging_label_dataset(extractor)
 
+                if self._identifier == "id":
+                    self._uris_by_external_id_by_type[core_node.resource_type].update(extractor._uriref_by_external_id)
+
+    def _extract_relationship_target_triples(self):
+        for id_, predicate, type_, external_id in self._relationship_subject_predicate_type_external_id:
+            try:
+                object_uri = self._uris_by_external_id_by_type[InstanceIdPrefix.from_str(type_)][external_id]
+            except KeyError:
+                warnings.warn(NeatValueWarning(f"Missing externalId {external_id} for {type_}"), stacklevel=2)
+            else:
+                yield id_, predicate, object_uri
+
     def _extract_labels(self):
         for chunk in self._chunk(list(self._labels), description="Extracting labels"):
             label_iterator = self._client.labels.retrieve(external_id=list(chunk), ignore_unknown_ids=True)
@@ -275,6 +337,11 @@ class ClassicGraphExtractor(KnowledgeGraphExtractor):
         for chunk in self._chunk(list(self._data_set_ids), description="Extracting data sets"):
             data_set_iterator = self._client.data_sets.retrieve_multiple(ids=list(chunk), ignore_unknown_ids=True)
             yield from DataSetExtractor(data_set_iterator, **self._extractor_args).extract()
+        for chunk in self._chunk(list(self._data_set_external_ids), description="Extracting data sets"):
+            data_set_iterator = self._client.data_sets.retrieve_multiple(
+                external_ids=list(chunk), ignore_unknown_ids=True
+            )
+            yield from DataSetExtractor(data_set_iterator, **self._extractor_args).extract()
 
     def _extract_with_logging_label_dataset(
         self, extractor: ClassicCDFBaseExtractor, resource_type: InstanceIdPrefix | None = None
@@ -285,9 +352,11 @@ class ClassicGraphExtractor(KnowledgeGraphExtractor):
             elif triple[1] == self._namespace.labels:
                 self._labels.add(remove_namespace_from_uri(triple[2]).removeprefix(InstanceIdPrefix.label))
             elif triple[1] == self._namespace.dataSetId:
-                self._data_set_ids.add(
-                    int(remove_namespace_from_uri(triple[2]).removeprefix(InstanceIdPrefix.data_set))
-                )
+                identifier = remove_namespace_from_uri(triple[2]).removeprefix(InstanceIdPrefix.data_set)
+                try:
+                    self._data_set_ids.add(int(identifier))
+                except ValueError:
+                    self._data_set_external_ids.add(identifier)
             yield triple
 
     @staticmethod
@@ -297,3 +366,15 @@ class ClassicGraphExtractor(KnowledgeGraphExtractor):
             return iterate_progress_bar(to_iterate, (len(items) // 1_000) + 1, description)
         else:
             return to_iterate
+
+    def _lookup_dataset(self, dataset_id: int) -> str:
+        if dataset_id not in self._dataset_external_ids_by_id:
+            try:
+                if (dataset := self._client.data_sets.retrieve(id=dataset_id)) and dataset.external_id:
+                    self._dataset_external_ids_by_id[dataset_id] = dataset.external_id
+                else:
+                    raise KeyError(f"Could not find dataset with id {dataset_id}.")
+            except CogniteAPIError as e:
+                warnings.warn(CDFAuthWarning("lookup dataset", str(e)), stacklevel=2)
+                return f"{InstanceIdPrefix.data_set}{dataset_id}"
+        return self._dataset_external_ids_by_id[dataset_id]

{cognite_neat-0.107.0 → cognite_neat-0.109.0}/cognite/neat/_graph/extractors/_classic_cdf/_relationships.py

@@ -1,11 +1,16 @@
+import typing
+import warnings
 from collections import defaultdict
 from collections.abc import Callable, Iterable, Set
 from pathlib import Path
+from typing import Any
 
 from cognite.client import CogniteClient
 from cognite.client.data_classes import Relationship, RelationshipList
-from rdflib import Namespace
+from rdflib import Namespace, URIRef
 
+from cognite.neat._issues.warnings import NeatValueWarning
+from cognite.neat._shared import Triple
 from cognite.neat._utils.auxiliary import create_sha256_hash
 
 from ._base import DEFAULT_SKIP_METADATA_VALUES, ClassicCDFBaseExtractor, InstanceIdPrefix, T_CogniteResource
@@ -29,9 +34,17 @@ class RelationshipsExtractor(ClassicCDFBaseExtractor[Relationship]):
         camel_case: bool = True,
         as_write: bool = False,
         prefix: str | None = None,
+        identifier: typing.Literal["id", "externalId"] = "id",
     ):
+        # This is used by the ClassicExtractor to log the target nodes, such
+        # that it can extract them.
+        # It is private to avoid exposing it to the user.
+        self._target_external_ids_by_type: dict[InstanceIdPrefix, set[str]] = defaultdict(set)
+        self._log_target_nodes = False
+        # Ensure that this becomes an iterator, even if it is a list.
+        to_iterate = (self._log_target_nodes_if_set(item) for item in items)
         super().__init__(
-            items,
+            to_iterate,
             namespace=namespace,
             to_type=to_type,
             total=total,
@@ -41,12 +54,39 @@ class RelationshipsExtractor(ClassicCDFBaseExtractor[Relationship]):
             camel_case=camel_case,
             as_write=as_write,
             prefix=prefix,
+            identifier=identifier,
         )
-        # This is used by the ClassicExtractor to log the target nodes, such
-        # that it can extract them.
-        # It is private to avoid exposing it to the user.
-        self._log_target_nodes = False
-        self._target_external_ids_by_type: dict[InstanceIdPrefix, set[str]] = defaultdict(set)
+        self._uri_by_external_id_by_by_type: dict[InstanceIdPrefix, dict[str, URIRef]] = defaultdict(dict)
+        self._target_triples: list[tuple[URIRef, URIRef, str, str]] = []
+
+    def _log_target_nodes_if_set(self, item: Relationship) -> Relationship:
+        if not self._log_target_nodes:
+            return item
+        if item.target_type and item.target_external_id:
+            self._target_external_ids_by_type[InstanceIdPrefix.from_str(item.target_type)].add(item.target_external_id)
+        return item
+
+    def _item2triples_special_cases(self, id_: URIRef, dumped: dict[str, Any]) -> list[Triple]:
+        if self.identifier == "externalId":
+            return []
+        triples: list[Triple] = []
+        if (source_external_id := dumped.pop("sourceExternalId")) and "sourceType" in dumped:
+            source_type = dumped["sourceType"]
+            try:
+                source_uri = self._uri_by_external_id_by_by_type[InstanceIdPrefix.from_str(source_type)][
+                    source_external_id
+                ]
+            except KeyError:
+                warnings.warn(
+                    NeatValueWarning(f"Missing externalId {source_external_id} for {source_type}"), stacklevel=2
+                )
+            else:
+                triples.append((id_, self.namespace["sourceExternalId"], source_uri))
+        if (target_external_id := dumped.pop("targetExternalId")) and "targetType" in dumped:
+            target_type = dumped["targetType"]
+            # We do not yet have the target nodes, so we log them for later extraction.
+            self._target_triples.append((id_, self.namespace["targetExternalId"], target_type, target_external_id))
+        return triples
 
     @classmethod
     def _from_dataset(
@@ -69,10 +109,6 @@ class RelationshipsExtractor(ClassicCDFBaseExtractor[Relationship]):
         return len(relationships), relationships
 
     def _fallback_id(self, item: Relationship) -> str | None:
-        if item.external_id and item.source_external_id and item.target_external_id:
-            if self._log_target_nodes and item.target_type and item.target_external_id:
-                self._target_external_ids_by_type[InstanceIdPrefix.from_str(item.target_type)].add(
-                    item.target_external_id
-                )
+        if item.external_id:
             return create_sha256_hash(item.external_id)
         return None

{cognite_neat-0.107.0 → cognite_neat-0.109.0}/cognite/neat/_graph/extractors/_classic_cdf/_sequences.py

@@ -1,5 +1,6 @@
 import itertools
 import json
+import typing
 from collections.abc import Callable, Iterable, Set
 from pathlib import Path
 from typing import Any
@@ -53,10 +54,21 @@ class SequencesExtractor(ClassicCDFBaseExtractor[NeatSequence]):
         camel_case: bool = True,
         as_write: bool = False,
         prefix: str | None = None,
+        identifier: typing.Literal["id", "externalId"] = "id",
         unpack_columns: bool = False,
     ):
         super().__init__(
-            items, namespace, to_type, total, limit, unpack_metadata, skip_metadata_values, camel_case, as_write, prefix
+            items,
+            namespace,
+            to_type,
+            total,
+            limit,
+            unpack_metadata,
+            skip_metadata_values,
+            camel_case,
+            as_write,
+            prefix,
+            identifier,
         )
         self.unpack_columns = unpack_columns
 
@@ -73,6 +85,7 @@ class SequencesExtractor(ClassicCDFBaseExtractor[NeatSequence]):
         camel_case: bool = True,
         as_write: bool = False,
         prefix: str | None = None,
+        identifier: typing.Literal["id", "externalId"] = "id",
         unpack_columns: bool = False,
     ):
         total, items = cls._handle_no_access(lambda: cls._from_dataset(client, data_set_external_id))
@@ -87,6 +100,7 @@ class SequencesExtractor(ClassicCDFBaseExtractor[NeatSequence]):
             camel_case,
             as_write,
             prefix,
+            identifier,
             unpack_columns,
         )
 
@@ -103,6 +117,7 @@ class SequencesExtractor(ClassicCDFBaseExtractor[NeatSequence]):
         camel_case: bool = True,
         as_write: bool = False,
         prefix: str | None = None,
+        identifier: typing.Literal["id", "externalId"] = "id",
         unpack_columns: bool = False,
     ):
         total, items = cls._handle_no_access(lambda: cls._from_hierarchy(client, root_asset_external_id))
@@ -117,6 +132,7 @@ class SequencesExtractor(ClassicCDFBaseExtractor[NeatSequence]):
             camel_case,
             as_write,
             prefix,
+            identifier,
             unpack_columns,
         )
 
@@ -132,6 +148,7 @@ class SequencesExtractor(ClassicCDFBaseExtractor[NeatSequence]):
         camel_case: bool = True,
         as_write: bool = False,
         prefix: str | None = None,
+        identifier: typing.Literal["id", "externalId"] = "id",
         unpack_columns: bool = False,
     ):
         total, items = cls._from_file(file_path)
@@ -146,6 +163,7 @@ class SequencesExtractor(ClassicCDFBaseExtractor[NeatSequence]):
             camel_case,
             as_write,
             prefix,
+            identifier,
             unpack_columns,
         )