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

cognite/extractorutils/unstable/core/errors.py

@@ -1,3 +1,7 @@
+"""
+This module defines the Error and ErrorLevel classes for reporting errors in extractors.
+"""
+
 import logging
 from enum import Enum
 from types import TracebackType
@@ -15,12 +19,19 @@ __all__ = ["Error", "ErrorLevel"]
 
 
 class ErrorLevel(Enum):
+    """
+    Enumeration of error levels for reporting errors in extractors.
+    """
+
     warning = "warning"
     error = "error"
     fatal = "fatal"
 
     @property
     def log_level(self) -> int:
+        """
+        Returns the corresponding logging level for the error level.
+        """
         match self:
             case ErrorLevel.warning:
                 return logging.WARNING
@@ -33,6 +44,20 @@ class ErrorLevel(Enum):
 
 
 class Error:
+    """
+    Represents an error that occurred during the run of an extractor.
+
+    This class should not be instantiated directly. Instead, use the ``CogniteLogger`` methods (either in the
+    TaskContext or the extractor base class) to create errors.
+
+    Args:
+        level: The severity level of the error.
+        description: A brief description of the error.
+        details: Additional details about the error, if any.
+        task_name: The name of the task during which the error occurred, if applicable.
+        extractor: The extractor instance that reported the error.
+    """
+
     def __init__(
         self,
         level: ErrorLevel,
@@ -55,6 +80,9 @@ class Error:
         self._extractor._report_error(self)
 
     def instant(self) -> None:
+        """
+        Make this error an instant error, meaning it does not have a duration.
+        """
         # Only end the error once
         if self.end_time is not None:
             return
@@ -65,6 +93,11 @@ class Error:
         self._extractor._report_error(self)
 
     def finish(self) -> None:
+        """
+        Mark the error as finished, setting the end time to the current time.
+
+        This method should be called when the error is resolved or no longer relevant.
+        """
         # Only end the error once
         if self.end_time is not None:
             return
@@ -75,6 +108,11 @@ class Error:
         self._extractor._report_error(self)
 
     def __enter__(self) -> "Error":
+        """
+        Start tracking an error as a context manager.
+
+        This allows the error to be automatically finished when exiting the context.
+        """
         return self
 
     def __exit__(
@@ -83,5 +121,8 @@ class Error:
         exc_val: BaseException | None,
         exc_tb: TracebackType | None,
     ) -> bool:
+        """
+        Finish the error context manager, marking the error as finished.
+        """
         self.finish()
         return exc_val is None

cognite/extractorutils/unstable/core/logger.py

@@ -1,3 +1,10 @@
+"""
+This module provides the ``CogniteLogger`` base class, which is an abstract base class for logging.
+
+This class is subclassed by both the ``TaskContext`` and the ``Extractor`` base classes, providing a unified interface
+for logging and error handling in extractors.
+"""
+
 from abc import ABC, abstractmethod
 from logging import Logger, getLogger
 from traceback import format_exception
@@ -9,6 +16,16 @@ from cognite.extractorutils.unstable.core.errors import Error, ErrorLevel
 
 
 class CogniteLogger(ABC):
+    """
+    Base class for logging and error handling in extractors.
+
+    This class provides methods to log messages at different levels (debug, info, warning, error, fatal) and to
+    create and manage errors that occur during the execution of an extractor.
+
+    If you use this class instead of a standard logger, you will get additional functionality such as reporting errors
+    back to CDF.
+    """
+
     def __init__(self) -> None:
         self._logger: Logger = getLogger()
 
@@ -24,9 +41,15 @@ class CogniteLogger(ABC):
         pass
 
     def debug(self, message: str) -> None:
+        """
+        Log a debug message.
+        """
         self._logger.debug(message)
 
     def info(self, message: str) -> None:
+        """
+        Log an information message.
+        """
         self._logger.info(message)
 
     def begin_warning(
@@ -36,6 +59,32 @@ class CogniteLogger(ABC):
         details: str | None = None,
         auto_log: bool = True,
     ) -> Error:
+        """
+        Begin a warning error.
+
+        This will both log the message and create an error object that can be used to track and report the error.
+
+        Args:
+            message: The message to log and include in the error.
+            details: Additional details about the error, if any.
+            auto_log: If True, the message will be logged to the standard logging framework as well. Defaults to True.
+
+        Returns:
+            An ``Error`` object representing the warning error, tied to the current extractor instance.
+
+        Examples:
+            To track and complete an error, you can keep a reference to the error object and call its ``finish``
+            method when the error is resolved, or use it in a context manager to automatically finish it:
+
+            ... code-block:: python
+                error = logger.begin_warning("This is a warning", details="Some details")
+                # Do something
+                error.finish()
+
+            ... code-block:: python
+                with logger.begin_warning("This is a warning", details="Some details")
+                    # Do something
+        """
         if auto_log:
             self._logger.warning(message)
         return self._new_error(
@@ -51,6 +100,32 @@ class CogniteLogger(ABC):
         details: str | None = None,
         auto_log: bool = True,
     ) -> Error:
+        """
+        Begin an error.
+
+        This will both log the message and create an error object that can be used to track and report the error.
+
+        Args:
+            message: The message to log and include in the error.
+            details: Additional details about the error, if any.
+            auto_log: If True, the message will be logged to the standard logging framework as well. Defaults to True.
+
+        Returns:
+            An ``Error`` object representing the error, tied to the current extractor instance.
+
+        Examples:
+            To track and complete an error, you can keep a reference to the error object and call its ``finish``
+            method when the error is resolved, or use it in a context manager to automatically finish it:
+
+            ... code-block:: python
+                error = logger.begin_error("This is an error", details="Some details")
+                # Do something
+                error.finish()
+
+            ... code-block:: python
+                with logger.begin_error("This is an error", details="Some details")
+                    # Do something
+        """
         if auto_log:
             self._logger.error(message)
         return self._new_error(
@@ -66,6 +141,32 @@ class CogniteLogger(ABC):
         details: str | None = None,
         auto_log: bool = True,
     ) -> Error:
+        """
+        Begin a fatal error.
+
+        This will both log the message and create an error object that can be used to track and report the error.
+
+        Args:
+            message: The message to log and include in the error.
+            details: Additional details about the error, if any.
+            auto_log: If True, the message will be logged to the standard logging framework as well. Defaults to True.
+
+        Returns:
+            An ``Error`` object representing the fatal error, tied to the current extractor instance.
+
+        Examples:
+            To track and complete an error, you can keep a reference to the error object and call its ``finish``
+            method when the error is resolved, or use it in a context manager to automatically finish it:
+
+            ... code-block:: python
+                error = logger.begin_fatal("This is a fatal error", details="Some details")
+                # Do something
+                error.finish()
+
+            ... code-block:: python
+                with logger.begin_fatal("This is a fatal error", details="Some details")
+                    # Do something
+        """
         if auto_log:
             self._logger.critical(message)
         return self._new_error(
@@ -81,6 +182,17 @@ class CogniteLogger(ABC):
         details: str | None = None,
         auto_log: bool = True,
     ) -> None:
+        """
+        Report an instant warning.
+
+        This will log the message and create an error object that is marked as instant, meaning it does not have a
+        duration.
+
+        Args:
+            message: The message to log and include in the error.
+            details: Additional details about the error, if any.
+            auto_log: If True, the message will be logged to the standard logging framework as well. Defaults to True.
+        """
         if auto_log:
             self._logger.warning(message)
         self._new_error(
@@ -96,6 +208,17 @@ class CogniteLogger(ABC):
         details: str | None = None,
         auto_log: bool = True,
     ) -> None:
+        """
+        Report an instant error.
+
+        This will log the message and create an error object that is marked as instant, meaning it does not have a
+        duration.
+
+        Args:
+            message: The message to log and include in the error.
+            details: Additional details about the error, if any.
+            auto_log: If True, the message will be logged to the standard logging framework as well. Defaults to True.
+        """
         if auto_log:
             self._logger.error(message)
         self._new_error(
@@ -111,6 +234,17 @@ class CogniteLogger(ABC):
         details: str | None = None,
         auto_log: bool = True,
     ) -> None:
+        """
+        Report an instant fatal.
+
+        This will log the message and create an error object that is marked as instant, meaning it does not have a
+        duration.
+
+        Args:
+            message: The message to log and include in the error.
+            details: Additional details about the error, if any.
+            auto_log: If True, the message will be logged to the standard logging framework as well. Defaults to True.
+        """
         if auto_log:
             self._logger.critical(message)
         self._new_error(
@@ -128,6 +262,21 @@ class CogniteLogger(ABC):
         include_details: Literal["stack_trace"] | Literal["exception_message"] | bool = "exception_message",
         auto_log: bool = True,
     ) -> None:
+        """
+        Report an exception as an error.
+
+        This will log the message and create an error object that is marked as instant, meaning it does not have a
+        duration. The exception details can be included in the error.
+
+        Args:
+            message: The message to log and include in the error.
+            exception: The exception to report.
+            level: The severity level of the error. Defaults to ``ErrorLevel.error``.
+            include_details: How to include details about the exception. Can be "stack_trace", "exception_message",
+                or True (equivalent to "exception_message"). If False, no details are included. Defaults to
+                "exception_message".
+            auto_log: If True, the message will be logged to the standard logging framework as well. Defaults to True.
+        """
         if auto_log:
             self._logger.log(level=level.log_level, msg=message, exc_info=exception)
 

cognite/extractorutils/unstable/core/restart_policy.py

@@ -1,3 +1,17 @@
+"""
+This module defines the restart policies for extractors.
+
+Is is used by the ``Runtime`` to determine whether an extractor should be restarted after a task failure.
+
+It provides three predefined restart policies:
+- ``NEVER``: The extractor will never be restarted.
+- ``WHEN_ANY_TASK_CRASHES``: The extractor will be restarted if any task crashes.
+- ``WHEN_CONTINUOUS_TASKS_CRASHES``: The extractor will be restarted only if a continuous task crashes.
+
+Users can also define their own restart policies by providing a callable that takes a `Task` and an `Exception`
+and returns a boolean indicating whether the extractor should be restarted.
+"""
+
 from collections.abc import Callable
 
 from cognite.extractorutils.unstable.core.tasks import ContinuousTask, Task
@@ -22,8 +36,8 @@ WHEN_CONTINUOUS_TASKS_CRASHES = _is_continuous
 WHEN_ANY_TASK_CRASHES = _true
 
 __all__ = [
-    "RestartPolicy",
     "NEVER",
-    "WHEN_CONTINUOUS_TASKS_CRASHES",
     "WHEN_ANY_TASK_CRASHES",
+    "WHEN_CONTINUOUS_TASKS_CRASHES",
+    "RestartPolicy",
 ]

cognite/extractorutils/unstable/core/runtime.py

@@ -1,3 +1,28 @@
+"""
+Module providing the runtime for an extractor.
+
+The runtime is responsible for starting the extractor in a separate process, managing its lifecycle, and handling
+configuration loading and updates. It also handles errors and restarts the extractor if necessary.
+
+It is the preferred way to run an extractor, as it provides a more robust and flexible way to manage the extractor's
+lifecycle compared to running it directly in the main process.
+
+The runtime also contains a command line interface (CLI) for starting the extractor, which allows users to specify
+the connection configuration and other parameters.
+
+.. code-block:: python
+
+    from cognite.extractorutils.unstable.core.runtime import Runtime
+    from my_extractor import MyExtractor
+
+    def main() -> None:
+        runtime = Runtime(MyExtractor)
+        runtime.run()
+
+    if __name__ == "__main__":
+        main()
+"""
+
 import logging
 import os
 import sys
@@ -9,7 +34,7 @@ from random import randint
 from typing import Any, Generic, TypeVar
 from uuid import uuid4
 
-from requests.exceptions import ConnectionError
+from requests.exceptions import ConnectionError as RequestsConnectionError
 from typing_extensions import assert_never
 
 from cognite.client import CogniteClient
@@ -32,12 +57,19 @@ from cognite.extractorutils.util import now
 from ._messaging import RuntimeMessage
 from .base import ConfigRevision, ConfigType, Extractor, FullConfig
 
-__all__ = ["Runtime", "ExtractorType"]
+__all__ = ["ExtractorType", "Runtime"]
 
 ExtractorType = TypeVar("ExtractorType", bound=Extractor)
 
 
 class Runtime(Generic[ExtractorType]):
+    """
+    The runtime for an extractor.
+
+    This class is responsible for starting the extractor in a separate process, managing its lifecycle, and handling
+    configuration loading and updates. It also handles errors and restarts the extractor if necessary.
+    """
+
     RETRY_CONFIG_INTERVAL = 30
 
     def __init__(
@@ -73,8 +105,8 @@ class Runtime(Generic[ExtractorType]):
             help="Connection parameters",
         )
         argparser.add_argument(
-            "-l",
-            "--local-override",
+            "-f",
+            "--force-local-config",
             nargs=1,
             type=Path,
             required=False,
@@ -248,11 +280,11 @@ class Runtime(Generic[ExtractorType]):
                 self.logger.critical(str(e.message))
 
             else:
-                self.logger.critical(f"Error while connecting to CDF {str(e)}")
+                self.logger.critical(f"Error while connecting to CDF {e!s}")
 
             return False
 
-        except ConnectionError as e:
+        except RequestsConnectionError as e:
             # This is sometime thrown, I've seen it when trying to get an auth token but it might happen elsewhere too
             self.logger.error(str(e))
             self.logger.critical("Could not initiate connection. Please check your configuration.")
@@ -261,6 +293,12 @@ class Runtime(Generic[ExtractorType]):
         return True
 
     def run(self) -> None:
+        """
+        Run the extractor runtime.
+
+        This is intended as the main entry point for the extractor runtime, and starts by parsing command line
+        arguments.
+        """
         argparser = self._create_argparser()
         args = argparser.parse_args()
 

cognite/extractorutils/unstable/core/tasks.py

@@ -1,3 +1,7 @@
+"""
+This module defines the base classes for tasks in the extractor framework.
+"""
+
 import logging
 from collections.abc import Callable
 from typing import TYPE_CHECKING
@@ -14,10 +18,16 @@ from cognite.extractorutils.unstable.core.logger import CogniteLogger
 if TYPE_CHECKING:
     from cognite.extractorutils.unstable.core.base import Extractor
 
-__all__ = ["ScheduledTask", "ContinuousTask", "StartupTask", "Task", "TaskContext"]
+__all__ = ["ContinuousTask", "ScheduledTask", "StartupTask", "Task", "TaskContext"]
 
 
 class TaskContext(CogniteLogger):
+    """
+    Context for a task execution.
+
+    This class is used to log errors and messages related to the task execution.
+    """
+
     def __init__(self, task: "Task", extractor: "Extractor"):
         super().__init__()
         self._task = task
@@ -58,6 +68,18 @@ class _Task:
 
 
 class ScheduledTask(_Task):
+    """
+    A task that is scheduled to run at specific intervals or according to a cron expression.
+
+    This class allows you to define tasks that can be scheduled using either an interval or a cron expression.
+
+    Args:
+        name: The name of the task.
+        target: A callable that takes a ``TaskContext`` and performs the task.
+        description: An optional description of the task.
+        schedule: A ``ScheduleConfig`` object that defines the scheduling configuration for the task.
+    """
+
     def __init__(
         self,
         *,
@@ -73,6 +95,15 @@ class ScheduledTask(_Task):
     def from_interval(
         cls, *, interval: str, name: str, target: TaskTarget, description: str | None = None
     ) -> "ScheduledTask":
+        """
+        Create a scheduled task that runs at regular intervals.
+
+        Args:
+            interval: A string representing the time interval (e.g., "5m" for 5 minutes).
+            name: The name of the task.
+            target: A callable that takes a ``TaskContext`` and performs the task.
+            description: An optional description of the task.
+        """
         return ScheduledTask(
             name=name,
             target=target,
@@ -82,6 +113,15 @@ class ScheduledTask(_Task):
 
     @classmethod
     def from_cron(cls, *, cron: str, name: str, target: TaskTarget, description: str | None = None) -> "ScheduledTask":
+        """
+        Create a scheduled task that runs according to a cron expression.
+
+        Args:
+            cron: A string representing the cron expression (e.g., "0 0 * * *" for daily at midnight).
+            name: The name of the task.
+            target: A callable that takes a ``TaskContext`` and performs the task.
+            description: An optional description of the task.
+        """
         return ScheduledTask(
             name=name,
             target=target,
@@ -91,6 +131,12 @@ class ScheduledTask(_Task):
 
 
 class ContinuousTask(_Task):
+    """
+    A task that runs continuously.
+
+    Continuous tasks are started when the extractor starts and are expected to run until the extractor stops.
+    """
+
     def __init__(
         self,
         *,
@@ -102,6 +148,12 @@ class ContinuousTask(_Task):
 
 
 class StartupTask(_Task):
+    """
+    A task that runs once at the startup of the extractor.
+
+    Startup tasks are executed before any continuous or scheduled tasks and are typically used for initialization.
+    """
+
     def __init__(
         self,
         *,

cognite/extractorutils/unstable/scheduling/__init__.py

@@ -1,3 +1,16 @@
+"""
+This module provides a task scheduler.
+
+It is inspired by the ``APScheduler`` library and is designed to manage the scheduling of tasks within the extractor
+framework. It differs from ``APScheduler`` in a few key ways:
+- It is designed to be used within the extractor framework, allowing for better integration with the extractor's
+  lifecycle and error handling. For example, it respects the extractor's cancellation token and will gracefully shut
+  down upon cancellation.
+- It has a simpler interface, focusing on the core functionality needed for scheduling tasks without the additional
+  complexity of a full-featured scheduler like ``APScheduler``.
+- It is fully typed, providing better type safety and autocompletion in IDEs.
+"""
+
 from ._scheduler import TaskScheduler
 
 __all__ = ["TaskScheduler"]

cognite/extractorutils/unstable/scheduling/_scheduler.py

@@ -51,7 +51,7 @@ class TaskScheduler:
             return []
         with self._jobs_lock:
             next_runs = sorted([(j.schedule.next(), j) for j in self._jobs.values()], key=lambda tup: tup[0])
-        return [job for (next, job) in next_runs if next == next_runs[0][0]] if next_runs else []
+        return [job for (scheduled_time, job) in next_runs if scheduled_time == next_runs[0][0]] if next_runs else []
 
     def _run_job(self, job: Job) -> bool:
         with self._running_lock:

cognite/extractorutils/uploader/__init__.py

@@ -13,8 +13,10 @@
 #  limitations under the License.
 
 """
-Module containing upload queue classes. The UploadQueue classes chunks together items and uploads them together to CDF,
-both to minimize the load on the API, and also to speed up uploading as requests can be slow.
+Module containing upload queue classes.
+
+The UploadQueue classes chunks together items and uploads them together to CDF,both to minimize the load on the API, and
+also to speed up uploading as requests can be slow.
 
 Each upload queue comes with some configurable conditions that, when met, automatically triggers an upload.
 
@@ -78,13 +80,13 @@ from .time_series import (
 
 __all__ = [
     "AssetUploadQueue",
-    "EventUploadQueue",
     "BytesUploadQueue",
+    "DataPoint",
+    "DataPointList",
+    "EventUploadQueue",
     "FileUploadQueue",
     "IOFileUploadQueue",
     "RawUploadQueue",
-    "DataPoint",
-    "DataPointList",
     "SequenceUploadQueue",
     "TimeSeriesUploadQueue",
     "default_time_series_factory",

cognite/extractorutils/uploader/_base.py

@@ -84,7 +84,7 @@ class AbstractUploadQueue(ABC):
 
     def _post_upload(self, uploaded: list[Any]) -> None:
         """
-        Perform post_upload_function to uploaded data, if applicable
+        Perform post_upload_function to uploaded data, if applicable.
 
         Args:
             uploaded: list of uploaded data
@@ -103,7 +103,7 @@ class AbstractUploadQueue(ABC):
 
     def _run(self) -> None:
         """
-        Internal run method for upload thread
+        Internal run method for upload thread.
         """
         while not self.cancellation_token.wait(timeout=self.max_upload_interval):
             try:
@@ -117,8 +117,7 @@ class AbstractUploadQueue(ABC):
 
     def start(self) -> None:
         """
-        Start upload thread if max_upload_interval is set, this called the upload method every max_upload_interval
-        seconds.
+        Start upload thread if max_upload_interval is set.
         """
         if self.max_upload_interval is not None:
             self.thread.start()
@@ -137,7 +136,7 @@ class AbstractUploadQueue(ABC):
 
     def __len__(self) -> int:
         """
-        The size of the upload queue
+        The size of the upload queue.
 
         Returns:
             Number of events in queue

cognite/extractorutils/uploader/assets.py

@@ -1,3 +1,7 @@
+"""
+Upload queue for (legacy) assets.
+"""
+
 #  Copyright 2023 Cognite AS
 #
 #  Licensed under the Apache License, Version 2.0 (the "License");
@@ -37,7 +41,7 @@ from cognite.extractorutils.util import cognite_exceptions, retry
 
 class AssetUploadQueue(AbstractUploadQueue):
     """
-    Upload queue for assets
+    Upload queue for assets.
 
     Args:
         cdf_client: Cognite Data Fusion client to use
@@ -77,8 +81,9 @@ class AssetUploadQueue(AbstractUploadQueue):
 
     def add_to_upload_queue(self, asset: Asset) -> None:
         """
-        Add asset to upload queue. The queue will be uploaded if the queue size is larger than the threshold
-        specified in the __init__.
+        Add asset to upload queue.
+
+        The queue will be uploaded if the queue size is larger than the threshold specified in the ``__init__``.
 
         Args:
             asset: Asset to add
@@ -92,7 +97,7 @@ class AssetUploadQueue(AbstractUploadQueue):
 
     def upload(self) -> None:
         """
-        Trigger an upload of the queue, clears queue afterwards
+        Trigger an upload of the queue, clears queue afterwards.
         """
 
         @retry(
@@ -107,8 +112,8 @@ class AssetUploadQueue(AbstractUploadQueue):
             try:
                 self.cdf_client.assets.create(self.upload_queue)
             except CogniteDuplicatedError as e:
-                duplicated_ids = set([dup["externalId"] for dup in e.duplicated if "externalId" in dup])
-                failed: list[Asset] = [e for e in e.failed]
+                duplicated_ids = {dup["externalId"] for dup in e.duplicated if "externalId" in dup}
+                failed: list[Asset] = list(e.failed)
                 to_create = []
                 to_update = []
                 for asset in failed:
@@ -138,7 +143,7 @@ class AssetUploadQueue(AbstractUploadQueue):
 
     def __enter__(self) -> "AssetUploadQueue":
         """
-        Wraps around start method, for use as context manager
+        Wraps around start method, for use as context manager.
 
         Returns:
             self
@@ -153,7 +158,7 @@ class AssetUploadQueue(AbstractUploadQueue):
         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