cognite-extractor-utils 7.5.14__py3-none-any.whl → 7.7.0__py3-none-any.whl

cognite/extractorutils/uploader/assets.py

@@ -1,3 +1,7 @@
+"""
+Upload queue for (legacy) assets.
+"""
+
 #  Copyright 2023 Cognite AS
 #
 #  Licensed under the Apache License, Version 2.0 (the "License");
@@ -37,7 +41,7 @@ from cognite.extractorutils.util import cognite_exceptions, retry
 
 class AssetUploadQueue(AbstractUploadQueue):
     """
-    Upload queue for assets
+    Upload queue for assets.
 
     Args:
         cdf_client: Cognite Data Fusion client to use
@@ -77,8 +81,9 @@ class AssetUploadQueue(AbstractUploadQueue):
 
     def add_to_upload_queue(self, asset: Asset) -> None:
         """
-        Add asset to upload queue. The queue will be uploaded if the queue size is larger than the threshold
-        specified in the __init__.
+        Add asset to upload queue.
+
+        The queue will be uploaded if the queue size is larger than the threshold specified in the ``__init__``.
 
         Args:
             asset: Asset to add
@@ -92,7 +97,7 @@ class AssetUploadQueue(AbstractUploadQueue):
 
     def upload(self) -> None:
         """
-        Trigger an upload of the queue, clears queue afterwards
+        Trigger an upload of the queue, clears queue afterwards.
         """
 
         @retry(
@@ -107,8 +112,8 @@ class AssetUploadQueue(AbstractUploadQueue):
             try:
                 self.cdf_client.assets.create(self.upload_queue)
             except CogniteDuplicatedError as e:
-                duplicated_ids = set([dup["externalId"] for dup in e.duplicated if "externalId" in dup])
-                failed: list[Asset] = [e for e in e.failed]
+                duplicated_ids = {dup["externalId"] for dup in e.duplicated if "externalId" in dup}
+                failed: list[Asset] = list(e.failed)
                 to_create = []
                 to_update = []
                 for asset in failed:
@@ -138,7 +143,7 @@ class AssetUploadQueue(AbstractUploadQueue):
 
     def __enter__(self) -> "AssetUploadQueue":
         """
-        Wraps around start method, for use as context manager
+        Wraps around start method, for use as context manager.
 
         Returns:
             self
@@ -153,7 +158,7 @@ class AssetUploadQueue(AbstractUploadQueue):
         exc_tb: TracebackType | None,
     ) -> None:
         """
-        Wraps around stop method, for use as context manager
+        Wraps around stop method, for use as context manager.
 
         Args:
             exc_type: Exception type

cognite/extractorutils/uploader/data_modeling.py

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

cognite/extractorutils/uploader/events.py

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

cognite/extractorutils/uploader/files.py

@@ -1,3 +1,7 @@
+"""
+Upload queue for files.
+"""
+
 #  Copyright 2023 Cognite AS
 #
 #  Licensed under the Apache License, Version 2.0 (the "License");
@@ -88,6 +92,9 @@ class ChunkedStream(RawIOBase, BinaryIO):
         self._current_chunk_size = -1
 
     def tell(self) -> int:
+        """
+        Get the current position of the stream.
+        """
         return self._pos
 
     # RawIOBase is (stupidly) incompatible with BinaryIO
@@ -97,12 +104,21 @@ class ChunkedStream(RawIOBase, BinaryIO):
     # required to satisfy mypy.
     # This may be solvable by changing the typing in the python SDK to use typing.Protocol.
     def writelines(self, __lines: Any) -> None:
+        """
+        Not supported for ChunkedStream.
+        """
         raise NotImplementedError()
 
     def write(self, __b: Any) -> int:
+        """
+        Not supported for ChunkedStream.
+        """
         raise NotImplementedError()
 
     def __enter__(self) -> "ChunkedStream":
+        """
+        Wraps around the inner stream's ``__enter__`` method, for use as context manager.
+        """
         return super().__enter__()
 
     def __exit__(
@@ -111,14 +127,23 @@ class ChunkedStream(RawIOBase, BinaryIO):
         exc_val: BaseException | None,
         exc_tb: TracebackType | None,
     ) -> None:
+        """
+        Wraps around the inner stream's ``__exit__`` method, for use as context manager.
+        """
         return super().__exit__(exc_type, exc_val, exc_tb)
 
     @property
     def chunk_count(self) -> int:
+        """
+        Number of chunks in the stream.
+        """
         return ceil(self._stream_length / self._max_chunk_size)
 
     @property
     def len(self) -> int:
+        """
+        Length of the current chunk, in bytes.
+        """
         return len(self)
 
     @property
@@ -129,12 +154,21 @@ class ChunkedStream(RawIOBase, BinaryIO):
         return self._chunk_index
 
     def __len__(self) -> int:
+        """
+        Length of the current chunk, in bytes.
+        """
         return self._current_chunk_size
 
     def readable(self) -> bool:
+        """
+        Check if the stream is readable. Always True for ChunkedStream.
+        """
         return True
 
     def read(self, size: int = -1) -> bytes:
+        """
+        Read bytes from the current chunk.
+        """
         if size < 0:
             size = self._current_chunk_size - self._pos
 
@@ -162,6 +196,10 @@ class ChunkedStream(RawIOBase, BinaryIO):
 
 
 class IOByteStream(SyncByteStream):
+    """
+    Wraps a BinaryIO stream to be used as a httpx SyncByteStream.
+    """
+
     CHUNK_SIZE = 65_536
 
     def __init__(self, stream: BinaryIO) -> None:
@@ -169,6 +207,9 @@ class IOByteStream(SyncByteStream):
         self._is_stream_consumed = False
 
     def __iter__(self) -> Iterator[bytes]:
+        """
+        Iterate over the stream, yielding chunks of data.
+        """
         if self._is_stream_consumed:
             raise StreamConsumed()
         chunk = self._stream.read(self.CHUNK_SIZE)
@@ -179,7 +220,7 @@ class IOByteStream(SyncByteStream):
 
 class IOFileUploadQueue(AbstractUploadQueue):
     """
-    Upload queue for files using BinaryIO
+    Upload queue for files using BinaryIO.
 
     Note that if the upload fails, the stream needs to be restarted, so
     the enqueued callback needs to produce a new IO object for each call.
@@ -264,6 +305,9 @@ class IOFileUploadQueue(AbstractUploadQueue):
             _QUEUES += 1
 
     def initialize_failure_logging(self) -> None:
+        """
+        Initialize the failure logging manager if a path is provided in the constructor.
+        """
         self._file_failure_manager: FileFailureManager | None = (
             FileFailureManager(path_to_file=self.failure_logging_path)
             if self.failure_logging_path is not None
@@ -271,14 +315,23 @@ class IOFileUploadQueue(AbstractUploadQueue):
         )
 
     def get_failure_logger(self) -> FileFailureManager | None:
+        """
+        Get the failure logger for this upload queue, if it exists.
+        """
         return self._file_failure_manager
 
     def add_entry_failure_logger(self, file_name: str, error: Exception) -> None:
+        """
+        Add an entry to the failure logger if it exists.
+        """
         if self._file_failure_manager is not None:
             error_reason = str(error)
             self._file_failure_manager.add(file_name=file_name, error_reason=error_reason)
 
     def flush_failure_logger(self) -> None:
+        """
+        Flush the failure logger if it exists, writing all failures to the file.
+        """
         if self._file_failure_manager is not None:
             self.logger.info("Flushing failure logs")
             self._file_failure_manager.write_to_file()
@@ -313,14 +366,14 @@ class IOFileUploadQueue(AbstractUploadQueue):
             )
             if any_unchaged:
                 update = FileMetadataUpdate(external_id=file_meta.external_id)
-                any = False
+                need_update = False
                 if file_meta.source:
-                    any = True
+                    need_update = True
                     update.source.set(file_meta.source)
                 if file_meta.directory:
-                    any = True
+                    need_update = True
                     update.directory.set(file_meta.directory)
-                if any:
+                if need_update:
                     self.cdf_client.files.update(update)
 
         return file_meta_response, url
@@ -373,7 +426,7 @@ class IOFileUploadQueue(AbstractUploadQueue):
     def _upload_multipart(self, size: int, file: BinaryIO, file_meta: FileMetadataOrCogniteExtractorFile) -> None:
         chunks = ChunkedStream(file, self.max_file_chunk_size, size)
         self.logger.debug(
-            f"File {file_meta.external_id} is larger than 5GiB ({size})" f", uploading in {chunks.chunk_count} chunks"
+            f"File {file_meta.external_id} is larger than 5GiB ({size}), uploading in {chunks.chunk_count} chunks"
         )
 
         returned_file_metadata = self._create_multi_part(file_meta, chunks)
@@ -428,13 +481,14 @@ class IOFileUploadQueue(AbstractUploadQueue):
         extra_retries: tuple[type[Exception], ...] | dict[type[Exception], Callable[[Any], bool]] | None = None,
     ) -> None:
         """
-        Add file to upload queue. The file will start uploading immedeately. If the size of the queue is larger than
-        the specified max size, this call will block until it's completed the upload.
+        Add file to upload queue.
+
+        The file will start uploading immediately. If the size of the queue is larger than the specified max size, this
+        call will block until it's completed the upload.
 
         Args:
             file_meta: File metadata-object
-            file_name: Path to file to be uploaded.
-                If none, the file object will still be created, but no data is uploaded
+            read_file: Callable that returns a BinaryIO stream to read the file from.
             extra_retries: Exception types that might be raised by ``read_file`` that should be retried
         """
         retries = cognite_exceptions()
@@ -568,7 +622,7 @@ class IOFileUploadQueue(AbstractUploadQueue):
 
     def upload(self, fail_on_errors: bool = True, timeout: float | None = None) -> None:
         """
-        Wait for all uploads to finish
+        Wait for all uploads to finish.
         """
         for future in self.upload_queue:
             future.result(timeout=timeout)
@@ -581,7 +635,7 @@ class IOFileUploadQueue(AbstractUploadQueue):
 
     def __enter__(self) -> "IOFileUploadQueue":
         """
-        Wraps around start method, for use as context manager
+        Wraps around start method, for use as context manager.
 
         Returns:
             self
@@ -598,7 +652,7 @@ class IOFileUploadQueue(AbstractUploadQueue):
         exc_tb: TracebackType | None,
     ) -> None:
         """
-        Wraps around stop method, for use as context manager
+        Wraps around stop method, for use as context manager.
 
         Args:
             exc_type: Exception type
@@ -610,7 +664,7 @@ class IOFileUploadQueue(AbstractUploadQueue):
 
     def __len__(self) -> int:
         """
