cognite-extractor-utils 7.1.5__tar.gz → 7.2.0__tar.gz

{cognite_extractor_utils-7.1.5 → cognite_extractor_utils-7.2.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: cognite-extractor-utils
-Version: 7.1.5
+Version: 7.2.0
 Summary: Utilities for easier development of extractors for CDF
 Home-page: https://github.com/cognitedata/python-extractor-utils
 License: Apache-2.0
@@ -18,10 +18,11 @@ Provides-Extra: experimental
 Requires-Dist: arrow (>=1.0.0,<2.0.0)
 Requires-Dist: azure-identity (>=1.14.0,<2.0.0)
 Requires-Dist: azure-keyvault-secrets (>=4.7.0,<5.0.0)
-Requires-Dist: cognite-sdk (>=7.41.0,<8.0.0)
+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.5 → cognite_extractor_utils-7.2.0}/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.5"
+__version__ = "7.2.0"
 from .base import Extractor

cognite_extractor_utils-7.2.0/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.0/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.0/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.5/cognite/extractorutils/statestore.py → cognite_extractor_utils-7.2.0/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.0/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.5 → cognite_extractor_utils-7.2.0}/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.5 → cognite_extractor_utils-7.2.0}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "cognite-extractor-utils"
-version = "7.1.5"
+version = "7.2.0"
 description = "Utilities for easier development of extractors for CDF"
 authors = ["Mathias Lohne <mathias.lohne@cognite.com>"]
 license = "Apache-2.0"
@@ -53,7 +53,7 @@ exclude = "tests/*"
 
 [tool.poetry.dependencies]
 python = "^3.8.0"
-cognite-sdk = "^7.41.0"
+cognite-sdk = "^7.43.3"
 prometheus-client = ">0.7.0, <=1.0.0"
 arrow = "^1.0.0"
 pyyaml = ">=5.3.0, <7"
@@ -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"]