infrahub-server 1.1.4__py3-none-any.whl → 1.1.6__py3-none-any.whl

infrahub/api/artifact.py

@@ -85,7 +85,10 @@ async def generate_artifact(
 
     service = request.app.state.service
     model = RequestArtifactDefinitionGenerate(
-        artifact_definition=artifact_definition.id, branch=branch_params.branch.name, limit=payload.nodes
+        artifact_definition_id=artifact_definition.id,
+        artifact_definition_name=artifact_definition.name.value,
+        branch=branch_params.branch.name,
+        limit=payload.nodes,
     )
 
     await service.workflow.submit_workflow(workflow=REQUEST_ARTIFACT_DEFINITION_GENERATE, parameters={"model": model})

infrahub/api/oidc.py

@@ -1,8 +1,9 @@
 from __future__ import annotations
 
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Any
 from urllib.parse import urljoin
 
+import jwt
 import ujson
 from authlib.integrations.httpx_client import AsyncOAuth2Client
 from fastapi import APIRouter, Depends, Request, Response
@@ -138,7 +139,7 @@ async def token(
 
     with trace.get_tracer(__name__).start_as_current_span("sso_token_request") as span:
         span.set_attribute("token_request_data", ujson.dumps(token_response.json()))
-        payload = token_response.json()
+        payload: dict[str, Any] = token_response.json()
 
     headers = {"Authorization": f"{payload.get('token_type')} {payload.get('access_token')}"}
 
@@ -148,8 +149,10 @@ async def token(
         userinfo_response = await service.http.post(str(oidc_config.userinfo_endpoint), headers=headers)
 
     _validate_response(response=userinfo_response)
-    user_info = userinfo_response.json()
-    sso_groups = user_info.get("groups", [])
+    user_info: dict[str, Any] = userinfo_response.json()
+    sso_groups = user_info.get("groups") or await _get_id_token_groups(
+        oidc_config=oidc_config, service=service, payload=payload, client_id=provider.client_id
+    )
 
     if not sso_groups and config.SETTINGS.security.sso_user_default_group:
         sso_groups = [config.SETTINGS.security.sso_user_default_group]
@@ -185,3 +188,29 @@ def _validate_response(response: httpx.Response) -> None:
         body=response.json(),
     )
     raise GatewayError(message="Invalid response from Authentication provider")
+
+
+async def _get_id_token_groups(
+    oidc_config: OIDCDiscoveryConfig, service: InfrahubServices, payload: dict[str, Any], client_id: str
+) -> list[str]:
+    id_token = payload.get("id_token")
+    if not id_token:
+        return []
+    jwks = await service.http.get(url=str(oidc_config.jwks_uri))
+
+    jwk_client = jwt.PyJWKClient(uri=str(oidc_config.jwks_uri), cache_jwk_set=True)
+    if jwk_client.jwk_set_cache:
+        jwk_client.jwk_set_cache.put(jwks.json())
+
+    signing_key = jwk_client.get_signing_key_from_jwt(id_token)
+
+    decoded_token: dict[str, Any] = jwt.decode(
+        jwt=id_token,
+        key=signing_key.key,
+        algorithms=oidc_config.id_token_signing_alg_values_supported,
+        audience=client_id,
+        issuer=str(oidc_config.issuer),
+        options={"verify_signature": False, "verify_aud": False, "verify_iss": False},
+    )
+
+    return decoded_token.get("groups", [])

infrahub/api/transformation.py

@@ -83,6 +83,7 @@ async def transform_python(
         commit=repository.commit.value,  # type: ignore[attr-defined]
         branch=branch_params.branch.name,
         transform_location=f"{transform.file_path.value}::{transform.class_name.value}",
+        timeout=transform.timeout.value,
         data=data,
     )
 
@@ -140,6 +141,7 @@ async def transform_jinja2(
         commit=repository.commit.value,  # type: ignore[attr-defined]
         branch=branch_params.branch.name,
         template_location=transform.template_path.value,
+        timeout=transform.timeout.value,
         data=data,
     )
 

infrahub/computed_attribute/tasks.py

@@ -4,7 +4,10 @@ from datetime import timedelta
 from typing import TYPE_CHECKING, Any
 
 import ujson
-from infrahub_sdk.protocols import CoreNode  # noqa: TC002
+from infrahub_sdk.protocols import (
+    CoreNode,  # noqa: TC002
+    CoreTransformPython,
+)
 from prefect import flow
 from prefect.automations import AutomationCore
 from prefect.client.orchestration import get_client
@@ -88,17 +91,18 @@ async def process_transform(
 
     for attribute_name, transform_attribute in transform_attributes.items():
         transform = await service.client.get(
-            kind="CoreTransformPython",
+            kind=CoreTransformPython,
             branch=branch_name,
             id=transform_attribute.transform,
             prefetch_relationships=True,
             populate_store=True,
         )
+
         if not transform:
             continue
 
         repo_node = await service.client.get(
-            kind=transform.repository.peer.typename,
+            kind=str(transform.repository.peer.typename),
             branch=branch_name,
             id=transform.repository.peer.id,
             raise_when_missing=True,
@@ -108,7 +112,7 @@ async def process_transform(
             repository_id=transform.repository.peer.id,
             name=transform.repository.peer.name.value,
             service=service,
-            repository_kind=transform.repository.peer.typename,
+            repository_kind=str(transform.repository.peer.typename),
             commit=repo_node.commit.value,
         )
 
@@ -120,7 +124,7 @@ async def process_transform(
             subscribers=[object_id],
         )
 
-        transformed_data = await repo.execute_python_transform(
+        transformed_data = await repo.execute_python_transform.with_options(timeout_seconds=transform.timeout.value)(
             branch_name=branch_name,
             commit=repo_node.commit.value,
             location=f"{transform.file_path.value}::{transform.class_name.value}",

infrahub/core/diff/calculator.py

@@ -1,20 +1,73 @@
+from dataclasses import dataclass, field
+
+from infrahub import config
 from infrahub.core import registry
 from infrahub.core.branch import Branch
 from infrahub.core.diff.query_parser import DiffQueryParser
-from infrahub.core.query.diff import DiffAllPathsQuery
+from infrahub.core.query.diff import (
+    DiffCalculationQuery,
+    DiffFieldPathsQuery,
+    DiffNodePathsQuery,
+    DiffPropertyPathsQuery,
+)
 from infrahub.core.timestamp import Timestamp
 from infrahub.database import InfrahubDatabase
 from infrahub.log import get_logger
 
-from .model.path import CalculatedDiffs, NodeFieldSpecifier
+from .model.path import CalculatedDiffs
 
 log = get_logger()
 
 
+@dataclass
+class DiffCalculationRequest:
+    base_branch: Branch
+    diff_branch: Branch
+    branch_from_time: Timestamp
+    from_time: Timestamp
+    to_time: Timestamp
+    current_node_field_specifiers: dict[str, set[str]] | None = field(default=None)
+    new_node_field_specifiers: dict[str, set[str]] | None = field(default=None)
+
+
 class DiffCalculator:
     def __init__(self, db: InfrahubDatabase) -> None:
         self.db = db
 
+    async def _run_diff_calculation_query(
+        self,
+        diff_parser: DiffQueryParser,
+        query_class: type[DiffCalculationQuery],
+        calculation_request: DiffCalculationRequest,
+        limit: int,
+    ) -> None:
+        has_more_data = True
+        offset = 0
+        while has_more_data:
+            diff_query = await query_class.init(
+                db=self.db,
+                branch=calculation_request.diff_branch,
+                base_branch=calculation_request.base_branch,
+                diff_branch_from_time=calculation_request.branch_from_time,
+                diff_from=calculation_request.from_time,
+                diff_to=calculation_request.to_time,
+                current_node_field_specifiers=calculation_request.current_node_field_specifiers,
+                new_node_field_specifiers=calculation_request.new_node_field_specifiers,
+                limit=limit,
+                offset=offset,
+            )
+            log.info(f"Beginning one diff calculation query {limit=}, {offset=}")
+            await diff_query.execute(db=self.db)
+            log.info("Diff calculation query complete")
+            last_result = None
+            for query_result in diff_query.get_results():
+                diff_parser.read_result(query_result=query_result)
+                last_result = query_result
+            has_more_data = False
+            if last_result:
+                has_more_data = last_result.get_as_type("has_more_data", bool)
+            offset += limit
+
     async def calculate_diff(
         self,
         base_branch: Branch,
@@ -22,7 +75,7 @@ class DiffCalculator:
         from_time: Timestamp,
         to_time: Timestamp,
         include_unchanged: bool = True,
-        previous_node_specifiers: set[NodeFieldSpecifier] | None = None,
+        previous_node_specifiers: dict[str, set[str]] | None = None,
     ) -> CalculatedDiffs:
         if diff_branch.name == registry.default_branch:
             diff_branch_from_time = from_time
@@ -36,46 +89,85 @@ class DiffCalculator:
             to_time=to_time,
             previous_node_field_specifiers=previous_node_specifiers,
         )
-        branch_diff_query = await DiffAllPathsQuery.init(
-            db=self.db,
-            branch=diff_branch,
+        node_limit = int(config.SETTINGS.database.query_size_limit / 10)
+        fields_limit = int(config.SETTINGS.database.query_size_limit / 3)
+        properties_limit = int(config.SETTINGS.database.query_size_limit)
+
+        calculation_request = DiffCalculationRequest(
             base_branch=base_branch,
-            diff_branch_from_time=diff_branch_from_time,
-            diff_from=from_time,
-            diff_to=to_time,
+            diff_branch=diff_branch,
+            branch_from_time=diff_branch_from_time,
+            from_time=from_time,
+            to_time=to_time,
         )
-        log.info("Beginning diff calculation query for branch")
-        await branch_diff_query.execute(db=self.db)
-        log.info("Diff calculation query for branch complete")
-        log.info("Reading results of query for branch")
-        for query_result in branch_diff_query.get_results():
-            diff_parser.read_result(query_result=query_result)
-        log.info("Results of query for branch read")
+
+        log.info("Beginning diff node-level calculation queries for branch")
+        await self._run_diff_calculation_query(
+            diff_parser=diff_parser,
+            query_class=DiffNodePathsQuery,
+            calculation_request=calculation_request,
+            limit=node_limit,
+        )
+        log.info("Diff node-level calculation queries for branch complete")
+
+        log.info("Beginning diff field-level calculation queries for branch")
+        await self._run_diff_calculation_query(
+            diff_parser=diff_parser,
+            query_class=DiffFieldPathsQuery,
+            calculation_request=calculation_request,
+            limit=fields_limit,
+        )
+        log.info("Diff field-level calculation queries for branch complete")
+
+        log.info("Beginning diff property-level calculation queries for branch")
+        await self._run_diff_calculation_query(
+            diff_parser=diff_parser,
+            query_class=DiffPropertyPathsQuery,
+            calculation_request=calculation_request,
+            limit=properties_limit,
+        )
+        log.info("Diff property-level calculation queries for branch complete")
 
         if base_branch.name != diff_branch.name:
-            new_node_field_specifiers = diff_parser.get_new_node_field_specifiers()
             current_node_field_specifiers = diff_parser.get_current_node_field_specifiers()
-
-            base_diff_query = await DiffAllPathsQuery.init(
-                db=self.db,
-                branch=base_branch,
+            new_node_field_specifiers = diff_parser.get_new_node_field_specifiers()
+            calculation_request = DiffCalculationRequest(
                 base_branch=base_branch,
-                diff_branch_from_time=diff_branch_from_time,
-                diff_from=from_time,
-                diff_to=to_time,
-                current_node_field_specifiers=[
-                    (nfs.node_uuid, nfs.field_name) for nfs in current_node_field_specifiers
-                ],
-                new_node_field_specifiers=[(nfs.node_uuid, nfs.field_name) for nfs in new_node_field_specifiers],
+                diff_branch=base_branch,
+                branch_from_time=diff_branch_from_time,
+                from_time=from_time,
+                to_time=to_time,
+                current_node_field_specifiers=current_node_field_specifiers,
+                new_node_field_specifiers=new_node_field_specifiers,
             )
 
-            log.info("Beginning diff calculation query for base")
-            await base_diff_query.execute(db=self.db)
-            log.info("Diff calculation query for base complete")
-            log.info("Reading results of query for base")
-            for query_result in base_diff_query.get_results():
-                diff_parser.read_result(query_result=query_result)
-            log.info("Results of query for branch read")
+            log.info("Beginning diff node-level calculation queries for base")
+            await self._run_diff_calculation_query(
+                diff_parser=diff_parser,
+                query_class=DiffNodePathsQuery,
+                calculation_request=calculation_request,
+                limit=node_limit,
+            )
+            log.info("Diff node-level calculation queries for base complete")
+
+            log.info("Beginning diff field-level calculation queries for base")
+            await self._run_diff_calculation_query(
+                diff_parser=diff_parser,
+                query_class=DiffFieldPathsQuery,
+                calculation_request=calculation_request,
+                limit=fields_limit,
+            )
+            log.info("Diff field-level calculation queries for base complete")
+
+            log.info("Beginning diff property-level calculation queries for base")
+            await self._run_diff_calculation_query(
+                diff_parser=diff_parser,
+                query_class=DiffPropertyPathsQuery,
+                calculation_request=calculation_request,
+                limit=properties_limit,
+            )
+            log.info("Diff property-level calculation queries for base complete")
+
         log.info("Parsing calculated diff")
         diff_parser.parse(include_unchanged=include_unchanged)
         log.info("Calculated diff parsed")

infrahub/core/diff/coordinator.py

@@ -15,7 +15,6 @@ from .model.path import (
     EnrichedDiffs,
     EnrichedDiffsMetadata,
     NameTrackingId,
-    NodeFieldSpecifier,
     TrackingId,
 )
 
@@ -44,7 +43,7 @@ class EnrichedDiffRequest:
     from_time: Timestamp
     to_time: Timestamp
     tracking_id: TrackingId | None = field(default=None)
-    node_field_specifiers: set[NodeFieldSpecifier] = field(default_factory=set)
+    node_field_specifiers: dict[str, set[str]] = field(default_factory=dict)
 
     def __repr__(self) -> str:
         return (

infrahub/core/diff/ipam_diff_parser.py

@@ -116,7 +116,7 @@ class IpamDiffParser:
                 rels = await node_from_db.ip_namespace.get_relationships(db=self.db)  # type: ignore[attr-defined]
                 if rels:
                     cnd.namespace_id = rels[0].get_peer_id()
-            if cnd.ip_value and cnd.namespace_id:
+            if cnd.ip_value and cnd.namespace_id and cnd.node_uuid in uuids_missing_data:
                 uuids_missing_data.remove(cnd.node_uuid)
 
     async def _add_missing_values(

infrahub/core/diff/merger/serializer.py

@@ -294,10 +294,6 @@ class DiffMergeSerializer:
         actions_and_peers = self._get_actions_and_peers(relationship_diff=relationship_diff)
         added_peer_ids = [peer_id for action, peer_id in actions_and_peers if action is DiffAction.ADDED]
         removed_peer_ids = [peer_id for action, peer_id in actions_and_peers if action is DiffAction.REMOVED]
-        if not added_peer_ids:
-            added_peer_ids = [relationship_diff.peer_id]
-        if not removed_peer_ids:
-            removed_peer_ids = [relationship_diff.peer_id]
         for action, peer_id in actions_and_peers:
             if (
                 peer_id
@@ -314,6 +310,7 @@ class DiffMergeSerializer:
             relationship_diff_properties=relationship_diff.properties,
             added_peer_ids=added_peer_ids,
             removed_peer_ids=removed_peer_ids,
+            unchanged_peer_id=relationship_diff.peer_id,
         )
         return relationship_dicts, relationship_property_dicts
 
@@ -324,9 +321,14 @@ class DiffMergeSerializer:
         relationship_diff_properties: set[EnrichedDiffProperty],
         added_peer_ids: list[str],
         removed_peer_ids: list[str],
+        unchanged_peer_id: str,
     ) -> list[RelationshipPropertyMergeDict]:
-        added_property_dicts = self._get_default_property_merge_dicts(action=DiffAction.ADDED)
-        removed_property_dicts = self._get_default_property_merge_dicts(action=DiffAction.REMOVED)
+        added_property_dicts = {}
+        removed_property_dicts = {}
+        if added_peer_ids:
+            added_property_dicts = self._get_default_property_merge_dicts(action=DiffAction.ADDED)
+        if removed_peer_ids:
+            removed_property_dicts = self._get_default_property_merge_dicts(action=DiffAction.REMOVED)
         for property_diff in relationship_diff_properties:
             if property_diff.property_type is DatabaseEdgeType.IS_RELATED:
                 # handled above
@@ -348,10 +350,15 @@ class DiffMergeSerializer:
                 elif action is DiffAction.REMOVED:
                     removed_property_dicts[property_diff.property_type] = property_dict
         relationship_property_dicts = []
+        peers_and_property_dicts: list[tuple[str, dict[DatabaseEdgeType, PropertyMergeDict]]] = []
         if added_property_dicts:
-            peers_and_property_dicts = [(peer_id, added_property_dicts) for peer_id in added_peer_ids]
+            peers_and_property_dicts += [
+                (peer_id, added_property_dicts) for peer_id in (added_peer_ids or [unchanged_peer_id])
+            ]
         if removed_property_dicts:
-            peers_and_property_dicts += [(peer_id, removed_property_dicts) for peer_id in removed_peer_ids]
+            peers_and_property_dicts += [
+                (peer_id, removed_property_dicts) for peer_id in (removed_peer_ids or [unchanged_peer_id])
+            ]
         for peer_id, property_dicts in peers_and_property_dicts:
             if (
                 peer_id

infrahub/core/diff/model/path.py

@@ -82,15 +82,6 @@ def deserialize_tracking_id(tracking_id_str: str) -> TrackingId:
     raise ValueError(f"{tracking_id_str} is not a valid TrackingId")
 
 
-@dataclass
-class NodeFieldSpecifier:
-    node_uuid: str
-    field_name: str
-
-    def __hash__(self) -> int:
-        return hash(f"{self.node_uuid}:{self.field_name}")
-
-
 @dataclass
 class NodeDiffFieldSummary:
     kind: str

infrahub/core/diff/query/all_conflicts.py

@@ -0,0 +1,64 @@
+from typing import Any, Generator
+
+from neo4j.graph import Node as Neo4jNode
+
+from infrahub.core.query import Query, QueryType
+from infrahub.database import InfrahubDatabase
+
+from ..model.path import TrackingId
+
+
+class EnrichedDiffAllConflictsQuery(Query):
+    name = "enriched_diff_all_conflicts"
+    type = QueryType.READ
+
+    def __init__(
+        self, diff_branch_name: str, tracking_id: TrackingId | None = None, diff_id: str | None = None, **kwargs: Any
+    ) -> None:
+        super().__init__(**kwargs)
+        if (diff_id is None and tracking_id is None) or (diff_id and tracking_id):
+            raise ValueError("EnrichedDiffAllConflictsQuery requires one and only one of `tracking_id` or `diff_id`")
+        self.diff_branch_name = diff_branch_name
+        self.tracking_id = tracking_id
+        self.diff_id = diff_id
+        if self.tracking_id is None and self.diff_id is None:
+            raise RuntimeError("tracking_id or diff_id is required")
+
+    async def query_init(self, db: InfrahubDatabase, **kwargs: Any) -> None:
+        self.params = {
+            "diff_branch_name": self.diff_branch_name,
+            "diff_id": self.diff_id,
+            "tracking_id": self.tracking_id.serialize() if self.tracking_id else None,
+        }
+        query = """
+MATCH (root:DiffRoot)
+WHERE ($diff_id IS NOT NULL AND root.uuid = $diff_id)
+OR ($tracking_id IS NOT NULL AND root.tracking_id = $tracking_id AND root.diff_branch = $diff_branch_name)
+CALL {
+    WITH root
+    MATCH (root)-[:DIFF_HAS_NODE]->(node:DiffNode)-[:DIFF_HAS_CONFLICT]->(node_conflict:DiffConflict)
+    RETURN node.path_identifier AS path_identifier, node_conflict AS conflict
+    UNION
+    WITH root
+    MATCH (root)-[:DIFF_HAS_NODE]->(node:DiffNode)-[:DIFF_HAS_ATTRIBUTE]->(:DiffAttribute)
+        -[:DIFF_HAS_PROPERTY]->(property:DiffProperty)-[:DIFF_HAS_CONFLICT]->(attr_property_conflict:DiffConflict)
+    RETURN property.path_identifier AS path_identifier, attr_property_conflict AS conflict
+    UNION
+    WITH root
+    MATCH (root)-[:DIFF_HAS_NODE]->(node:DiffNode)-[:DIFF_HAS_RELATIONSHIP]->(:DiffRelationship)
+        -[:DIFF_HAS_ELEMENT]->(element:DiffRelationshipElement)-[:DIFF_HAS_CONFLICT]->(rel_element_conflict:DiffConflict)
+    RETURN element.path_identifier AS path_identifier, rel_element_conflict AS conflict
+    UNION
+    WITH root
+    MATCH (root)-[:DIFF_HAS_NODE]->(node:DiffNode)-[:DIFF_HAS_RELATIONSHIP]->(:DiffRelationship)
+        -[:DIFF_HAS_ELEMENT]->(:DiffRelationshipElement)-[:DIFF_HAS_PROPERTY]->(property:DiffProperty)
+        -[:DIFF_HAS_CONFLICT]->(rel_property_conflict:DiffConflict)
+    RETURN property.path_identifier AS path_identifier, rel_property_conflict AS conflict
+}
+"""
+        self.return_labels = ["path_identifier", "conflict"]
+        self.add_to_query(query=query)
+
+    def get_conflict_paths_and_nodes(self) -> Generator[tuple[str, Neo4jNode], None, None]:
+        for result in self.get_results():
+            yield (result.get_as_type("path_identifier", str), result.get_node("conflict"))

infrahub/core/diff/query/diff_get.py

@@ -17,8 +17,6 @@ QUERY_MATCH_NODES = """
     AND ($to_time IS NULL OR diff_root.to_time <= $to_time)
     AND ($tracking_id IS NULL OR diff_root.tracking_id = $tracking_id)
     AND ($diff_ids IS NULL OR diff_root.uuid IN $diff_ids)
-    WITH diff_root
-    ORDER BY diff_root.base_branch, diff_root.diff_branch, diff_root.from_time, diff_root.to_time
     // get all the nodes attached to the diffs
     OPTIONAL MATCH (diff_root)-[:DIFF_HAS_NODE]->(diff_node:DiffNode)
     """
@@ -79,26 +77,21 @@ class EnrichedDiffGetQuery(Query):
             self.add_to_query(query=query_filters)
 
         query_2 = """
-        // group by diff node uuid for pagination
-        WITH diff_node.uuid AS diff_node_uuid, diff_node.kind AS diff_node_kind, collect([diff_root, diff_node]) AS node_root_tuples
-        // order by kind and latest label for each diff_node uuid
-        CALL {
-            WITH node_root_tuples
-            UNWIND node_root_tuples AS nrt
-            WITH nrt[0] AS diff_root, nrt[1] AS diff_node
-            ORDER BY diff_root.from_time DESC
-            RETURN diff_node.label AS latest_node_label
-            LIMIT 1
-        }
-        WITH diff_node_kind, node_root_tuples, latest_node_label
-        ORDER BY diff_node_kind, latest_node_label
+        WITH diff_root, diff_node
+        ORDER BY diff_root.base_branch, diff_root.diff_branch, diff_root.from_time, diff_root.to_time, diff_node.uuid
+        // -------------------------------------
+        // Limit number of results
+        // -------------------------------------
         SKIP COALESCE($offset, 0)
         LIMIT $limit
-        UNWIND node_root_tuples AS nrt
-        WITH nrt[0] AS diff_root, nrt[1] AS diff_node
-        WITH diff_root, diff_node
+        // -------------------------------------
+        // Check if more data after this limited group
+        // -------------------------------------
+        WITH collect([diff_root, diff_node]) AS limited_results
+        WITH limited_results, size(limited_results) = $limit AS has_more_data
+        UNWIND limited_results AS one_result
+        WITH one_result[0] AS diff_root, one_result[1] AS diff_node, has_more_data
         // if depth limit, make sure not to exceed it when traversing linked nodes
-        WITH diff_root, diff_node
         // -------------------------------------
         // Retrieve Parents
         // -------------------------------------
@@ -109,12 +102,12 @@ class EnrichedDiffGetQuery(Query):
             ORDER BY size(nodes(parents_path)) DESC
             LIMIT 1
         }
-        WITH diff_root, diff_node, parents_path
+        WITH diff_root, diff_node, has_more_data, parents_path
         // -------------------------------------
         // Retrieve conflicts
         // -------------------------------------
         OPTIONAL MATCH (diff_node)-[:DIFF_HAS_CONFLICT]->(diff_node_conflict:DiffConflict)
-        WITH diff_root, diff_node, parents_path, diff_node_conflict
+        WITH diff_root, diff_node, has_more_data, parents_path, diff_node_conflict
         // -------------------------------------
         // Retrieve Attributes
         // -------------------------------------
@@ -128,7 +121,7 @@ class EnrichedDiffGetQuery(Query):
             RETURN diff_attribute, diff_attr_property, diff_attr_property_conflict
             ORDER BY diff_attribute.name, diff_attr_property.property_type
         }
-        WITH diff_root, diff_node, parents_path, diff_node_conflict, collect([diff_attribute, diff_attr_property, diff_attr_property_conflict]) as diff_attributes
+        WITH diff_root, diff_node, has_more_data, parents_path, diff_node_conflict, collect([diff_attribute, diff_attr_property, diff_attr_property_conflict]) as diff_attributes
         // -------------------------------------
         // Retrieve Relationships
         // -------------------------------------
@@ -150,6 +143,7 @@ class EnrichedDiffGetQuery(Query):
         WITH
             diff_root,
             diff_node,
+            has_more_data,
             parents_path,
             diff_node_conflict,
             diff_attributes,
@@ -161,6 +155,7 @@ class EnrichedDiffGetQuery(Query):
         self.return_labels = [
             "diff_root",
             "diff_node",
+            "has_more_data",
             "parents_path",
             "diff_node_conflict",
             "diff_attributes",