cognite-extractor-utils 6.4.1__tar.gz → 7.0.0__tar.gz

{cognite_extractor_utils-6.4.1 → cognite_extractor_utils-7.0.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: cognite-extractor-utils
-Version: 6.4.1
+Version: 7.0.0
 Summary: Utilities for easier development of extractors for CDF
 Home-page: https://github.com/cognitedata/python-extractor-utils
 License: Apache-2.0
@@ -21,8 +21,6 @@ Requires-Dist: azure-keyvault-secrets (>=4.7.0,<5.0.0)
 Requires-Dist: cognite-sdk (>=7,<8)
 Requires-Dist: dacite (>=1.6.0,<2.0.0)
 Requires-Dist: decorator (>=5.1.1,<6.0.0)
-Requires-Dist: jq (>=1.3.0,<2.0.0) ; sys_platform == "darwin"
-Requires-Dist: jq (>=1.3.0,<2.0.0) ; sys_platform == "linux"
 Requires-Dist: more-itertools (>=10.0.0,<11.0.0)
 Requires-Dist: prometheus-client (>0.7.0,<=1.0.0)
 Requires-Dist: psutil (>=5.7.0,<6.0.0)

{cognite_extractor_utils-6.4.1 → cognite_extractor_utils-7.0.0}/cognite/extractorutils/__init__.py

@@ -16,5 +16,5 @@
 Cognite extractor utils is a Python package that simplifies the development of new extractors.
 """
 
-__version__ = "6.4.1"
+__version__ = "7.0.0"
 from .base import Extractor

{cognite_extractor_utils-6.4.1 → cognite_extractor_utils-7.0.0}/cognite/extractorutils/_inner_util.py

@@ -34,7 +34,7 @@ class _DecimalEncoder(json.JSONEncoder):
 
 class _DecimalDecoder(json.JSONDecoder):
     def __init__(self, *args: Any, **kwargs: Any) -> None:
-        json.JSONDecoder.__init__(self, object_hook=self.object_hook, *args, **kwargs)
+        json.JSONDecoder.__init__(self, *args, object_hook=self.object_hook, **kwargs)
 
     def object_hook(self, obj_dict: Dict[str, str]) -> Union[Dict[str, str], Decimal]:
         if obj_dict.get("type") == "decimal_encoded":

{cognite_extractor_utils-6.4.1 → cognite_extractor_utils-7.0.0}/cognite/extractorutils/base.py

@@ -17,7 +17,7 @@ import os
 import sys
 from dataclasses import is_dataclass
 from enum import Enum
-from threading import Event, Thread
+from threading import Thread
 from types import TracebackType
 from typing import Any, Callable, Dict, Generic, Optional, Type, TypeVar
 
@@ -29,7 +29,7 @@ from cognite.extractorutils.configtools import BaseConfig, ConfigResolver, State
 from cognite.extractorutils.exceptions import InvalidConfigError
 from cognite.extractorutils.metrics import BaseMetrics
 from cognite.extractorutils.statestore import AbstractStateStore, LocalStateStore, NoStateStore
-from cognite.extractorutils.util import set_event_on_interrupt
+from cognite.extractorutils.threading import CancellationToken
 
 
 class ReloadConfigAction(Enum):
@@ -77,11 +77,13 @@ class Extractor(Generic[CustomConfigClass]):
         name: str,
         description: str,
         version: Optional[str] = None,
-        run_handle: Optional[Callable[[CogniteClient, AbstractStateStore, CustomConfigClass, Event], None]] = None,
+        run_handle: Optional[
+            Callable[[CogniteClient, AbstractStateStore, CustomConfigClass, CancellationToken], None]
+        ] = None,
         config_class: Type[CustomConfigClass],
         metrics: Optional[BaseMetrics] = None,
         use_default_state_store: bool = True,
-        cancellation_token: Event = Event(),
+        cancellation_token: Optional[CancellationToken] = None,
         config_file_path: Optional[str] = None,
         continuous_extractor: bool = False,
         heartbeat_waiting_time: int = 600,
@@ -95,7 +97,7 @@ class Extractor(Generic[CustomConfigClass]):
         self.config_class = config_class
         self.use_default_state_store = use_default_state_store
         self.version = version or "unknown"
-        self.cancellation_token = cancellation_token
+        self.cancellation_token = cancellation_token.create_child_token() if cancellation_token else CancellationToken()
         self.config_file_path = config_file_path
         self.continuous_extractor = continuous_extractor
         self.heartbeat_waiting_time = heartbeat_waiting_time
@@ -136,7 +138,7 @@ class Extractor(Generic[CustomConfigClass]):
         Extractor._config_singleton = self.config  # type: ignore
 
         def config_refresher() -> None:
-            while not self.cancellation_token.is_set():
+            while not self.cancellation_token.is_cancelled:
                 self.cancellation_token.wait(self.reload_config_interval)
                 if self.config_resolver.has_changed:
                     self._reload_config()
@@ -158,7 +160,7 @@ class Extractor(Generic[CustomConfigClass]):
 
         elif self.reload_config_action == ReloadConfigAction.SHUTDOWN:
             self.logger.info("Shutting down, expecting to be restarted")
-            self.cancellation_token.set()
+            self.cancellation_token.cancel()
 
         elif self.reload_config_action == ReloadConfigAction.CALLBACK:
             self.logger.info("Loading in new config file")
@@ -265,7 +267,7 @@ class Extractor(Generic[CustomConfigClass]):
         self.logger.info(f"Loaded {'remote' if self.config_resolver.is_remote else 'local'} config file")
 
         if self.handle_interrupts:
-            set_event_on_interrupt(self.cancellation_token)
+            self.cancellation_token.cancel_on_interrupt()
 
         self.cognite_client = self.config.cognite.get_cognite_client(self.name)
         self._load_state_store()
@@ -279,10 +281,10 @@ class Extractor(Generic[CustomConfigClass]):
             pass
 
         def heartbeat_loop() -> None:
-            while not self.cancellation_token.is_set():
+            while not self.cancellation_token.is_cancelled:
                 self.cancellation_token.wait(self.heartbeat_waiting_time)
 
-                if not self.cancellation_token.is_set():
+                if not self.cancellation_token.is_cancelled:
                     self.logger.info("Reporting new heartbeat")
                     try:
                         self.cognite_client.extraction_pipelines.runs.create(
@@ -329,7 +331,7 @@ class Extractor(Generic[CustomConfigClass]):
         Returns:
             True if the extractor shut down cleanly, False if the extractor was shut down due to an unhandled error
         """
-        self.cancellation_token.set()
+        self.cancellation_token.cancel()
 
         if self.state_store:
             self.state_store.synchronize()

{cognite_extractor_utils-6.4.1 → cognite_extractor_utils-7.0.0}/cognite/extractorutils/configtools/elements.py

@@ -18,7 +18,6 @@ from dataclasses import dataclass, field
 from datetime import timedelta
 from enum import Enum
 from logging.handlers import TimedRotatingFileHandler
-from threading import Event
 from time import sleep
 from typing import Any, Dict, List, Optional, Tuple, Union
 from urllib.parse import urljoin
@@ -46,6 +45,7 @@ from cognite.extractorutils.statestore import (
     NoStateStore,
     RawStateStore,
 )
+from cognite.extractorutils.threading import CancellationToken
 from cognite.extractorutils.util import EitherId
 
 _logger = logging.getLogger(__name__)
@@ -481,7 +481,7 @@ class MetricsConfig:
     cognite: Optional[_CogniteMetricsConfig]
     server: Optional[_PromServerConfig]
 
-    def start_pushers(self, cdf_client: CogniteClient, cancellation_token: Event = Event()) -> None:
+    def start_pushers(self, cdf_client: CogniteClient, cancellation_token: Optional[CancellationToken] = None) -> None:
         self._pushers: List[AbstractMetricsPusher] = []
         self._clear_on_stop: Dict[PrometheusPusher, int] = {}
 

{cognite_extractor_utils-6.4.1 → cognite_extractor_utils-7.0.0}/cognite/extractorutils/configtools/loaders.py

@@ -20,6 +20,7 @@ import re
 import sys
 from enum import Enum
 from hashlib import sha256
+from pathlib import Path
 from typing import Any, Callable, Dict, Generic, Iterable, Optional, TextIO, Type, TypeVar, Union
 
 import dacite
@@ -111,7 +112,7 @@ class KeyVaultLoader:
         try:
             return self.client.get_secret(node.value).value  # type: ignore  # _init_client guarantees not None
         except (ResourceNotFoundError, ServiceRequestError, HttpResponseError) as e:
-            raise InvalidConfigError(str(e))
+            raise InvalidConfigError(str(e)) from e
 
 
 class _EnvLoader(yaml.SafeLoader):
@@ -188,11 +189,15 @@ def _load_yaml(
 
     try:
         config = dacite.from_dict(
-            data=config_dict, data_class=config_type, config=dacite.Config(strict=True, cast=[Enum, TimeIntervalConfig])
+            data=config_dict,
+            data_class=config_type,
+            config=dacite.Config(strict=True, cast=[Enum, TimeIntervalConfig, Path]),
         )
     except dacite.UnexpectedDataError as e:
         unknowns = [f'"{k.replace("_", "-") if case_style == "hyphen" else k}"' for k in e.keys]
-        raise InvalidConfigError(f"Unknown config parameter{'s' if len(unknowns) > 1 else ''} {', '.join(unknowns)}")
+        raise InvalidConfigError(
+            f"Unknown config parameter{'s' if len(unknowns) > 1 else ''} {', '.join(unknowns)}"
+        ) from e
 
     except (dacite.WrongTypeError, dacite.MissingValueError, dacite.UnionMatchError) as e:
         if e.field_path:
@@ -212,11 +217,11 @@ def _load_yaml(
 
             raise InvalidConfigError(
                 f'Wrong type for field "{path}" - got "{e.value}" of type {got_type} instead of {need_type}'
-            )
-        raise InvalidConfigError(f'Missing mandatory field "{path}"')
+            ) from e
+        raise InvalidConfigError(f'Missing mandatory field "{path}"') from e
 
     except dacite.ForwardReferenceError as e:
-        raise ValueError(f"Invalid config class: {str(e)}")
+        raise ValueError(f"Invalid config class: {str(e)}") from e
 
     config._file_hash = sha256(json.dumps(config_dict).encode("utf-8")).hexdigest()
 

{cognite_extractor_utils-6.4.1 → cognite_extractor_utils-7.0.0}/cognite/extractorutils/metrics.py

@@ -41,7 +41,6 @@ import logging
 import os
 import threading
 from abc import ABC, abstractmethod
-from threading import Event
 from time import sleep
 from types import TracebackType
 from typing import Any, Callable, Dict, List, Optional, Tuple, Type, TypeVar, Union
@@ -55,6 +54,7 @@ from prometheus_client.exposition import basic_auth_handler, delete_from_gateway
 from cognite.client import CogniteClient
 from cognite.client.data_classes import Asset, Datapoints, DatapointsArray, TimeSeries
 from cognite.client.exceptions import CogniteDuplicatedError
+from cognite.extractorutils.threading import CancellationToken
 from cognite.extractorutils.util import EitherId
 
 from .util import ensure_time_series
@@ -179,14 +179,14 @@ class AbstractMetricsPusher(ABC):
         self,
         push_interval: Optional[int] = None,
         thread_name: Optional[str] = None,
-        cancellation_token: Event = Event(),
+        cancellation_token: Optional[CancellationToken] = None,
     ):
         self.push_interval = push_interval
         self.thread_name = thread_name
 
         self.thread: Optional[threading.Thread] = None
         self.thread_name = thread_name
-        self.cancellation_token = cancellation_token
+        self.cancellation_token = cancellation_token.create_child_token() if cancellation_token else CancellationToken()
 
         self.logger = logging.getLogger(__name__)
 
@@ -201,7 +201,7 @@ class AbstractMetricsPusher(ABC):
         """
         Run push loop.
         """
-        while not self.cancellation_token.is_set():
+        while not self.cancellation_token.is_cancelled:
             self._push_to_server()
             self.cancellation_token.wait(self.push_interval)
 
@@ -210,7 +210,6 @@ class AbstractMetricsPusher(ABC):
         Starts a thread that pushes the default registry to the configured gateway at certain intervals.
 
         """
-        self.cancellation_token.clear()
         self.thread = threading.Thread(target=self._run, daemon=True, name=self.thread_name)
         self.thread.start()
 
@@ -220,7 +219,7 @@ class AbstractMetricsPusher(ABC):
         """
         # Make sure everything is pushed
         self._push_to_server()
-        self.cancellation_token.set()
+        self.cancellation_token.cancel()
 
     def __enter__(self) -> "AbstractMetricsPusher":
         """
@@ -268,7 +267,7 @@ class PrometheusPusher(AbstractMetricsPusher):
         username: Optional[str] = None,
         password: Optional[str] = None,
         thread_name: Optional[str] = None,
-        cancellation_token: Event = Event(),
+        cancellation_token: Optional[CancellationToken] = None,
     ):
         super(PrometheusPusher, self).__init__(push_interval, thread_name, cancellation_token)
 
@@ -344,7 +343,7 @@ class CognitePusher(AbstractMetricsPusher):
         asset: Optional[Asset] = None,
         data_set: Optional[EitherId] = None,
         thread_name: Optional[str] = None,
-        cancellation_token: Event = Event(),
+        cancellation_token: Optional[CancellationToken] = None,
     ):
         super(CognitePusher, self).__init__(push_interval, thread_name, cancellation_token)
 

{cognite_extractor_utils-6.4.1 → cognite_extractor_utils-7.0.0}/cognite/extractorutils/statestore.py

@@ -92,18 +92,17 @@ from abc import ABC, abstractmethod
 from types import TracebackType
 from typing import Any, Callable, Dict, Iterator, List, Optional, Tuple, Type, Union
 
-from requests.exceptions import ConnectionError
-
 from cognite.client import CogniteClient
-from cognite.client.exceptions import CogniteAPIError, CogniteException
+from cognite.client.exceptions import CogniteAPIError
+from cognite.extractorutils.threading import CancellationToken
 from cognite.extractorutils.uploader import DataPointList
 
 from ._inner_util import _DecimalDecoder, _DecimalEncoder, _resolve_log_level
-from .util import retry
+from .util import cognite_exceptions, retry
 
 RETRY_BACKOFF_FACTOR = 1.5
-RETRY_MAX_DELAY = 15
-RETRY_DELAY = 5
+RETRY_MAX_DELAY = 60
+RETRY_DELAY = 1
 RETRIES = 10
 
 
@@ -124,7 +123,7 @@ class AbstractStateStore(ABC):
         save_interval: Optional[int] = None,
         trigger_log_level: str = "DEBUG",
         thread_name: Optional[str] = None,
-        cancellation_token: threading.Event = threading.Event(),
+        cancellation_token: Optional[CancellationToken] = None,
     ):
         self._initialized = False
         self._local_state: Dict[str, Dict[str, Any]] = {}
@@ -135,7 +134,7 @@ class AbstractStateStore(ABC):
 
         self.thread = threading.Thread(target=self._run, daemon=True, name=thread_name)
         self.lock = threading.RLock()
-        self.cancellation_token: threading.Event = cancellation_token
+        self.cancellation_token = cancellation_token.create_child_token() if cancellation_token else CancellationToken()
 
         self._deleted: List[str] = []
 
@@ -145,7 +144,6 @@ class AbstractStateStore(ABC):
         This calls the synchronize method every save_interval seconds.
         """
         if self.save_interval is not None:
-            self.cancellation_token.clear()
             self.thread.start()
 
     def stop(self, ensure_synchronize: bool = True) -> None:
@@ -155,7 +153,7 @@ class AbstractStateStore(ABC):
         Args:
             ensure_synchronize (bool): (Optional). Call synchronize one last time after shutting down thread.
         """
-        self.cancellation_token.set()
+        self.cancellation_token.cancel()
         if ensure_synchronize:
             self.synchronize()
 
@@ -337,7 +335,7 @@ class RawStateStore(AbstractStateStore):
         save_interval: Optional[int] = None,
         trigger_log_level: str = "DEBUG",
         thread_name: Optional[str] = None,
-        cancellation_token: threading.Event = threading.Event(),
+        cancellation_token: Optional[CancellationToken] = None,
     ):
         super().__init__(save_interval, trigger_log_level, thread_name, cancellation_token)
 
@@ -347,79 +345,87 @@ class RawStateStore(AbstractStateStore):
 
         self._ensure_table()
 
-    @retry(
-        exceptions=(CogniteException, ConnectionError),
-        tries=RETRIES,
-        delay=RETRY_DELAY,
-        max_delay=RETRY_MAX_DELAY,
-        backoff=RETRY_BACKOFF_FACTOR,
-    )
     def _ensure_table(self) -> None:
-        try:
-            self._cdf_client.raw.databases.create(self.database)
-        except CogniteAPIError as e:
-            if not e.code == 400:
-                raise e
-        try:
-            self._cdf_client.raw.tables.create(self.database, self.table)
-        except CogniteAPIError as e:
-            if not e.code == 400:
-                raise e
-
-    def initialize(self, force: bool = False) -> None:
-        self._initialize_implementation(force)
-
-    @retry(
-        exceptions=(CogniteException, ConnectionError),
-        tries=RETRIES,
-        delay=RETRY_DELAY,
-        max_delay=RETRY_MAX_DELAY,
-        backoff=RETRY_BACKOFF_FACTOR,
-    )
-    def _initialize_implementation(self, force: bool = False) -> None:
-        """
-        Get all known states.
-
-        Args:
-            force: Enable re-initialization, ie overwrite when called multiple times
-        """
-        if self._initialized and not force:
-            return
-
-        # ignore type since list _is_ optional, sdk types are wrong
-        rows = self._cdf_client.raw.rows.list(db_name=self.database, table_name=self.table, limit=None)  # type: ignore
+        @retry(
+            exceptions=cognite_exceptions(),
+            cancellation_token=self.cancellation_token,
+            tries=RETRIES,
+            delay=RETRY_DELAY,
+            max_delay=RETRY_MAX_DELAY,
+            backoff=RETRY_BACKOFF_FACTOR,
+        )
+        def impl() -> None:
+            try:
+                self._cdf_client.raw.databases.create(self.database)
+            except CogniteAPIError as e:
+                if not e.code == 400:
+                    raise e
+            try:
+                self._cdf_client.raw.tables.create(self.database, self.table)
+            except CogniteAPIError as e:
+                if not e.code == 400:
+                    raise e
 
-        with self.lock:
-            self._local_state.clear()
-            for row in rows:
-                if row.key is None or row.columns is None:
-                    self.logger.warning(f"None encountered in row: {str(row)}")
-                    # should never happen, but type from sdk is optional
-                    continue
-                self._local_state[row.key] = row.columns
+        impl()
 
-        self._initialized = True
+    def initialize(self, force: bool = False) -> None:
+        @retry(
+            exceptions=cognite_exceptions(),
+            cancellation_token=self.cancellation_token,
+            tries=RETRIES,
+            delay=RETRY_DELAY,
+            max_delay=RETRY_MAX_DELAY,
+            backoff=RETRY_BACKOFF_FACTOR,
+        )
+        def impl() -> None:
+            """
+            Get all known states.
+
+            Args:
+                force: Enable re-initialization, ie overwrite when called multiple times
+            """
+            if self._initialized and not force:
+                return
+
+            # ignore type since list _is_ optional, sdk types are wrong
+            rows = self._cdf_client.raw.rows.list(db_name=self.database, table_name=self.table, limit=None)  # type: ignore
+
+            with self.lock:
+                self._local_state.clear()
+                for row in rows:
+                    if row.key is None or row.columns is None:
+                        self.logger.warning(f"None encountered in row: {str(row)}")
+                        # should never happen, but type from sdk is optional
+                        continue
+                    self._local_state[row.key] = row.columns
+
+            self._initialized = True
+
+        impl()
 
     def synchronize(self) -> None:
-        self._synchronize_implementation()
-
-    @retry(
-        exceptions=(CogniteException, ConnectionError),
-        tries=RETRIES,
-        delay=RETRY_DELAY,
-        max_delay=RETRY_MAX_DELAY,
-        backoff=RETRY_BACKOFF_FACTOR,
-    )
-    def _synchronize_implementation(self) -> None:
-        """
-        Upload local state store to CDF
-        """
-        self._cdf_client.raw.rows.insert(db_name=self.database, table_name=self.table, row=self._local_state)
-        # Create a copy of deleted to facilitate testing (mock library stores list, and as it changes, the assertions
-        # fail)
-        self._cdf_client.raw.rows.delete(db_name=self.database, table_name=self.table, key=[k for k in self._deleted])
-        with self.lock:
-            self._deleted.clear()
+        @retry(
+            exceptions=cognite_exceptions(),
+            cancellation_token=self.cancellation_token,
+            tries=RETRIES,
+            delay=RETRY_DELAY,
+            max_delay=RETRY_MAX_DELAY,
+            backoff=RETRY_BACKOFF_FACTOR,
+        )
+        def impl() -> None:
+            """
+            Upload local state store to CDF
+            """
+            self._cdf_client.raw.rows.insert(db_name=self.database, table_name=self.table, row=self._local_state)
+            # Create a copy of deleted to facilitate testing (mock library stores list, and as it changes, the
+            # assertions fail)
+            self._cdf_client.raw.rows.delete(
+                db_name=self.database, table_name=self.table, key=[k for k in self._deleted]
+            )
+            with self.lock:
+                self._deleted.clear()
+
+        impl()
 
     def __enter__(self) -> "RawStateStore":
         """
@@ -464,7 +470,7 @@ class LocalStateStore(AbstractStateStore):
         save_interval: Optional[int] = None,
         trigger_log_level: str = "DEBUG",
         thread_name: Optional[str] = None,
-        cancellation_token: threading.Event = threading.Event(),
+        cancellation_token: Optional[CancellationToken] = None,
     ):
         super().__init__(save_interval, trigger_log_level, thread_name, cancellation_token)
 

