cognite-extractor-utils 7.8.1__tar.gz → 7.9.0__tar.gz

{cognite_extractor_utils-7.8.1 → cognite_extractor_utils-7.9.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: cognite-extractor-utils
-Version: 7.8.1
+Version: 7.9.0
 Summary: Utilities for easier development of extractors for CDF
 Project-URL: repository, https://github.com/cognitedata/python-extractor-utils
 Author-email: Mathias Lohne <mathias.lohne@cognite.com>

{cognite_extractor_utils-7.8.1 → cognite_extractor_utils-7.9.0}/cognite/extractorutils/__init__.py

@@ -16,7 +16,7 @@
 Cognite extractor utils is a Python package that simplifies the development of new extractors.
 """
 
-__version__ = "7.8.1"
+__version__ = "7.9.0"
 from .base import Extractor
 
 __all__ = ["Extractor"]

{cognite_extractor_utils-7.8.1 → cognite_extractor_utils-7.9.0}/cognite/extractorutils/unstable/configuration/models.py

@@ -8,9 +8,11 @@ from collections.abc import Iterator
 from datetime import timedelta
 from enum import Enum
 from pathlib import Path
+from time import sleep
 from typing import Annotated, Any, Literal, TypeVar
 
 from humps import kebabize
+from prometheus_client import REGISTRY, start_http_server
 from pydantic import BaseModel, ConfigDict, Field, GetCoreSchemaHandler
 from pydantic_core import CoreSchema, core_schema
 from typing_extensions import assert_never
@@ -22,8 +24,10 @@ from cognite.client.credentials import (
     OAuthClientCertificate,
     OAuthClientCredentials,
 )
+from cognite.client.data_classes import Asset
 from cognite.extractorutils.configtools._util import _load_certificate_data
 from cognite.extractorutils.exceptions import InvalidConfigError
+from cognite.extractorutils.metrics import AbstractMetricsPusher, CognitePusher, PrometheusPusher
 from cognite.extractorutils.statestore import (
     AbstractStateStore,
     LocalStateStore,
@@ -31,6 +35,7 @@ from cognite.extractorutils.statestore import (
     RawStateStore,
 )
 from cognite.extractorutils.threading import CancellationToken
+from cognite.extractorutils.util import EitherId
 
 __all__ = [
     "AuthenticationConfig",
@@ -43,6 +48,7 @@ __all__ = [
     "LogFileHandlerConfig",
     "LogHandlerConfig",
     "LogLevel",
+    "MetricsConfig",
     "ScheduleConfig",
     "TimeIntervalConfig",
 ]
@@ -405,6 +411,15 @@ class LogLevel(Enum):
     INFO = "INFO"
     DEBUG = "DEBUG"
 
+    @classmethod
+    def _missing_(cls, value: object) -> "LogLevel":
+        if not isinstance(value, str):
+            raise ValueError(f"{value} is not a valid log level")
+        for member in cls:
+            if member.value == value.upper():
+                return member
+        raise ValueError(f"{value} is not a valid log level")
+
 
 class LogFileHandlerConfig(ConfigModel):
     """
@@ -429,11 +444,339 @@ class LogConsoleHandlerConfig(ConfigModel):
 LogHandlerConfig = Annotated[LogFileHandlerConfig | LogConsoleHandlerConfig, Field(discriminator="type")]
 
 
+class EitherIdConfig(ConfigModel):
+    """
+    Configuration parameter representing an ID in CDF, which can either be an external or internal ID.
+
+    An EitherId can only hold one ID type, not both.
+    """
+
+    id: int | None = None
+    external_id: str | None = None
+
+    @property
+    def either_id(self) -> EitherId:
+        """
+        Returns an EitherId object based on the current configuration.
+
+        Raises:
+            TypeError: If both id and external_id are None, or if both are set.
+        """
+        return EitherId(id=self.id, external_id=self.external_id)
+
+
+class _PushGatewayConfig(ConfigModel):
+    """
+    Configuration for pushing metrics to a Prometheus Push Gateway.
+    """
+
+    host: str
+    job_name: str
+    username: str | None = None
+    password: str | None = None
+
+    clear_after: TimeIntervalConfig | None = None
+    push_interval: TimeIntervalConfig = Field(default_factory=lambda: TimeIntervalConfig("30s"))
+
+
+class _PromServerConfig(ConfigModel):
+    """
+    Configuration for pushing metrics to a Prometheus server.
+    """
+
+    port: int = 9000
+    host: str = "0.0.0.0"
+
+
+class _CogniteMetricsConfig(ConfigModel):
+    """
+    Configuration for pushing metrics to Cognite Data Fusion.
+    """
+
+    external_id_prefix: str
+    asset_name: str | None = None
+    asset_external_id: str | None = None
+    data_set: EitherIdConfig | None = None
+
+    push_interval: TimeIntervalConfig = Field(default_factory=lambda: TimeIntervalConfig("30s"))
+
+
+class MetricsPushManager:
+    """
+    Manages the pushing of metrics to various backends.
+
+    Starts and stops pushers based on a given configuration.
+
+    Args:
+        metrics_config: Configuration for the metrics to be pushed.
+        cdf_client: The CDF tenant to upload time series to
+        cancellation_token: Event object to be used as a thread cancelation event
+    """
+
+    def __init__(
+        self,
+        metrics_config: "MetricsConfig",
+        cdf_client: CogniteClient,
+        cancellation_token: CancellationToken | None = None,
+    ) -> None:
+        """
+        Initialize the MetricsPushManager.
+        """
+        self.metrics_config = metrics_config
+        self.cdf_client = cdf_client
+        self.cancellation_token = cancellation_token
+        self.pushers: list[AbstractMetricsPusher] = []
+        self.clear_on_stop: dict[AbstractMetricsPusher, int] = {}
+
+    def start(self) -> None:
+        """
+        Start all metric pushers.
+        """
+        push_gateways = self.metrics_config.push_gateways or []
+        for counter, push_gateway in enumerate(push_gateways):
+            prometheus_pusher = PrometheusPusher(
+                job_name=push_gateway.job_name,
+                username=push_gateway.username,
+                password=push_gateway.password,
+                url=push_gateway.host,
+                push_interval=push_gateway.push_interval.seconds,
+                thread_name=f"MetricsPusher_{counter}",
+                cancellation_token=self.cancellation_token,
+            )
+            prometheus_pusher.start()
+            self.pushers.append(prometheus_pusher)
+            if push_gateway.clear_after is not None:
+                self.clear_on_stop[prometheus_pusher] = push_gateway.clear_after.seconds
+
+        if self.metrics_config.cognite:
+            asset = None
+            if self.metrics_config.cognite.asset_name and self.metrics_config.cognite.asset_external_id:
+                asset = Asset(
+                    name=self.metrics_config.cognite.asset_name,
+                    external_id=self.metrics_config.cognite.asset_external_id,
+                )
+            cognite_pusher = CognitePusher(
+                cdf_client=self.cdf_client,
+                external_id_prefix=self.metrics_config.cognite.external_id_prefix,
+                push_interval=self.metrics_config.cognite.push_interval.seconds,
+                asset=asset,
+                data_set=self.metrics_config.cognite.data_set.either_id
+                if self.metrics_config.cognite.data_set
+                else None,
+                thread_name="CogniteMetricsPusher",
+                cancellation_token=self.cancellation_token,
+            )
+            cognite_pusher.start()
+            self.pushers.append(cognite_pusher)
+
+        if self.metrics_config.server:
+            start_http_server(self.metrics_config.server.port, self.metrics_config.server.host, registry=REGISTRY)
+
+    def stop(self) -> None:
+        """
+        Stop all metric pushers.
+        """
+        for pusher in self.pushers:
+            pusher.stop()
+
+        # Clear Prometheus pushers gateways if required
+        if self.clear_on_stop:
+            wait_time = max(self.clear_on_stop.values())
+            sleep(wait_time)
+            for pusher in (p for p in self.clear_on_stop if isinstance(p, PrometheusPusher)):
+                pusher.clear_gateway()
+
+
+class MetricsConfig(ConfigModel):
+    """
+    Destination(s) for metrics.
+
+    Including options for one or several Prometheus push gateways, and pushing as CDF Time Series.
+    """
+
+    push_gateways: list[_PushGatewayConfig] | None = None
+    cognite: _CogniteMetricsConfig | None = None
+    server: _PromServerConfig | None = None
+
+    def create_manager(
+        self, cdf_client: CogniteClient, cancellation_token: CancellationToken | None = None
+    ) -> MetricsPushManager:
+        """
+        Create a MetricsPushManager based on the current configuration.
+
+        Args:
+            cdf_client: An instance of CogniteClient to interact with CDF.
+            cancellation_token: Optional token to signal cancellation of metric pushing.
+
+        Returns:
+            MetricsPushManager: An instance of MetricsPushManager configured with the provided parameters.
+        """
+        return MetricsPushManager(self, cdf_client, cancellation_token)
+
+
 # Mypy BS
 def _log_handler_default() -> list[LogHandlerConfig]:
     return [LogConsoleHandlerConfig(type="console", level=LogLevel.INFO)]
 
 
+class FileSizeConfig:
+    """
+    Configuration parameter for setting a file size.
+    """
+
+    def __init__(self, expression: str) -> None:
+        self._bytes, self._expression = FileSizeConfig._parse_expression(expression)
+
+    @classmethod
+    def __get_pydantic_core_schema__(cls, source_type: Any, handler: GetCoreSchemaHandler) -> CoreSchema:  # noqa: ANN401
+        """
+        Pydantic hook to define how this class should be serialized/deserialized.
+
+        This allows the class to be used as a field in Pydantic models.
+        """
+        return core_schema.no_info_after_validator_function(cls, handler(str | int))
+
+    def __eq__(self, other: object) -> bool:
+        """
+        Two FileSizeConfig objects are equal if they have the same number of bytes.
+        """
+        if not isinstance(other, FileSizeConfig):
+            return NotImplemented
+        return self._bytes == other._bytes
+
+    def __hash__(self) -> int:
+        """
+        Hash function for FileSizeConfig based on the number of bytes.
+        """
+        return hash(self._bytes)
+
+    @classmethod
+    def _parse_expression(cls, expression: str) -> tuple[int, str]:
+        sizes = {
+            "kb": 1000,
+            "mb": 1_000_000,
+            "gb": 1_000_000_000,
+            "tb": 1_000_000_000_000,
+            "kib": 1024,
+            "mib": 1_048_576,
+            "gib": 1_073_741_824,
+            "tib": 1_099_511_627_776,
+        }
+
+        expression_normalized = expression.strip().lower()
+        try:
+            num_value = float(expression_normalized)
+            return int(num_value), expression
+        except ValueError:
+            pass
+
+        match = re.match(r"^([0-9]*\.?[0-9]+)\s*([a-zA-Z]*)$", expression_normalized)
+        if not match:
+            raise InvalidConfigError(f"Invalid file size format: '{expression}'. Must start with a number.")
+
+        num_str, unit_str = match.groups()
+        try:
+            num_value = float(num_str)
+        except ValueError as e:
+            raise InvalidConfigError(f"Invalid numeric value in file size: '{num_str}'") from e
+
+        if not unit_str:
+            return int(num_value), expression
+
+        unit_lower = unit_str.lower()
+        if unit_lower in sizes:
+            return int(num_value * sizes[unit_lower]), expression
+
+        raise InvalidConfigError(f"Invalid unit for file size: '{unit_str}'. Valid units: {list(sizes.keys())}")
+
+    @property
+    def bytes(self) -> int:
+        """
+        File size in bytes.
+        """
+        return self._bytes
+
+    @property
+    def kilobytes(self) -> float:
+        """
+        File size in kilobytes.
+        """
+        return self._bytes / 1000
+
+    @property
+    def megabytes(self) -> float:
+        """
+        File size in megabytes.
+        """
+        return self._bytes / 1_000_000
+
+    @property
+    def gigabytes(self) -> float:
+        """
+        File size in gigabytes.
+        """
+        return self._bytes / 1_000_000_000
+
+    @property
+    def terabytes(self) -> float:
+        """
+        File size in terabytes.
+        """
+        return self._bytes / 1_000_000_000_000
+
+    @property
+    def kibibytes(self) -> float:
+        """
+        File size in kibibytes (1024 bytes).
+        """
+        return self._bytes / 1024
+
+    @property
+    def mebibytes(self) -> float:
+        """
+        File size in mebibytes (1024 kibibytes).
+        """
+        return self._bytes / 1_048_576
+
+    @property
+    def gibibytes(self) -> float:
+        """
+        File size in gibibytes (1024 mebibytes).
+        """
+        return self._bytes / 1_073_741_824
+
+    @property
+    def tebibytes(self) -> float:
+        """
+        File size in tebibytes (1024 gibibytes).
+        """
+        return self._bytes / 1_099_511_627_776
+
+    def __int__(self) -> int:
+        """
+        Returns the file size as bytes.
+        """
+        return int(self._bytes)
+
+    def __float__(self) -> float:
+        """
+        Returns the file size as bytes.
+        """
+        return float(self._bytes)
+
+    def __str__(self) -> str:
+        """
+        Returns the file size as a human readable string.
+        """
+        return self._expression
+
+    def __repr__(self) -> str:
+        """
+        Returns the file size as a human readable string.
+        """
+        return self._expression
+
+
 class RawDestinationConfig(ConfigModel):
     """
     Configuration parameters for using Raw.
@@ -523,6 +866,7 @@ class ExtractorConfig(ConfigModel):
     """
 
     state_store: StateStoreConfig | None = None
+    metrics: MetricsConfig | None = None
     log_handlers: list[LogHandlerConfig] = Field(default_factory=_log_handler_default)
     retry_startup: bool = True
 

{cognite_extractor_utils-7.8.1 → cognite_extractor_utils-7.9.0}/cognite/extractorutils/unstable/core/base.py

@@ -59,6 +59,7 @@ from humps import pascalize
 from typing_extensions import Self, assert_never
 
 from cognite.extractorutils._inner_util import _resolve_log_level
+from cognite.extractorutils.metrics import BaseMetrics
 from cognite.extractorutils.statestore import (
     AbstractStateStore,
     LocalStateStore,
@@ -147,7 +148,9 @@ class Extractor(Generic[ConfigType], CogniteLogger):
 
     cancellation_token: CancellationToken
 
-    def __init__(self, config: FullConfig[ConfigType], checkin_worker: CheckinWorker) -> None:
+    def __init__(
+        self, config: FullConfig[ConfigType], checkin_worker: CheckinWorker, metrics: BaseMetrics | None = None
+    ) -> None:
         self._logger = logging.getLogger(f"{self.EXTERNAL_ID}.main")
         self._checkin_worker = checkin_worker
 
@@ -156,6 +159,7 @@ class Extractor(Generic[ConfigType], CogniteLogger):
 
         self.connection_config = config.connection_config
         self.application_config = config.application_config
+        self.metrics_config = config.application_config.metrics
         self.current_config_revision: ConfigRevision = config.current_config_revision
         self.log_level_override = config.log_level_override
 
@@ -170,6 +174,13 @@ class Extractor(Generic[ConfigType], CogniteLogger):
 
         self._tasks: list[Task] = []
         self._start_time: datetime
+        self._metrics: BaseMetrics | None = metrics
+
+        self.metrics_push_manager = (
+            self.metrics_config.create_manager(self.cognite_client, cancellation_token=self.cancellation_token)
+            if self.metrics_config
+            else None
+        )
 
         self.__init_tasks__()
 
@@ -367,8 +378,10 @@ class Extractor(Generic[ConfigType], CogniteLogger):
         self.cancellation_token.cancel()
 
     @classmethod
-    def _init_from_runtime(cls, config: FullConfig[ConfigType], checkin_worker: CheckinWorker) -> Self:
-        return cls(config, checkin_worker)
+    def _init_from_runtime(
+        cls, config: FullConfig[ConfigType], checkin_worker: CheckinWorker, metrics: BaseMetrics
+    ) -> Self:
+        return cls(config, checkin_worker, metrics)
 
     def add_task(self, task: Task) -> None:
         """
@@ -438,6 +451,8 @@ class Extractor(Generic[ConfigType], CogniteLogger):
         self.state_store.start()
 
         Thread(target=self._run_checkin, name="ExtractorCheckin", daemon=True).start()
+        if self.metrics_push_manager:
+            self.metrics_push_manager.start()
 
     def stop(self) -> None:
         """
@@ -446,6 +461,8 @@ class Extractor(Generic[ConfigType], CogniteLogger):
         Instead of calling this method directly, it is recommended to use the context manager interface by using the
         ``with`` statement, which ensures proper cleanup on exit.
         """
+        if self.metrics_push_manager:
+            self.metrics_push_manager.stop()
         self.cancellation_token.cancel()
 
     def __enter__(self) -> Self:

{cognite_extractor_utils-7.8.1 → cognite_extractor_utils-7.9.0}/cognite/extractorutils/unstable/core/runtime.py

@@ -47,6 +47,7 @@ from cognite.client.exceptions import (
     CogniteAuthError,
     CogniteConnectionError,
 )
+from cognite.extractorutils.metrics import BaseMetrics
 from cognite.extractorutils.threading import CancellationToken
 from cognite.extractorutils.unstable.configuration.exceptions import InvalidArgumentError, InvalidConfigError
 from cognite.extractorutils.unstable.configuration.loaders import (
@@ -78,6 +79,7 @@ def _extractor_process_entrypoint(
     controls: _RuntimeControls,
     config: FullConfig,
     checkin_worker: CheckinWorker,
+    metrics: BaseMetrics | None = None,
 ) -> None:
     logger = logging.getLogger(f"{extractor_class.EXTERNAL_ID}.runtime")
     checkin_worker.active_revision = config.current_config_revision
@@ -85,7 +87,9 @@ def _extractor_process_entrypoint(
     checkin_worker.set_on_revision_change_handler(lambda _: on_revision_changed(controls))
     if config.application_config.retry_startup:
         checkin_worker.set_retry_startup(config.application_config.retry_startup)
-    extractor = extractor_class._init_from_runtime(config, checkin_worker)
+    if not metrics:
+        metrics = BaseMetrics(extractor_name=extractor_class.NAME, extractor_version=extractor_class.VERSION)
+    extractor = extractor_class._init_from_runtime(config, checkin_worker, metrics)
     extractor._attach_runtime_controls(
         cancel_event=controls.cancel_event,
         message_queue=controls.message_queue,
@@ -135,11 +139,13 @@ class Runtime(Generic[ExtractorType]):
     def __init__(
         self,
         extractor: type[ExtractorType],
+        metrics: BaseMetrics | None = None,
     ) -> None:
         self._extractor_class = extractor
         self._cancellation_token = CancellationToken()
         self._cancellation_token.cancel_on_interrupt()
         self._message_queue: Queue[RuntimeMessage] = Queue()
+        self._metrics = metrics
         self.logger = logging.getLogger(f"{self._extractor_class.EXTERNAL_ID}.runtime")
         self._setup_logging()
         self._cancel_event: MpEvent | None = None
@@ -268,7 +274,7 @@ class Runtime(Generic[ExtractorType]):
 
         process = Process(
             target=_extractor_process_entrypoint,
-            args=(self._extractor_class, controls, config, checkin_worker),
+            args=(self._extractor_class, controls, config, checkin_worker, self._metrics),
         )
 
         process.start()

{cognite_extractor_utils-7.8.1 → cognite_extractor_utils-7.9.0}/cognite/extractorutils/uploader/time_series.py

@@ -53,7 +53,7 @@ from cognite.extractorutils.uploader._metrics import (
 from cognite.extractorutils.util import EitherId, cognite_exceptions, retry
 
 MIN_DATAPOINT_TIMESTAMP = -2208988800000
-MAX_DATAPOINT_STRING_LENGTH = 255
+MAX_DATAPOINT_STRING_BYTES = 1023
 MAX_DATAPOINT_VALUE = 1e100
 MIN_DATAPOINT_VALUE = -1e100
 
@@ -154,7 +154,7 @@ class BaseTimeSeriesUploadQueue(AbstractUploadQueue, Generic[IdType]):
                 math.isnan(value) or math.isinf(value) or value > MAX_DATAPOINT_VALUE or value < MIN_DATAPOINT_VALUE
             )
         elif isinstance(value, str):
-            return len(value) <= MAX_DATAPOINT_STRING_LENGTH
+            return len(value.encode("utf-8")) <= MAX_DATAPOINT_STRING_BYTES
         return not isinstance(value, datetime)
 
     def _is_datapoint_valid(

{cognite_extractor_utils-7.8.1 → cognite_extractor_utils-7.9.0}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "cognite-extractor-utils"
-version = "7.8.1"
+version = "7.9.0"
 description = "Utilities for easier development of extractors for CDF"
 authors = [
     {name = "Mathias Lohne", email = "mathias.lohne@cognite.com"}