Fix minor issues

Author: Hmnt39

cognite/extractorutils/unstable/core/base.py

@@ -22,6 +22,11 @@ class MyConfig(ExtractorConfig):
         another_parameter: int
         schedule: ScheduleConfig
 
+    class MyMetrics(BaseMetrics):
+        def __init__(self, extractor_name: str, extractor_version: str):
+            super().__init__(extractor_name, extractor_version)
+            self.custom_counter = Counter("custom_counter", "A custom counter")
+
     class MyExtractor(Extractor[MyConfig]):
         NAME = "My Extractor"
         EXTERNAL_ID = "my-extractor"
@@ -30,6 +35,9 @@ class MyExtractor(Extractor[MyConfig]):
 
         CONFIG_TYPE = MyConfig
 
+        # Override metrics type annotation for IDE support
+        metrics: MyMetrics
+
         def __init_tasks__(self) -> None:
             self.add_task(
                 ScheduledTask(
@@ -42,6 +50,8 @@ def __init_tasks__(self) -> None:
 
         def my_task_function(self, task_context: TaskContext) -> None:
             task_context.logger.info("Running my task")
+            # IDE will now autocomplete custom_counter
+            self.metrics.custom_counter.inc()
 """
 
 import logging
@@ -123,7 +133,7 @@ def __init__(
         self.application_config = application_config
         self.current_config_revision: ConfigRevision = current_config_revision
         self.log_level_override = log_level_override
-        self.metrics_class = metrics_class
+        self.metrics_class: type[BaseMetrics] | None = metrics_class
 
 
 class Extractor(Generic[ConfigType], CogniteLogger):
@@ -148,7 +158,7 @@ class Extractor(Generic[ConfigType], CogniteLogger):
     RESTART_POLICY: RestartPolicy = WHEN_CONTINUOUS_TASKS_CRASHES
     USE_DEFAULT_STATE_STORE: bool = True
     _statestore_singleton: AbstractStateStore | None = None
-    _metrics_singleton: BaseMetrics
+    _metrics_singleton: BaseMetrics | None = None
 
     cancellation_token: CancellationToken
 
@@ -271,16 +281,15 @@ def _load_metrics(self, metrics_class: type[BaseMetrics] | None = None) -> BaseM
         Reuses existing singleton if available to avoid Prometheus registry conflicts.
         """
         if Extractor._metrics_singleton is not None:
-            self.metrics = Extractor._metrics_singleton
-            return self.metrics
+            return Extractor._metrics_singleton
 
-        if metrics_class:
-            self.metrics = safe_get(metrics_class)
+        if metrics_class and issubclass(metrics_class, BaseMetrics):
+            metrics_instance = safe_get(metrics_class)
         else:
-            self.metrics = BaseMetrics(extractor_name=self.EXTERNAL_ID, extractor_version=self.VERSION)
+            metrics_instance = safe_get(BaseMetrics, extractor_name=self.EXTERNAL_ID, extractor_version=self.VERSION)
 
-        Extractor._metrics_singleton = self.metrics
-        return self.metrics
+        Extractor._metrics_singleton = metrics_instance
+        return metrics_instance
 
     @classmethod
     def get_current_metrics(cls) -> BaseMetrics:

tests/test_unstable/conftest.py

@@ -16,6 +16,7 @@
 from cognite.client import CogniteClient
 from cognite.client.config import ClientConfig
 from cognite.client.credentials import OAuthClientCredentials
+from cognite.extractorutils import metrics
 from cognite.extractorutils.metrics import BaseMetrics
 from cognite.extractorutils.unstable.configuration.models import (
     ConnectionConfig,
@@ -39,6 +40,7 @@ class variables are reset, and Prometheus collectors are unregistered,
     # Clean up before test
     Extractor._statestore_singleton = None
     Extractor._metrics_singleton = None
+    metrics._metrics_singularities.clear()
 
     # Unregister all collectors to prevent "Duplicated timeseries" errors
     collectors = list(REGISTRY._collector_to_names.keys())
@@ -51,6 +53,7 @@ class variables are reset, and Prometheus collectors are unregistered,
     # Clean up after test
     Extractor._statestore_singleton = None
     Extractor._metrics_singleton = None
+    metrics._metrics_singularities.clear()
 
     # Unregister all collectors again
     collectors = list(REGISTRY._collector_to_names.keys())