cognite_extractor_utils-7.0.0/cognite/extractorutils/threading.py

@@ -0,0 +1,90 @@
+import logging
+import signal
+from threading import Condition
+from typing import Any, Optional
+
+
+class CancellationToken:
+    """
+    Abstraction for a hierarchical cancellation token.
+
+    Using this you can create hierarchies of cancellation tokens, to cancel a part of the extractor
+    without cancelling the whole process. Use ``create_child_token`` to create a token that will be
+    cancelled if the parent is cancelled, but can be canceled alone without affecting the parent token.
+    """
+
+    def __init__(self, condition: Optional[Condition] = None) -> None:
+        self._cv: Condition = condition or Condition()
+        self._is_cancelled_int: bool = False
+        self._parent: Optional["CancellationToken"] = None
+
+    def __repr__(self) -> str:
+        cls = self.__class__
+        status = "cancelled" if self.is_cancelled else "not cancelled"
+        return f"<{cls.__module__}.{cls.__qualname__} at {id(self):#x}: {status}>"
+
+    @property
+    def is_cancelled(self) -> bool:
+        """
+        ``True`` if the token has been cancelled, or if some parent token has been cancelled.
+        """
+        return self._is_cancelled_int or self._parent is not None and self._parent.is_cancelled
+
+    def is_set(self) -> bool:
+        """
+        Deprecated, use ``is_cancelled`` instead.
+
+        ``True`` if the token has been cancelled, or if some parent token has been cancelled.
+        """
+        return self.is_cancelled
+
+    def cancel(self) -> None:
+        """
+        Cancel the token, notifying any waiting threads.
+        """
+        # No point in cancelling if a parent token is already canceled.
+        if self.is_cancelled:
+            return
+
+        with self._cv:
+            self._is_cancelled_int = True
+            self._cv.notify_all()
+
+    def set(self) -> None:
+        """
+        Deprecated, use ``cancel`` instead. This will be removed in the next major release.
+
+        Cancel the token, notifying any waiting threads.
+        """
+        self.cancel()
+
+    def wait(self, timeout: Optional[float] = None) -> bool:
+        while not self.is_cancelled:
+            with self._cv:
+                did_not_time_out = self._cv.wait(timeout)
+                if not did_not_time_out:
+                    return False
+        return True
+
+    def create_child_token(self) -> "CancellationToken":
+        child = CancellationToken(self._cv)
+        child._parent = self
+        return child
+
+    def cancel_on_interrupt(self) -> None:
+        """
+        Register an interrupt handler to capture SIGINT (Ctrl-C) and cancel this token,
+        instead of throwing a KeyboardInterrupt exception.
+        """
+
+        def sigint_handler(sig_num: int, frame: Any) -> None:
+            logger = logging.getLogger(__name__)
+            logger.warning("Interrupt signal received, stopping extractor gracefully")
+            self.cancel()
+            logger.info("Waiting for threads to complete. Send another interrupt to force quit.")
+            signal.signal(signal.SIGINT, signal.default_int_handler)
+
+        try:
+            signal.signal(signal.SIGINT, sigint_handler)
+        except ValueError as e:
+            logging.getLogger(__name__).warning(f"Could not register handler for interrupt signals: {str(e)}")

{cognite_extractor_utils-6.4.1 → cognite_extractor_utils-7.0.0}/cognite/extractorutils/uploader/_base.py

@@ -22,6 +22,7 @@ from arrow import Arrow
 
 from cognite.client import CogniteClient
 from cognite.extractorutils._inner_util import _resolve_log_level
+from cognite.extractorutils.threading import CancellationToken
 
 
 class AbstractUploadQueue(ABC):
@@ -47,7 +48,7 @@ class AbstractUploadQueue(ABC):
         max_upload_interval: Optional[int] = None,
         trigger_log_level: str = "DEBUG",
         thread_name: Optional[str] = None,
-        cancellation_token: threading.Event = threading.Event(),
+        cancellation_token: Optional[CancellationToken] = None,
     ):
         self.cdf_client = cdf_client
 
@@ -59,7 +60,9 @@ class AbstractUploadQueue(ABC):
 
         self.thread = threading.Thread(target=self._run, daemon=True, name=thread_name)
         self.lock = threading.RLock()
-        self.cancellation_token: threading.Event = cancellation_token
+        self.cancellation_token: CancellationToken = (
+            cancellation_token.create_child_token() if cancellation_token else CancellationToken()
+        )
 
         self.max_upload_interval = max_upload_interval
 
@@ -117,7 +120,6 @@ class AbstractUploadQueue(ABC):
         seconds.
         """
         if self.max_upload_interval is not None:
-            self.cancellation_token.clear()
             self.thread.start()
 
     def stop(self, ensure_upload: bool = True) -> None:
@@ -128,7 +130,7 @@ class AbstractUploadQueue(ABC):
             ensure_upload (bool): (Optional). Call upload one last time after shutting down thread to ensure empty
                 upload queue.
         """
-        self.cancellation_token.set()
+        self.cancellation_token.cancel()
         if ensure_upload:
             self.upload()
 
@@ -148,7 +150,7 @@ class TimestampedObject:
     created: Arrow
 
 
-RETRY_BACKOFF_FACTOR = 1.5
-RETRY_MAX_DELAY = 15
-RETRY_DELAY = 5
+RETRY_BACKOFF_FACTOR = 2
+RETRY_MAX_DELAY = 60
+RETRY_DELAY = 1
 RETRIES = 10