Merge branch 'master' into ann-rule

Author: mathialo

CHANGELOG.md

@@ -12,12 +12,32 @@ Changes are grouped as follows
 - `Fixed` for any bug fixes.
 - `Security` in case of vulnerabilities.
 
-## Next
+## 7.7.0
+
+### Added
+
+* Added the `CDMTimeSeriesUploadQueue` class to upload time series datapoints to instances in the Cognite Data Model (CDM). This new uploader uses `NodeId` for identification and can automatically create `CogniteExtractorTimeSeriesApply` instances if they are missing.
+
+### Changed
+
+* Refactored the time series uploaders for better code reuse and maintainability. A new generic base class, `BaseTimeSeriesUploadQueue`, has been introduced to contain common logic for datapoint validation and sanitization. Both the existing `TimeSeriesUploadQueue` and the new `CDMTimeSeriesUploadQueue` now inherit from this base class.
+
+
+## 7.6.0
+
+### Added
+
+* Added a common method `_report_run` that can be called from any extractors using this library to report a extractor run into Extractor Pipelines. This automatically handles truncating the message to less than 1000 characters if required.
+
+### Fixed
+
+* Fixed `_report_success` calls failing if the `success_message` was longer than 1000 characters
 
 ### Changed
 
 * In the `unstable` package: `-l` / `--local-override` is renamed to `-f` / `--force-local-config`
 
+
 ## 7.5.14
 
 ### Added

cognite/extractorutils/__init__.py

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

cognite/extractorutils/base.py

@@ -1,6 +1,7 @@
 """
 Module containing the base class for extractors.
 """
+
 #  Copyright 2021 Cognite AS
 #
 #  Licensed under the Apache License, Version 2.0 (the "License");
@@ -22,20 +23,32 @@
 from collections.abc import Callable
 from dataclasses import is_dataclass
 from enum import Enum
+from textwrap import shorten
 from threading import Thread
 from types import TracebackType
-from typing import Any, Generic, TypeVar
+from typing import Any, Generic, Literal, TypeVar
 
 from dotenv import load_dotenv
 
 from cognite.client import CogniteClient
 from cognite.client.data_classes import ExtractionPipeline, ExtractionPipelineRun
-from cognite.extractorutils.configtools import BaseConfig, ConfigResolver, StateStoreConfig
+from cognite.client.exceptions import CogniteAPIError
+from cognite.extractorutils.configtools import (
+    BaseConfig,
+    ConfigResolver,
+    StateStoreConfig,
+)
 from cognite.extractorutils.exceptions import InvalidConfigError
 from cognite.extractorutils.metrics import BaseMetrics
-from cognite.extractorutils.statestore import AbstractStateStore, LocalStateStore, NoStateStore
+from cognite.extractorutils.statestore import (
+    AbstractStateStore,
+    LocalStateStore,
+    NoStateStore,
+)
 from cognite.extractorutils.threading import CancellationToken
 
+ReportStatus = Literal["success", "failure", "seen"]
+
 
 class ReloadConfigAction(Enum):
     """
@@ -227,19 +240,53 @@ def recursive_find_state_store(d: dict[str, Any]) -> StateStoreConfig | None:
 
         Extractor._statestore_singleton = self.state_store
 
-    def _report_success(self) -> None:
+    def _report_run(self, status: ReportStatus, message: str) -> None:
         """
