Fixed the channel implementation to be non-lazy

Author: angelosantos4

nodestream/databases/query_executor_with_statistics.py

@@ -1,12 +1,12 @@
 from typing import Iterable
 
 from ..metrics import (
-    Metric,
-    Metrics,
+    INGEST_HOOKS_EXECUTED,
     NODES_UPSERTED,
     RELATIONSHIPS_UPSERTED,
     TIME_TO_LIVE_OPERATIONS,
-    INGEST_HOOKS_EXECUTED,
+    Metric,
+    Metrics,
 )
 from ..model import IngestionHook, Node, RelationshipWithNodes, TimeToLiveConfiguration
 from .query_executor import (
@@ -82,9 +82,7 @@ async def upsert_relationships_in_bulk_of_same_operation(
 
         # Increment metrics in bulk
         metrics = Metrics.get()
-        metrics.increment(
-            RELATIONSHIPS_UPSERTED, total_relationships
-        )
+        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)

nodestream/metrics.py

@@ -149,7 +149,6 @@ def get_gauge(self, metric: Metric) -> Gauge:
             return self.instruments_by_metric[metric]
 
         def increment(self, metric: Metric, value: Number):
-            print(f"Incrementing {metric.name} by {value}")
             self.get_gauge(metric).inc(value)
 
         def decrement(self, metric: Metric, value: Number):

nodestream/pipeline/channel.py

@@ -37,32 +37,22 @@ class Channel:
     prevent one step from overwhelming another step with too many records.
     """
 
-    __slots__ = ("queue", "input_dropped", "_metric", "input_name", "output_name")
+    __slots__ = ("queue", "input_dropped", "metric")
 
-    def __init__(self, size: int) -> None:
+    def __init__(self, size: int, metric: Metric) -> None:
         self.queue = Queue(maxsize=size)
         self.input_dropped = False
-        self.input_name = "Void"
-        self.output_name = "Void"
-        self._metric = None
-
-    @property
-    def metric(self) -> Metric:
-        """Get the metric for the channel."""
-        if self._metric is None:
-            self._metric = Metric(
-                f"buffered_{self.output_name}_to_{self.input_name}",
-                f"Records buffered: {self.output_name} → {self.input_name}",
-            )
-        return self._metric
-
-    def register_input(self, name: str) -> None:
-        """Register the name of the step that will consume from this channel."""
-        self.input_name = name
-
-    def register_output(self, name: str) -> None:
-        """Register the name of the step that will produce to this channel."""
-        self.output_name = name
+        self.metric = metric
+
+    @classmethod
+    def create_with_naming(
+        cls, size: int, input_name: str = "Void", output_name: str = "Void"
+    ) -> "Channel":
+        metric = Metric(
+            f"buffered_{input_name}_to_{output_name}",
+            f"Records buffered: {input_name} → {output_name}",
+        )
+        return cls(size, metric)
 
     async def get(self):
         """Get an object from the channel.
@@ -160,10 +150,6 @@ class StepInput:
     def __init__(self, channel: Channel) -> None:
         self.channel = channel
 
-    def register(self, name: str) -> None:
-        """Register the name of the step that will consume from this channel."""
-        self.channel.register_input(name)
-
     async def get(self) -> Optional[object]:
         """Get an object from the input channel.
 
@@ -187,13 +173,15 @@ def done(self):
         self.channel.input_dropped = True
 
 
-def channel(size: int) -> Tuple[StepInput, StepOutput]:
+def channel(
+    size: int, input_name: str | None = None, output_name: str | None = None
+) -> Tuple[StepInput, StepOutput]:
     """Create a new input and output channel.
 
     Args:
         size: The size of the channel.
         input_name: The name of the input step.
         output_name: The name of the output step.
     """
-    channel = Channel(size)
+    channel = Channel.create_with_naming(size, input_name, output_name)
     return StepInput(channel), StepOutput(channel)

nodestream/pipeline/pipeline.py

@@ -3,9 +3,9 @@
 from typing import Iterable, List, Tuple
 
 from ..metrics import (
-    Metrics,
     RECORDS,
     STEPS_RUNNING,
+    Metrics,
 )
 from ..schema import ExpandsSchema, ExpandsSchemaFromChildren
 from .channel import StepInput, StepOutput, channel
@@ -62,7 +62,6 @@ async def emit_record(self, record):
     async def drive_step(self):
         try:
             while (next_record := await self.input.get()) is not None:
-                print(f"Driving step with record {next_record}")
                 results = self.step.process_record(next_record, self.context)
                 async for record in results:
                     if not await self.emit_record(record):
@@ -179,8 +178,12 @@ async def run(self, reporter: PipelineProgressReporter):
         # the steps in the pipeline. The channels have a fixed size to control
         # the flow of records between the steps.
         executors: List[StepExecutor] = []
-        current_input, current_output = channel(self.step_outbox_size)
-        print(current_input, current_output)
+        current_input_name = None
+        current_output_name = self.steps[-1].__class__.__name__ + f"_{len(self.steps)}"
+
+        current_input, current_output = channel(
+            self.step_outbox_size, current_output_name, current_input_name
+        )
         pipeline_output = PipelineOutput(current_input, reporter)
 
         # Create the executors for the steps in the pipeline. The executors
@@ -191,9 +194,16 @@ async def run(self, reporter: PipelineProgressReporter):
             index = len(self.steps) - reversed_index - 1
             storage = self.object_store.namespaced(str(index))
             context = StepContext(step.__class__.__name__, index, reporter, storage)
-            current_input, next_output = channel(self.step_outbox_size)
-            current_input.register(step.__class__.__name__)
-            current_output.register(step.__class__.__name__)
+            current_output_name = (
+                self.steps[reversed_index - 1].__class__.__name__
+                + f"_{reversed_index - 1}"
+                if reversed_index - 1 >= 0
+                else None
+            )
+            current_input_name = step.__class__.__name__ + f"_{reversed_index}"
+            current_input, next_output = channel(
+                self.step_outbox_size, current_output_name, current_input_name
+            )
             exec = StepExecutor(step, current_input, current_output, context)
             current_output = next_output
             executors.append(exec)

nodestream/pipeline/step.py

@@ -1,9 +1,9 @@
 from typing import AsyncGenerator, Optional
 
 from ..metrics import (
-    Metrics,
     FATAL_ERRORS,
     NON_FATAL_ERRORS,
+    Metrics,
 )
 from .object_storage import ObjectStore
 from .progress_reporter import PipelineProgressReporter

tests/unit/databases/test_query_executor_with_statistics.py

@@ -6,11 +6,11 @@
     QueryExecutorWithStatistics,
 )
 from nodestream.metrics import (
-    Metrics,
+    INGEST_HOOKS_EXECUTED,
     NODES_UPSERTED,
     RELATIONSHIPS_UPSERTED,
     TIME_TO_LIVE_OPERATIONS,
-    INGEST_HOOKS_EXECUTED,
+    Metrics,
 )
 from nodestream.model import Node, Relationship, RelationshipWithNodes
 
@@ -24,44 +24,42 @@ def query_executor_with_statistics(mocker):
 async def test_upsert_nodes_in_bulk_with_same_operation_increments_counter_by_size_of_list(
     query_executor_with_statistics, mocker
 ):
+    nodes = [
+        Node("node_type", "node1", "id1"),
+        Node("node_type", "node2", "id2"),
+    ]
     with Metrics.capture() as metrics:
         metrics.increment = mocker.Mock()
         await query_executor_with_statistics.upsert_nodes_in_bulk_with_same_operation(
             "operation",
-            [Node("node_type", "node1", "id1"), Node("node_type", "node2", "id2")],
+            nodes,
         )
         query_executor_with_statistics.inner.upsert_nodes_in_bulk_with_same_operation.assert_awaited_once_with(
             "operation",
-            [Node("node_type", "node1", "id1"), Node("node_type", "node2", "id2")],
+            nodes,
         )
 
         assert "node_type" in query_executor_with_statistics.node_metric_by_type
         assert (
             call(query_executor_with_statistics.node_metric_by_type["node_type"], 2)
             in metrics.increment.call_args_list
         )
-        assert (
-            call(NODES_UPSERTED, 2)
-            in metrics.increment.call_args_list
-        )
+        assert call(NODES_UPSERTED, 2) in metrics.increment.call_args_list
 
 
 @pytest.mark.asyncio
 async def test_upsert_relationships_in_bulk_of_same_operation_increments_counter_by_size_of_list(
     query_executor_with_statistics, mocker
 ):
+    relationships = [
+        RelationshipWithNodes("node1", "node2", Relationship("relationship_type")),
+        RelationshipWithNodes("node3", "node4", Relationship("relationship_type")),
+    ]
     with Metrics.capture() as metrics:
         metrics.increment = mocker.Mock()
         await query_executor_with_statistics.upsert_relationships_in_bulk_of_same_operation(
             "operation",
-            [
-                RelationshipWithNodes(
-                    "node1", "node2", Relationship("relationship_type")
-                ),
-                RelationshipWithNodes(
-                    "node3", "node4", Relationship("relationship_type")
-                ),
-            ],
+            relationships,
         )
         query_executor_with_statistics.inner.upsert_relationships_in_bulk_of_same_operation.assert_awaited_once_with(
             "operation",
@@ -87,10 +85,7 @@ async def test_upsert_relationships_in_bulk_of_same_operation_increments_counter
             )
             in metrics.increment.call_args_list
         )
-        assert (
-            call(RELATIONSHIPS_UPSERTED, 2)
-            in metrics.increment.call_args_list
-        )
+        assert call(RELATIONSHIPS_UPSERTED, 2) in metrics.increment.call_args_list
 
 
 @pytest.mark.asyncio
@@ -103,9 +98,7 @@ async def test_perform_ttl_op_increments_counter_by_one(
         query_executor_with_statistics.inner.perform_ttl_op.assert_awaited_once_with(
             "config"
         )
-        metrics.increment.assert_called_once_with(
-            TIME_TO_LIVE_OPERATIONS
-        )
+        metrics.increment.assert_called_once_with(TIME_TO_LIVE_OPERATIONS)
 
 
 @pytest.mark.asyncio
@@ -118,9 +111,7 @@ async def test_execute_hook_increments_counter_by_one(
         query_executor_with_statistics.inner.execute_hook.assert_awaited_once_with(
             "hook"
         )
-        metrics.increment.assert_called_once_with(
-            INGEST_HOOKS_EXECUTED
-        )
+        metrics.increment.assert_called_once_with(INGEST_HOOKS_EXECUTED)
 
 
 @pytest.mark.asyncio

tests/unit/pipeline/test_pipeline.py

@@ -1,10 +1,6 @@
 import pytest
 
-from nodestream.pipeline import Extractor
-from nodestream.pipeline.channel import Channel
-from nodestream.pipeline.object_storage import NullObjectStore
 from nodestream.pipeline.pipeline import (
-    Pipeline,
     PipelineOutput,
     PipelineProgressReporter,
     Step,
@@ -140,43 +136,3 @@ def on_start_callback():
 
     output.call_handling_errors(output.reporter.on_start_callback)
     output.reporter.logger.exception.assert_called_once()
-
-
-class TestExtractor(Extractor):
-    async def extract_records(self):
-        return None
-
-
-class TestStep2(Step):
-    async def process_record(self, record, context: StepContext):
-        yield record
-
-
-@pytest.mark.asyncio
-async def test_pipeline_channels_obtain_input_and_output_names_from_steps(mocker):
-    channel = mocker.Mock(Channel)
-    mock_input = mocker.Mock(StepInput, channel=channel)
-    mock_output = mocker.Mock(StepOutput, channel=channel)
-    mock_input.get.return_value = None
-    mock_output.put.return_value = True
-    with mocker.patch(
-        "nodestream.pipeline.pipeline.channel", return_value=(mock_input, mock_output)
-    ):
-        step1 = TestExtractor()
-        step2 = TestStep2()
-        pipeline = Pipeline(
-            steps=[step1, step2], step_outbox_size=1, object_store=NullObjectStore()
-        )
-        await pipeline.run(
-            reporter=PipelineProgressReporter(
-                on_start_callback=lambda: None, logger=mocker.Mock()
-            )
-        )
-        assert mock_input.register.call_args_list == [
-            mocker.call("TestStep2"),
-            mocker.call("TestExtractor"),
-        ]
-        assert mock_output.register.call_args_list == [
-            mocker.call("TestStep2"),
-            mocker.call("TestExtractor"),
-        ]

tests/unit/test_metrics.py

@@ -3,16 +3,12 @@
 import pytest
 
 from nodestream.metrics import (
+    RECORDS,
     AggregateHandler,
     ConsoleMetricHandler,
     JsonLogMetricHandler,
     Metric,
     Metrics,
-    RECORDS,
-    NODES_UPSERTED,
-    RELATIONSHIPS_UPSERTED,
-    TIME_TO_LIVE_OPERATIONS,
-    INGEST_HOOKS_EXECUTED,
     NullMetricHandler,
     PrometheusMetricHandler,
 )