Merge branch 'main' into splunk_extractor

Author: AGhafaryy

.gitignore

@@ -160,4 +160,5 @@ cython_debug/
 #  be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
 #  and can be added to the global gitignore or merged into this file.  For a more nuclear
 #  option (not recommended) you can uncomment the following to ignore the entire idea folder.
-.idea/
+.idea/
+pytest.xml

Makefile

@@ -29,11 +29,13 @@ venv: poetry.lock
 format: venv
 	poetry run black nodestream tests
 	poetry run isort nodestream tests
+	poetry run ruff check nodestream tests --fix
 
 .PHONY: lint
 lint: venv
 	poetry run black nodestream tests --check
 	poetry run ruff check nodestream tests
+	poetry run isort nodestream tests --check-only
 
 .PHONY: test-unit
 test-unit: venv

nodestream/cli/operations/run_copy.py

@@ -3,6 +3,7 @@
 from ...databases import Copier, GraphDatabaseWriter
 from ...pipeline import Pipeline
 from ...pipeline.object_storage import ObjectStore
+from ...pipeline.progress_reporter import PipelineProgressReporter
 from ...project import Project, Target
 from ..commands.nodestream_command import NodestreamCommand
 from .operation import Operation
@@ -20,12 +21,13 @@ def __init__(
         self.from_target = from_target
         self.to_target = to_target
         self.project = project
+        self.schema = self.project.get_schema()
         self.node_types = node_types
         self.relationship_types = relationship_types
 
     async def perform(self, command: NodestreamCommand):
         pipeline = self.build_pipeline()
-        await pipeline.run()
+        await pipeline.run(reporter=PipelineProgressReporter())
 
     def build_pipeline(self) -> Pipeline:
         copier = self.build_copier()
@@ -37,7 +39,7 @@ def build_pipeline(self) -> Pipeline:
     def build_copier(self) -> Copier:
         return Copier(
             self.from_target.make_type_retriever(),
-            self.project,
+            self.schema,
             self.node_types,
             self.relationship_types,
         )

nodestream/cli/operations/run_pipeline.py

@@ -1,3 +1,4 @@
+from logging import getLogger
 from typing import Iterable, Optional
 
 from yaml import safe_dump
@@ -113,7 +114,7 @@ def get_progress_indicator(
         self, command: NodestreamCommand, pipeline_name: str
     ) -> "ProgressIndicator":
         if command.has_json_logging_set:
-            return ProgressIndicator(command, pipeline_name)
+            return JsonProgressIndicator(command, pipeline_name)
 
         return SpinnerProgressIndicator(command, pipeline_name)
 
@@ -160,14 +161,15 @@ def on_start(self):
         self.progress = self.command.progress_indicator()
         self.progress.start(f"Running pipeline: '{self.pipeline_name}'")
 
-    def progress_callback(self, index, _):
+    def progress_callback(self, index, metrics: Metrics):
         self.progress.set_message(
             f"Currently processing record at index: <info>{index}</info>"
         )
+        metrics.tick()
 
     def on_finish(self, metrics: Metrics):
         self.progress.finish(f"Finished running pipeline: '{self.pipeline_name}'")
-
+        metrics.tick()
         if self.exception:
             raise self.exception
 
@@ -176,3 +178,27 @@ def on_fatal_error(self, exception: Exception):
             "<error>Encountered a fatal error while running pipeline</error>"
         )
         self.exception = exception
+
+
+class JsonProgressIndicator(ProgressIndicator):
+    def __init__(self, command: NodestreamCommand, pipeline_name: str) -> None:
+        super().__init__(command, pipeline_name)
+        self.logger = getLogger()
+        self.exception = None
+
+    def on_start(self):
+        self.logger.info("Starting Pipeline")
+
+    def progress_callback(self, index, metrics: Metrics):
+        self.logger.info("Processing Record", extra={"index": index})
+        metrics.tick()
+
+    def on_finish(self, metrics: Metrics):
+        self.logger.info("Pipeline Completed")
+        metrics.tick()
+        if self.exception:
+            raise self.exception
+
+    def on_fatal_error(self, exception: Exception):
+        self.logger.error("Pipeline Failed", exc_info=exception)
+        self.exception = exception

nodestream/databases/query_executor_with_statistics.py

@@ -1,6 +1,13 @@
 from typing import Iterable
 