-        Called on a successful exit of the extractor.
+        Report the status of the extractor run to the extraction pipeline.
+
+        Args:
+            status: Status of the run, either success or failure or seen
+            message: Message to report to the extraction pipeline
         """
+        MAX_MESSAGE_LENGTH_FOR_EXTRACTION_PIPELINE_RUN = 1000
         if self.extraction_pipeline:
-            self.logger.info("Reporting new successful run")
-            self.cognite_client.extraction_pipelines.runs.create(
-                ExtractionPipelineRun(
+            try:
+                message = message or ""
+                shortened_message = shorten(
+                    message,
+                    width=MAX_MESSAGE_LENGTH_FOR_EXTRACTION_PIPELINE_RUN,
+                    placeholder="...",
+                )
+
+                run = ExtractionPipelineRun(
                     extpipe_external_id=self.extraction_pipeline.external_id,
-                    status="success",
-                    message=self.success_message,
+                    status=status,
+                    message=shortened_message,
                 )
-            )
+
+                self.logger.info(f"Reporting new {status} run: {message}")
+                self.cognite_client.extraction_pipelines.runs.create(run)
+            except CogniteAPIError as e:
+                self.logger.exception(f"Error while reporting run - status {status} - message {message} . Error: {e!s}")
+
+    def _report_success(self, message: str | None = None) -> None:
+        """
+        Called on a successful exit of the extractor.
+
+        Args:
+            message: Message to report to the extraction pipeline. If not provided, Extractor.success_message is taken.
+        """
+        message = message or self.success_message
+        self._report_run("success", message)
+
+    def _report_failure(self, message: str) -> None:
+        """
+        Called on an unsuccessful exit of the extractor.
+
+        Args:
+            message: Message to report to the extraction pipeline
+        """
+        self._report_run("failure", message)
 
     def _report_error(self, exception: BaseException) -> None:
         """
@@ -248,16 +295,8 @@ def _report_error(self, exception: BaseException) -> None:
         Args:
             exception: Exception object that caused the extractor to fail
         """
-        self.logger.error("Unexpected error during extraction", exc_info=exception)
-        if self.extraction_pipeline:
-            message = f"{type(exception).__name__}: {exception!s}"[:1000]
-
-            self.logger.info(f"Reporting new failed run: {message}")
-            self.cognite_client.extraction_pipelines.runs.create(
-                ExtractionPipelineRun(
-                    extpipe_external_id=self.extraction_pipeline.external_id, status="failure", message=message
-                )
-            )
+        message = f"{type(exception).__name__}: {exception!s}"
+        self._report_run("failure", message)
 
     def __enter__(self) -> "Extractor":
         """
@@ -279,7 +318,9 @@ def __enter__(self) -> "Extractor":
         try:
             self._initial_load_config(override_path=self.config_file_path)
         except InvalidConfigError as e:
-            print("Critical error: Could not read config file", file=sys.stderr)  # noqa: T201
+            print(  # noqa: T201
+                "Critical error: Could not read config file", file=sys.stderr
+            )
             print(str(e), file=sys.stderr)  # noqa: T201
             sys.exit(1)
 
@@ -338,7 +379,10 @@ def heartbeat_loop() -> None:
         return self
 
     def __exit__(
-        self, exc_type: type[BaseException] | None, exc_val: BaseException | None, exc_tb: TracebackType | None
+        self,
+        exc_type: type[BaseException] | None,
+        exc_val: BaseException | None,
+        exc_tb: TracebackType | None,
     ) -> bool:
         """
         Shuts down the extractor. Makes sure states are preserved, that all uploads of data and metrics are done, etc.
@@ -381,7 +425,12 @@ def run(self) -> None:
         if not self.started:
             raise ValueError("You must run the extractor in a context manager")
         if self.run_handle:
-            self.run_handle(self.cognite_client, self.state_store, self.config, self.cancellation_token)
+            self.run_handle(
+                self.cognite_client,
+                self.state_store,
+                self.config,
+                self.cancellation_token,
+            )
         else:
             raise ValueError("No run_handle defined")
 

cognite/extractorutils/metrics.py

@@ -54,6 +54,7 @@ def __init__(self):
 
 from cognite.client import CogniteClient
 from cognite.client.data_classes import Asset, Datapoints, DatapointsArray, TimeSeries
+from cognite.client.data_classes.data_modeling import NodeId
 from cognite.client.exceptions import CogniteDuplicatedError
 from cognite.extractorutils.threading import CancellationToken
 from cognite.extractorutils.util import EitherId
@@ -413,7 +414,7 @@ def _push_to_server(self) -> None:
         """
         timestamp = int(arrow.get().float_timestamp * 1000)
 
-        datapoints: list[dict[str, str | int | list[Any] | Datapoints | DatapointsArray]] = []
+        datapoints: list[dict[str, str | int | list[Any] | Datapoints | DatapointsArray | NodeId]] = []
 
         for metric in REGISTRY.collect():
             if isinstance(metric, Metric) and metric.type in ["gauge", "counter"]:

cognite/extractorutils/unstable/configuration/models.py

@@ -4,6 +4,7 @@
 
 import os
 import re
+from collections.abc import Iterator
 from datetime import timedelta
 from enum import Enum
 from pathlib import Path
@@ -53,23 +54,44 @@ class ConfigModel(BaseModel):
     )
 
 
-class _ClientCredentialsConfig(ConfigModel):
-    type: Literal["client-credentials"]
+class Scopes(str):
+    def __init__(self, scopes: str) -> None:
+        self._scopes = list(scopes.split(" "))
+
+    @classmethod
+    def __get_pydantic_core_schema__(cls, source_type: Any, handler: GetCoreSchemaHandler) -> CoreSchema:
+        return core_schema.no_info_after_validator_function(cls, handler(str))
+
+    def __eq__(self, other: object) -> bool:
+        if not isinstance(other, Scopes):
+            return NotImplemented
+        return self._scopes == other._scopes
+
+    def __hash__(self) -> int:
+        return hash(self._scopes)
+
+    def __iter__(self) -> Iterator[str]:
+        return iter(self._scopes)
+
+
+class BaseCredentialsConfig(ConfigModel):
     client_id: str
+    scopes: Scopes
+
+
+class _ClientCredentialsConfig(BaseCredentialsConfig):
+    type: Literal["client-credentials"]
     client_secret: str
     token_url: str
-    scopes: list[str]
     resource: str | None = None
     audience: str | None = None
 
 
-class _ClientCertificateConfig(ConfigModel):
+class _ClientCertificateConfig(BaseCredentialsConfig):
     type: Literal["client-certificate"]
-    client_id: str
     path: Path
     password: str | None = None
     authority_url: str
-    scopes: list[str]
 
 
 AuthenticationConfig = Annotated[_ClientCredentialsConfig | _ClientCertificateConfig, Field(discriminator="type")]
@@ -191,18 +213,26 @@ def __repr__(self) -> str:
         return self._expression
 
 
-class _ConnectionParameters(ConfigModel):
-    gzip_compression: bool = False
-    status_forcelist: list[int] = Field(default_factory=lambda: [429, 502, 503, 504])
-    max_retries: int = 10
-    max_retries_connect: int = 3
-    max_retry_backoff: TimeIntervalConfig = Field(default_factory=lambda: TimeIntervalConfig("30s"))
-    max_connection_pool_size: int = 50
-    ssl_verify: bool = True
-    proxies: dict[str, str] = Field(default_factory=dict)
+class RetriesConfig(ConfigModel):
+    max_retries: int = Field(default=10, ge=-1)
+    max_backoff: TimeIntervalConfig = Field(default_factory=lambda: TimeIntervalConfig("30s"))
     timeout: TimeIntervalConfig = Field(default_factory=lambda: TimeIntervalConfig("30s"))
 
 
+class SslCertificatesConfig(ConfigModel):
+    verify: bool = True
+    allow_list: list[str] | None = None
+
+
+class ConnectionParameters(ConfigModel):
+    retries: RetriesConfig = Field(default_factory=RetriesConfig)
+    ssl_certificates: SslCertificatesConfig = Field(default_factory=SslCertificatesConfig)
+
+
+class IntegrationConfig(ConfigModel):
+    external_id: str
+
+
 class ConnectionConfig(ConfigModel):
     """
     Configuration for connecting to a Cognite Data Fusion project.
@@ -216,11 +246,11 @@ class ConnectionConfig(ConfigModel):
     project: str
     base_url: str
 
-    integration: str
+    integration: IntegrationConfig
 
     authentication: AuthenticationConfig
 
-    connection: _ConnectionParameters = Field(default_factory=_ConnectionParameters)
+    connection: ConnectionParameters = Field(default_factory=ConnectionParameters)
 
     def get_cognite_client(self, client_name: str) -> CogniteClient:
         """