-        The size of the upload queue
+        The size of the upload queue.
 
         Returns:
             Number of events in queue
@@ -620,7 +674,7 @@ class IOFileUploadQueue(AbstractUploadQueue):
 
 class FileUploadQueue(IOFileUploadQueue):
     """
-    Upload queue for files
+    Upload queue for files.
 
     Args:
         cdf_client: Cognite Data Fusion client to use
@@ -661,8 +715,10 @@ class FileUploadQueue(IOFileUploadQueue):
         file_name: str | PathLike,
     ) -> None:
         """
-        Add file to upload queue. The queue will be uploaded if the queue size is larger than the threshold
-        specified in the __init__.
+        Add file to upload queue.
+
+        The file will start uploading immediately. If the size of the queue is larger than the specified max size, this
+        call will block until it's completed the upload.
 
         Args:
             file_meta: File metadata-object
@@ -678,7 +734,7 @@ class FileUploadQueue(IOFileUploadQueue):
 
 class BytesUploadQueue(IOFileUploadQueue):
     """
-    Upload queue for bytes
+    Upload queue for bytes.
 
     Args:
         cdf_client: Cognite Data Fusion client to use
@@ -714,11 +770,14 @@ class BytesUploadQueue(IOFileUploadQueue):
 
     def add_to_upload_queue(self, content: bytes, file_meta: FileMetadataOrCogniteExtractorFile) -> None:
         """
