cognite-extractor-utils 7.1.6__tar.gz → 7.2.1__tar.gz

{cognite_extractor_utils-7.1.6 → cognite_extractor_utils-7.2.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: cognite-extractor-utils
-Version: 7.1.6
+Version: 7.2.1
 Summary: Utilities for easier development of extractors for CDF
 Home-page: https://github.com/cognitedata/python-extractor-utils
 License: Apache-2.0
@@ -22,6 +22,7 @@ Requires-Dist: cognite-sdk (>=7.43.3,<8.0.0)
 Requires-Dist: dacite (>=1.6.0,<2.0.0)
 Requires-Dist: decorator (>=5.1.1,<6.0.0)
 Requires-Dist: more-itertools (>=10.0.0,<11.0.0)
+Requires-Dist: orjson (>=3.10.3,<4.0.0)
 Requires-Dist: prometheus-client (>0.7.0,<=1.0.0)
 Requires-Dist: psutil (>=5.7.0,<6.0.0)
 Requires-Dist: python-dotenv (>=1.0.0,<2.0.0)

{cognite_extractor_utils-7.1.6 → cognite_extractor_utils-7.2.1}/cognite/extractorutils/__init__.py

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

{cognite_extractor_utils-7.1.6 → cognite_extractor_utils-7.2.1}/cognite/extractorutils/configtools/loaders.py

@@ -139,11 +139,10 @@ def _env_constructor(_: yaml.SafeLoader, node: yaml.Node) -> bool:
     return bool_values.get(expanded_value.lower(), expanded_value)
 
 
-def _load_yaml_dict(
+def _load_yaml_dict_raw(
     source: Union[TextIO, str],
-    case_style: str = "hyphen",
     expand_envvars: bool = True,
-    dict_manipulator: Callable[[Dict[str, Any]], Dict[str, Any]] = lambda x: x,
+    keyvault_loader: Optional[KeyVaultLoader] = None,
 ) -> Dict[str, Any]:
     loader = _EnvLoader if expand_envvars else yaml.SafeLoader
 
@@ -159,11 +158,14 @@ def _load_yaml_dict(
     if not isinstance(source, str):
         source.seek(0)
 
-    keyvault_config = initial_load.get("azure-keyvault", initial_load.get("key-vault"))
+    if keyvault_loader:
+        _EnvLoader.add_constructor("!keyvault", keyvault_loader)
+    else:
+        keyvault_config = initial_load.get("azure-keyvault", initial_load.get("key-vault"))
+        _EnvLoader.add_constructor("!keyvault", KeyVaultLoader(keyvault_config))
 
     _EnvLoader.add_implicit_resolver("!env", re.compile(r"\$\{([^}^{]+)\}"), None)
     _EnvLoader.add_constructor("!env", _env_constructor)
-    _EnvLoader.add_constructor("!keyvault", KeyVaultLoader(keyvault_config))
 
     try:
         config_dict = yaml.load(source, Loader=loader)  # noqa: S506
@@ -173,12 +175,26 @@ def _load_yaml_dict(
         cause = e.problem or e.context
         raise InvalidConfigError(f"Invalid YAML{formatted_location}: {cause or ''}") from e
 
-    if "azure-keyvault" in config_dict:
-        config_dict.pop("azure-keyvault")
+    return config_dict
+
+
+def _load_yaml_dict(
+    source: Union[TextIO, str],
+    case_style: str = "hyphen",
+    expand_envvars: bool = True,
+    dict_manipulator: Callable[[Dict[str, Any]], Dict[str, Any]] = lambda x: x,
+    keyvault_loader: Optional[KeyVaultLoader] = None,
+) -> Dict[str, Any]:
+    config_dict = _load_yaml_dict_raw(source, expand_envvars, keyvault_loader)
 
     config_dict = dict_manipulator(config_dict)
     config_dict = _to_snake_case(config_dict, case_style)
 
+    if "azure_keyvault" in config_dict:
+        config_dict.pop("azure_keyvault")
+    if "key_vault" in config_dict:
+        config_dict.pop("key_vault")
+
     return config_dict
 
 
@@ -188,9 +204,14 @@ def _load_yaml(
     case_style: str = "hyphen",
     expand_envvars: bool = True,
     dict_manipulator: Callable[[Dict[str, Any]], Dict[str, Any]] = lambda x: x,
+    keyvault_loader: Optional[KeyVaultLoader] = None,
 ) -> CustomConfigClass:
     config_dict = _load_yaml_dict(
-        source, case_style=case_style, expand_envvars=expand_envvars, dict_manipulator=dict_manipulator
+        source,
+        case_style=case_style,
+        expand_envvars=expand_envvars,
+        dict_manipulator=dict_manipulator,
+        keyvault_loader=keyvault_loader,
     )
 
     try:
@@ -239,6 +260,7 @@ def load_yaml(
     config_type: Type[CustomConfigClass],
     case_style: str = "hyphen",
     expand_envvars: bool = True,
+    keyvault_loader: Optional[KeyVaultLoader] = None,
 ) -> CustomConfigClass:
     """
     Read a YAML file, and create a config object based on its contents.
@@ -249,6 +271,7 @@ def load_yaml(
         case_style: Casing convention of config file. Valid options are 'snake', 'hyphen' or 'camel'. Should be
             'hyphen'.
         expand_envvars: Substitute values with the pattern ${VAR} with the content of the environment variable VAR
+        keyvault_loader: Pre-built loader for keyvault tags. Will be loaded from config if not set.
 
     Returns:
         An initialized config object.
@@ -256,13 +279,20 @@ def load_yaml(
     Raises:
         InvalidConfigError: If any config field is given as an invalid type, is missing or is unknown
     """
-    return _load_yaml(source=source, config_type=config_type, case_style=case_style, expand_envvars=expand_envvars)
+    return _load_yaml(
+        source=source,
+        config_type=config_type,
+        case_style=case_style,
+        expand_envvars=expand_envvars,
+        keyvault_loader=keyvault_loader,
+    )
 
 
 def load_yaml_dict(
     source: Union[TextIO, str],
     case_style: str = "hyphen",
     expand_envvars: bool = True,
+    keyvault_loader: Optional[KeyVaultLoader] = None,
 ) -> Dict[str, Any]:
     """
     Read a YAML file and return a dictionary from its contents
@@ -272,6 +302,7 @@ def load_yaml_dict(
         case_style: Casing convention of config file. Valid options are 'snake', 'hyphen' or 'camel'. Should be
             'hyphen'.
         expand_envvars: Substitute values with the pattern ${VAR} with the content of the environment variable VAR
+        keyvault_loader: Pre-built loader for keyvault tags. Will be loaded from config if not set.
 
     Returns:
         A raw dict with the contents of the config file.
@@ -279,7 +310,9 @@ def load_yaml_dict(
     Raises:
         InvalidConfigError: If any config field is given as an invalid type, is missing or is unknown
     """
-    return _load_yaml_dict(source=source, case_style=case_style, expand_envvars=expand_envvars)
+    return _load_yaml_dict(
+        source=source, case_style=case_style, expand_envvars=expand_envvars, keyvault_loader=keyvault_loader
+    )
 
 
 class ConfigResolver(Generic[CustomConfigClass]):
@@ -374,6 +407,10 @@ class ConfigResolver(Generic[CustomConfigClass]):
             and tmp_config.cognite.idp_authentication == self._config.cognite.idp_authentication
         )
 
+    def _get_keyvault_loader(self) -> KeyVaultLoader:
+        temp_config = _load_yaml_dict_raw(self._config_text)
+        return KeyVaultLoader(temp_config.get("azure-keyvault", temp_config.get("key-vault")))
+
     def _resolve_config(self) -> None:
         self._reload_file()
 
@@ -400,6 +437,7 @@ class ConfigResolver(Generic[CustomConfigClass]):
                 source=response.config,
                 config_type=self.config_type,
                 dict_manipulator=lambda d: self._inject_cognite(tmp_config, d),
+                keyvault_loader=self._get_keyvault_loader(),
             )
 
         else:

cognite_extractor_utils-7.2.1/cognite/extractorutils/statestore/__init__.py

@@ -0,0 +1,12 @@
+from .hashing import AbstractHashStateStore, LocalHashStateStore, RawHashStateStore
+from .watermark import AbstractStateStore, LocalStateStore, NoStateStore, RawStateStore
+
+__all__ = [
+    "AbstractStateStore",
+    "RawStateStore",
+    "LocalStateStore",
+    "NoStateStore",
+    "AbstractHashStateStore",
+    "RawHashStateStore",
+    "LocalHashStateStore",
+]

cognite_extractor_utils-7.2.1/cognite/extractorutils/statestore/_base.py

@@ -0,0 +1,81 @@
+import logging
+import threading
+from abc import ABC, abstractmethod
+from typing import Optional
+
+from cognite.extractorutils._inner_util import _resolve_log_level
+from cognite.extractorutils.threading import CancellationToken
+
+RETRY_BACKOFF_FACTOR = 1.5
+RETRY_MAX_DELAY = 60
+RETRY_DELAY = 1
+RETRIES = 10
+
+
+class _BaseStateStore(ABC):
+    def __init__(
+        self,
+        save_interval: Optional[int] = None,
+        trigger_log_level: str = "DEBUG",
+        thread_name: Optional[str] = None,
+        cancellation_token: Optional[CancellationToken] = None,
+    ) -> None:
+        self._initialized = False
+
+        self.logger = logging.getLogger(__name__)
+        self.trigger_log_level = _resolve_log_level(trigger_log_level)
+        self.save_interval = save_interval
+
+        self.thread = threading.Thread(target=self._run, daemon=cancellation_token is None, name=thread_name)
+        self.lock = threading.RLock()
+        self.cancellation_token = cancellation_token.create_child_token() if cancellation_token else CancellationToken()
+
+    def start(self, initialize: bool = True) -> None:
+        """
+        Start saving state periodically if save_interval is set.
+        This calls the synchronize method every save_interval seconds.
+        """
+        if initialize and not self._initialized:
+            self.initialize()
+        if self.save_interval is not None:
+            self.thread.start()
+
+    def stop(self, ensure_synchronize: bool = True) -> None:
+        """
+        Stop synchronize thread if running, and ensure state is saved if ensure_synchronize is True.
+
+        Args:
+            ensure_synchronize (bool): (Optional). Call synchronize one last time after shutting down thread.
+        """
+        self.cancellation_token.cancel()
+        if ensure_synchronize:
+            self.synchronize()
+
+    def _run(self) -> None:
+        """
+        Internal run method for synchronize thread
+        """
+        self.initialize()
+        while not self.cancellation_token.wait(timeout=self.save_interval):
+            try:
+                self.logger.log(self.trigger_log_level, "Triggering scheduled state store synchronization")
+                self.synchronize()
+            except Exception as e:
+                self.logger.error("Unexpected error while synchronizing state store: %s.", str(e))
+
+        # trigger stop event explicitly to drain the queue
+        self.stop(ensure_synchronize=True)
+
+    @abstractmethod
+    def initialize(self, force: bool = False) -> None:
+        """
+        Get states from remote store
+        """
+        pass
+
+    @abstractmethod
+    def synchronize(self) -> None:
+        """
+        Upload states to remote store
+        """
+        pass

cognite_extractor_utils-7.2.1/cognite/extractorutils/statestore/hashing.py

@@ -0,0 +1,258 @@
+import hashlib
+import json
+from abc import ABC
+from types import TracebackType
+from typing import Any, Dict, Iterable, Iterator, Optional, Set, Type
+
+import orjson
+
+from cognite.client import CogniteClient
+from cognite.client.data_classes import Row
+from cognite.client.exceptions import CogniteAPIError
+from cognite.extractorutils._inner_util import _DecimalDecoder, _DecimalEncoder
+from cognite.extractorutils.threading import CancellationToken
+from cognite.extractorutils.util import cognite_exceptions, retry
+
+from ._base import RETRIES, RETRY_BACKOFF_FACTOR, RETRY_DELAY, RETRY_MAX_DELAY, _BaseStateStore
+
+
+class AbstractHashStateStore(_BaseStateStore, ABC):
+    def __init__(
+        self,
+        save_interval: Optional[int] = None,
+        trigger_log_level: str = "DEBUG",
+        thread_name: Optional[str] = None,
+        cancellation_token: Optional[CancellationToken] = None,
+    ) -> None:
+        super().__init__(
+            save_interval=save_interval,
+            trigger_log_level=trigger_log_level,
+            thread_name=thread_name,
+            cancellation_token=cancellation_token,
+        )
+
+        self._local_state: Dict[str, Dict[str, str]] = {}
+        self._seen: Set[str] = set()
+
+    def get_state(self, external_id: str) -> Optional[str]:
+        with self.lock:
+            return self._local_state.get(external_id, {}).get("digest")
+
+    def _hash_row(self, data: Dict[str, Any]) -> str:
+        return hashlib.sha256(orjson.dumps(data, option=orjson.OPT_SORT_KEYS)).hexdigest()
+
+    def set_state(self, external_id: str, data: Dict[str, Any]) -> None:
+        with self.lock:
+            self._local_state[external_id] = {"digest": self._hash_row(data)}
+
+    def has_changed(self, external_id: str, data: Dict[str, Any]) -> bool:
+        with self.lock:
+            if external_id not in self._local_state:
+                return True
+
+            return self._hash_row(data) != self._local_state[external_id]["digest"]
+
+    def __getitem__(self, external_id: str) -> Optional[str]:
+        return self.get_state(external_id)
+
+    def __setitem__(self, key: str, value: Dict[str, Any]) -> None:
+        self.set_state(external_id=key, data=value)
+
+    def __contains__(self, external_id: str) -> bool:
+        return external_id in self._local_state
+
+    def __len__(self) -> int:
+        return len(self._local_state)
+
+    def __iter__(self) -> Iterator[str]:
+        with self.lock:
+            for key in self._local_state:
+                yield key
+
+
+class RawHashStateStore(AbstractHashStateStore):
+    def __init__(
+        self,
+        cdf_client: CogniteClient,
+        database: str,
+        table: str,
+        save_interval: Optional[int] = None,
+        trigger_log_level: str = "DEBUG",
+        thread_name: Optional[str] = None,
+        cancellation_token: Optional[CancellationToken] = None,
+    ) -> None:
+        super().__init__(
+            save_interval=save_interval,
+            trigger_log_level=trigger_log_level,
+            thread_name=thread_name,
+            cancellation_token=cancellation_token,
+        )
+        self._cdf_client = cdf_client
+        self.database = database
+        self.table = table
+
+    def synchronize(self) -> 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:
+            """
+            Upload local state store to CDF
+            """
+            with self.lock:
+                self._cdf_client.raw.rows.insert(
+                    db_name=self.database,
+                    table_name=self.table,
+                    row=self._local_state,
+                    ensure_parent=True,
+                )
+
+        impl()
+
+    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
+
+            rows: Iterable[Row]
+            try:
+                rows = self._cdf_client.raw.rows.list(db_name=self.database, table_name=self.table, limit=None)
+            except CogniteAPIError as e:
+                if e.code == 404:
+                    rows = []
+                else:
+                    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
+                    state = row.columns.get("digest")
+                    if state:
+                        self._local_state[row.key] = {"digest": state}
+
+            self._initialized = True
+
+        impl()
+
+    def __enter__(self) -> "RawHashStateStore":
+        """
+        Wraps around start method, for use as context manager
+
+        Returns:
+            self
+        """
+        self.start()
+        return self
+
+    def __exit__(
+        self,
+        exc_type: Optional[Type[BaseException]],
+        exc_val: Optional[BaseException],
+        exc_tb: Optional[TracebackType],
+    ) -> None:
+        """
+        Wraps around stop method, for use as context manager
+
+        Args:
+            exc_type: Exception type
+            exc_val: Exception value
+            exc_tb: Traceback
+        """
+        self.stop()
+
+
+class LocalHashStateStore(AbstractHashStateStore):
+    def __init__(
+        self,
+        file_path: str,
+        save_interval: Optional[int] = None,
+        trigger_log_level: str = "DEBUG",
+        thread_name: Optional[str] = None,
+        cancellation_token: Optional[CancellationToken] = None,
+    ) -> None:
+        super().__init__(
+            save_interval=save_interval,
+            trigger_log_level=trigger_log_level,
+            thread_name=thread_name,
+            cancellation_token=cancellation_token,
+        )
+
+        self._file_path = file_path
+
+    def initialize(self, force: bool = False) -> None:
+        """
+        Load states from specified JSON file
+
+        Args:
+            force: Enable re-initialization, ie overwrite when called multiple times
+        """
+        if self._initialized and not force:
+            return
+
+        with self.lock:
+            try:
+                with open(self._file_path, "r") as f:
+                    self._local_state = json.load(f, cls=_DecimalDecoder)
+            except FileNotFoundError:
+                pass
+            except json.decoder.JSONDecodeError as e:
+                raise ValueError(f"Invalid JSON in state store file: {str(e)}") from e
+
+        self._initialized = True
+
+    def synchronize(self) -> None:
+        """
+        Save states to specified JSON file
+        """
+        with self.lock:
+            with open(self._file_path, "w") as f:
+                json.dump(self._local_state, f, cls=_DecimalEncoder)
+
+    def __enter__(self) -> "LocalHashStateStore":
+        """
+        Wraps around start method, for use as context manager
+
+        Returns:
+            self
+        """
+        self.start()
+        return self
+
+    def __exit__(
+        self,
+        exc_type: Optional[Type[BaseException]],
+        exc_val: Optional[BaseException],
+        exc_tb: Optional[TracebackType],
+    ) -> None:
+        """
+        Wraps around stop method, for use as context manager
+
+        Args:
+            exc_type: Exception type
+            exc_val: Exception value
+            exc_tb: Traceback
+        """
+        self.stop()

cognite_extractor_utils-7.1.6/cognite/extractorutils/statestore.py → cognite_extractor_utils-7.2.1/cognite/extractorutils/statestore/watermark.py

@@ -86,27 +86,21 @@ You can set a state store to automatically update on upload triggers from an upl
 """
 
 import json
-import logging
-import threading
-from abc import ABC, abstractmethod
+from abc import ABC
 from types import TracebackType
 from typing import Any, Callable, Dict, Iterator, List, Optional, Tuple, Type, Union
 
 from cognite.client import CogniteClient
 from cognite.client.exceptions import CogniteAPIError
+from cognite.extractorutils._inner_util import _DecimalDecoder, _DecimalEncoder
 from cognite.extractorutils.threading import CancellationToken
 from cognite.extractorutils.uploader import DataPointList
+from cognite.extractorutils.util import cognite_exceptions, retry
 
-from ._inner_util import _DecimalDecoder, _DecimalEncoder, _resolve_log_level
-from .util import cognite_exceptions, retry
+from ._base import RETRIES, RETRY_BACKOFF_FACTOR, RETRY_DELAY, RETRY_MAX_DELAY, _BaseStateStore
 
-RETRY_BACKOFF_FACTOR = 1.5
-RETRY_MAX_DELAY = 60
-RETRY_DELAY = 1
-RETRIES = 10
 
-
-class AbstractStateStore(ABC):
+class AbstractStateStore(_BaseStateStore, ABC):
     """
     Base class for a state store.
 
@@ -125,69 +119,16 @@ class AbstractStateStore(ABC):
         thread_name: Optional[str] = None,
         cancellation_token: Optional[CancellationToken] = None,
     ):
-        self._initialized = False
-        self._local_state: Dict[str, Dict[str, Any]] = {}
-        self.save_interval = save_interval
-        self.trigger_log_level = _resolve_log_level(trigger_log_level)
-
-        self.logger = logging.getLogger(__name__)
-
-        self.thread = threading.Thread(target=self._run, daemon=cancellation_token is None, name=thread_name)
-        self.lock = threading.RLock()
-        self.cancellation_token = cancellation_token.create_child_token() if cancellation_token else CancellationToken()
+        super().__init__(
+            save_interval=save_interval,
+            trigger_log_level=trigger_log_level,
+            thread_name=thread_name,
+            cancellation_token=cancellation_token,
+        )
 
+        self._local_state: Dict[str, Dict[str, Any]] = {}
         self._deleted: List[str] = []
 
-    def start(self, initialize: bool = True) -> None:
-        """
-        Start saving state periodically if save_interval is set.
-        This calls the synchronize method every save_interval seconds.
-        """
-        if initialize and not self._initialized:
-            self.initialize()
-        if self.save_interval is not None:
-            self.thread.start()
-
-    def stop(self, ensure_synchronize: bool = True) -> None:
-        """
-        Stop synchronize thread if running, and ensure state is saved if ensure_synchronize is True.
-
-        Args:
-            ensure_synchronize (bool): (Optional). Call synchronize one last time after shutting down thread.
-        """
-        self.cancellation_token.cancel()
-        if ensure_synchronize:
-            self.synchronize()
-
-    def _run(self) -> None:
-        """
-        Internal run method for synchronize thread
-        """
-        self.initialize()
-        while not self.cancellation_token.wait(timeout=self.save_interval):
-            try:
-                self.logger.log(self.trigger_log_level, "Triggering scheduled state store synchronization")
-                self.synchronize()
-            except Exception as e:
-                self.logger.error("Unexpected error while synchronizing state store: %s.", str(e))
-
-        # trigger stop event explicitly to drain the queue
-        self.stop(ensure_synchronize=True)
-
-    @abstractmethod
-    def initialize(self, force: bool = False) -> None:
-        """
-        Get states from remote store
-        """
-        pass
-
-    @abstractmethod
-    def synchronize(self) -> None:
-        """
-        Upload states to remote store
-        """
-        pass
-
     def get_state(self, external_id: Union[str, List[str]]) -> Union[Tuple[Any, Any], List[Tuple[Any, Any]]]:
         """
         Get state(s) for external ID(s)

cognite_extractor_utils-7.2.1/cognite/extractorutils/uploader/data_modeling.py

@@ -0,0 +1,115 @@
+from types import TracebackType
+from typing import Any, Callable, List, Optional, Type
+
+from cognite.client import CogniteClient
+from cognite.client.data_classes.data_modeling import EdgeApply, NodeApply
+from cognite.extractorutils.threading import CancellationToken
+from cognite.extractorutils.uploader._base import (
+    RETRIES,
+    RETRY_BACKOFF_FACTOR,
+    RETRY_DELAY,
+    RETRY_MAX_DELAY,
+    AbstractUploadQueue,
+)
+from cognite.extractorutils.util import cognite_exceptions, retry
+
+
+class InstanceUploadQueue(AbstractUploadQueue):
+    def __init__(
+        self,
+        cdf_client: CogniteClient,
+        post_upload_function: Optional[Callable[[List[Any]], None]] = None,
+        max_queue_size: Optional[int] = None,
+        max_upload_interval: Optional[int] = None,
+        trigger_log_level: str = "DEBUG",
+        thread_name: Optional[str] = None,
+        cancellation_token: Optional[CancellationToken] = None,
+        auto_create_start_nodes: bool = True,
+        auto_create_end_nodes: bool = True,
+        auto_create_direct_relations: bool = True,
+    ):
+        super().__init__(
+            cdf_client,
+            post_upload_function,
+            max_queue_size,
+            max_upload_interval,
+            trigger_log_level,
+            thread_name,
+            cancellation_token,
+        )
+
+        self.auto_create_start_nodes = auto_create_start_nodes
+        self.auto_create_end_nodes = auto_create_end_nodes
+        self.auto_create_direct_relations = auto_create_direct_relations
+
+        self.node_queue: List[NodeApply] = []
+        self.edge_queue: List[EdgeApply] = []
+
+    def add_to_upload_queue(
+        self,
+        *,
+        node_data: Optional[List[NodeApply]] = None,
+        edge_data: Optional[List[EdgeApply]] = None,
+    ) -> None:
+        if node_data:
+            with self.lock:
+                self.node_queue.extend(node_data)
+                self.upload_queue_size += len(node_data)
+
+        if edge_data:
+            with self.lock:
+                self.edge_queue.extend(edge_data)
+                self.upload_queue_size += len(edge_data)
+
+        with self.lock:
+            self._check_triggers()
+
+    def upload(self) -> 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 upload_batch() -> None:
+            self.cdf_client.data_modeling.instances.apply(
+                nodes=self.node_queue,
+                edges=self.edge_queue,
+                auto_create_start_nodes=self.auto_create_start_nodes,
+                auto_create_end_nodes=self.auto_create_end_nodes,
+                auto_create_direct_relations=self.auto_create_direct_relations,
+            )
+            self.node_queue.clear()
+            self.edge_queue.clear()
+            self.upload_queue_size = 0
+
+        with self.lock:
+            upload_batch()
+
+    def __enter__(self) -> "InstanceUploadQueue":
+        """
+        Wraps around start method, for use as context manager
+
+        Returns:
+            self
+        """
+        self.start()
+        return self
+
+    def __exit__(
+        self,
+        exc_type: Optional[Type[BaseException]],
+        exc_val: Optional[BaseException],
+        exc_tb: Optional[TracebackType],
+    ) -> None:
+        """
+        Wraps around stop method, for use as context manager
+
+        Args:
+            exc_type: Exception type
+            exc_val: Exception value
+            exc_tb: Traceback
+        """
+        self.stop()

{cognite_extractor_utils-7.1.6 → cognite_extractor_utils-7.2.1}/cognite/extractorutils/uploader/time_series.py

@@ -264,9 +264,9 @@ class TimeSeriesUploadQueue(AbstractUploadQueue):
                         missing = [id_dict for id_dict in ex.not_found if id_dict.get("externalId") not in retry_these]
                         missing_num = len(ex.not_found) - len(create_these_ids)
                         self.logger.error(
-                            f"{missing_num} time series not found, and could not be created automatically:\n"
+                            f"{missing_num} time series not found, and could not be created automatically: "
                             + str(missing)
-                            + "\nData will be dropped"
+                            + " Data will be dropped"
                         )
 
                 # Remove entries with non-existing time series from upload queue

{cognite_extractor_utils-7.1.6 → cognite_extractor_utils-7.2.1}/cognite/extractorutils/util.py

@@ -319,11 +319,14 @@ def _retry_internal(
 ) -> _T2:
     logger = logging.getLogger(__name__)
 
-    while tries and not cancellation_token.is_cancelled:
+    while tries:
         try:
             return f()
 
         except Exception as e:
+            if cancellation_token.is_cancelled:
+                break
+
             if isinstance(exceptions, tuple):
                 for ex_type in exceptions:
                     if isinstance(e, ex_type):

{cognite_extractor_utils-7.1.6 → cognite_extractor_utils-7.2.1}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "cognite-extractor-utils"
-version = "7.1.6"
+version = "7.2.1"
 description = "Utilities for easier development of extractors for CDF"
 authors = ["Mathias Lohne <mathias.lohne@cognite.com>"]
 license = "Apache-2.0"
@@ -65,6 +65,7 @@ typing-extensions = ">=3.7.4, <5"
 python-dotenv = "^1.0.0"
 azure-identity = "^1.14.0"
 azure-keyvault-secrets = "^4.7.0"
+orjson = "^3.10.3"
 
 [tool.poetry.extras]
 experimental = ["cognite-sdk-experimental"]
@@ -84,6 +85,7 @@ parameterized = "*"
 requests = "^2.31.0"
 types-requests = "^2.31.0.20240125"
 httpx = "^0.27.0"
+faker = "^25.2.0"
 
 [build-system]
 requires = ["poetry-core>=1.0.0"]