@@ -235,14 +265,9 @@ def get_cognite_client(self, client_name: str) -> CogniteClient:
         from cognite.client.config import global_config
 
         global_config.disable_pypi_version_check = True
-        global_config.disable_gzip = not self.connection.gzip_compression
-        global_config.status_forcelist = set(self.connection.status_forcelist)
-        global_config.max_retries = self.connection.max_retries
-        global_config.max_retries_connect = self.connection.max_retries_connect
-        global_config.max_retry_backoff = self.connection.max_retry_backoff.seconds
-        global_config.max_connection_pool_size = self.connection.max_connection_pool_size
-        global_config.disable_ssl = not self.connection.ssl_verify
-        global_config.proxies = self.connection.proxies
+        global_config.max_retries = self.connection.retries.max_retries
+        global_config.max_retry_backoff = self.connection.retries.max_backoff.seconds
+        global_config.disable_ssl = not self.connection.ssl_certificates.verify
 
         credential_provider: CredentialProvider
         match self.authentication:
@@ -270,7 +295,7 @@ def get_cognite_client(self, client_name: str) -> CogniteClient:
                     client_id=client_certificate.client_id,
                     cert_thumbprint=str(thumbprint),
                     certificate=str(key),
-                    scopes=client_certificate.scopes,
+                    scopes=list(client_certificate.scopes),
                 )
 
             case _:
@@ -280,7 +305,7 @@ def get_cognite_client(self, client_name: str) -> CogniteClient:
             project=self.project,
             base_url=self.base_url,
             client_name=client_name,
-            timeout=self.connection.timeout.seconds,
+            timeout=self.connection.retries.timeout.seconds,
             credentials=credential_provider,
         )
 
@@ -315,7 +340,9 @@ def from_environment(cls) -> "ConnectionConfig":
                 client_id=os.environ["COGNITE_CLIENT_ID"],
                 client_secret=os.environ["COGNITE_CLIENT_SECRET"],
                 token_url=os.environ["COGNITE_TOKEN_URL"],
-                scopes=os.environ["COGNITE_TOKEN_SCOPES"].split(","),
+                scopes=Scopes(
+                    os.environ["COGNITE_TOKEN_SCOPES"],
+                ),
             )
         elif "COGNITE_CLIENT_CERTIFICATE_PATH" in os.environ:
             auth = _ClientCertificateConfig(
@@ -324,15 +351,17 @@ def from_environment(cls) -> "ConnectionConfig":
                 path=Path(os.environ["COGNITE_CLIENT_CERTIFICATE_PATH"]),
                 password=os.environ.get("COGNITE_CLIENT_CERTIFICATE_PATH"),
                 authority_url=os.environ["COGNITE_AUTHORITY_URL"],
-                scopes=os.environ["COGNITE_TOKEN_SCOPES"].split(","),
+                scopes=Scopes(
+                    os.environ["COGNITE_TOKEN_SCOPES"],
+                ),
             )
         else:
             raise KeyError("Missing auth, either COGNITE_CLIENT_SECRET or COGNITE_CLIENT_CERTIFICATE_PATH must be set")
 
         return ConnectionConfig(
             project=os.environ["COGNITE_PROJECT"],
             base_url=os.environ["COGNITE_BASE_URL"],
-            integration=os.environ["COGNITE_INTEGRATION"],
+            integration=IntegrationConfig(external_id=os.environ["COGNITE_INTEGRATION"]),
             authentication=auth,
         )