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

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
@@ -5,27 +30,48 @@ import time
 from argparse import ArgumentParser, Namespace
 from multiprocessing import Process, Queue
 from pathlib import Path
+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.exceptions import CogniteAPIError, CogniteAuthError, CogniteConnectionError
+from cognite.client import CogniteClient
+from cognite.client.exceptions import (
+    CogniteAPIError,
+    CogniteAuthError,
+    CogniteConnectionError,
+)
 from cognite.extractorutils.threading import CancellationToken
 from cognite.extractorutils.unstable.configuration.exceptions import InvalidConfigError
-from cognite.extractorutils.unstable.configuration.loaders import load_file, load_from_cdf
+from cognite.extractorutils.unstable.configuration.loaders import (
+    load_file,
+    load_from_cdf,
+)
 from cognite.extractorutils.unstable.configuration.models import ConnectionConfig
 from cognite.extractorutils.unstable.core._dto import Error
+from cognite.extractorutils.unstable.core.errors import ErrorLevel
+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__(
         self,
         extractor: type[ExtractorType],
@@ -37,6 +83,8 @@ class Runtime(Generic[ExtractorType]):
         self.logger = logging.getLogger(f"{self._extractor_class.EXTERNAL_ID}.runtime")
         self._setup_logging()
 
+        self._cognite_client: CogniteClient
+
     def _create_argparser(self) -> ArgumentParser:
         argparser = ArgumentParser(
             prog=sys.argv[0],
@@ -57,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,
@@ -121,7 +169,7 @@ class Runtime(Generic[ExtractorType]):
         self.logger.info(f"Started extractor with PID {process.pid}")
         return process
 
-    def _get_application_config(
+    def _try_get_application_config(
         self,
         args: Namespace,
         connection_config: ConnectionConfig,
@@ -143,39 +191,65 @@ class Runtime(Generic[ExtractorType]):
 
         else:
             self.logger.info("Loading application config from CDF")
-            client = connection_config.get_cognite_client(
-                f"{self._extractor_class.EXTERNAL_ID}-{self._extractor_class.VERSION}"
+
+            application_config, current_config_revision = load_from_cdf(
+                self._cognite_client,
+                connection_config.integration,
+                self._extractor_class.CONFIG_TYPE,
             )
 
-            errors: list[Error] = []
+        return application_config, current_config_revision
+
+    def _safe_get_application_config(
+        self,
+        args: Namespace,
+        connection_config: ConnectionConfig,
+    ) -> tuple[ConfigType, ConfigRevision] | None:
+        prev_error: str | None = None
 
+        while not self._cancellation_token.is_cancelled:
             try:
-                application_config, current_config_revision = load_from_cdf(
-                    client,
-                    connection_config.integration,
-                    self._extractor_class.CONFIG_TYPE,
+                return self._try_get_application_config(args, connection_config)
+
+            except Exception as e:
+                error_message = str(e)
+                if error_message == prev_error:
+                    # Same error as before, no need to log it again
+                    self._cancellation_token.wait(randint(1, self.RETRY_CONFIG_INTERVAL))
+                    continue
+                prev_error = error_message
+
+                ts = now()
+                error = Error(
+                    external_id=str(uuid4()),
+                    level=ErrorLevel.fatal.value,
+                    start_time=ts,
+                    end_time=ts,
+                    description=error_message,
+                    details=None,
+                    task=None,
                 )
 
-            finally:
-                if errors:
-                    client.post(
-                        f"/api/v1/projects/{client.config.project}/odin/checkin",
-                        json={
-                            "externalId": connection_config.integration,
-                            "errors": [e.model_dump() for e in errors],
-                        },
-                        headers={"cdf-version": "alpha"},
-                    )
+                self._cognite_client.post(
+                    f"/api/v1/projects/{self._cognite_client.config.project}/odin/checkin",
+                    json={
+                        "externalId": connection_config.integration,
+                        "errors": [error.model_dump()],
+                    },
+                    headers={"cdf-version": "alpha"},
+                )
 
-        return application_config, current_config_revision
+                self._cancellation_token.wait(randint(1, self.RETRY_CONFIG_INTERVAL))
+
+        return None
 
     def _verify_connection_config(self, connection_config: ConnectionConfig) -> bool:
-        client = connection_config.get_cognite_client(
+        self._cognite_client = connection_config.get_cognite_client(
             f"{self._extractor_class.EXTERNAL_ID}-{self._extractor_class.VERSION}"
         )
         try:
-            client.post(
-                f"/api/v1/projects/{client.config.project}/odin/checkin",
+            self._cognite_client.post(
+                f"/api/v1/projects/{self._cognite_client.config.project}/odin/checkin",
                 json={
                     "externalId": connection_config.integration,
                 },
@@ -206,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.")
@@ -219,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()
 
@@ -234,16 +314,19 @@ class Runtime(Generic[ExtractorType]):
         if not args.skip_init_checks and not self._verify_connection_config(connection_config):
             sys.exit(1)
 
-        # This has to be Any. We don't know the type of the extractors' config at type checking since the sel doesn't
+        # This has to be Any. We don't know the type of the extractors' config at type checking since the self doesn't
         # exist yet, and I have not found a way to represent it in a generic way that isn't just an Any in disguise.
         application_config: Any
+        config: tuple[Any, ConfigRevision] | None
+
         while not self._cancellation_token.is_cancelled:
-            try:
-                application_config, current_config_revision = self._get_application_config(args, connection_config)
+            config = self._safe_get_application_config(args, connection_config)
+            if config is None:
+                if self._cancellation_token.is_cancelled:
+                    break
+                continue
 
-            except InvalidConfigError:
-                self.logger.critical("Could not get a valid application config file. Shutting down")
-                sys.exit(1)
+            application_config, current_config_revision = config
 
             # Start extractor in separate process, and wait for it to end
             process = self._spawn_extractor(

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

cognite/extractorutils/uploader/data_modeling.py

@@ -1,3 +1,7 @@
+"""
+Module for uploading data modeling instances to CDF.
+"""
+
 from collections.abc import Callable
 from types import TracebackType
 from typing import Any
@@ -16,6 +20,24 @@ from cognite.extractorutils.util import cognite_exceptions, retry
 
 
 class InstanceUploadQueue(AbstractUploadQueue):
+    """
+    Upload queue for data modeling instances (nodes and edges).
+
+    Args:
+        cdf_client: Cognite Data Fusion client to use.
+        post_upload_function: A function that will be called after each upload. The function will be given one argument:
+            A list of the nodes and edges that were uploaded.
+        max_queue_size: Maximum size of upload queue. Defaults to no max size.
+        max_upload_interval: Automatically trigger an upload on an interval when run as a thread (use start/stop
+            methods). Unit is seconds.
+        trigger_log_level: Log level to log upload triggers to.
+        thread_name: Thread name of uploader thread.
+        cancellation_token: Cancellation token for managing thread cancellation.
+        auto_create_start_nodes: Automatically create start nodes if they do not exist.
+        auto_create_end_nodes: Automatically create end nodes if they do not exist.
+        auto_create_direct_relations: Automatically create direct relations if they do not exist.
+    """
+
     def __init__(
         self,
         cdf_client: CogniteClient,
@@ -52,6 +74,15 @@ class InstanceUploadQueue(AbstractUploadQueue):
         node_data: list[NodeApply] | None = None,
         edge_data: list[EdgeApply] | None = None,
     ) -> None:
+        """
+        Add instances to the upload queue.
+
+        The queue will be uploaded if the queue size is larger than the threshold specified in the ``__init__``.
+
+        Args:
+            node_data: List of nodes to add to the upload queue.
+            edge_data: List of edges to add to the upload queue.
+        """
         if node_data:
             with self.lock:
                 self.node_queue.extend(node_data)
@@ -66,6 +97,10 @@ class InstanceUploadQueue(AbstractUploadQueue):
             self._check_triggers()
 
     def upload(self) -> None:
+        """
+        Trigger an upload of the queue, clears queue afterwards.
+        """
+
         @retry(
             exceptions=cognite_exceptions(),
             cancellation_token=self.cancellation_token,
@@ -91,7 +126,7 @@ class InstanceUploadQueue(AbstractUploadQueue):
 
     def __enter__(self) -> "InstanceUploadQueue":
         """
-        Wraps around start method, for use as context manager
+        Wraps around start method, for use as context manager.
 
         Returns:
             self
@@ -106,7 +141,7 @@ class InstanceUploadQueue(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

cognite/extractorutils/uploader/events.py

@@ -1,3 +1,7 @@
+"""
+Upload queue for (legacy) events.
+"""
+
 #  Copyright 2023 Cognite AS
 #
 #  Licensed under the Apache License, Version 2.0 (the "License");
@@ -36,7 +40,7 @@ from cognite.extractorutils.util import cognite_exceptions, retry
 
 class EventUploadQueue(AbstractUploadQueue):
     """
-    Upload queue for events
+    Upload queue for events.
 
     Args:
         cdf_client: Cognite Data Fusion client to use
@@ -78,8 +82,9 @@ class EventUploadQueue(AbstractUploadQueue):
 
     def add_to_upload_queue(self, event: Event) -> None:
         """
-        Add event to upload queue. The queue will be uploaded if the queue size is larger than the threshold
-        specified in the __init__.
+        Add event to upload queue.
+
+        The queue will be uploaded if the queue size is larger than the threshold specified in the ``__init__``.
 
         Args:
             event: Event to add
@@ -94,7 +99,7 @@ class EventUploadQueue(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,10 +112,10 @@ class EventUploadQueue(AbstractUploadQueue):
         )
         def _upload_batch() -> None:
             try:
-                self.cdf_client.events.create([e for e in self.upload_queue])
+                self.cdf_client.events.create(list(self.upload_queue))
             except CogniteDuplicatedError as e:
-                duplicated_ids = set([dup["externalId"] for dup in e.duplicated if "externalId" in dup])
-                failed: list[Event] = [e for e in e.failed]
+                duplicated_ids = {dup["externalId"] for dup in e.duplicated if "externalId" in dup}
+                failed: list[Event] = list(e.failed)
                 to_create = []
                 to_update = []
                 for evt in failed:
@@ -142,7 +147,7 @@ class EventUploadQueue(AbstractUploadQueue):
 
     def __enter__(self) -> "EventUploadQueue":
         """
-        Wraps around start method, for use as context manager
+        Wraps around start method, for use as context manager.
 
         Returns:
             self
@@ -154,7 +159,7 @@ class EventUploadQueue(AbstractUploadQueue):
         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