Merge pull request #34 from nodestream-proj/RecordMetrics

Recording Database Runtime Metrics

Author: angelosantos4

Makefile

@@ -27,13 +27,15 @@ venv: poetry.lock
 
 .PHONY: format
 format: venv
-	poetry run black nodestream_plugin_neo4j tests
+	poetry run ruff format nodestream_plugin_neo4j tests
 	poetry run isort nodestream_plugin_neo4j tests
+	poetry run black nodestream_plugin_neo4j tests
 
 .PHONY: lint
-lint: venv
+lint: venv	
+	poetry run ruff check nodestream_plugin_neo4j tests
+	poetry run isort nodestream_plugin_neo4j tests --check
 	poetry run black nodestream_plugin_neo4j tests --check
-	poetry run ruff nodestream_plugin_neo4j tests
 
 .PHONY: test-unit
 test-unit: venv

nodestream_plugin_neo4j/database_connector.py

@@ -26,7 +26,7 @@ def from_file_data(
         execute_chunks_in_parallel: bool = True,
         retries_per_chunk: int = 3,
         _experimental_set_first_ingested_at: bool = False,
-        **connection_args
+        **connection_args,
     ):
         database_connection = Neo4jDatabaseConnection.from_configuration(
             **connection_args

nodestream_plugin_neo4j/extractor.py

@@ -15,7 +15,7 @@ def from_file_data(
         query: str,
         parameters: Optional[Dict[str, Any]] = None,
         limit: int = 100,
-        **connection_args
+        **connection_args,
     ):
         connector = Neo4jDatabaseConnection.from_configuration(**connection_args)
         return cls(query, connector, parameters, limit)

nodestream_plugin_neo4j/neo4j_database.py

@@ -2,7 +2,14 @@
 from logging import getLogger
 from typing import Awaitable, Iterable, Tuple, Union
 
-from neo4j import AsyncDriver, AsyncGraphDatabase, AsyncSession, Record, RoutingControl
+from neo4j import (
+    AsyncDriver,
+    AsyncGraphDatabase,
+    AsyncSession,
+    EagerResult,
+    Record,
+    RoutingControl,
+)
 from neo4j.auth_management import AsyncAuthManagers
 from neo4j.exceptions import (
     AuthError,
@@ -13,6 +20,7 @@
 from nodestream.file_io import LazyLoadedArgument
 
 from .query import Query
+from .result import Neo4jQueryStatistics, Neo4jResult
 
 RETRYABLE_EXCEPTIONS = (TransientError, ServiceUnavailable, SessionExpired, AuthError)
 
@@ -74,6 +82,7 @@ def __init__(
         self.max_retry_attempts = max_retry_attempts
         self.retry_factor = retry_factor
         self._driver = None
+        self.query_set: set[Query] = set()
 
     def acquire_driver(self) -> AsyncDriver:
         self._driver = self.driver_factory()
@@ -85,38 +94,39 @@ def driver(self):
         return self._driver
 
     def log_query_start(self, query: Query):
-        self.logger.info(
-            "Executing Cypher Query to Neo4j",
-            extra={
-                "query": query.query_statement,
-                "uri": self.driver._pool.address.host,
-            },
-        )
-
-    def log_record(self, record: Record):
-        self.logger.debug(
-            "Gathered Query Results",
-            extra=dict(**record, uri=self.driver._pool.address.host),
-        )
+        if query.query_statement not in self.query_set:
+            self.logger.info(
+                "Executing Cypher Query to Neo4j.",
+                extra={
+                    "query": query.query_statement,
+                    "uri": self.driver._pool.address.host,
+                },
+            )
+            self.query_set.add(query.query_statement)
 
     async def _execute_query(
         self,
         query: Query,
         log_result: bool = False,
         routing_=RoutingControl.WRITE,
-    ) -> Record:
-        result = await self.driver.execute_query(
+    ) -> Iterable[Record]:
+        result: EagerResult = await self.driver.execute_query(
             query.query_statement,
             query.parameters,
             database_=self.database_name,
             routing_=routing_,
         )
-        records = result.records
+        neo4j_result = Neo4jResult(query, result)
         if log_result:
-            for record in records:
-                self.log_record(record)
+            statistics = neo4j_result.obtain_query_statistics()
+            self.log_error_messages_from_statistics(statistics)
+            statistics.update_metrics_from_summary()
+
+        return neo4j_result.records
 
-        return records
+    def log_error_messages_from_statistics(self, statistics: Neo4jQueryStatistics):
+        for error in statistics.error_messages:
+            self.logger.error("Query Error Occurred.", extra={"error": error})
 
     async def execute(
         self,
@@ -132,7 +142,8 @@ async def execute(
                 return await self._execute_query(query, log_result, routing_)
             except RETRYABLE_EXCEPTIONS as e:
                 self.logger.warning(
-                    f"Error executing query, retrying. Attempt {attempts + 1}",
+                    "Error executing query, retrying. Attempt %s",
+                    attempts,
                     exc_info=e,
                 )
                 await asyncio.sleep(self.retry_factor * attempts)

nodestream_plugin_neo4j/query.py

@@ -1,18 +1,58 @@
-from dataclasses import dataclass
-from typing import Any, Dict, List
+from dataclasses import dataclass, field
+from typing import Any
+
+
+@dataclass
+class ApocUpdateStatistics:
+    """Type definition for APOC update statistics."""
+
+    relationshipsDeleted: int = 0  # Number of relationships deleted
+    relationshipsCreated: int = 0  # Number of relationships created
+    nodesDeleted: int = 0  # Number of nodes deleted
+    nodesCreated: int = 0  # Number of nodes created
+    labelsRemoved: int = 0  # Number of labels removed
+    labelsAdded: int = 0  # Number of labels added
+    propertiesSet: int = 0  # Number of properties set
+
+
+@dataclass
+class ApocBatchResponse:
+    """Type definition for APOC periodic.iterate response fields."""
+
+    batches: int = 0  # Number of batches processed
+    total: int = 0  # Total number of operations
+    timeTaken: int = 0  # Time taken in milliseconds
+    committedOperations: int = 0  # Number of operations committed
+    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] = field(default_factory=dict)  # Map of error messages
+    wasTerminated: bool = False  # Whether the operation was terminated
+    updateStatistics: ApocUpdateStatistics = field(default_factory=ApocUpdateStatistics)
+
+
+# Get fields from ApocBatchResponse dataclass
+APOC_BATCH_QUERY_RESPONSE_FIELDS = [
+    field.name for field in ApocBatchResponse.__dataclass_fields__.values()
+]
 
 UNWIND_QUERY = "UNWIND $batched_parameter_sets as params RETURN params"
-COMMIT_QUERY = """
+
+# Build the YIELD and RETURN clauses using the standardized fields
+YIELD_CLAUSE = f"YIELD {', '.join(APOC_BATCH_QUERY_RESPONSE_FIELDS)}"
+RETURN_CLAUSE = f"RETURN {', '.join(APOC_BATCH_QUERY_RESPONSE_FIELDS)}"
+
+COMMIT_QUERY = f"""
 CALL apoc.periodic.iterate(
     $iterable_query,
     $batched_query,
-    {batchSize: $chunk_size, parallel: $execute_chunks_in_parallel, retries: $retries_per_chunk, params: $iterate_params}
+    {{batchSize: $chunk_size, parallel: $execute_chunks_in_parallel, retries: $retries_per_chunk, params: $iterate_params}}
 )
-YIELD batches, committedOperations, failedOperations, errorMessages
-RETURN batches, committedOperations, failedOperations, errorMessages
+{YIELD_CLAUSE}
+{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
@@ -23,7 +63,8 @@
 @dataclass(slots=True, frozen=True)
 class Query:
     query_statement: str
-    parameters: Dict[str, Any]
+    parameters: dict[str, Any]
+    is_apoc: bool = False  # Indicates if this is an APOC query
 
     @classmethod
     def from_statement(cls, query_statement: str, **parameters: Any) -> "Query":
@@ -47,13 +88,15 @@ def feed_batched_query(
                 "chunk_size": chunk_size,
                 "retries_per_chunk": retries_per_chunk,
             },
+            is_apoc=True,  # This is an APOC query
         )
 
 
 @dataclass(slots=True, frozen=True)
 class QueryBatch:
     query_statement: str
-    batched_parameter_sets: List[Dict[str, Any]]
+    batched_parameter_sets: list[dict[str, Any]]
+    is_apoc: bool = False  # Indicates if this is an APOC query
 
     def as_query(
         self,
@@ -63,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
@@ -74,4 +117,5 @@ def as_query(
                 "chunk_size": chunk_size,
                 "retries_per_chunk": retries_per_chunk,
             },
+            is_apoc=apoc_iterate,  # Set is_apoc based on apoc_iterate parameter
         )