cognite-extractor-utils 7.5.14__py3-none-any.whl → 7.6.0__py3-none-any.whl

cognite/extractorutils/configtools/loaders.py

@@ -1,3 +1,6 @@
+"""
+Module containing functions and classes for loading configuration files.
+"""
 #  Copyright 2023 Cognite AS
 #
 #  Licensed under the Apache License, Version 2.0 (the "License");
@@ -53,19 +56,25 @@ CustomConfigClass = TypeVar("CustomConfigClass", bound=BaseConfig)
 
 
 class KeyVaultAuthenticationMethod(Enum):
+    """
+    Enum representing the authentication methods for Azure KeyVault.
+    """
+
     DEFAULT = "default"
     CLIENTSECRET = "client-secret"
 
 
 class KeyVaultLoader:
     """
-    Class responsible for configuring keyvault for clients using Azure
+    Class responsible for configuring keyvault for clients using Azure.
+
+    Args:
+        config: A dictionary containing the configuration for the keyvault.
     """
 
     def __init__(self, config: dict | None):
         self.config = config
 
-        self.credentials: TokenCredential | None = None
         self.client: SecretClient | None = None
 
     def _init_client(self) -> None:
@@ -89,9 +98,10 @@ class KeyVaultLoader:
 
         vault_url = f"https://{keyvault_name}.vault.azure.net"
 
+        credentials: TokenCredential
         if self.config["authentication-method"] == KeyVaultAuthenticationMethod.DEFAULT.value:
             _logger.info("Using Azure DefaultCredentials to access KeyVault")
-            self.credentials = DefaultAzureCredential()
+            credentials = DefaultAzureCredential()
 
         elif self.config["authentication-method"] == KeyVaultAuthenticationMethod.CLIENTSECRET.value:
             auth_parameters = ("client-id", "tenant-id", "secret")
@@ -104,11 +114,11 @@ class KeyVaultLoader:
                 _logger.info(f"Local environment file not found at {Path.cwd() / '.env'}")
 
             if all(param in self.config for param in auth_parameters):
-                tenant_id = os.path.expandvars(self.config.get("tenant-id", None))
-                client_id = os.path.expandvars(self.config.get("client-id", None))
-                secret = os.path.expandvars(self.config.get("secret", None))
+                tenant_id = os.path.expandvars(self.config["tenant-id"])
+                client_id = os.path.expandvars(self.config["client-id"])
+                secret = os.path.expandvars(self.config["secret"])
 
