Saving

Author: angelosantos4

nodestream_plugin_neo4j/neo4j_database.py

@@ -106,7 +106,7 @@ async def _execute_query(
         query: Query,
         log_result: bool = False,
         routing_=RoutingControl.WRITE,
-    ) -> Record:
+    ) -> Iterable[Record]:
         result: EagerResult = await self.driver.execute_query(
             query.query_statement,
             query.parameters,
@@ -119,11 +119,12 @@ async def _execute_query(
             self.log_error_messages_from_statistics(statistics)
             statistics.update_metrics_from_summary()
 
-        return neo4j_result.records
+        for record in neo4j_result.records:
+            yield record
 
     def log_error_messages_from_statistics(self, statistics: Neo4jQueryStatistics):
         for error in statistics.error_messages:
-            self.logger.error(f"Query Errors Occured: {error}")
+            self.logger.error("Query Error Occurred: %s", error)
 
     async def execute(
         self,

nodestream_plugin_neo4j/query.py

@@ -1,4 +1,4 @@
-from dataclasses import dataclass
+from dataclasses import dataclass, field
 from typing import Any
 
 
@@ -26,15 +26,9 @@ class ApocBatchResponse:
     failedOperations: int = 0  # Number of operations that failed
     failedBatches: int = 0  # Number of batches that failed
     retries: int = 0  # Number of retries performed
-    errorMessages: dict[str, Any] = None  # Map of error messages
+    errorMessages: dict[str, Any] = field(default_factory=dict)  # Map of error messages
     wasTerminated: bool = False  # Whether the operation was terminated
-    updateStatistics: ApocUpdateStatistics = None  # Statistics about updates
-
-    def __post_init__(self):
-        if self.errorMessages is None:
-            self.errorMessages = {}
-        if self.updateStatistics is None:
-            self.updateStatistics = ApocUpdateStatistics()
+    updateStatistics: ApocUpdateStatistics = field(default_factory=ApocUpdateStatistics)
 
 
 # Get fields from ApocBatchResponse dataclass
@@ -58,7 +52,7 @@ def __post_init__(self):
 {RETURN_CLAUSE}
 """
 
-NON_APOCH_COMMIT_QUERY = """
+NON_APOC_COMMIT_QUERY = """
 UNWIND $iterate_params.batched_parameter_sets AS param
 CALL apoc.cypher.doIt($batched_query, {params: param})
 YIELD value
@@ -112,7 +106,7 @@ def as_query(
         retries_per_chunk: int = 3,
     ) -> Query:
         return Query(
-            {True: COMMIT_QUERY, False: NON_APOCH_COMMIT_QUERY}[apoc_iterate],
+            {True: COMMIT_QUERY, False: NON_APOC_COMMIT_QUERY}[apoc_iterate],
             {
                 "iterate_params": {
                     "batched_parameter_sets": self.batched_parameter_sets

nodestream_plugin_neo4j/result.py

@@ -207,34 +207,29 @@ def from_result(
 
     def update_metrics_from_summary(self):
         metrics = Metrics.get()
-        # Timing metrics
-        metrics.increment(PLANNING_TIME, self.timing.planning_time_ms)
-        metrics.increment(PROCESSING_TIME, self.timing.processing_time_ms)
-        metrics.increment(TOTAL_TIME, self.timing.total_time_ms)
-        metrics.increment(APOC_TIME, self.timing.apoc_time_ms)
-
-        # Write metrics
-        metrics.increment(NODES_CREATED, self.write_metrics.nodes_created)
-        metrics.increment(NODES_DELETED, self.write_metrics.nodes_deleted)
-        metrics.increment(
-            RELATIONSHIPS_CREATED, self.write_metrics.relationships_created
-        )
-        metrics.increment(
-            RELATIONSHIPS_DELETED, self.write_metrics.relationships_deleted
-        )
-        metrics.increment(PROPERTIES_SET, self.write_metrics.properties_set)
-        metrics.increment(LABELS_ADDED, self.write_metrics.labels_added)
-        metrics.increment(LABELS_REMOVED, self.write_metrics.labels_removed)
-        metrics.increment(CONSTRAINTS_ADDED, self.write_metrics.constraints_added)
-        metrics.increment(CONSTRAINTS_REMOVED, self.write_metrics.constraints_removed)
-        metrics.increment(INDEXES_ADDED, self.write_metrics.indexes_added)
-
-        # APOC specific metrics
-        metrics.increment(WAS_TERMINATED, int(self.was_terminated))
-        metrics.increment(RETRIES, self.retries)
-
-        # Error tracking
-        metrics.increment(ERROR_MESSAGES, len(self.error_messages))
+
+        metric_updates: list[tuple[Metric, int]] = [
+            (PLANNING_TIME, self.timing.planning_time_ms),
+            (PROCESSING_TIME, self.timing.processing_time_ms),
+            (TOTAL_TIME, self.timing.total_time_ms + self.timing.processing_time_ms),
+            (APOC_TIME, self.timing.apoc_time_ms),
+            (NODES_CREATED, self.write_metrics.nodes_created),
+            (NODES_DELETED, self.write_metrics.nodes_deleted),
+            (RELATIONSHIPS_CREATED, self.write_metrics.relationships_created),
+            (RELATIONSHIPS_DELETED, self.write_metrics.relationships_deleted),
+            (PROPERTIES_SET, self.write_metrics.properties_set),
+            (LABELS_ADDED, self.write_metrics.labels_added),
+            (LABELS_REMOVED, self.write_metrics.labels_removed),
+            (CONSTRAINTS_ADDED, self.write_metrics.constraints_added),
+            (CONSTRAINTS_REMOVED, self.write_metrics.constraints_removed),
+            (INDEXES_ADDED, self.write_metrics.indexes_added),
+            (INDEXES_REMOVED, self.write_metrics.indexes_removed),
+            (WAS_TERMINATED, int(self.was_terminated)),
+            (RETRIES, self.retries),
+            (ERROR_MESSAGES, len(self.error_messages)),
+        ]
+
+        metrics.increment(*metric_updates)
 
 
 class Neo4jResult: