cognite-extractor-utils 7.5.14__py3-none-any.whl → 7.7.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
@@ -13,32 +17,69 @@ from cognite.extractorutils.exceptions import InvalidConfigError as OldInvalidCo
 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
@@ -67,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:
@@ -95,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)
 
@@ -119,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]:

cognite/extractorutils/unstable/core/base.py

@@ -1,5 +1,50 @@
+"""
+This module provides the base class for extractors.
+
+It includes functionality for task management, logging, error handling, and configuration management.
+
+Extractors should subclass the `Extractor` class and implement the `__init_tasks__` method to define their tasks.
+The subclass should also define several class attributes:
+- ``NAME``: A human-readable name for the extractor.
+- ``EXTERNAL_ID``: A unique identifier for the extractor, used when reporting to CDF Integrations.
+- ``DESCRIPTION``: A brief description of the extractor.
+- ``VERSION``: The version of the extractor, used when reporting to CDF Integrations. This should follow semantic
+   versioning.
+- ``CONFIG_TYPE``: The type of the application configuration for the extractor, which should be a subclass of
+  ``ExtractorConfig``. This should be the same class as the one used for the generic type parameter of the
+  ``Extractor`` class.
+
+
+.. code-block:: python
+
+    class MyConfig(ExtractorConfig):
+        parameter: str
+        another_parameter: int
+        schedule: ScheduleConfig
+
+    class MyExtractor(Extractor[MyConfig]):
+        NAME = "My Extractor"
+        EXTERNAL_ID = "my-extractor"
+        DESCRIPTION = "An example extractor"
+        VERSION = "1.0.0"
+
+        CONFIG_TYPE = MyConfig
+
+        def __init_tasks__(self) -> None:
+            self.add_task(
+                ScheduledTask(
+                    name="my_task",
+                    description="An example task",
+                    schedule=self.application_config.schedule,
+                    target=self.my_task_function,
+                )
+            )
+
+        def my_task_function(self, task_context: TaskContext) -> None:
+            task_context.logger.info("Running my task")
+"""
+
 import logging
-import logging.config
 import time
 from concurrent.futures import ThreadPoolExecutor
 from functools import partial
@@ -30,7 +75,7 @@ from cognite.extractorutils.unstable.core.tasks import ContinuousTask, Scheduled
 from cognite.extractorutils.unstable.scheduling import TaskScheduler
 from cognite.extractorutils.util import now
 
-__all__ = ["ConfigType", "ConfigRevision", "Extractor"]
+__all__ = ["ConfigRevision", "ConfigType", "Extractor"]
 
 ConfigType = TypeVar("ConfigType", bound=ExtractorConfig)
 ConfigRevision = Literal["local"] | int
@@ -40,6 +85,13 @@ _T = TypeVar("_T", bound=ExtractorConfig)
 
 
 class FullConfig(Generic[_T]):
+    """
+    A class that holds the full configuration for an extractor.
+
+    This includes the connection configuration, application configuration, and which revision of the application
+    configuration is currently active.
+    """
+
     def __init__(
         self,
         connection_config: ConnectionConfig,
@@ -52,6 +104,16 @@ class FullConfig(Generic[_T]):
 
 
 class Extractor(Generic[ConfigType], CogniteLogger):
+    """
+    Base class for all extractors.
+
+    This class provides the basic functionality for running an extractor, including task management, logging,
+    error handling, and configuration management.
+
+    It designed to be subclassed by specific extractors, which should implement the `__init_tasks__` method
+    to define their tasks.
+    """
+
     NAME: str
     EXTERNAL_ID: str
     DESCRIPTION: str
@@ -128,6 +190,13 @@ class Extractor(Generic[ConfigType], CogniteLogger):
                     root.addHandler(fh)
 
     def __init_tasks__(self) -> None:
+        """
+        This method should be overridden by subclasses to define their tasks.
+
+        It is called automatically when the extractor is initialized.
+
+        Subclasses should call ``self.add_task(...)`` to add tasks to the extractor.
+        """
         pass
 
     def _set_runtime_message_queue(self, queue: Queue) -> None:
@@ -200,6 +269,9 @@ class Extractor(Generic[ConfigType], CogniteLogger):
         )
 
     def restart(self) -> None:
+        """
+        Trigger a restart of the extractor.
+        """
         self._logger.info("Restarting extractor")
         if self._runtime_messages:
             self._runtime_messages.put(RuntimeMessage.RESTART)
@@ -210,12 +282,20 @@ class Extractor(Generic[ConfigType], CogniteLogger):
         return cls(config)
 
     def add_task(self, task: Task) -> None:
+        """
+        Add a task to the extractor.
+
+        This method wraps the task's target function to include error handling and task tracking.
+
+        Args:
+            task: The task to add. It should be an instance of ``StartupTask``, ``ContinuousTask``, or ``ScheduledTask``
+        """
         # Store this for later, since we'll override it with the wrapped version
         target = task.target
 
         def run_task(task_context: TaskContext) -> None:
             """
-            A wrapped version of the task's target, with tracking and error handling
+            A wrapped version of the task's target, with tracking and error handling.
             """
             # Record a task start
             with self._checkin_lock:
@@ -275,7 +355,7 @@ class Extractor(Generic[ConfigType], CogniteLogger):
                     {
                         "name": t.name,
                         "type": "continuous" if isinstance(t, ContinuousTask) else "batch",
-                        "action": True if isinstance(t, ScheduledTask) else False,
+                        "action": bool(isinstance(t, ScheduledTask)),
                         "description": t.description,
                     }
                     for t in self._tasks
@@ -285,14 +365,29 @@ class Extractor(Generic[ConfigType], CogniteLogger):
         )
 
     def start(self) -> None:
+        """
+        Start the extractor.
+
+        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.
+        """
         self._setup_logging()
         self._report_extractor_info()
         Thread(target=self._run_checkin, name="ExtractorCheckin", daemon=True).start()
 
     def stop(self) -> None:
+        """
+        Stop the extractor.
+
+        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.
+        """
         self.cancellation_token.cancel()
 
     def __enter__(self) -> Self:
+        """
+        Start the extractor in a context manager.
+        """
         self.start()
         return self
 
@@ -302,6 +397,9 @@ class Extractor(Generic[ConfigType], CogniteLogger):
         exc_val: BaseException | None,
         exc_tb: TracebackType | None,
     ) -> bool:
+        """
+        Stop the extractor when exiting the context manager.
+        """
         self.stop()
         with self._checkin_lock:
             self._checkin()
@@ -310,6 +408,17 @@ class Extractor(Generic[ConfigType], CogniteLogger):
         return exc_val is None
 
     def run(self) -> None:
+        """
+        Run the extractor. This method starts the extractor and runs all tasks that have been added.
+
+        This method assumes ``self.start()`` has been called first. The recommended way to use this method is
+        to use the context manager interface, which ensures that the extractor is started and stopped properly.
+
+        .. code-block:: python
+
+            with extractor:
+                extractor.run()
+        """
         has_scheduled = False
 
         startup: list[StartupTask] = []