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

cognite/extractorutils/unstable/configuration/__init__.py

@@ -0,0 +1,3 @@
+"""
+New version of ``configtools`` based on pydantic instead of dataclasses.
+"""

cognite/extractorutils/unstable/configuration/exceptions.py

@@ -1,10 +1,15 @@
+"""
+Exceptions representing invalid configurations.
+"""
+
+
 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):
@@ -15,7 +20,13 @@ class InvalidConfigError(Exception):
         self.attempted_revision: int | None = None
 
     def __str__(self) -> str:
+        """
+        Underlying message prefixed with 'Invalid config:'.
+        """
         return f"Invalid config: {self.message}"
 
     def __repr__(self) -> str:
+        """
+        Underlying message prefixed with 'Invalid config:'.
+        """
         return self.__str__()

cognite/extractorutils/unstable/configuration/loaders.py

@@ -1,3 +1,7 @@
+"""
+Module containing functions and classes for loading configuration files.
+"""
+
 import json
 from enum import Enum
 from io import StringIO
@@ -7,46 +11,89 @@ from typing import TextIO, TypeVar
 from pydantic import ValidationError
 
 from cognite.client import CogniteClient
+from cognite.client.exceptions import CogniteAPIError
 from cognite.extractorutils.configtools.loaders import _load_yaml_dict_raw
 from cognite.extractorutils.exceptions import InvalidConfigError as OldInvalidConfigError
 from cognite.extractorutils.unstable.configuration.exceptions import InvalidConfigError
 from cognite.extractorutils.unstable.configuration.models import ConfigModel
 
-__all__ = ["ConfigFormat", "load_file", "load_from_cdf", "load_io", "load_dict"]
+__all__ = ["ConfigFormat", "load_dict", "load_file", "load_from_cdf", "load_io"]
 
 
 _T = TypeVar("_T", bound=ConfigModel)
 
 
 class ConfigFormat(Enum):
+    """
+    Enumeration of supported configuration file formats.
+
+    Attributes:
+        JSON: Represents the JSON configuration file format.
+        YAML: Represents the YAML configuration file format.
+    """
+
     JSON = "json"
     YAML = "yaml"
 
 
 def load_file(path: Path, schema: type[_T]) -> _T:
+    """
+    Load a configuration file from the given path and parse it into the specified schema.
+
+    Args:
+        path: Path to the configuration file.
+        schema: The schema class to parse the configuration into.
+
+    Returns:
+        An instance of the schema populated with the configuration data.
+
+    Raises:
+        InvalidConfigError: If the file type is unknown or the configuration is invalid.
+    """
     if path.suffix in [".yaml", ".yml"]:
-        format = ConfigFormat.YAML
+        file_format = ConfigFormat.YAML
     elif path.suffix == ".json":
-        format = ConfigFormat.JSON
+        file_format = ConfigFormat.JSON
     else:
         raise InvalidConfigError(f"Unknown file type {path.suffix}")
 
     with open(path) as stream:
-        return load_io(stream, format, schema)
+        return load_io(stream, file_format, schema)
 
 
 def load_from_cdf(
     cognite_client: CogniteClient, external_id: str, schema: type[_T], revision: int | None = None
 ) -> tuple[_T, int]:
+    """
+    Load a configuration from a CDF integration using the provided external ID and schema.
+
+    Args:
+        cognite_client: An instance of CogniteClient to interact with CDF.
+        external_id: The external ID of the integration to load configuration from.
+        schema: The schema class to parse the configuration into.
+        revision: the specific revision of the configuration to load, otherwise get the latest.
+
+    Returns:
+        A tuple containing the parsed configuration instance and the revision number.
+
+    Raises:
+        InvalidConfigError: If the configuration is invalid or not found.
+        CogniteAPIError: If there is an unexpected error communicating with CDF.
+    """
     params: dict[str, str | int] = {"integration": external_id}
     if revision:
         params["revision"] = revision
