Undoing Registry.

Author: angelosantos4

nodestream/databases/query_executor_with_statistics.py

@@ -3,7 +3,10 @@
 from ..metrics import (
     Metric,
     Metrics,
-    NodestreamMetricRegistry,
+    NODES_UPSERTED,
+    RELATIONSHIPS_UPSERTED,
+    TIME_TO_LIVE_OPERATIONS,
+    INGEST_HOOKS_EXECUTED,
 )
 from ..model import IngestionHook, Node, RelationshipWithNodes, TimeToLiveConfiguration
 from .query_executor import (
@@ -53,7 +56,7 @@ async def upsert_nodes_in_bulk_with_same_operation(
 
         # Increment metrics in bulk
         metrics = Metrics.get()
-        metrics.increment(NodestreamMetricRegistry.NODES_UPSERTED, total_nodes)
+        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)
@@ -80,19 +83,19 @@ async def upsert_relationships_in_bulk_of_same_operation(
         # Increment metrics in bulk
         metrics = Metrics.get()
         metrics.increment(
-            NodestreamMetricRegistry.RELATIONSHIPS_UPSERTED, total_relationships
+            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(NodestreamMetricRegistry.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(NodestreamMetricRegistry.INGEST_HOOKS_EXECUTED)
+        Metrics.get().increment(INGEST_HOOKS_EXECUTED)
 
     async def finish(self):
         await self.inner.finish()

nodestream/metrics.py

@@ -73,63 +73,26 @@ def tick(self):
         pass
 
 
-class MetricRegistry:
-    """A registry of metrics that collects metrics from class variables in subclasses.
-
-    This class provides a way to collect metrics defined as class attributes in subclasses.
-    Metrics can be declared as class attributes in subclasses. For example:
-
-    class MyMetricRegistry(MetricRegistry):
-        MY_METRIC = Metric("my_metric", "Description of my metric")
-        ANOTHER_METRIC = Metric("another_metric", "Description of another metric")
-    """
-
-    _subclasses: list[type["MetricRegistry"]] = []
-
-    def __init_subclass__(cls, **kwargs):
-        """Register subclasses for metric collection."""
-        super().__init_subclass__(**kwargs)
-        MetricRegistry._subclasses.append(cls)
-
-    @classmethod
-    def get_all_metrics(cls) -> dict[str, Metric]:
-        """Get all metrics from all subclasses.
-
-        Returns a dictionary mapping metric names to their Metric instances,
-        collected from all MetricRegistry subclasses.
-        """
-        all_metrics = {}
-        for subclass in cls._subclasses:
-            # Collect metrics from class variables
-            for attr_name, attr_value in subclass.__dict__.items():
-                if isinstance(attr_value, Metric):
-                    all_metrics[attr_name] = attr_value
-        return all_metrics
-
-
-class NodestreamMetricRegistry(MetricRegistry):
-    """A registry of metrics for the Nodestream project."""
-
-    # Core metrics
-    RECORDS = Metric("records", "Number of records processed")
-    NON_FATAL_ERRORS = Metric("non_fatal_errors", "Number of non-fatal errors")
-    FATAL_ERRORS = Metric("fatal_errors", "Number of fatal errors")
-    NODES_UPSERTED = Metric("nodes_upserted", "Number of nodes upserted to the graph")
-    RELATIONSHIPS_UPSERTED = Metric(
-        "relationships_upserted", "Number of relationships upserted to the graph"
-    )
-    TIME_TO_LIVE_OPERATIONS = Metric(
-        "time_to_live_operations", "Number of time-to-live operations executed"
-    )
-    INGEST_HOOKS_EXECUTED = Metric(
-        "ingest_hooks_executed", "Number of ingest hooks executed to the graph"
-    )
-    BUFFERED_RECORDS = Metric(
-        "buffered_records", "Number of records buffered between steps"
-    )
-    STEPS_RUNNING = Metric(
-        "steps_running", "Number of steps currently running in the pipeline"
-    )
+# Core metrics
+RECORDS = Metric("records", "Number of records processed")
+NON_FATAL_ERRORS = Metric("non_fatal_errors", "Number of non-fatal errors")
+FATAL_ERRORS = Metric("fatal_errors", "Number of fatal errors")
+NODES_UPSERTED = Metric("nodes_upserted", "Number of nodes upserted to the graph")
+RELATIONSHIPS_UPSERTED = Metric(
+    "relationships_upserted", "Number of relationships upserted to the graph"
+)
+TIME_TO_LIVE_OPERATIONS = Metric(
+    "time_to_live_operations", "Number of time-to-live operations executed"
+)
+INGEST_HOOKS_EXECUTED = Metric(
+    "ingest_hooks_executed", "Number of ingest hooks executed to the graph"
+)
+BUFFERED_RECORDS = Metric(
+    "buffered_records", "Number of records buffered between steps"
+)
+STEPS_RUNNING = Metric(
+    "steps_running", "Number of steps currently running in the pipeline"
+)
 
 
 try:
@@ -223,9 +186,7 @@ class ConsoleMetricHandler(MetricHandler):
     """A metric handler that prints metrics to the console."""
 
     def __init__(self, command: Command):
-        self.metrics: dict[Metric, Number] = {
-            metric: 0 for metric in MetricRegistry.get_all_metrics().values()
-        }
+        self.metrics: dict[Metric, Number] = {}
         self.command = command
 
     def increment(self, metric: Metric, value: Number):
@@ -250,9 +211,7 @@ class JsonLogMetricHandler(MetricHandler):
     """A metric handler that logs metrics in JSON format."""
 
     def __init__(self):
-        self.metrics: dict[Metric, Number] = {
-            metric: 0 for metric in MetricRegistry.get_all_metrics().values()
-        }
+        self.metrics: dict[Metric, Number] = {}
         self.logger = getLogger(__name__)
 
     def increment(self, metric: Metric, value: Number):

nodestream/pipeline/pipeline.py

@@ -2,7 +2,11 @@
 from logging import getLogger
 from typing import Iterable, List, Tuple
 
-from ..metrics import Metrics, NodestreamMetricRegistry
+from ..metrics import (
+    Metrics,
+    RECORDS,
+    STEPS_RUNNING,
+)
 from ..schema import ExpandsSchema, ExpandsSchemaFromChildren
 from .channel import StepInput, StepOutput, channel
 from .object_storage import ObjectStore
@@ -34,14 +38,14 @@ def __init__(
 
     async def start_step(self):
         try:
-            Metrics.get().increment(NodestreamMetricRegistry.STEPS_RUNNING)
+            Metrics.get().increment(STEPS_RUNNING)
             await self.step.start(self.context)
         except Exception as e:
             self.context.report_error("Error starting step", e)
 
     async def stop_step(self):
         try:
-            Metrics.get().decrement(NodestreamMetricRegistry.STEPS_RUNNING)
+            Metrics.get().decrement(STEPS_RUNNING)
             await self.step.finish(self.context)
         except Exception as e:
             self.context.report_error("Error stopping step", e)
@@ -116,7 +120,7 @@ async def run(self):
 
         index = 0
         while (record := await self.input.get()) is not None:
-            metrics.increment(NodestreamMetricRegistry.RECORDS)
+            metrics.increment(RECORDS)
             self.call_handling_errors(self.reporter.report, index, metrics)
             if self.observe_results:
                 self.call_handling_errors(self.reporter.observe, record)

nodestream/pipeline/step.py

@@ -1,6 +1,10 @@
 from typing import AsyncGenerator, Optional
 
-from ..metrics import Metrics, NodestreamMetricRegistry
+from ..metrics import (
+    Metrics,
+    FATAL_ERRORS,
+    NON_FATAL_ERRORS,
+)
 from .object_storage import ObjectStore
 from .progress_reporter import PipelineProgressReporter
 
@@ -57,9 +61,9 @@ def report_error(
         )
         if fatal:
             self.reporter.on_fatal_error(exception)
-            Metrics.get().increment(NodestreamMetricRegistry.FATAL_ERRORS)
+            Metrics.get().increment(FATAL_ERRORS)
         else:
-            Metrics.get().increment(NodestreamMetricRegistry.NON_FATAL_ERRORS)
+            Metrics.get().increment(NON_FATAL_ERRORS)
 
     def debug(self, message: str, **extras):
         """Log a debug message.

tests/unit/databases/test_query_executor_with_statistics.py

@@ -5,7 +5,13 @@
 from nodestream.databases.query_executor_with_statistics import (
     QueryExecutorWithStatistics,
 )
-from nodestream.metrics import Metrics, NodestreamMetricRegistry
+from nodestream.metrics import (
+    Metrics,
+    NODES_UPSERTED,
+    RELATIONSHIPS_UPSERTED,
+    TIME_TO_LIVE_OPERATIONS,
+    INGEST_HOOKS_EXECUTED,
+)
 from nodestream.model import Node, Relationship, RelationshipWithNodes
 
 
@@ -35,7 +41,7 @@ async def test_upsert_nodes_in_bulk_with_same_operation_increments_counter_by_si
             in metrics.increment.call_args_list
         )
         assert (
-            call(NodestreamMetricRegistry.NODES_UPSERTED, 2)
+            call(NODES_UPSERTED, 2)
             in metrics.increment.call_args_list
         )
 
@@ -82,7 +88,7 @@ async def test_upsert_relationships_in_bulk_of_same_operation_increments_counter
             in metrics.increment.call_args_list
         )
         assert (
-            call(NodestreamMetricRegistry.RELATIONSHIPS_UPSERTED, 2)
+            call(RELATIONSHIPS_UPSERTED, 2)
             in metrics.increment.call_args_list
         )
 
@@ -98,7 +104,7 @@ async def test_perform_ttl_op_increments_counter_by_one(
             "config"
         )
         metrics.increment.assert_called_once_with(
-            NodestreamMetricRegistry.TIME_TO_LIVE_OPERATIONS
+            TIME_TO_LIVE_OPERATIONS
         )
 
 
@@ -113,7 +119,7 @@ async def test_execute_hook_increments_counter_by_one(
             "hook"
         )
         metrics.increment.assert_called_once_with(
-            NodestreamMetricRegistry.INGEST_HOOKS_EXECUTED
+            INGEST_HOOKS_EXECUTED
         )
 
 

tests/unit/test_metrics.py

@@ -8,27 +8,26 @@
     JsonLogMetricHandler,
     Metric,
     Metrics,
-    NodestreamMetricRegistry,
+    RECORDS,
+    NODES_UPSERTED,
+    RELATIONSHIPS_UPSERTED,
+    TIME_TO_LIVE_OPERATIONS,
+    INGEST_HOOKS_EXECUTED,
     NullMetricHandler,
     PrometheusMetricHandler,
 )
 
 
-def test_metric_registry_contains_subclasses_in_all_metrics():
-    for metric in NodestreamMetricRegistry.get_all_metrics().values():
-        assert metric in NodestreamMetricRegistry.get_all_metrics().values()
-
-
 def test_metric_increment_on_handler_increments_metric_on_handler(mocker):
     handler = mocker.Mock()
-    metric = NodestreamMetricRegistry.RECORDS
+    metric = RECORDS
     metric.increment_on(handler, 1)
     handler.increment.assert_called_once_with(metric, 1)
 
 
 def test_metric_decrement_on_handler_decrements_metric_on_handler(mocker):
     handler = mocker.Mock()
-    metric = NodestreamMetricRegistry.RECORDS
+    metric = RECORDS
     metric.decrement_on(handler, 1)
     handler.decrement.assert_called_once_with(metric, 1)
 
@@ -42,8 +41,8 @@ def test_prometheus_metric_handler_registers_new_metrics(mocker):
 
 def test_null_metric_handler():
     handler = NullMetricHandler()
-    handler.increment(NodestreamMetricRegistry.RECORDS, 1)
-    handler.decrement(NodestreamMetricRegistry.RECORDS, 1)
+    handler.increment(RECORDS, 1)
+    handler.decrement(RECORDS, 1)
     # No assertions needed as NullMetricHandler does nothing
 
 
@@ -53,16 +52,16 @@ def test_prometheus_metric_handler(mocker):
     handler = PrometheusMetricHandler()
     handler.start()
     mock_start_http_server.assert_called_once()
-    handler.increment(NodestreamMetricRegistry.RECORDS, 1)
-    handler.decrement(NodestreamMetricRegistry.RECORDS, 1)
+    handler.increment(RECORDS, 1)
+    handler.decrement(RECORDS, 1)
     handler.stop()
 
 
 def test_console_metric_handler(mocker):
     mock_command = mocker.Mock()
     handler = ConsoleMetricHandler(mock_command)
-    handler.increment(NodestreamMetricRegistry.RECORDS, 1)
-    handler.decrement(NodestreamMetricRegistry.RECORDS, 1)
+    handler.increment(RECORDS, 1)
+    handler.decrement(RECORDS, 1)
     handler.tick()
     handler.stop()
     mock_command.table.assert_called_once()
@@ -71,8 +70,8 @@ def test_console_metric_handler(mocker):
 def test_json_log_metric_handler(mocker):
     mock_logger = mocker.patch("nodestream.metrics.getLogger").return_value
     handler = JsonLogMetricHandler()
-    handler.increment(NodestreamMetricRegistry.RECORDS, 1)
-    handler.decrement(NodestreamMetricRegistry.RECORDS, 1)
+    handler.increment(RECORDS, 1)
+    handler.decrement(RECORDS, 1)
     handler.tick()
     handler.stop()
     mock_logger.info.assert_called_once()
@@ -85,12 +84,12 @@ def test_aggregate_handler(mocker):
     handler.start()
     mock_handler1.start.assert_called_once()
     mock_handler2.start.assert_called_once()
-    handler.increment(NodestreamMetricRegistry.RECORDS, 1)
-    mock_handler1.increment.assert_called_once_with(NodestreamMetricRegistry.RECORDS, 1)
-    mock_handler2.increment.assert_called_once_with(NodestreamMetricRegistry.RECORDS, 1)
-    handler.decrement(NodestreamMetricRegistry.RECORDS, 1)
-    mock_handler1.decrement.assert_called_once_with(NodestreamMetricRegistry.RECORDS, 1)
-    mock_handler2.decrement.assert_called_once_with(NodestreamMetricRegistry.RECORDS, 1)
+    handler.increment(RECORDS, 1)
+    mock_handler1.increment.assert_called_once_with(RECORDS, 1)
+    mock_handler2.increment.assert_called_once_with(RECORDS, 1)
+    handler.decrement(RECORDS, 1)
+    mock_handler1.decrement.assert_called_once_with(RECORDS, 1)
+    mock_handler2.decrement.assert_called_once_with(RECORDS, 1)
     handler.stop()
     mock_handler1.stop.assert_called_once()
     mock_handler2.stop.assert_called_once()
@@ -99,8 +98,8 @@ def test_aggregate_handler(mocker):
 def test_metrics_context(mocker):
     handler = mocker.Mock()
     with Metrics.capture(handler) as metrics:
-        metrics.increment(NodestreamMetricRegistry.RECORDS, 1)
-        handler.increment.assert_called_once_with(NodestreamMetricRegistry.RECORDS, 1)
+        metrics.increment(RECORDS, 1)
+        handler.increment.assert_called_once_with(RECORDS, 1)
     handler.stop.assert_called_once()