cognite-extractor-utils 6.1.1__tar.gz → 6.2.1__tar.gz

{cognite_extractor_utils-6.1.1 → cognite_extractor_utils-6.2.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: cognite-extractor-utils
-Version: 6.1.1
+Version: 6.2.1
 Summary: Utilities for easier development of extractors for CDF
 Home-page: https://github.com/cognitedata/python-extractor-utils
 License: Apache-2.0

{cognite_extractor_utils-6.1.1 → cognite_extractor_utils-6.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__ = "6.1.1"
+__version__ = "6.2.1"
 from .base import Extractor

{cognite_extractor_utils-6.1.1 → cognite_extractor_utils-6.2.1}/cognite/extractorutils/uploader/__init__.py

@@ -66,7 +66,7 @@ instead. If both are used, the condition being met first will trigger the upload
 
 from .assets import AssetUploadQueue
 from .events import EventUploadQueue
-from .files import BytesUploadQueue, FileUploadQueue
+from .files import BytesUploadQueue, FileUploadQueue, IOFileUploadQueue
 from .raw import RawUploadQueue
 from .time_series import (
     DataPoint,

{cognite_extractor_utils-6.1.1 → cognite_extractor_utils-6.2.1}/cognite/extractorutils/uploader/files.py

@@ -13,15 +13,16 @@
 #  limitations under the License.
 
 import threading
-from concurrent.futures import ThreadPoolExecutor
+from concurrent.futures import Future, ThreadPoolExecutor
+from io import BytesIO
 from os import PathLike
 from types import TracebackType
-from typing import Any, Callable, List, Optional, Tuple, Type, Union
+from typing import BinaryIO, Callable, List, Optional, Tuple, Type, Union
 
 from requests import ConnectionError
 
 from cognite.client import CogniteClient
-from cognite.client.data_classes import Event, FileMetadata
+from cognite.client.data_classes import FileMetadata
 from cognite.client.exceptions import CogniteAPIError
 from cognite.extractorutils.uploader._base import (
     RETRIES,
@@ -31,9 +32,6 @@ from cognite.extractorutils.uploader._base import (
     AbstractUploadQueue,
 )
 from cognite.extractorutils.uploader._metrics import (
-    BYTES_UPLOADER_QUEUE_SIZE,
-    BYTES_UPLOADER_QUEUED,
-    BYTES_UPLOADER_WRITTEN,
     FILES_UPLOADER_QUEUE_SIZE,
     FILES_UPLOADER_QUEUED,
     FILES_UPLOADER_WRITTEN,
@@ -41,9 +39,12 @@ from cognite.extractorutils.uploader._metrics import (
 from cognite.extractorutils.util import retry
 
 
-class FileUploadQueue(AbstractUploadQueue):
+class IOFileUploadQueue(AbstractUploadQueue):
     """
-    Upload queue for files
+    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.
 
     Args:
         cdf_client: Cognite Data Fusion client to use
@@ -54,18 +55,23 @@ class FileUploadQueue(AbstractUploadQueue):
             methods).
         trigger_log_level: Log level to log upload triggers to.
         thread_name: Thread name of uploader thread.
+        max_parallelism: Maximum number of parallel uploads. If this is greater than 0,
+            the largest of this and client.config.max_workers is used to limit the number
+            of parallel uploads. This may be important if the IO objects being processed
+            also load data from an external system.
     """
 
     def __init__(
         self,
         cdf_client: CogniteClient,
-        post_upload_function: Optional[Callable[[List[Event]], None]] = None,
+        post_upload_function: Optional[Callable[[List[FileMetadata]], None]] = None,
         max_queue_size: Optional[int] = None,
         max_upload_interval: Optional[int] = None,
         trigger_log_level: str = "DEBUG",
         thread_name: Optional[str] = None,
         overwrite_existing: bool = False,
         cancellation_token: threading.Event = threading.Event(),
+        max_parallelism: int = 0,
     ):
         # Super sets post_upload and threshold
         super().__init__(
@@ -78,14 +84,20 @@ class FileUploadQueue(AbstractUploadQueue):
             cancellation_token,
         )
 
-        self.upload_queue: List[Tuple[FileMetadata, Union[str, PathLike]]] = []
+        self.upload_queue: List[Tuple[FileMetadata, Union[str, Callable[[], BinaryIO]]]] = []
         self.overwrite_existing = overwrite_existing
 
+        self.parallelism = self.cdf_client.config.max_workers
+        if max_parallelism > 0 and max_parallelism < self.parallelism:
+            self.parallelism = max_parallelism
+        if self.parallelism <= 0:
+            self.parallelism = 4
+
         self.files_queued = FILES_UPLOADER_QUEUED
         self.files_written = FILES_UPLOADER_WRITTEN
         self.queue_size = FILES_UPLOADER_QUEUE_SIZE
 
-    def add_to_upload_queue(self, file_meta: FileMetadata, file_name: Union[str, PathLike]) -> None:
+    def add_io_to_upload_queue(self, file_meta: FileMetadata, read_file: Callable[[], BinaryIO]) -> None:
         """
         Add file to upload queue. The queue will be uploaded if the queue size is larger than the threshold
         specified in the __init__.
@@ -96,7 +108,7 @@ class FileUploadQueue(AbstractUploadQueue):
                 If none, the file object will still be created, but no data is uploaded
         """
         with self.lock:
-            self.upload_queue.append((file_meta, file_name))
+            self.upload_queue.append((file_meta, read_file))
             self.upload_queue_size += 1
             self.files_queued.inc()
             self.queue_size.set(self.upload_queue_size)
@@ -116,7 +128,7 @@ class FileUploadQueue(AbstractUploadQueue):
             self.files_written.inc(self.upload_queue_size)
 
             try:
-                self._post_upload(self.upload_queue)
+                self._post_upload([el[0] for el in self.upload_queue])
             except Exception as e:
                 self.logger.error("Error in upload callback: %s", str(e))
             self.upload_queue.clear()
@@ -131,21 +143,41 @@ class FileUploadQueue(AbstractUploadQueue):
         max_delay=RETRY_MAX_DELAY,
         backoff=RETRY_BACKOFF_FACTOR,
     )
-    def _upload_single(self, index: int, file_name: Union[str, PathLike], file_meta: FileMetadata) -> None:
+    def _upload_single(self, index: int, read_file: Callable[[], BinaryIO], file_meta: FileMetadata) -> None:
         # Upload file
-        file_meta = self.cdf_client.files.upload(str(file_name), overwrite=self.overwrite_existing, **file_meta.dump())  # type: ignore
+        with read_file() as file:
+            file_meta = self.cdf_client.files.upload_bytes(
+                file,
+                file_meta.name if file_meta.name is not None else "",
+                overwrite=self.overwrite_existing,
+                external_id=file_meta.external_id,
+                source=file_meta.source,
+                mime_type=file_meta.mime_type,
+                metadata=file_meta.metadata,
+                directory=file_meta.directory,
+                asset_ids=file_meta.asset_ids,
+                data_set_id=file_meta.data_set_id,
+                labels=file_meta.labels,
+                geo_location=file_meta.geo_location,
+                source_created_time=file_meta.source_created_time,
+                source_modified_time=file_meta.source_modified_time,
+                security_categories=file_meta.security_categories,
+            )
 
         # Update meta-object in queue
-        self.upload_queue[index] = (file_meta, file_name)
+        self.upload_queue[index] = (file_meta, read_file)
 
     def _upload_batch(self) -> None:
         # Concurrently execute file-uploads
 
-        with ThreadPoolExecutor(self.cdf_client.config.max_workers) as pool:
+        futures: List[Future] = []
+        with ThreadPoolExecutor(self.parallelism) as pool:
             for i, (file_meta, file_name) in enumerate(self.upload_queue):
-                pool.submit(self._upload_single, i, file_name, file_meta)
+                futures.append(pool.submit(self._upload_single, i, file_name, file_meta))
+        for fut in futures:
+            fut.result(0.0)
 
-    def __enter__(self) -> "FileUploadQueue":
+    def __enter__(self) -> "IOFileUploadQueue":
         """
         Wraps around start method, for use as context manager
 
@@ -178,9 +210,9 @@ class FileUploadQueue(AbstractUploadQueue):
         return self.upload_queue_size
 
 
-class BytesUploadQueue(AbstractUploadQueue):
+class FileUploadQueue(IOFileUploadQueue):
     """
-    Upload queue for bytes
+    Upload queue for files
 
     Args:
         cdf_client: Cognite Data Fusion client to use
@@ -191,20 +223,20 @@ class BytesUploadQueue(AbstractUploadQueue):
             methods).
         trigger_log_level: Log level to log upload triggers to.
         thread_name: Thread name of uploader thread.
-        overwrite_existing: If 'overwrite' is set to true, fields for the files found for externalIds can be overwritten
     """
 
     def __init__(
         self,
         cdf_client: CogniteClient,
-        post_upload_function: Optional[Callable[[List[Any]], None]] = None,
+        post_upload_function: Optional[Callable[[List[FileMetadata]], None]] = None,
         max_queue_size: Optional[int] = None,
         max_upload_interval: Optional[int] = None,
         trigger_log_level: str = "DEBUG",
         thread_name: Optional[str] = None,
         overwrite_existing: bool = False,
         cancellation_token: threading.Event = threading.Event(),
-    ) -> None:
+    ):
+        # Super sets post_upload and threshold
         super().__init__(
             cdf_client,
             post_upload_function,
@@ -212,105 +244,75 @@ class BytesUploadQueue(AbstractUploadQueue):
             max_upload_interval,
             trigger_log_level,
             thread_name,
+            overwrite_existing,
             cancellation_token,
         )
-        self.upload_queue: List[Tuple[bytes, FileMetadata]] = []
-        self.overwrite_existing = overwrite_existing
-        self.upload_queue_size = 0
 
-        self.bytes_queued = BYTES_UPLOADER_QUEUED
-        self.queue_size = BYTES_UPLOADER_QUEUE_SIZE
-        self.bytes_written = BYTES_UPLOADER_WRITTEN
-
-    def add_to_upload_queue(self, content: bytes, metadata: FileMetadata) -> None:
+    def add_to_upload_queue(self, file_meta: FileMetadata, file_name: Union[str, PathLike]) -> None:
         """
-        Add object to upload queue. The queue will be uploaded if the queue size is larger than the threshold
+        Add file to upload queue. The queue will be uploaded if the queue size is larger than the threshold
         specified in the __init__.
-        Args:
-            content: bytes object to upload
-            metadata: metadata for the given bytes object
-        """
-        with self.lock:
-            self.upload_queue.append((content, metadata))
-            self.upload_queue_size += 1
-            self.bytes_queued.inc()
-            self.queue_size.set(self.upload_queue_size)
 
-    def upload(self) -> None:
-        """
-        Trigger an upload of the queue, clears queue afterwards
+        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
         """
-        if len(self.upload_queue) == 0:
-            return
 
-        with self.lock:
-            # Upload frames in batches
-            self._upload_batch()
+        def load_file_from_path() -> BinaryIO:
+            return open(file_name, "rb")
 
-            # Log stats
-            self.bytes_written.inc(self.upload_queue_size)
+        self.add_io_to_upload_queue(file_meta, load_file_from_path)
 
-            try:
-                self._post_upload(self.upload_queue)
-            except Exception as e:
-                self.logger.error("Error in upload callback: %s", str(e))
 
-            # Clear queue
-            self.upload_queue.clear()
-            self.upload_queue_size = 0
-            self.logger.info(f"Uploaded {self.upload_queue_size} files")
-            self.queue_size.set(self.upload_queue_size)
+class BytesUploadQueue(IOFileUploadQueue):
+    """
+    Upload queue for bytes
 
-    def _upload_batch(self) -> None:
-        # Concurrently execute bytes-uploads
-        with ThreadPoolExecutor(self.cdf_client.config.max_workers) as pool:
-            for i, (frame, metadata) in enumerate(self.upload_queue):
-                pool.submit(self._upload_single, i, frame, metadata)
+    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 events that were uploaded.
+        max_queue_size: Maximum size of upload queue. Defaults to no max size.
+        max_upload_interval: Automatically trigger an upload each m seconds when run as a thread (use start/stop
+            methods).
+        trigger_log_level: Log level to log upload triggers to.
+        thread_name: Thread name of uploader thread.
+        overwrite_existing: If 'overwrite' is set to true, fields for the files found for externalIds can be overwritten
+    """
 
-    @retry(
-        exceptions=(CogniteAPIError, ConnectionError),
-        tries=RETRIES,
-        delay=RETRY_DELAY,
-        max_delay=RETRY_MAX_DELAY,
-        backoff=RETRY_BACKOFF_FACTOR,
-    )
-    def _upload_single(self, index: int, content: bytes, metadata: FileMetadata) -> None:
-        # Upload object
-        file_meta_data: FileMetadata = self.cdf_client.files.upload_bytes(
-            content, overwrite=self.overwrite_existing, **metadata.dump()
+    def __init__(
+        self,
+        cdf_client: CogniteClient,
+        post_upload_function: Optional[Callable[[List[FileMetadata]], None]] = None,
+        max_queue_size: Optional[int] = None,
+        max_upload_interval: Optional[int] = None,
+        trigger_log_level: str = "DEBUG",
+        thread_name: Optional[str] = None,
+        overwrite_existing: bool = False,
+        cancellation_token: threading.Event = threading.Event(),
+    ) -> None:
+        super().__init__(
+            cdf_client,
+            post_upload_function,
+            max_queue_size,
+            max_upload_interval,
+            trigger_log_level,
+            thread_name,
+            overwrite_existing,
+            cancellation_token,
         )
 
-        # Update meta-object in queue
-        self.upload_queue[index] = (content, file_meta_data)
-
-    def __enter__(self) -> "BytesUploadQueue":
-        """
-        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:
+    def add_to_upload_queue(self, content: bytes, metadata: FileMetadata) -> None:
         """
-        Wraps around stop method, for use as context manager
-
+        Add object to upload queue. The queue will be uploaded if the queue size is larger than the threshold
+        specified in the __init__.
         Args:
-            exc_type: Exception type
-            exc_val: Exception value
-            exc_tb: Traceback
+            content: bytes object to upload
+            metadata: metadata for the given bytes object
         """
-        self.stop()
 
-    def __len__(self) -> int:
-        """
-        The size of the upload queue
+        def get_byte_io() -> BinaryIO:
+            return BytesIO(content)
 
-        Returns:
-            Number of events in queue
-        """
-        return self.upload_queue_size
+        self.add_io_to_upload_queue(metadata, get_byte_io)

{cognite_extractor_utils-6.1.1 → cognite_extractor_utils-6.2.1}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "cognite-extractor-utils"
-version = "6.1.1"
+version = "6.2.1"
 description = "Utilities for easier development of extractors for CDF"
 authors = ["Mathias Lohne <mathias.lohne@cognite.com>"]
 license = "Apache-2.0"
@@ -67,7 +67,7 @@ jq = [{version = "^1.3.0", platform = "darwin"}, {version = "^1.3.0", platform =
 experimental = ["cognite-sdk-experimental"]
 
 [tool.poetry.group.dev.dependencies]
-mypy = "1.7.1"
+mypy = "1.8.0"
 ruff = "^0.1.0"
 pytest = "^7.0.0"
 pytest-cov = "^4.0.0"