cognite-extractor-utils 7.5.13__tar.gz → 7.6.0__tar.gz

{cognite_extractor_utils-7.5.13 → cognite_extractor_utils-7.6.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: cognite-extractor-utils
-Version: 7.5.13
+Version: 7.6.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.5.13 → cognite_extractor_utils-7.6.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.5.13"
+__version__ = "7.6.0"
 from .base import Extractor
 
 __all__ = ["Extractor"]

{cognite_extractor_utils-7.5.13 → cognite_extractor_utils-7.6.0}/cognite/extractorutils/_inner_util.py

@@ -13,7 +13,7 @@
 #  limitations under the License.
 
 """
-A module containing utilities meant for use inside the extractor-utils package
+A module containing utilities meant for use inside the extractor-utils package.
 """
 
 import json

{cognite_extractor_utils-7.5.13 → cognite_extractor_utils-7.6.0}/cognite/extractorutils/base.py

@@ -1,3 +1,7 @@
+"""
+Module containing the base class for extractors.
+"""
+
 #  Copyright 2021 Cognite AS
 #
 #  Licensed under the Apache License, Version 2.0 (the "License");
@@ -12,28 +16,45 @@
 #  See the License for the specific language governing permissions and
 #  limitations under the License.
 
+import contextlib
 import logging
 import os
 import sys
 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):
+    """
+    Enum for actions to take when a config file is reloaded.
+    """
+
     DO_NOTHING = 1
     REPLACE_ATTRIBUTE = 2
     SHUTDOWN = 3
@@ -125,8 +146,10 @@ class Extractor(Generic[CustomConfigClass]):
 
     def _initial_load_config(self, override_path: str | None = None) -> None:
         """
-        Load a configuration file, either from the specified path, or by a path specified by the user in a command line
-        arg. Will quit further execution of no path is given.
+        Load a configuration file.
+
+        Either from the specified path, or from a path specified by the user in a command line arg. Will quit further
+        execution if no path is specified.
 
         Args:
             override_path: Optional override for file path, ie don't parse command line arguments
@@ -149,6 +172,11 @@ class Extractor(Generic[CustomConfigClass]):
             Thread(target=config_refresher, name="ConfigReloader", daemon=True).start()
 
     def reload_config_callback(self) -> None:
+        """
+        If the reload_config_action was set to CALLBACK, this method will be called when the config file is reloaded.
+
+        This method should be overridden in subclasses to provide custom behavior when the config file is reloaded.
+        """
         self.logger.error("Method for reloading configs has not been overridden in subclass")
 
     def _reload_config(self) -> None:
@@ -172,9 +200,11 @@ class Extractor(Generic[CustomConfigClass]):
 
     def _load_state_store(self) -> None:
         """
-        Searches through the config object for a StateStoreConfig. If found, it will use that configuration to generate
-        a state store, if no such config is found it will either create a LocalStateStore or a NoStateStore depending
-        on whether the  ``use_default_state_store`` argument to the constructor was true or false.
+        Searches through the config object for a StateStoreConfig.
+
+        If found, it will use that configuration to generate a state store, if no such config is found it will either
+        create a LocalStateStore or a NoStateStore depending on whether the ``use_default_state_store`` argument to the
+        constructor was true or false.
 
         Either way, the state_store attribute is guaranteed to be set after calling this method.
         """
@@ -210,37 +240,63 @@ class Extractor(Generic[CustomConfigClass]):
 
         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:
         """
-        Called on an unsuccessful exit of the extractor
+        Called on an unsuccessful exit of the extractor.
 
         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__}: {str(exception)}"[: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":
         """
@@ -249,7 +305,6 @@ class Extractor(Generic[CustomConfigClass]):
         Returns:
             self
         """
-
         if str(os.getenv("COGNITE_FUNCTION_RUNTIME", False)).lower() != "true":
             # Environment Variables
             env_file_found = load_dotenv(dotenv_path="./.env", override=True)
@@ -263,7 +318,9 @@ class Extractor(Generic[CustomConfigClass]):
         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)
 
@@ -284,10 +341,8 @@ class Extractor(Generic[CustomConfigClass]):
 
         self.extraction_pipeline = self.config.cognite.get_extraction_pipeline(self.cognite_client)
 
-        try:
+        with contextlib.suppress(AttributeError):
             self.config.metrics.start_pushers(self.cognite_client)  # type: ignore
-        except AttributeError:
-            pass
 
         def heartbeat_loop() -> None:
             while not self.cancellation_token.is_cancelled:
@@ -324,7 +379,10 @@ class Extractor(Generic[CustomConfigClass]):
         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.
@@ -345,10 +403,8 @@ class Extractor(Generic[CustomConfigClass]):
         if self.state_store:
             self.state_store.synchronize()
 
-        try:
+        with contextlib.suppress(AttributeError):
             self.config.metrics.stop_pushers()  # type: ignore
-        except AttributeError:
-            pass
 
         if exc_val:
             self._report_error(exc_val)
@@ -360,26 +416,50 @@ class Extractor(Generic[CustomConfigClass]):
 
     def run(self) -> None:
         """
-        Run the extractor. Ensures that the Extractor is set up correctly (``run`` called within a ``with``) and calls
-        the ``run_handle``.
+        Run the extractor.
+
+        Ensures that the Extractor is set up correctly (``run`` called within a ``with``) and calls the ``run_handle``.
 
-        Can be overrided in subclasses.
+        Can be overriden in subclasses.
         """
         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")
 
     @classmethod
     def get_current_config(cls) -> CustomConfigClass:
+        """
+        Get the current configuration singleton.
+
+        Returns:
+            The current configuration singleton
+
+        Raises:
+            ValueError: If no configuration singleton has been created, meaning no config file has been loaded.
+        """
         if Extractor._config_singleton is None:  # type: ignore
             raise ValueError("No config singleton created. Have a config file been loaded?")
         return Extractor._config_singleton  # type: ignore
 
     @classmethod
     def get_current_statestore(cls) -> AbstractStateStore:
+        """
+        Get the current state store singleton.
+
+        Returns:
+            The current state store singleton
+
+        Raises:
+            ValueError: If no state store singleton has been created, meaning no state store has been loaded.
+        """
         if Extractor._statestore_singleton is None:
             raise ValueError("No state store singleton created. Have a state store been loaded?")
         return Extractor._statestore_singleton

{cognite_extractor_utils-7.5.13 → cognite_extractor_utils-7.6.0}/cognite/extractorutils/configtools/__init__.py

@@ -13,8 +13,7 @@
 #  limitations under the License.
 
 """
-Module containing tools for loading and verifying config files, and a YAML loader to automatically serialize these
-dataclasses from a config file.
+Module containing tools for loading and verifying config files.
 
 Configs are described as ``dataclass`` es, and use the ``BaseConfig`` class as a superclass to get a few things
 built-in: config version, Cognite project and logging. Use type hints to specify types, use the ``Optional`` type to
@@ -114,10 +113,13 @@ __all__ = [
     "CastableInt",
     "CertificateConfig",
     "CogniteConfig",
+    "ConfigResolver",
     "ConfigType",
     "ConnectionConfig",
     "EitherIdConfig",
     "FileSizeConfig",
+    "KeyVaultAuthenticationMethod",
+    "KeyVaultLoader",
     "LocalStateStoreConfig",
     "LoggingConfig",
     "MetricsConfig",
@@ -126,9 +128,6 @@ __all__ = [
     "RawStateStoreConfig",
     "StateStoreConfig",
     "TimeIntervalConfig",
-    "ConfigResolver",
-    "KeyVaultAuthenticationMethod",
-    "KeyVaultLoader",
     "load_yaml",
     "load_yaml_dict",
 ]

{cognite_extractor_utils-7.5.13 → cognite_extractor_utils-7.6.0}/cognite/extractorutils/configtools/_util.py

@@ -27,8 +27,9 @@ from cognite.extractorutils.exceptions import InvalidConfigError
 
 def _to_snake_case(dictionary: dict[str, Any], case_style: str) -> dict[str, Any]:
     """
-    Ensure that all keys in the dictionary follows the snake casing convention (recursively, so any sub-dictionaries are
-    changed too).
+    Ensure that all keys in the dictionary follows the snake casing convention.
+
+    This function will recursively fix any list or dictionaries. The input dictionary is never modified in place.
 
     Args:
         dictionary: Dictionary to update.