-                self.credentials = ClientSecretCredential(
+                credentials = ClientSecretCredential(
                     tenant_id=tenant_id,
                     client_id=client_id,
                     client_secret=secret,
@@ -122,9 +132,12 @@ class KeyVaultLoader:
                 "Invalid KeyVault authentication method. Possible values : default or client-secret"
             )
 
-        self.client = SecretClient(vault_url=vault_url, credential=self.credentials)  # type: ignore
+        self.client = SecretClient(vault_url=vault_url, credential=credentials)
 
     def __call__(self, _: yaml.SafeLoader, node: yaml.Node) -> str:
+        """
+        Method to be called when the !keyvault tag is encountered in the YAML file.
+        """
         self._init_client()
         try:
             return self.client.get_secret(node.value).value  # type: ignore  # _init_client guarantees not None
@@ -137,7 +150,14 @@ class _EnvLoader(yaml.SafeLoader):
 
 
 class SafeLoaderIgnoreUnknown(yaml.SafeLoader):
+    """
+    Variant of PyYAML's SafeLoader that ignores unknown tags.
+    """
+
     def ignore_unknown(self, node: yaml.Node) -> None:
+        """
+        Constructor for unknown tags that does nothing.
+        """
         return None
 
 
@@ -185,7 +205,9 @@ def _load_yaml_dict_raw(
         config_dict = yaml.load(source, Loader=loader)  # noqa: S506
     except ScannerError as e:
         location = e.problem_mark or e.context_mark
-        formatted_location = f" at line {location.line+1}, column {location.column+1}" if location is not None else ""
+        formatted_location = (
+            f" at line {location.line + 1}, column {location.column + 1}" if location is not None else ""
+        )
         cause = e.problem or e.context
         raise InvalidConfigError(f"Invalid YAML{formatted_location}: {cause or ''}") from e
 
@@ -241,10 +263,7 @@ def _load_yaml(
         ) from e
 
     except (dacite.WrongTypeError, dacite.MissingValueError, dacite.UnionMatchError) as e:
-        if e.field_path:
-            path = e.field_path.replace("_", "-") if case_style == "hyphen" else e.field_path
-        else:
-            path = None
+        path = (e.field_path.replace("_", "-") if case_style == "hyphen" else e.field_path) if e.field_path else None
 
         def name(type_: type) -> str:
             return type_.__name__ if hasattr(type_, "__name__") else str(type_)
@@ -262,7 +281,7 @@ def _load_yaml(
         raise InvalidConfigError(f'Missing mandatory field "{path}"') from e
 
     except dacite.ForwardReferenceError as e:
-        raise ValueError(f"Invalid config class: {str(e)}") from e
+        raise ValueError(f"Invalid config class: {e!s}") from e
 
     config._file_hash = sha256(json.dumps(config_dict).encode("utf-8")).hexdigest()
 
@@ -309,7 +328,7 @@ def load_yaml_dict(
     keyvault_loader: KeyVaultLoader | None = None,
 ) -> dict[str, Any]:
     """
-    Read a YAML file and return a dictionary from its contents
+    Read a YAML file and return a dictionary from its contents.
 
     Args:
         source: Input stream (as returned by open(...)) or string containing YAML.
@@ -331,7 +350,7 @@ def load_yaml_dict(
 
 def compile_patterns(ignore_patterns: list[str | IgnorePattern]) -> list[re.Pattern[str]]:
     """
-    list of patterns to compile
+    List of patterns to compile.
 
     Args:
         ignore_patterns: A list of strings or IgnorePattern to be compiled.
@@ -349,6 +368,12 @@ def compile_patterns(ignore_patterns: list[str | IgnorePattern]) -> list[re.Patt
 
 
 class ConfigResolver(Generic[CustomConfigClass]):
+    """
+    Class for resolving configuration files, either from a local file or a remote CDF extraction pipeline.
+
+    Automatically reloads the configuration file if it has changed
+    """
+
     def __init__(self, config_path: str, config_type: type[CustomConfigClass]):
         self.config_path = config_path
         self.config_type = config_type
@@ -364,6 +389,9 @@ class ConfigResolver(Generic[CustomConfigClass]):
 
     @property
     def cognite_client(self) -> CogniteClient | None:
+        """
+        Returns a CogniteClient instance based on the configuration.
+        """
         if self._cognite_client is None and self._config is not None:
             self._cognite_client = self._config.cognite.get_cognite_client("config_resolver")
         return self._cognite_client
@@ -376,6 +404,9 @@ class ConfigResolver(Generic[CustomConfigClass]):
 
     @property
     def is_remote(self) -> bool:
+        """
+        Returns True if the configuration is a remote CDF extraction pipeline config, False if it is a local file.
+        """
         raw_config_type = load_yaml_dict(self._config_text).get("type")
         if raw_config_type is None:
             _logger.warning("No config type specified, default to local")
@@ -385,6 +416,9 @@ class ConfigResolver(Generic[CustomConfigClass]):
 
     @property
     def has_changed(self) -> bool:
+        """
+        Returns True if the configuration file has changed since the last accepted configuration.
+        """
         try:
             self._resolve_config()
         except Exception:
@@ -394,18 +428,36 @@ class ConfigResolver(Generic[CustomConfigClass]):
 
     @property
     def config(self) -> CustomConfigClass:
+        """
+        Returns the current configuration object. If it has not been resolved yet, it will resolve it first.
+        """
         if self._config is None:
             self._resolve_config()
             self.accept_new_config()
         return self._config  # type: ignore
 
     def accept_new_config(self) -> None:
+        """
+        Accepts the new configuration, making it the current configuration.
+        """
         self._config = self._next_config
 
     @classmethod
     def from_cli(
         cls, name: str, description: str, version: str, config_type: type[CustomConfigClass]
     ) -> "ConfigResolver":
+        """
+        Creates a ConfigResolver instance from command line arguments.
+
+        Args:
+            name: The name of the extractor.
+            description: A description of the extractor.
+            version: The version of the extractor.
+            config_type: The type of the configuration class to be used.
+
+        Returns:
+            A ConfigResolver instance initialized with the configuration file path.
+        """
         argument_parser = argparse.ArgumentParser(sys.argv[0], description=description)
         argument_parser.add_argument(
             "config", nargs=1, type=str, help="The YAML file containing configuration for the extractor."

cognite/extractorutils/configtools/validators.py

@@ -1,3 +1,7 @@
+"""
+Module containing utility functions for validating config values.
+"""
+
 import logging
 import re
 
@@ -15,7 +19,7 @@ def matches_patterns(patterns: list[str | re.Pattern[str]], string: str) -> bool
     Returns:
         boolean value indicating whether string matches any of the patterns.
     """
-    return any([matches_pattern(pattern, string) for pattern in patterns])
+    return any(matches_pattern(pattern, string) for pattern in patterns)
 
 
 def matches_pattern(pattern: str | re.Pattern[str], string: str) -> bool:

cognite/extractorutils/exceptions.py

@@ -1,3 +1,6 @@
+"""
+This module defines custom exceptions for the extractorutils package.
+"""
 #  Copyright 2020 Cognite AS
 #
 #  Licensed under the Apache License, Version 2.0 (the "License");
@@ -15,11 +18,11 @@
 
 class InvalidConfigError(Exception):
     """
-    Exception thrown from ``load_yaml`` and ``load_yaml_dict`` if config file is invalid. This can be due to
+    Exception thrown from ``load_yaml`` and ``load_yaml_dict`` if config file is invalid. This can be due to.
 
       * Missing fields
       * Incompatible types
-      * Unkown fields
+      * Unknown fields
     """
 
     def __init__(self, message: str, details: list[str] | None = None):
@@ -28,7 +31,13 @@ class InvalidConfigError(Exception):
         self.details = details
 
     def __str__(self) -> str:
+        """
+        Returns a string representation of the error.
+        """
         return f"Invalid config: {self.message}"
 
     def __repr__(self) -> str:
+        """
+        Returns a string representation of the error.
+        """
         return self.__str__()

cognite/extractorutils/metrics.py

@@ -54,6 +54,7 @@ from prometheus_client.exposition import basic_auth_handler, delete_from_gateway
 
 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
@@ -84,6 +85,8 @@ def safe_get(cls: type[T], *args: Any, **kwargs: Any) -> T:
 
     Args:
         cls: Metrics class to either create or get a cached version of
+        args: Arguments passed as-is to the class constructor
+        kwargs: Keyword arguments passed as-is to the class constructor
 
     Returns:
         An instance of given class
@@ -98,8 +101,10 @@ def safe_get(cls: type[T], *args: Any, **kwargs: Any) -> T:
 
 class BaseMetrics:
     """
-    Base collection of extractor metrics. The class also spawns a collector thread on init that regularly fetches
-    process information and update the ``process_*`` gauges.
+    Base collection of extractor metrics.
+
+    The class also spawns a collector thread on init that regularly fetches process information and update the
+    ``process_*`` gauges.
 
     To create a set of metrics for an extractor, create a subclass of this class.
 
@@ -144,7 +149,7 @@ class BaseMetrics:
 
     def _proc_collect(self) -> None:
         """
-        Collect values for process metrics
+        Collect values for process metrics.
         """
         total_memory_available = psutil.virtual_memory().total
         while True:
@@ -157,7 +162,7 @@ class BaseMetrics:
 
     def _start_proc_collector(self) -> None:
         """
-        Start a thread that collects process metrics at a regular interval
+        Start a thread that collects process metrics at a regular interval.
         """
         thread = threading.Thread(target=self._proc_collect, name="ProcessMetricsCollector", daemon=True)
         thread.start()
@@ -165,8 +170,9 @@ class BaseMetrics:
 
 class AbstractMetricsPusher(ABC):
     """
-    Base class for metric pushers. Metric pushers spawns a thread that routinely pushes metrics to a configured
-    destination.
+    Base class for metric pushers.
+
+    Metric pushers spawns a thread that routinely pushes metrics to a configured destination.
 
     Contains all the logic for starting and running threads.
 
@@ -194,7 +200,7 @@ class AbstractMetricsPusher(ABC):
     @abstractmethod
     def _push_to_server(self) -> None:
         """
-        Push metrics to a remote server, to be overrided in subclasses.
+        Push metrics to a remote server, to be overridden in subclasses.
         """
         pass
 
@@ -209,7 +215,6 @@ class AbstractMetricsPusher(ABC):
     def start(self) -> None:
         """
         Starts a thread that pushes the default registry to the configured gateway at certain intervals.
-
         """
         self.thread = threading.Thread(target=self._run, daemon=True, name=self.thread_name)
         self.thread.start()
@@ -224,7 +229,7 @@ class AbstractMetricsPusher(ABC):
 
     def __enter__(self) -> "AbstractMetricsPusher":
         """
-        Wraps around start method, for use as context manager
+        Wraps around start method, for use as context manager.
 
         Returns:
             self
@@ -236,7 +241,7 @@ class AbstractMetricsPusher(ABC):
         self, exc_type: type[BaseException] | None, exc_val: BaseException | None, exc_tb: TracebackType | None
     ) -> None:
         """
-        Wraps around stop method, for use as context manager
+        Wraps around stop method, for use as context manager.
 
         Args:
             exc_type: Exception type
@@ -403,11 +408,11 @@ class CognitePusher(AbstractMetricsPusher):
 
     def _push_to_server(self) -> None:
         """
-        Create datapoints an push them to their respective time series
+        Create datapoints an push them to their respective time series.
         """
         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/statestore/__init__.py

@@ -1,12 +1,86 @@
+"""
+Module containing state stores for extractors.
+
+The ``statestore`` module contains classes for keeping track of the extraction state of individual items, facilitating
+incremental load and speeding up startup times.
+
+At the beginning of a run the extractor typically calls the ``initialize`` method, which loads the states from the
+remote store (which can either be a local JSON file or a table in CDF RAW), and during and/or at the end of a run, the
+``synchronize`` method is called, which saves the current states to the remote store.
+
+You can choose the back-end for your state store with which class you're instantiating:
+
+.. code-block:: python
+
+    # A state store using a JSON file as remote storage:
+    states = LocalStateStore("state.json")
+    states.initialize()
+
+    # A state store using a RAW table as remote storage:
+    states = RawStateStore(
+        cdf_client = CogniteClient(),
+        database = "extractor_states",
+        table = "my_extractor_deployment"
+    )
+    states.initialize()
+
+You can now use this state store to get states:
+
+.. code-block:: python
+
+    low, high = states.get_state(external_id = "my-id")
+
+You can set states:
+
+.. code-block:: python
+
+    states.set_state(external_id = "another-id", high=100)
+
+and similar for ``low``. The ``set_state(...)`` method will always overwrite the current state. Some times you might
+want to only set state *if larger* than the previous state, in that case consider ``expand_state(...)``:
+
+.. code-block:: python
+
+    # High watermark of another-id is already 100, nothing happens in this call:
+    states.expand_state(external_id = "another-id", high=50)
+
+    # This will set high to 150 as it is larger than the previous state
+    states.expand_state(external_id = "another-id", high=150)
+
+To store the state to the remote store, use the ``synchronize()`` method:
+
+.. code-block:: python
+
+    states.synchronize()
+
+You can set a state store to automatically update on upload triggers from an upload queue by using the
+``post_upload_function`` in the upload queue:
+
+.. code-block:: python
+
+    states = LocalStateStore("state.json")
+    states.initialize()
+
+    uploader = TimeSeriesUploadQueue(
+        cdf_client = CogniteClient(),
+        max_upload_interval = 10
+        post_upload_function = states.post_upload_handler()
+    )
+
+    # The state store is now updated automatically!
+
+    states.synchronize()
+"""
+
 from .hashing import AbstractHashStateStore, LocalHashStateStore, RawHashStateStore
 from .watermark import AbstractStateStore, LocalStateStore, NoStateStore, RawStateStore
 
 __all__ = [
+    "AbstractHashStateStore",
     "AbstractStateStore",
-    "RawStateStore",
+    "LocalHashStateStore",
     "LocalStateStore",
     "NoStateStore",
-    "AbstractHashStateStore",
     "RawHashStateStore",
-    "LocalHashStateStore",
+    "RawStateStore",
 ]

cognite/extractorutils/statestore/_base.py

@@ -32,7 +32,11 @@ class _BaseStateStore(ABC):
     def start(self, initialize: bool = True) -> None:
         """
         Start saving state periodically if save_interval is set.
+
         This calls the synchronize method every save_interval seconds.
+
+        Args:
+            initialize (bool): (Optional). If True, call initialize method before starting the thread.
         """
         if initialize and not self._initialized:
             self.initialize()
@@ -52,7 +56,7 @@ class _BaseStateStore(ABC):
 
     def _run(self) -> None:
         """
-        Internal run method for synchronize thread
+        Internal run method for synchronize thread.
         """
         self.initialize()
         while not self.cancellation_token.wait(timeout=self.save_interval):
@@ -68,13 +72,13 @@ class _BaseStateStore(ABC):
     @abstractmethod
     def initialize(self, force: bool = False) -> None:
         """
-        Get states from remote store
+        Get states from remote store.
         """
         pass
 
     @abstractmethod
     def synchronize(self) -> None:
         """
-        Upload states to remote store
+        Upload states to remote store.
         """
         pass