-    response = cognite_client.get(
-        f"/api/v1/projects/{cognite_client.config.project}/odin/config",
-        params=params,
-        headers={"cdf-version": "alpha"},
-    )
-    response.raise_for_status()
+    try:
+        response = cognite_client.get(
+            f"/api/v1/projects/{cognite_client.config.project}/odin/config",
+            params=params,
+            headers={"cdf-version": "alpha"},
+        )
+    except CogniteAPIError as e:
+        if e.code == 404:
+            raise InvalidConfigError("No configuration found for the given integration") from e
+        raise e
+
     data = response.json()
 
     try:
@@ -61,11 +108,25 @@ def load_from_cdf(
         raise new_e from e
 
 
-def load_io(stream: TextIO, format: ConfigFormat, schema: type[_T]) -> _T:
-    if format == ConfigFormat.JSON:
+def load_io(stream: TextIO, file_format: ConfigFormat, schema: type[_T]) -> _T:
+    """
+    Load a configuration from a stream (e.g., file or string) and parse it into the specified schema.
+
+    Args:
+        stream: A text stream containing the configuration data.
+        file_format: The format of the configuration data.
+        schema: The schema class to parse the configuration into.
+
+    Returns:
+        An instance of the schema populated with the configuration data.
+
+    Raises:
+        InvalidConfigError: If the file format is unknown or the configuration is invalid.
+    """
+    if file_format == ConfigFormat.JSON:
         data = json.load(stream)
 
-    elif format == ConfigFormat.YAML:
+    elif file_format == ConfigFormat.YAML:
         data = _load_yaml_dict_raw(stream)
 
         if "azure-keyvault" in data:
@@ -89,15 +150,25 @@ def _make_loc_str(loc: tuple) -> str:
             loc_str = f"{loc_str}{lo}"
             needs_sep = True
         else:
-            if isinstance(lo, int):
-                loc_str = f"{loc_str}[{lo}]"
-            else:
-                loc_str = f"{loc_str}.{lo}"
+            loc_str = f"{loc_str}[{lo}]" if isinstance(lo, int) else f"{loc_str}.{lo}"
 
     return loc_str
 
 
 def load_dict(data: dict, schema: type[_T]) -> _T:
+    """
+    Load a configuration from a dictionary and parse it into the specified schema.
+
+    Args:
+        data: A dictionary containing the configuration data.
+        schema: The schema class to parse the configuration into.
+
+    Returns:
+        An instance of the schema populated with the configuration data.
+
+    Raises:
+        InvalidConfigError: If the configuration is invalid.
+    """
     try:
         return schema.model_validate(data)
 
@@ -113,8 +184,8 @@ def load_dict(data: dict, schema: type[_T]) -> _T:
 
             if "ctx" in err and "error" in err["ctx"]:
                 exc = err["ctx"]["error"]
-                if isinstance(exc, ValueError) or isinstance(exc, AssertionError):
-                    messages.append(f"{str(exc)}: {loc_str}")
+                if isinstance(exc, ValueError | AssertionError):
+                    messages.append(f"{exc!s}: {loc_str}")
                     continue
 
             messages.append(f"{err.get('msg')}: {loc_str}")

cognite/extractorutils/unstable/configuration/models.py

@@ -1,3 +1,7 @@
+"""
+Module containing pre-built models for common extractor configuration.
+"""
+
 import os
 import re
 from datetime import timedelta
@@ -21,22 +25,26 @@ from cognite.extractorutils.configtools._util import _load_certificate_data
 from cognite.extractorutils.exceptions import InvalidConfigError
 
 __all__ = [
-    "ConfigModel",
     "AuthenticationConfig",
-    "TimeIntervalConfig",
+    "ConfigModel",
     "ConnectionConfig",
     "CronConfig",
+    "ExtractorConfig",
     "IntervalConfig",
-    "ScheduleConfig",
-    "LogLevel",
-    "LogFileHandlerConfig",
     "LogConsoleHandlerConfig",
+    "LogFileHandlerConfig",
     "LogHandlerConfig",
-    "ExtractorConfig",
+    "LogLevel",
+    "ScheduleConfig",
+    "TimeIntervalConfig",
 ]
 
 
 class ConfigModel(BaseModel):
+    """
+    Base model for configuration objects, setting the correct pydantic options for extractor config.
+    """
+
     model_config = ConfigDict(
         alias_generator=kebabize,
         populate_by_name=True,
@@ -69,7 +77,7 @@ AuthenticationConfig = Annotated[_ClientCredentialsConfig | _ClientCertificateCo
 
 class TimeIntervalConfig:
     """
-    Configuration parameter for setting a time interval
+    Configuration parameter for setting a time interval.
     """
 
     def __init__(self, expression: str) -> None:
@@ -77,14 +85,25 @@ class TimeIntervalConfig:
 
     @classmethod
     def __get_pydantic_core_schema__(cls, source_type: Any, handler: GetCoreSchemaHandler) -> CoreSchema:
+        """
+        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 TimeIntervalConfig objects are equal if they have the same number of seconds in their interval.
+        """
         if not isinstance(other, TimeIntervalConfig):
             return NotImplemented
         return self._interval == other._interval
 
     def __hash__(self) -> int:
+        """
+        Hash function for TimeIntervalConfig based on the number of seconds in the interval.
+        """
         return hash(self._interval)
 
     @classmethod
@@ -106,36 +125,69 @@ class TimeIntervalConfig:
 
     @property
     def seconds(self) -> int:
+        """
+        Time interval as number of seconds.
+        """
         return self._interval
 
     @property
     def minutes(self) -> float:
+        """
+        Time interval as number of minutes.
+
+        This is a float since the underlying interval is in seconds.
+        """
         return self._interval / 60
 
     @property
     def hours(self) -> float:
+        """
+        Time interval as number of hours.
+
+        This is a float since the underlying interval is in seconds.
+        """
         return self._interval / (60 * 60)
 
     @property
     def days(self) -> float:
+        """
+        Time interval as number of days.
+
+        This is a float since the underlying interval is in seconds.
+        """
         return self._interval / (60 * 60 * 24)
 
     @property
     def timedelta(self) -> timedelta:
+        """
+        Time interval as a timedelta object.
+        """
         days = self._interval // (60 * 60 * 24)
         seconds = self._interval % (60 * 60 * 24)
         return timedelta(days=days, seconds=seconds)
 
     def __int__(self) -> int:
+        """
+        Returns the time interval as a number of seconds.
+        """
         return int(self._interval)
 
     def __float__(self) -> float:
+        """
+        Returns the time interval as a number of seconds.
+        """
         return float(self._interval)
 
     def __str__(self) -> str:
+        """
+        Returns the time interval as a human readable string.
+        """
         return self._expression
 
     def __repr__(self) -> str:
+        """
+        Returns the time interval as a human readable string.
+        """
         return self._expression
 
 
@@ -152,6 +204,15 @@ class _ConnectionParameters(ConfigModel):
 
 
 class ConnectionConfig(ConfigModel):
+    """
+    Configuration for connecting to a Cognite Data Fusion project.
+
+    This configuration includes the project name, base URL, integration name, and authentication details, as well as
+    optional connection parameters.
+
+    This configuration is common for all extractors.
+    """
+
     project: str
     base_url: str
 
@@ -162,6 +223,15 @@ class ConnectionConfig(ConfigModel):
     connection: _ConnectionParameters = Field(default_factory=_ConnectionParameters)
 
     def get_cognite_client(self, client_name: str) -> CogniteClient:
+        """
+        Create a CogniteClient instance using the configuration parameters.
+
+        Args:
+            client_name: Name of the client, set as the x-cdp-app header in the requests
+
+        Returns:
+            CogniteClient: An instance of CogniteClient configured with the provided parameters.
+        """
         from cognite.client.config import global_config
 
         global_config.disable_pypi_version_check = True
@@ -218,6 +288,26 @@ class ConnectionConfig(ConfigModel):
 
     @classmethod
     def from_environment(cls) -> "ConnectionConfig":
+        """
+        Create a ConnectionConfig instance from environment variables.
+
+        Environment variables should be set as follows:
+        - COGNITE_PROJECT: The name of the Cognite Data Fusion project.
+        - COGNITE_BASE_URL: The base URL of the Cognite Data Fusion instance.
+        - COGNITE_INTEGRATION: The external ID of the corresponding integration in CDF.
+        - COGNITE_CLIENT_ID: The client ID for authentication.
+        - COGNITE_TOKEN_SCOPES: The scopes for the token.
+        - COGNITE_CLIENT_SECRET: The client secret for authentication (if using client credentials).
+        - COGNITE_TOKEN_URL: The token URL for authentication (if using client credentials).
+        - COGNITE_CLIENT_CERTIFICATE_PATH: The path to the client certificate (if using client certificate).
+        - COGNITE_AUTHORITY_URL: The authority URL for authentication (if using client certificate).
+
+        Returns:
+            ConnectionConfig: An instance of ConnectionConfig populated with the environment variables.
+
+        Raises:
+            KeyError: If any of the required environment variables are missing.
+        """
         auth: AuthenticationConfig
         if "COGNITE_CLIENT_SECRET" in os.environ:
             auth = _ClientCredentialsConfig(
@@ -248,11 +338,19 @@ class ConnectionConfig(ConfigModel):
 
 
 class CronConfig(ConfigModel):
+    """
+    Configuration parameter for setting a cron schedule.
+    """
+
     type: Literal["cron"]
     expression: str
 
 
 class IntervalConfig(ConfigModel):
+    """
+    Configuration parameter for setting an interval schedule.
+    """
+
     type: Literal["interval"]
     expression: TimeIntervalConfig
 
@@ -261,6 +359,10 @@ ScheduleConfig = Annotated[CronConfig | IntervalConfig, Field(discriminator="typ
 
 
 class LogLevel(Enum):
+    """
+    Enumeration of log levels for the extractor.
+    """
+
     CRITICAL = "CRITICAL"
     ERROR = "ERROR"
     WARNING = "WARNING"
@@ -269,6 +371,10 @@ class LogLevel(Enum):
 
 
 class LogFileHandlerConfig(ConfigModel):
+    """
+    Configuration for a log handler that writes to a file, with daily rotation.
+    """
+
     type: Literal["file"]
     path: Path
     level: LogLevel
@@ -276,6 +382,10 @@ class LogFileHandlerConfig(ConfigModel):
 
 
 class LogConsoleHandlerConfig(ConfigModel):
+    """
+    Configuration for a log handler that writes to standard output.
+    """
+
     type: Literal["console"]
     level: LogLevel
 
@@ -289,4 +399,8 @@ def _log_handler_default() -> list[LogHandlerConfig]:
 
 
 class ExtractorConfig(ConfigModel):
+    """
+    Base class for application configuration for extractors.
+    """
+
     log_handlers: list[LogHandlerConfig] = Field(default_factory=_log_handler_default)

cognite/extractorutils/unstable/core/__init__.py

@@ -0,0 +1,5 @@
+"""
+The ``core`` package contains the core functionality for defining and managing extractors.
+
+It contains the base class for extractors, the runtime for running extractors, and classes for tasks and errors.
+"""

cognite/extractorutils/unstable/core/_dto.py

@@ -1,5 +1,5 @@
 """
-Temporary holding place for DTOs against Extraction Pipelines 2.0 until it's in the SDK
+Temporary holding place for DTOs against Extraction Pipelines 2.0 until it's in the SDK.
 """
 
 from typing import Any, Literal
@@ -10,9 +10,11 @@ from pydantic import BaseModel, ConfigDict
 
 class CogniteModel(BaseModel):
     """
-    Base class for DTO classes based on pydantic, but with a few tweaks to make it inline with the CDF API guidelines:
+    Base class for DTO classes based on pydantic.
+
+    With a few tweaks to make it inline with the CDF API guidelines:
       * camelCase instead of snake_case when serializing/deserializing into/from JSON
-      * exclude Nones from serialized JSON instead of having nulls in the response text
+      * exclude Nones from serialized JSON instead of having nulls in the response text.
     """
 
     def model_dump(self, *args: Any, **kwargs: Any) -> dict[str, Any]: