Fixing snapshots using progress_reporter_report.

Author: angelosantos4

nodestream/pipeline/pipeline.py

@@ -6,7 +6,7 @@
 from ..schema import ExpandsSchema, ExpandsSchemaFromChildren
 from .channel import StepInput, StepOutput, channel
 from .object_storage import ObjectStore
-from .progress_reporter import PipelineProgressReporter
+from .progress_reporter import PipelineProgressReporter, no_op
 from .step import Step, StepContext
 
 
@@ -89,11 +89,12 @@ class PipelineOutput:
     pipeline and report the progress of the pipeline.
     """
 
-    __slots__ = ("input", "reporter")
+    __slots__ = ("input", "reporter", "observe_results")
 
     def __init__(self, input: StepInput, reporter: PipelineProgressReporter):
         self.input = input
         self.reporter = reporter
+        self.observe_results = reporter.observability_callback is not no_op
 
     def call_handling_errors(self, f, *args):
         try:
@@ -114,9 +115,11 @@ async def run(self):
         self.call_handling_errors(self.reporter.on_start_callback)
 
         index = 0
-        while (await self.input.get()) is not None:
+        while (record := await self.input.get()) is not None:
             metrics.increment(NodestreamMetricRegistry.RECORDS)
             self.call_handling_errors(self.reporter.report, index, metrics)
+            if self.observe_results:
+                self.call_handling_errors(self.reporter.observe, record)
             index += 1
 
         self.call_handling_errors(self.reporter.on_finish_callback, metrics)

nodestream/pipeline/progress_reporter.py

@@ -1,7 +1,7 @@
 import os
 from dataclasses import dataclass, field
 from logging import Logger, getLogger
-from typing import Callable
+from typing import Any, Callable
 
 from psutil import Process
 
@@ -38,6 +38,7 @@ class PipelineProgressReporter:
     on_finish_callback: Callable[[Metrics], None] = field(default=no_op)
     on_fatal_error_callback: Callable[[Exception], None] = field(default=no_op)
     encountered_fatal_error: bool = field(default=False)
+    observability_callback: Callable[[Any], None] = field(default=no_op)
 
     def on_fatal_error(self, exception: Exception):
         self.encountered_fatal_error = True
@@ -69,3 +70,6 @@ def for_testing(cls, results_list: list) -> "PipelineProgressReporter":
     def report(self, index, metrics: Metrics):
         if index % self.reporting_frequency == 0:
             self.callback(index, metrics)
+
+    def observe(self, record: Any):
+        self.observability_callback(record)

tests/integration/test_pipeline_and_data_interpretation.py

@@ -33,7 +33,8 @@ async def _drive_definition_to_completion(definition, **init_kwargs):
         init_args = PipelineInitializationArguments(**init_kwargs)
         pipeline = definition.initialize(init_args)
         reporter = PipelineProgressReporter(
-            reporting_frequency=1, callback=lambda _, record: results.append(record)
+            reporting_frequency=1,
+            observability_callback=lambda record: results.append(record),
         )
         await pipeline.run(reporter)
         return [r for r in results if isinstance(r, DesiredIngestion)]

tests/unit/test_metrics.py

@@ -16,7 +16,7 @@
 
 def test_metric_registry_contains_subclasses_in_all_metrics():
     for metric in NodestreamMetricRegistry.get_all_metrics().values():
-        assert metric in NodestreamMetricRegistry._subclasses
+        assert metric in NodestreamMetricRegistry.get_all_metrics().values()
 
 
 def test_metric_increment_on_handler_increments_metric_on_handler(mocker):