cognite-neat 0.110.0__py3-none-any.whl → 0.111.0__py3-none-any.whl

cognite/neat/_graph/loaders/_rdf2dms.py

@@ -6,7 +6,7 @@ from collections import defaultdict
 from collections.abc import Iterable, Sequence
 from dataclasses import dataclass
 from pathlib import Path
-from typing import Any, Literal, cast, get_args
+from typing import Any, Literal, cast, get_args, overload
 
 import yaml
 from cognite.client import CogniteClient
@@ -17,14 +17,20 @@ from cognite.client.data_classes.data_modeling.ids import InstanceId
 from cognite.client.data_classes.data_modeling.views import SingleEdgeConnection
 from cognite.client.exceptions import CogniteAPIError
 from pydantic import BaseModel, ValidationInfo, create_model, field_validator
-from rdflib import RDF
+from rdflib import RDF, URIRef
 
 from cognite.neat._client import NeatClient
 from cognite.neat._client._api_client import SchemaAPI
 from cognite.neat._constants import DMS_DIRECT_RELATION_LIST_LIMIT, is_readonly_property
-from cognite.neat._issues import IssueList, NeatIssue, catch_issues
-from cognite.neat._issues.errors import ResourceCreationError, ResourceDuplicatedError, ResourceNotFoundError
+from cognite.neat._issues import IssueList, NeatError, NeatIssue, catch_issues
+from cognite.neat._issues.errors import (
+    AuthorizationError,
+    ResourceCreationError,
+    ResourceDuplicatedError,
+    ResourceNotFoundError,
+)
 from cognite.neat._issues.warnings import (
+    NeatValueWarning,
     PropertyDirectRelationLimitWarning,
     PropertyMultipleValueWarning,
     PropertyTypeNotSupportedWarning,
@@ -39,11 +45,11 @@ from cognite.neat._shared import InstanceType
 from cognite.neat._store import NeatGraphStore
 from cognite.neat._utils.auxiliary import create_sha256_hash
 from cognite.neat._utils.collection_ import iterate_progress_bar_if_above_config_threshold
-from cognite.neat._utils.rdf_ import remove_namespace_from_uri
-from cognite.neat._utils.text import humanize_collection
+from cognite.neat._utils.rdf_ import namespace_as_space, remove_namespace_from_uri, split_uri
+from cognite.neat._utils.text import NamingStandardization, humanize_collection
 from cognite.neat._utils.upload import UploadResult
 
-from ._base import _END_OF_CLASS, CDFLoader
+from ._base import _END_OF_CLASS, _START_OF_CLASS, CDFLoader
 
 
 @dataclass
@@ -89,6 +95,9 @@ class DMSLoader(CDFLoader[dm.InstanceApply]):
         client (NeatClient | None): This is used to lookup containers such that the loader
             creates instances in accordance with required constraints. Defaults to None.
         unquote_external_ids (bool): If True, the loader will unquote external ids before creating the instances.
+        neat_prefix_by_predicate_uri (dict[URIRef, str] | None): A dictionary that maps a predicate URIRef to a
+            prefix that Neat added to the object upon extraction. This is used to remove the prefix from the
+            object before creating the instance.
     """
 
     def __init__(
@@ -98,16 +107,23 @@ class DMSLoader(CDFLoader[dm.InstanceApply]):
         graph_store: NeatGraphStore,
         instance_space: str,
         space_property: str | None = None,
+        use_source_space: bool = False,
         client: NeatClient | None = None,
         create_issues: Sequence[NeatIssue] | None = None,
         unquote_external_ids: bool = False,
+        neat_prefix_by_predicate_uri: dict[URIRef, str] | None = None,
+        neat_prefix_by_type_uri: dict[URIRef, str] | None = None,
     ):
         super().__init__(graph_store)
         self.dms_rules = dms_rules
         self.info_rules = info_rules
+        self.neat_prefix_by_predicate_uri = neat_prefix_by_predicate_uri or {}
+        self.neat_prefix_by_type_uri = neat_prefix_by_type_uri or {}
         self._instance_space = instance_space
         self._space_property = space_property
-        self._space_by_uri: dict[str, str] = defaultdict(lambda: instance_space)
+        self._use_source_space = use_source_space
+        self._space_by_instance_uri: dict[URIRef, str] = defaultdict(lambda: instance_space)
+        self._external_id_by_uri: dict[URIRef, str] = {}
         self._issues = IssueList(create_issues or [])
         self._client = client
         self._unquote_external_ids = unquote_external_ids
@@ -132,7 +148,9 @@ class DMSLoader(CDFLoader[dm.InstanceApply]):
             else:
                 yaml.safe_dump(dumped, f, sort_keys=False)
 
-    def _load(self, stop_on_exception: bool = False) -> Iterable[dm.InstanceApply | NeatIssue | type[_END_OF_CLASS]]:
+    def _load(
+        self, stop_on_exception: bool = False
+    ) -> Iterable[dm.InstanceApply | NeatIssue | type[_END_OF_CLASS] | _START_OF_CLASS]:
         if self._issues.has_errors and stop_on_exception:
             raise self._issues.as_exception()
         elif self._issues.has_errors:
@@ -143,6 +161,16 @@ class DMSLoader(CDFLoader[dm.InstanceApply]):
         if self._space_property:
             yield from self._lookup_space_by_uri(view_iterations, stop_on_exception)
 
+            if self._client:
+                space_creation = self._create_instance_space_if_not_exists()
+                yield from space_creation.warnings
+                if space_creation.has_errors and stop_on_exception:
+                    raise space_creation.as_exception()
+                yield from space_creation.errors
+
+        if self.neat_prefix_by_type_uri:
+            self._lookup_identifier_by_uri()
+
         for it in view_iterations:
             view = it.view
             if view is None:
@@ -151,10 +179,15 @@ class DMSLoader(CDFLoader[dm.InstanceApply]):
             projection, issues = self._create_projection(view)
             yield from issues
             query = it.query
-            reader = self.graph_store.read(query.rdf_type, property_renaming_config=query.property_renaming_config)
+            reader = self.graph_store.read(
+                query.rdf_type,
+                property_renaming_config=query.property_renaming_config,
+                remove_uri_namespace=False,
+            )
             instance_iterable = iterate_progress_bar_if_above_config_threshold(
                 reader, it.instance_count, f"Loading {it.view_id!r}"
             )
+            yield _START_OF_CLASS(view.external_id)
             for identifier, properties in instance_iterable:
                 yield from self._create_instances(
                     identifier, properties, projection, stop_on_exception, exclude=it.hierarchical_properties
@@ -162,7 +195,10 @@ class DMSLoader(CDFLoader[dm.InstanceApply]):
             if it.hierarchical_properties:
                 # Force the creation of instances, before we create the hierarchical properties.
                 yield _END_OF_CLASS
+                yield _START_OF_CLASS(f"{view.external_id} ({humanize_collection(it.hierarchical_properties)})")
                 yield from self._create_hierarchical_properties(it, projection, stop_on_exception)
+            if reader is not instance_iterable:
+                print(f"Loaded {it.instance_count} instances for {it.view_id!r}")
 
             yield _END_OF_CLASS
 
@@ -188,6 +224,12 @@ class DMSLoader(CDFLoader[dm.InstanceApply]):
             views = self._client.data_modeling.views.retrieve(
                 list(iterations_by_view_id.keys()), include_inherited_properties=True
             )
+            if missing := set(iterations_by_view_id) - {view.as_id() for view in views}:
+                for missing_view in missing:
+                    issues.append(ResourceNotFoundError(missing_view, "view", more="The view is not found in CDF."))
+                return [], issues
+            # Todo: Remove if this turns out to be unnecessary.
+            hierarchical_properties_by_view_id: dict[dm.ViewId, set[str]] = {}
         else:
             views = dm.ViewList([])
             with catch_issues() as issues:
@@ -195,15 +237,17 @@ class DMSLoader(CDFLoader[dm.InstanceApply]):
                 views.extend(read_model.views)
             if issues.has_errors:
                 return [], issues
+            hierarchical_properties_by_view_id = {}
         views_by_id = {view.as_id(): view for view in views}
-        hierarchical_properties_by_view_id = SchemaAPI.get_hierarchical_properties(views)
 
         def sort_by_instance_type(id_: dm.ViewId) -> int:
             if id_ not in views_by_id:
                 return 0
-            return {"node": 1, "all": 2, "edge": 3}.get(views_by_id[id_].used_for, 0)
+            return {"node": 1, "all": 1, "edge": 3}.get(views_by_id[id_].used_for, 0)
 
-        ordered_view_ids = sorted(iterations_by_view_id.keys(), key=sort_by_instance_type)
+        ordered_view_ids = SchemaAPI.get_view_order_by_direct_relation_constraints(views)
+        # Sort is stable in Python, so we will keep the order of the views:
+        ordered_view_ids.sort(key=sort_by_instance_type)
         view_iterations: list[_ViewIterator] = []
         for view_id in ordered_view_ids:
             if view_id not in iterations_by_view_id:
@@ -245,8 +289,59 @@ class DMSLoader(CDFLoader[dm.InstanceApply]):
         instance_iterable = iterate_progress_bar_if_above_config_threshold(
             instance_iterable, total, f"Looking up spaces for {total} instances..."
         )
+        neat_prefix = self.neat_prefix_by_predicate_uri.get(space_property_uri)
+        warned_spaces: set[str] = set()
         for instance, space in instance_iterable:
-            self._space_by_uri[remove_namespace_from_uri(instance)] = space
+            if neat_prefix:
+                space = space.removeprefix(neat_prefix)
+
+            clean_space = NamingStandardization.standardize_space_str(space)
+            if clean_space != space and space not in warned_spaces:
+                issues.append(
+                    NeatValueWarning(
+                        f"Invalid space in property {self._space_property}: {space}. Fixed to {clean_space}"
+                    )
+                )
+                warned_spaces.add(space)
+
+            self._space_by_instance_uri[instance] = clean_space
+        return issues
+
+    def _lookup_identifier_by_uri(self) -> None:
+        if not self.neat_prefix_by_type_uri:
+            return
+
+        count = sum(count for _, count in self.graph_store.queries.summarize_instances())
+        instance_iterable = self.graph_store.queries.list_instances_ids()
+        instance_iterable = iterate_progress_bar_if_above_config_threshold(
+            instance_iterable, count, f"Looking up identifiers for {count} instances..."
+        )
+        count_by_identifier: dict[str, list[URIRef]] = defaultdict(list)
+        for instance_uri, type in instance_iterable:
+            if type not in self.neat_prefix_by_type_uri:
+                continue
+            prefix = self.neat_prefix_by_type_uri[type]
+            identifier = remove_namespace_from_uri(instance_uri)
+            if self._unquote_external_ids:
+                identifier = urllib.parse.unquote(identifier)
+            count_by_identifier[identifier.removeprefix(prefix)].append(instance_uri)
+
+        for identifier, uris in count_by_identifier.items():
+            if len(uris) == 1:
+                self._external_id_by_uri[uris[0]] = identifier
+
+    def _create_instance_space_if_not_exists(self) -> IssueList:
+        issues = IssueList()
+        if not self._client:
+            return issues
+
+        instance_spaces = set(self._space_by_instance_uri.values()) - {self._instance_space}
+        existing_spaces = {space.space for space in self._client.data_modeling.spaces.retrieve(list(instance_spaces))}
+        if missing_spaces := (instance_spaces - existing_spaces):
+            try:
+                self._client.data_modeling.spaces.apply([dm.SpaceApply(space=space) for space in missing_spaces])
+            except CogniteAPIError as e:
+                issues.append(AuthorizationError(f"Creating {len(missing_spaces)} instance spaces.", str(e)))
         return issues
 
     def _create_projection(self, view: dm.View) -> tuple[_Projection, IssueList]:
@@ -346,8 +441,8 @@ class DMSLoader(CDFLoader[dm.InstanceApply]):
             def parse_direct_relation(cls, value: list, info: ValidationInfo) -> dict | list[dict]:
                 # We validate above that we only get one value for single direct relations.
                 if list.__name__ in _get_field_value_types(cls, info):
-                    external_ids = (remove_namespace_from_uri(v) for v in value)
-                    result = [{"space": self._space_by_uri[e], "externalId": e} for e in external_ids]
+                    ids = (self._create_instance_id(v, "node", stop_on_exception=True) for v in value)
+                    result = [id_.dump(camel_case=True, include_instance_type=False) for id_ in ids]
                     # Todo: Account for max_list_limit
                     if len(result) <= DMS_DIRECT_RELATION_LIST_LIMIT:
                         return result
@@ -364,8 +459,9 @@ class DMSLoader(CDFLoader[dm.InstanceApply]):
                     result.sort(key=lambda x: (x["space"], x["externalId"]))
                     return result[:DMS_DIRECT_RELATION_LIST_LIMIT]
                 elif value:
-                    external_id = remove_namespace_from_uri(value[0])
-                    return {"space": self._space_by_uri[external_id], "externalId": external_id}
+                    return self._create_instance_id(value[0], "node", stop_on_exception=True).dump(
+                        camel_case=True, include_instance_type=False
+                    )
                 return {}
 
             validators["parse_direct_relation"] = field_validator(*direct_relation_by_property.keys(), mode="before")(  # type: ignore[assignment]
@@ -376,35 +472,50 @@ class DMSLoader(CDFLoader[dm.InstanceApply]):
 
             def parse_direct_relation_to_unit(cls, value: Any, info: ValidationInfo) -> dict | list[dict]:
                 if value:
-                    return {"space": "cdf_cdm_units", "externalId": remove_namespace_from_uri(value[0])}
+                    external_id = remove_namespace_from_uri(value[0])
+                    if self._unquote_external_ids:
+                        external_id = urllib.parse.unquote(external_id)
+                    return {"space": "cdf_cdm_units", "externalId": external_id}
                 return {}
 
             validators["parse_direct_relation_to_unit"] = field_validator(*unit_properties, mode="before")(  # type: ignore[assignment]
                 parse_direct_relation_to_unit  # type: ignore[arg-type]
             )
 
+        if text_fields:
+
+            def parse_text(cls, value: Any, info: ValidationInfo) -> str | list[str]:
+                if isinstance(value, list):
+                    return [remove_namespace_from_uri(v) if isinstance(v, URIRef) else str(v) for v in value]
+                return remove_namespace_from_uri(value) if isinstance(value, URIRef) else str(value)
+
+            validators["parse_text"] = field_validator(*text_fields, mode="before")(parse_text)  # type: ignore[assignment, arg-type]
+
         pydantic_cls = create_model(view.external_id, __validators__=validators, **field_definitions)  # type: ignore[arg-type, call-overload]
         return _Projection(view.as_id(), view.used_for, pydantic_cls, edge_by_type, edge_by_prop_id), issues
 
     def _create_instances(
         self,
-        identifier: str,
-        properties: dict[str | InstanceType, list[str]],
+        instance_uri: URIRef,
+        properties: dict[str | InstanceType, list[Any]],
         projection: _Projection,
-        stop_on_exception: bool = False,
+        stop_on_exception: Literal[True, False] = False,
         exclude: set[str] | None = None,
         include: set[str] | None = None,
     ) -> Iterable[dm.InstanceApply | NeatIssue]:
-        if self._unquote_external_ids:
-            identifier = urllib.parse.unquote(identifier)
+        instance_id = self._create_instance_id(instance_uri, "node", stop_on_exception)
+        if not isinstance(instance_id, InstanceId):
+            yield instance_id
+            return
+        space, external_id = instance_id.space, instance_id.external_id
         start_node, end_node = self._pop_start_end_node(properties)
         is_edge = start_node and end_node
         instance_type = "edge" if is_edge else "node"
         if (projection.used_for == "node" and is_edge) or (projection.used_for == "edge" and not is_edge):
             creation_error = ResourceCreationError(
-                identifier,
+                external_id,
                 instance_type,
-                f"View used for {projection.used_for} instance {identifier!s} but is {instance_type}",
+                f"View used for {projection.used_for} instance {external_id!s} but is {instance_type}",
             )
             if stop_on_exception:
                 raise creation_error from None
@@ -412,22 +523,18 @@ class DMSLoader(CDFLoader[dm.InstanceApply]):
             return
 
         if RDF.type not in properties:
-            error = ResourceCreationError(identifier, instance_type, "No rdf:type found")
+            error = ResourceCreationError(external_id, instance_type, "No rdf:type found")
             if stop_on_exception:
                 raise error from None
             yield error
             return
         _ = properties.pop(RDF.type)[0]
-        if start_node and self._unquote_external_ids:
-            start_node = urllib.parse.unquote(start_node)
-        if end_node and self._unquote_external_ids:
-            end_node = urllib.parse.unquote(end_node)
-
         if exclude:
             properties = {k: v for k, v in properties.items() if k not in exclude}
         if include:
             properties = {k: v for k, v in properties.items() if k in include}
 
+        sources = []
         with catch_issues() as property_issues:
             sources = [
                 dm.NodeOrEdgeData(
@@ -437,32 +544,46 @@ class DMSLoader(CDFLoader[dm.InstanceApply]):
             ]
         for issue in property_issues:
             if isinstance(issue, ResourceNeatWarning):
-                issue.identifier = identifier
+                issue.identifier = external_id
 
         if property_issues.has_errors and stop_on_exception:
             raise property_issues.as_exception()
         yield from property_issues
+        if not sources:
+            return
 
         if start_node and end_node:
-            yield dm.EdgeApply(
-                space=self._space_by_uri[identifier],
-                external_id=identifier,
-                type=(projection.view_id.space, projection.view_id.external_id),
-                start_node=(self._space_by_uri[start_node], start_node),
-                end_node=(self._space_by_uri[end_node], end_node),
-                sources=sources,
-            )
+            start = self._create_instance_id(start_node, "edge", stop_on_exception)
+            end = self._create_instance_id(end_node, "edge", stop_on_exception)
+            if isinstance(start, NeatError):
+                yield start
+            if isinstance(end, NeatError):
+                yield end
+            if isinstance(start, InstanceId) and isinstance(end, InstanceId):
+                yield dm.EdgeApply(
+                    space=space,
+                    external_id=external_id,
+                    type=(projection.view_id.space, projection.view_id.external_id),
+                    start_node=start.as_tuple(),
+                    end_node=end.as_tuple(),
+                    sources=sources,
+                )
         else:
             yield dm.NodeApply(
-                space=self._space_by_uri[identifier],
-                external_id=identifier,
+                space=space,
+                external_id=external_id,
                 type=(projection.view_id.space, projection.view_id.external_id),
                 sources=sources,
             )
-        yield from self._create_edges_without_properties(identifier, properties, projection)
+        yield from self._create_edges_without_properties(space, external_id, properties, projection, stop_on_exception)
 
     def _create_edges_without_properties(
-        self, identifier: str, properties: dict[str | InstanceType, list[str]], projection: _Projection
+        self,
+        space: str,
+        identifier: str,
+        properties: dict[str | InstanceType, list[str] | list[URIRef]],
+        projection: _Projection,
+        stop_on_exception: Literal[True, False],
     ) -> Iterable[dm.EdgeApply | NeatIssue]:
         for predicate, values in properties.items():
             if predicate in projection.edge_by_type:
@@ -480,15 +601,22 @@ class DMSLoader(CDFLoader[dm.InstanceApply]):
                 yield error
                 continue
             for target in values:
+                target_id = self._create_instance_id(target, "edge", stop_on_exception)  # type: ignore[call-overload]
+                if not isinstance(target_id, InstanceId):
+                    yield target_id
+                    continue
+                if isinstance(target, URIRef):
+                    target = remove_namespace_from_uri(target)
                 external_id = f"{identifier}.{prop_id}.{target}"
+
                 start_node, end_node = (
-                    (self._space_by_uri[identifier], identifier),
-                    (self._space_by_uri[target], target),
+                    (space, identifier),
+                    target_id.as_tuple(),
                 )
                 if edge.direction == "inwards":
                     start_node, end_node = end_node, start_node
                 yield dm.EdgeApply(
-                    space=self._space_by_uri[identifier],
+                    space=space,
                     external_id=(external_id if len(external_id) < 256 else create_sha256_hash(external_id)),
                     type=edge.type,
                     start_node=start_node,
@@ -496,7 +624,9 @@ class DMSLoader(CDFLoader[dm.InstanceApply]):
                 )
 
     @staticmethod
-    def _pop_start_end_node(properties: dict[str | InstanceType, list[str]]) -> tuple[str, str] | tuple[None, None]:
+    def _pop_start_end_node(
+        properties: dict[str | InstanceType, list[str] | list[URIRef]],
+    ) -> tuple[URIRef, URIRef] | tuple[None, None]:
         start_node = properties.pop("startNode", [None])[0]
         if not start_node:
             start_node = properties.pop("start_node", [None])[0]
@@ -504,9 +634,47 @@ class DMSLoader(CDFLoader[dm.InstanceApply]):
         if not end_node:
             end_node = properties.pop("end_node", [None])[0]
         if start_node and end_node:
-            return start_node, end_node
+            return start_node, end_node  # type: ignore[return-value]
         return None, None
 
+    @overload
+    def _create_instance_id(
+        self, uri: URIRef, instance_type: str, stop_on_exception: Literal[False] = False
+    ) -> InstanceId | NeatError: ...
+
+    @overload
+    def _create_instance_id(
+        self, uri: URIRef, instance_type: str, stop_on_exception: Literal[True] = True
+    ) -> InstanceId: ...
+
+    def _create_instance_id(
+        self, uri: URIRef, instance_type: str, stop_on_exception: bool = False
+    ) -> InstanceId | NeatError:
+        space: str | None = None
+        external_id: str | None = None
+        error: NeatError | None = None
+        if self._use_source_space:
+            namespace, external_id = split_uri(uri)
+            space = namespace_as_space(namespace)
+            if space is None:
+                error = ResourceCreationError(uri, instance_type, f"Could not find space for {uri!s}.")
+        else:
+            space = self._space_by_instance_uri[uri]
+            if uri in self._external_id_by_uri:
+                external_id = self._external_id_by_uri[uri]
+            else:
+                external_id = remove_namespace_from_uri(uri)
+
+        if external_id and self._unquote_external_ids:
+            external_id = urllib.parse.unquote(external_id)
+        if space and external_id:
+            return InstanceId(space, external_id)
+        if error is None:
+            raise ValueError(f"Bug in neat. Failed to create instance ID and determine error for {uri!r}")
+        if stop_on_exception:
+            raise error
+        return error
+
     def _get_required_capabilities(self) -> list[Capability]:
         return [
             DataModelInstancesAcl(
@@ -525,24 +693,11 @@ class DMSLoader(CDFLoader[dm.InstanceApply]):
         items: list[dm.InstanceApply],
         dry_run: bool,
         read_issues: IssueList,
+        class_name: str | None = None,
     ) -> Iterable[UploadResult]:
-        nodes: list[dm.NodeApply] = []
-        edges: list[dm.EdgeApply] = []
-        source_by_node_id: dict[dm.NodeId, str] = {}
-        source_by_edge_id: dict[dm.EdgeId, str] = {}
-        for item in items:
-            if isinstance(item, dm.NodeApply):
-                nodes.append(item)
-                if item.sources:
-                    source_by_node_id[item.as_id()] = item.sources[0].source.external_id
-                else:
-                    source_by_node_id[item.as_id()] = "node"
-            elif isinstance(item, dm.EdgeApply):
-                edges.append(item)
-                if item.sources:
-                    source_by_edge_id[item.as_id()] = item.sources[0].source.external_id
-                else:
-                    source_by_edge_id[item.as_id()] = "edge"
+        name = class_name or "Instances"
+        nodes = [item for item in items if isinstance(item, dm.NodeApply)]
+        edges = [item for item in items if isinstance(item, dm.EdgeApply)]
         try:
             upserted = client.data_modeling.instances.apply(
                 nodes,
@@ -552,29 +707,28 @@ class DMSLoader(CDFLoader[dm.InstanceApply]):
                 skip_on_version_conflict=True,
             )
         except CogniteAPIError as e:
-            result = UploadResult[InstanceId](name="Instances", issues=read_issues)
-            result.error_messages.append(str(e))
-            result.failed_upserted.update(item.as_id() for item in e.failed + e.unknown)
-            result.created.update(item.as_id() for item in e.successful)
-            yield result
+            if len(items) == 1:
+                yield UploadResult(
+                    name=name,
+                    issues=read_issues,
+                    failed_items=items,
+                    error_messages=[str(e)],
+                    failed_upserted={item.as_id() for item in items},  # type: ignore[attr-defined]
+                )
+            else:
+                half = len(items) // 2
+                yield from self._upload_to_cdf(client, items[:half], dry_run, read_issues, class_name)
+                yield from self._upload_to_cdf(client, items[half:], dry_run, read_issues, class_name)
         else:
-            for instances, ids_by_source in [
-                (upserted.nodes, source_by_node_id),
-                (upserted.edges, source_by_edge_id),
-            ]:
-                for name, subinstances in itertools.groupby(
-                    sorted(instances, key=lambda i: ids_by_source.get(i.as_id(), "")),  # type: ignore[call-overload, index, attr-defined]
-                    key=lambda i: ids_by_source.get(i.as_id(), ""),  # type: ignore[index, attr-defined]
-                ):
-                    result = UploadResult(name=name, issues=read_issues)
-                    for instance in subinstances:  # type: ignore[attr-defined]
-                        if instance.was_modified and instance.created_time == instance.last_updated_time:
-                            result.created.add(instance.as_id())
-                        elif instance.was_modified:
-                            result.changed.add(instance.as_id())
-                        else:
-                            result.unchanged.add(instance.as_id())
-                    yield result
+            result = UploadResult(name=name, issues=read_issues)  # type: ignore[var-annotated]
+            for instance in itertools.chain(upserted.nodes, upserted.edges):  # type: ignore[attr-defined]
+                if instance.was_modified and instance.created_time == instance.last_updated_time:
+                    result.created.add(instance.as_id())
+                elif instance.was_modified:
+                    result.changed.add(instance.as_id())
+                else:
+                    result.unchanged.add(instance.as_id())
+            yield result
 
 
 def _get_field_value_types(cls, info):