-        Add object to upload queue. The queue will be uploaded if the queue size is larger than the threshold
-        specified in the __init__.
+        Add file to upload queue.
+
+        The file will start uploading immediately. If the size of the queue is larger than the specified max size, this
+        call will block until it's completed the upload.
+
         Args:
             content: bytes object to upload
-            metadata: metadata for the given bytes object
+            file_meta: File metadata-object
         """
 
         def get_byte_io() -> BinaryIO:

cognite/extractorutils/uploader/raw.py

@@ -1,3 +1,7 @@
+"""
+Upload queue for RAW.
+"""
+
 #  Copyright 2023 Cognite AS
 #
 #  Licensed under the Apache License, Version 2.0 (the "License");
@@ -41,7 +45,7 @@ from cognite.extractorutils.util import cognite_exceptions, retry
 
 class RawUploadQueue(AbstractUploadQueue):
     """
-    Upload queue for RAW
+    Upload queue for RAW.
 
     Args:
         cdf_client: Cognite Data Fusion client to use
@@ -84,8 +88,9 @@ class RawUploadQueue(AbstractUploadQueue):
 
     def add_to_upload_queue(self, database: str, table: str, raw_row: Row) -> None:
         """
-        Adds a row to the upload queue. The queue will be uploaded if the queue size is larger than the threshold
-        specified in the __init__.
+        Adds a row to the upload queue.
+
+        The queue will be uploaded if the queue size is larger than the threshold specified in the ``__init__``.
 
         Args:
             database: The database to upload the Raw object to
@@ -109,7 +114,7 @@ class RawUploadQueue(AbstractUploadQueue):
 
     def upload(self) -> None:
         """
-        Trigger an upload of the queue, clears queue afterwards
+        Trigger an upload of the queue, clears queue afterwards.
         """
 
         @retry(
@@ -154,7 +159,7 @@ class RawUploadQueue(AbstractUploadQueue):
 
     def __enter__(self) -> "RawUploadQueue":
         """
-        Wraps around start method, for use as context manager
+        Wraps around start method, for use as context manager.
 
         Returns:
             self
@@ -166,7 +171,7 @@ class RawUploadQueue(AbstractUploadQueue):
         self, exc_type: type[BaseException] | None, exc_val: BaseException | None, exc_tb: TracebackType | None
     ) -> None:
         """
-        Wraps around stop method, for use as context manager
+        Wraps around stop method, for use as context manager.
 
         Args:
             exc_type: Exception type
@@ -177,7 +182,7 @@ class RawUploadQueue(AbstractUploadQueue):
 
     def __len__(self) -> int:
         """
-        The size of the upload queue
+        The size of the upload queue.
 
         Returns:
             Number of elements in queue