-from ..metrics import Metric, Metrics
+from ..metrics import (
+    INGEST_HOOKS_EXECUTED,
+    NODES_UPSERTED,
+    RELATIONSHIPS_UPSERTED,
+    TIME_TO_LIVE_OPERATIONS,
+    Metric,
+    Metrics,
+)
 from ..model import IngestionHook, Node, RelationshipWithNodes, TimeToLiveConfiguration
 from .query_executor import (
     OperationOnNodeIdentity,
@@ -14,12 +21,45 @@ class QueryExecutorWithStatistics(QueryExecutor):
 
     def __init__(self, inner: QueryExecutor) -> None:
         self.inner = inner
+        self.node_metric_by_type: dict[str, Metric] = {}
+        self.relationship_metric_by_relationship_type: dict[str, Metric] = {}
+
+    def _get_or_create_node_metric(self, node_type: str) -> Metric:
+        """Get or create a metric for a node type."""
+        if node_type not in self.node_metric_by_type:
+            self.node_metric_by_type[node_type] = Metric(
+                f"nodes_upserted_by_type_{node_type}",
+                f"Number of nodes upserted by type {node_type}",
+            )
+        return self.node_metric_by_type[node_type]
+
+    def _get_or_create_relationship_metric(self, relationship_type: str) -> Metric:
+        """Get or create a metric for a relationship type."""
+        if relationship_type not in self.relationship_metric_by_relationship_type:
+            self.relationship_metric_by_relationship_type[relationship_type] = Metric(
+                f"relationships_upserted_by_relationship_type_{relationship_type}",
+                f"Number of relationships upserted by relationship type {relationship_type}",
+            )
+        return self.relationship_metric_by_relationship_type[relationship_type]
 
     async def upsert_nodes_in_bulk_with_same_operation(
         self, operation: OperationOnNodeIdentity, nodes: Iterable[Node]
     ):
         await self.inner.upsert_nodes_in_bulk_with_same_operation(operation, nodes)
-        Metrics.get().increment(Metric.NODES_UPSERTED, len(nodes))
+
+        # Tally node types
+        node_type_counts: dict[str, int] = {}
+        total_nodes = 0
+        for node in nodes:
+            node_type_counts[node.type] = node_type_counts.get(node.type, 0) + 1
+            total_nodes += 1
+
+        # Increment metrics in bulk
+        metrics = Metrics.get()
+        metrics.increment(NODES_UPSERTED, total_nodes)
+        for node_type, count in node_type_counts.items():
+            metric = self._get_or_create_node_metric(node_type)
+            metrics.increment(metric, count)
 
     async def upsert_relationships_in_bulk_of_same_operation(
         self,
@@ -29,15 +69,31 @@ async def upsert_relationships_in_bulk_of_same_operation(
         await self.inner.upsert_relationships_in_bulk_of_same_operation(
             shape, relationships
         )
-        Metrics.get().increment(Metric.RELATIONSHIPS_UPSERTED, len(relationships))
+
+        # Tally relationship types
+        relationship_type_counts: dict[str, int] = {}
+        total_relationships = 0
+        for relationship in relationships:
+            rel_type = relationship.relationship.type
+            relationship_type_counts[rel_type] = (
+                relationship_type_counts.get(rel_type, 0) + 1
+            )
+            total_relationships += 1
+
+        # Increment metrics in bulk
+        metrics = Metrics.get()
+        metrics.increment(RELATIONSHIPS_UPSERTED, total_relationships)
+        for rel_type, count in relationship_type_counts.items():
+            metric = self._get_or_create_relationship_metric(rel_type)
+            metrics.increment(metric, count)
 
     async def perform_ttl_op(self, config: TimeToLiveConfiguration):
         await self.inner.perform_ttl_op(config)
-        Metrics.get().increment(Metric.TIME_TO_LIVE_OPERATIONS)
+        Metrics.get().increment(TIME_TO_LIVE_OPERATIONS)
 
     async def execute_hook(self, hook: IngestionHook):
         await self.inner.execute_hook(hook)
-        Metrics.get().increment(Metric.INGEST_HOOKS_EXECUTED)
+        Metrics.get().increment(INGEST_HOOKS_EXECUTED)
 
     async def finish(self):
         await self.inner.finish()

nodestream/databases/writer.py

@@ -14,7 +14,7 @@ def from_file_data(
         ),
         collect_stats: bool = True,
         batch_size: int = 1000,
-        **database_args
+        **database_args,
     ):
         connector = DatabaseConnector.from_database_args(
             database=database, **database_args