Saving relevant files

Author: angelosantos4

nodestream/cli/operations/initialize_logger.py

@@ -4,8 +4,9 @@
 
 from pythonjsonlogger.jsonlogger import JsonFormatter
 
-from ...metrics import Metrics
-from ..commands.nodestream_command import NodestreamCommand
+from nodestream.cli.commands.nodestream_command import BaseCommand
+from nodestream.metrics import Metrics
+from nodestream.logging_metrics import MetricsLoggingHandler
 from .operation import Operation
 
 
@@ -32,9 +33,13 @@ def record_factory(*args, **kwargs):
     )
     logger = logging.getLogger()  # Configure the root logger.
     logger.handlers[0].setFormatter(formatter)
+    
+    # Add metrics handler to track log level counts
+    metrics_handler = MetricsLoggingHandler()
+    logger.addHandler(metrics_handler)
 
 
 class InitializeLogger(Operation):
-    async def perform(self, command: NodestreamCommand) -> Any:
+    async def perform(self, command: BaseCommand) -> Any:
         if command.has_json_logging_set:
             configure_logging_with_json_defaults()

nodestream/logging_metrics.py

@@ -0,0 +1,36 @@
+"""
+Logging handler that increments metrics based on log levels.
+"""
+
+import logging
+from typing import Optional
+
+from .metrics import (
+    LOG_DEBUG_COUNT,
+    LOG_ERROR_COUNT,
+    LOG_INFO_COUNT,
+    LOG_WARNING_COUNT,
+    Metrics,
+)
+
+
+class MetricsLoggingHandler(logging.Handler):
+    """A logging handler that increments metrics based on log levels."""
+
+    def __init__(self, metrics: Optional[Metrics] = None):
+        super().__init__()
+        self.metrics = metrics or Metrics.get()
+
+    def emit(self, record: logging.LogRecord) -> None:
+        """Increment the appropriate metric based on log level."""
+        try:
+            if record.levelno >= logging.ERROR:
+                self.metrics.increment(LOG_ERROR_COUNT)
+            elif record.levelno >= logging.WARNING:
+                self.metrics.increment(LOG_WARNING_COUNT)
+            elif record.levelno >= logging.INFO:
+                self.metrics.increment(LOG_INFO_COUNT)
+            elif record.levelno >= logging.DEBUG:
+                self.metrics.increment(LOG_DEBUG_COUNT)
+        except Exception:
+            self.handleError(record)  # Handles any errors during metric increment

nodestream/metrics.py

@@ -92,6 +92,13 @@ def tick(self):
     "non_fatal_errors", "Number of non-fatal errors", accumulate=True
 )
 FATAL_ERRORS = Metric("fatal_errors", "Number of fatal errors", accumulate=True)
+
+# Logging metrics
+LOG_ERROR_COUNT = Metric("log_error_count", "Number of error log messages", accumulate=True)
+LOG_WARNING_COUNT = Metric("log_warning_count", "Number of warning log messages", accumulate=True)
+LOG_INFO_COUNT = Metric("log_info_count", "Number of info log messages", accumulate=True)
+LOG_DEBUG_COUNT = Metric("log_debug_count", "Number of debug log messages", accumulate=True)
+
 NODES_UPSERTED = Metric(
     "nodes_upserted", "Number of nodes upserted to the graph", accumulate=True
 )