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

cognite/extractorutils/statestore/hashing.py

@@ -1,3 +1,12 @@
+"""
+State store implementations that use hashing to track changes.
+
+This module provides two main classes for state management:
+- ``RawHashStateStore``: A state store that uses CDF RAW to store and persist states based on a hash of the data.
+- ``LocalHashStateStore``: A state store that uses a local JSON file to store and persist states based on a hash of the
+  data.
+"""
+
 import hashlib
 import json
 from abc import ABC
@@ -18,6 +27,12 @@ from ._base import RETRIES, RETRY_BACKOFF_FACTOR, RETRY_DELAY, RETRY_MAX_DELAY,
 
 
 class AbstractHashStateStore(_BaseStateStore, ABC):
+    """
+    Base class for state stores that use hashing to track changes.
+
+    This class is thread-safe.
+    """
+
     def __init__(
         self,
         save_interval: int | None = None,
@@ -36,6 +51,15 @@ class AbstractHashStateStore(_BaseStateStore, ABC):
         self._seen: set[str] = set()
 
     def get_state(self, external_id: str) -> str | None:
+        """
+        Get the state for a given external ID as a hash digest.
+
+        Args:
+            external_id: The external ID for which to retrieve the state.
+
+        Returns:
+            The hash digest of the state if it exists, otherwise None.
+        """
         with self.lock:
             return self._local_state.get(external_id, {}).get("digest")
 
@@ -43,10 +67,29 @@ class AbstractHashStateStore(_BaseStateStore, ABC):
         return hashlib.sha256(orjson.dumps(data, option=orjson.OPT_SORT_KEYS)).hexdigest()
 
     def set_state(self, external_id: str, data: dict[str, Any]) -> None:
+        """
+        Set the state for a given external ID based on a hash of the provided data.
+
+        Args:
+            external_id: The external ID for which to set the state.
+            data: The data to hash and store as the state.
+        """
         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:
+        """
+        Check if the provided data is different from the stored state for the given external ID.
+
+        This is done by comparing the hash of the provided data with the stored hash.
+
+        Args:
+            external_id: The external ID for which to check the state.
+            data: The data to hash and compare against the stored state.
+
+        Returns:
+            True if the data has changed (i.e., the hash is different or not present), otherwise False.
+        """
         with self.lock:
             if external_id not in self._local_state:
                 return True
@@ -54,23 +97,65 @@ class AbstractHashStateStore(_BaseStateStore, ABC):
             return self._hash_row(data) != self._local_state[external_id]["digest"]
 
     def __getitem__(self, external_id: str) -> str | None:
+        """
+        Get the state for a given external ID as a hash digest.
+
+        Args:
+            external_id: The external ID for which to retrieve the state.
+
+        Returns:
+            The hash digest of the state if it exists, otherwise None.
+        """
         return self.get_state(external_id)
 
     def __setitem__(self, key: str, value: dict[str, Any]) -> None:
+        """
+        Set the state for a given external ID based on a hash of the provided data.
+
+        Args:
+            key: The external ID for which to set the state.
+            value: The data to hash and store as the state.
+        """
         self.set_state(external_id=key, data=value)
 
     def __contains__(self, external_id: str) -> bool:
+        """
+        Check if the given external ID exists in the state store.
+        """
         return external_id in self._local_state
 
     def __len__(self) -> int:
+        """
+        Get the number of external IDs stored in the state store.
+        """
         return len(self._local_state)
 
     def __iter__(self) -> Iterator[str]:
+        """
+        Iterate over the external IDs stored in the state store.
+        """
         with self.lock:
             yield from self._local_state
 
 
 class RawHashStateStore(AbstractHashStateStore):
+    """
+    A version of AbstractHashStateStore that uses CDF RAW to store and persist states.
+
+    All states are stored in a CDF RAW table, where each row is identified by an external ID.
+
+    This class is thread-safe.
+
+    Args:
+        cdf_client: The CogniteClient instance to use for ingesting to/reading from RAW.
+        database: The name of the CDF RAW database.
+        table: The name of the CDF RAW table.
+        save_interval: If set, the state store will periodically synchronize with CDF RAW.
+        trigger_log_level: The logging level to use for synchronization triggers.
+        thread_name: Name of the thread used for synchronization.
+        cancellation_token: A CancellationToken to control the lifecycle of the state store.
+    """
+
     def __init__(
         self,
         cdf_client: CogniteClient,
@@ -92,6 +177,10 @@ class RawHashStateStore(AbstractHashStateStore):
         self.table = table
 
     def synchronize(self) -> None:
+        """
+        Upload local state store to CDF.
+        """
+
         @retry(
             exceptions=cognite_exceptions(),
             cancellation_token=self.cancellation_token,
@@ -101,9 +190,6 @@ class RawHashStateStore(AbstractHashStateStore):
             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,
@@ -115,6 +201,16 @@ class RawHashStateStore(AbstractHashStateStore):
         impl()
 
     def initialize(self, force: bool = False) -> None:
+        """
+        Initialize the state store by loading all known states from CDF RAW.
+
+        Unless ``force`` is set to True, this will not re-initialize the state store if it has already been initialized.
+        Subsequent calls to this method will be noop unless ``force`` is set to True.
+
+        Args:
+            force: Enable re-initialization, ie overwrite when called multiple times
+        """
+
         @retry(
             exceptions=cognite_exceptions(),
             cancellation_token=self.cancellation_token,
@@ -146,7 +242,7 @@ class RawHashStateStore(AbstractHashStateStore):
                 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)}")
+                        self.logger.warning(f"None encountered in row: {row!s}")
                         # should never happen, but type from sdk is optional
                         continue
                     state = row.columns.get("digest")
@@ -159,7 +255,7 @@ class RawHashStateStore(AbstractHashStateStore):
 
     def __enter__(self) -> "RawHashStateStore":
         """
-        Wraps around start method, for use as context manager
+        Wraps around start method, for use as context manager.
 
         Returns:
             self
@@ -174,7 +270,7 @@ class RawHashStateStore(AbstractHashStateStore):
         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
@@ -185,6 +281,22 @@ class RawHashStateStore(AbstractHashStateStore):
 
 
 class LocalHashStateStore(AbstractHashStateStore):
+    """
+    A version of AbstractHashStateStore that uses a local JSON file to store and persist states.
+
+    All states are stored in a JSON file, where each key is an external ID and the value is a dictionary containing
+    the hash digest of the data.
+
+    This class is thread-safe.
+
+    Args:
+        file_path: The path to the JSON file where states will be stored.
+        save_interval: If set, the state store will periodically synchronize with the JSON file.
+        trigger_log_level: The logging level to use for synchronization triggers.
+        thread_name: Name of the thread used for synchronization.
+        cancellation_token: A CancellationToken to control the lifecycle of the state store.
+    """
+
     def __init__(
         self,
         file_path: str,
@@ -204,10 +316,13 @@ class LocalHashStateStore(AbstractHashStateStore):
 
     def initialize(self, force: bool = False) -> None:
         """
-        Load states from specified JSON file
+        Load states from specified JSON file.
+
+        Unless ``force`` is set to True, this will not re-initialize the state store if it has already been initialized.
+        Subsequent calls to this method will be noop unless ``force`` is set to True.
 
         Args:
-            force: Enable re-initialization, ie overwrite when called multiple times
+            force: Enable re-initialization, i.e. overwrite when called multiple times
         """
         if self._initialized and not force:
             return
@@ -219,21 +334,20 @@ class LocalHashStateStore(AbstractHashStateStore):
             except FileNotFoundError:
                 pass
             except json.decoder.JSONDecodeError as e:
-                raise ValueError(f"Invalid JSON in state store file: {str(e)}") from e
+                raise ValueError(f"Invalid JSON in state store file: {e!s}") from e
 
         self._initialized = True
 
     def synchronize(self) -> None:
         """
-        Save states to specified JSON file
+        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)
+        with self.lock, 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
+        Wraps around start method, for use as context manager.
 
         Returns:
             self
@@ -248,7 +362,7 @@ class LocalHashStateStore(AbstractHashStateStore):
         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/statestore/watermark.py

@@ -11,78 +11,23 @@
 #  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 #  See the License for the specific language governing permissions and
 #  limitations under the License.
-
 """
-The ``statestore`` module contains classes for keeping track of the extraction state of individual items, facilitating
-incremental load and speeding up startup times.
-
-At the beginning of a run the extractor typically calls the ``initialize`` method, which loads the states from the
-remote store (which can either be a local JSON file or a table in CDF RAW), and during and/or at the end of a run, the
-``synchronize`` method is called, which saves the current states to the remote store.
-
-You can choose the back-end for your state store with which class you're instantiating:
-
-.. code-block:: python
-
-    # A state store using a JSON file as remote storage:
-    states = LocalStateStore("state.json")
-    states.initialize()
-
-    # A state store using a RAW table as remote storage:
-    states = RawStateStore(
-        cdf_client = CogniteClient(),
-        database = "extractor_states",
-        table = "my_extractor_deployment"
-    )
-    states.initialize()
-
-You can now use this state store to get states:
-
-.. code-block:: python
-
-    low, high = states.get_state(external_id = "my-id")
-
-You can set states:
-
-.. code-block:: python
-
-    states.set_state(external_id = "another-id", high=100)
-
-and similar for ``low``. The ``set_state(...)`` method will always overwrite the current state. Some times you might
-want to only set state *if larger* than the previous state, in that case consider ``expand_state(...)``:
-
-.. code-block:: python
+State store implementation that uses watermarks to track changes.
 
-    # High watermark of another-id is already 100, nothing happens in this call:
-    states.expand_state(external_id = "another-id", high=50)
+Watermarks are either low and high values, or just high values, that represent the known range of data that has been
+processed for a given external ID. This allows for incremental processing of data, where only new or changed data
+is processed in subsequent runs.
 
-    # This will set high to 150 as it is larger than the previous state
-    states.expand_state(external_id = "another-id", high=150)
+For example, if a time series has a low watermark of 100 and a high watermark of 200, the extractor can start processing
+new data from 201 onwards when starting up, and can begin backfilling historical data from 100 and backwards.
 
-To store the state to the remote store, use the ``synchronize()`` method:
-
-.. code-block:: python
-
-    states.synchronize()
-
-You can set a state store to automatically update on upload triggers from an upload queue by using the
-``post_upload_function`` in the upload queue:
-
-.. code-block:: python
-
-    states = LocalStateStore("state.json")
-    states.initialize()
-
-    uploader = TimeSeriesUploadQueue(
-        cdf_client = CogniteClient(),
-        max_upload_interval = 10
-        post_upload_function = states.post_upload_handler()
-    )
-
-    # The state store is now updated automatically!
-
-    states.synchronize()
+Or if a file has a high watermark of 1000, and the extractor receives a new file with a high watermark of 1500, the
+extractor will know that this the file has indeed changed.
 
+This module provides the following state store implementations:
+- `RawStateStore`: A state store that uses a CDF RAW table to store states.
+- `LocalStateStore`: A state store that uses a local JSON file to store states.
+- `NoStateStore`: A state store that does not persist states between runs, but keeps the state in memory only.
 """
 
 import json
@@ -105,6 +50,8 @@ class AbstractStateStore(_BaseStateStore, ABC):
     """
     Base class for a state store.
 
+    This class is thread-safe.
+
     Args:
         save_interval: Automatically trigger synchronize each m seconds when run as a thread (use start/stop
             methods).
@@ -132,7 +79,7 @@ class AbstractStateStore(_BaseStateStore, ABC):
 
     def get_state(self, external_id: str | list[str]) -> tuple[Any, Any] | list[tuple[Any, Any]]:
         """
-        Get state(s) for external ID(s)
+        Get state(s) for external ID(s).
 
         Args:
             external_id: An external ID or list of external IDs to get states for
@@ -157,6 +104,9 @@ class AbstractStateStore(_BaseStateStore, ABC):
         """
         Set/update state of a singe external ID.
 
+        Consider using `expand_state` instead, since this method will overwrite the current state no matter if it is
+        actually outside the current state.
+
         Args:
             external_id: External ID of e.g. time series to store state of
             low: Low watermark
@@ -169,8 +119,10 @@ class AbstractStateStore(_BaseStateStore, ABC):
 
     def expand_state(self, external_id: str, low: Any | None = None, high: Any | None = None) -> None:
         """
-        Like set_state, but only sets state if the proposed state is outside the stored state. That is if e.g. low is
-        lower than the stored low.
+        Only set/update state if the proposed state is outside the stored state.
+
+        Only updates the low watermark if the proposed low is lower than the stored low, and only updates the high
+        watermark if the proposed high is higher than the stored high.
 
         Args:
             external_id: External ID of e.g. time series to store state of
@@ -195,7 +147,9 @@ class AbstractStateStore(_BaseStateStore, ABC):
 
     def post_upload_handler(self) -> Callable[[list[dict[str, str | DataPointList]]], None]:
         """
-        Get a callable suitable for passing to a time series upload queue as post_upload_function, that will
+        Get a callback function to handle post-upload events.
+
+        This callable is suitable for passing to a time series upload queue as ``post_upload_function``, that will
         automatically update the states in this state store when that upload queue is uploading.
 
         Returns:
@@ -234,24 +188,38 @@ class AbstractStateStore(_BaseStateStore, ABC):
 
         if high is not None and new_state > high:
             return True
-        if low is not None and new_state < low:
-            return True
-
-        return False
+        return bool(low is not None and new_state < low)
 
     def __getitem__(self, external_id: str) -> tuple[Any, Any]:
+        """
+        Get state for a single external ID.
+        """
         return self.get_state(external_id)  # type: ignore  # will not be list if input is single str
 
     def __setitem__(self, key: str, value: tuple[Any, Any]) -> None:
+        """
+        Set state for a single external ID.
+
+        This will always overwrite the current state, so use with care.
+        """
         self.set_state(external_id=key, low=value[0], high=value[1])
 
     def __contains__(self, external_id: str) -> bool:
+        """
+        Check if an external ID is in the state store.
+        """
         return external_id in self._local_state
 
     def __len__(self) -> int:
+        """
+        Get the number of external IDs in the state store.
+        """
         return len(self._local_state)
 
     def __iter__(self) -> Iterator[str]:
+        """
+        Iterate over external IDs in the state store.
+        """
         yield from self._local_state
 
 
@@ -259,6 +227,8 @@ class RawStateStore(AbstractStateStore):
     """
     An extractor state store based on CDF RAW.
 
+    This class is thread-safe.
+
     Args:
         cdf_client: Cognite client to use
         database: Name of CDF database
@@ -312,6 +282,16 @@ class RawStateStore(AbstractStateStore):
         impl()
 
     def initialize(self, force: bool = False) -> None:
+        """
+        Initialize the state store by loading all known states from CDF RAW.
+
+        Unless ``force`` is set to True, this will not re-initialize the state store if it has already been initialized.
+        Subsequent calls to this method will be noop unless ``force`` is set to True.
+
+        Args:
+            force: Enable re-initialization, ie overwrite when called multiple times
+        """
+
         @retry(
             exceptions=cognite_exceptions(),
             cancellation_token=self.cancellation_token,
@@ -336,7 +316,7 @@ class RawStateStore(AbstractStateStore):
                 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)}")
+                        self.logger.warning(f"None encountered in row: {row!s}")
                         # should never happen, but type from sdk is optional
                         continue
                     self._local_state[row.key] = row.columns
@@ -346,6 +326,10 @@ class RawStateStore(AbstractStateStore):
         impl()
 
     def synchronize(self) -> None:
+        """
+        Upload the contents of the state store to CDF RAW.
+        """
+
         @retry(
             exceptions=cognite_exceptions(),
             cancellation_token=self.cancellation_token,
@@ -356,22 +340,20 @@ class RawStateStore(AbstractStateStore):
         )
         def impl() -> None:
             """
-            Upload local state store to CDF
+            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)
                 # 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]
-                )
+                self._cdf_client.raw.rows.delete(db_name=self.database, table_name=self.table, key=list(self._deleted))
                 self._deleted.clear()
 
         impl()
 
     def __enter__(self) -> "RawStateStore":
         """
-        Wraps around start method, for use as context manager
+        Wraps around start method, for use as context manager.
 
         Returns:
             self
@@ -383,7 +365,7 @@ class RawStateStore(AbstractStateStore):
         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
@@ -420,7 +402,7 @@ class LocalStateStore(AbstractStateStore):
 
     def initialize(self, force: bool = False) -> None:
         """
-        Load states from specified JSON file
+        Load states from specified JSON file.
 
         Args:
             force: Enable re-initialization, ie overwrite when called multiple times
@@ -435,13 +417,13 @@ class LocalStateStore(AbstractStateStore):
             except FileNotFoundError:
                 pass
             except json.decoder.JSONDecodeError as e:
-                raise ValueError(f"Invalid JSON in state store file: {str(e)}") from e
+                raise ValueError(f"Invalid JSON in state store file: {e!s}") from e
 
         self._initialized = True
 
     def synchronize(self) -> None:
         """
-        Save states to specified JSON file
+        Save states to specified JSON file.
         """
         with self.lock:
             with open(self._file_path, "w") as f:
@@ -450,7 +432,7 @@ class LocalStateStore(AbstractStateStore):
 
     def __enter__(self) -> "LocalStateStore":
         """
-        Wraps around start method, for use as context manager
+        Wraps around start method, for use as context manager.
 
         Returns:
             self
@@ -465,7 +447,7 @@ class LocalStateStore(AbstractStateStore):
         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
@@ -478,13 +460,21 @@ class LocalStateStore(AbstractStateStore):
 class NoStateStore(AbstractStateStore):
     """
     A state store that only keeps states in memory and never stores or initializes from external sources.
+
+    This class is thread-safe.
     """
 
     def __init__(self) -> None:
         super().__init__()
 
     def initialize(self, force: bool = False) -> None:
+        """
+        Does nothing.
+        """
         pass
 
     def synchronize(self) -> None:
+        """
+        Does nothing.
+        """
         pass

cognite/extractorutils/threading.py

@@ -1,3 +1,7 @@
+"""
+Module that provides additional threading utilities.
+"""
+
 import logging
 import signal
 from threading import Condition
@@ -20,6 +24,9 @@ class CancellationToken:
         self._parent: CancellationToken | None = None
 
     def __repr__(self) -> str:
+        """
+        Return a string representation of the CancellationToken instance.
+        """
         cls = self.__class__
         status = "cancelled" if self.is_cancelled else "not cancelled"
         return f"<{cls.__module__}.{cls.__qualname__} at {id(self):#x}: {status}>"
@@ -29,7 +36,7 @@ class CancellationToken:
         """
         ``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
+        return self._is_cancelled_int or (self._parent is not None and self._parent.is_cancelled)
 
     def is_set(self) -> bool:
         """
@@ -60,6 +67,19 @@ class CancellationToken:
         self.cancel()
 
     def wait(self, timeout: float | None = None) -> bool:
+        """
+        Wait for the token to be cancelled, or until the timeout expires.
+
+        This can also be used as a drop-in replacement for sleep if you want to wait for a certain amount of time. A
+        call to sleep will not be interrupted by a cancellation, but a call to wait will return immediately if the token
+        is cancelled.
+
+        Args:
+            timeout: The maximum time to wait in seconds. If None, wait indefinitely.
+
+        Returns:
+            ``True`` if the token was cancelled, ``False`` if the timeout expired before cancellation.
+        """
         endtime = None
         if timeout is not None:
             endtime = time() + timeout
@@ -78,14 +98,20 @@ class CancellationToken:
         return True
 
     def create_child_token(self) -> "CancellationToken":
+        """
+        Create a child cancellation token of this token.
+
+        The child token will be cancelled if this token is cancelled, but can also be cancelled independently.
+        """
         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.
+        Register an interrupt handler to capture SIGINT (Ctrl-C) and cancel this token.
+
+        This will set the cancellation token instead of throwing a KeyboardInterrupt exception.
         """
 
         def sigint_handler(sig_num: int, frame: Any) -> None:
@@ -98,4 +124,4 @@ class CancellationToken:
         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)}")
+            logging.getLogger(__name__).warning(f"Could not register handler for interrupt signals: {e!s}")

cognite/extractorutils/unstable/__init__.py

@@ -1,8 +1,8 @@
 """
-The unstable package contains experimental functions and classes currently
-deemed unstable. The contents of this package is subject to change without
-notice, even in minor or patch releases.
+The unstable package contains experimental functions and classes currently deemed unstable.
 
-Whenever you import anything from the unstable package, you should make sure to
-run a type checker such as mypy to help catch these changes.
+The contents of this package is subject to change without notice, even in minor or patch releases. Whenever you import
+anything from the unstable package, you should make sure to run a type checker such as mypy to help catch these changes.
+
+Parts of ``unstable`` might be promoted to the main library at some point.
 """