cognite-extractor-utils 6.4.0__py3-none-any.whl → 7.0.0__py3-none-any.whl

cognite/extractorutils/uploader/assets.py

@@ -12,12 +12,12 @@
 #  See the License for the specific language governing permissions and
 #  limitations under the License.
 
-from threading import Event
 from typing import Any, Callable, List, Optional, Type
 
 from cognite.client import CogniteClient
 from cognite.client.data_classes.assets import Asset
-from cognite.client.exceptions import CogniteAPIError, CogniteDuplicatedError
+from cognite.client.exceptions import CogniteDuplicatedError
+from cognite.extractorutils.threading import CancellationToken
 from cognite.extractorutils.uploader._base import (
     RETRIES,
     RETRY_BACKOFF_FACTOR,
@@ -30,7 +30,7 @@ from cognite.extractorutils.uploader._metrics import (
     ASSETS_UPLOADER_QUEUED,
     ASSETS_UPLOADER_WRITTEN,
 )
-from cognite.extractorutils.util import retry
+from cognite.extractorutils.util import cognite_exceptions, retry
 
 
 class AssetUploadQueue(AbstractUploadQueue):
@@ -57,7 +57,7 @@ class AssetUploadQueue(AbstractUploadQueue):
         max_upload_interval: Optional[int] = None,
         trigger_log_level: str = "DEBUG",
         thread_name: Optional[str] = None,
-        cancellation_token: Event = Event(),
+        cancellation_token: Optional[CancellationToken] = None,
     ):
         super().__init__(
             cdf_client,
@@ -92,9 +92,36 @@ class AssetUploadQueue(AbstractUploadQueue):
         """
         Trigger an upload of the queue, clears queue afterwards
         """
+
+        @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:
+            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]
+                to_create = []
+                to_update = []
+                for asset in failed:
+                    if asset.external_id is not None and asset.external_id in duplicated_ids:
+                        to_update.append(asset)
+                    else:
+                        to_create.append(asset)
+                if to_create:
+                    self.cdf_client.assets.create(to_create)
+                if to_update:
+                    self.cdf_client.assets.update(to_update)
+
         if len(self.upload_queue) > 0:
             with self.lock:
-                self._upload_batch()
+                _upload_batch()
 
                 try:
                     self._post_upload(self.upload_queue)
@@ -107,31 +134,6 @@ class AssetUploadQueue(AbstractUploadQueue):
                 self.upload_queue.clear()
                 self.queue_size.set(self.upload_queue_size)
 
-    @retry(
-        exceptions=(CogniteAPIError, ConnectionError),
-        tries=RETRIES,
-        delay=RETRY_DELAY,
-        max_delay=RETRY_MAX_DELAY,
-        backoff=RETRY_BACKOFF_FACTOR,
-    )
-    def _upload_batch(self) -> None:
-        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]
-            to_create = []
-            to_update = []
-            for asset in failed:
-                if asset.external_id is not None and asset.external_id in duplicated_ids:
-                    to_update.append(asset)
-                else:
-                    to_create.append(asset)
-            if to_create:
-                self.cdf_client.assets.create(to_create)
-            if to_update:
-                self.cdf_client.assets.update(to_update)
-
     def __enter__(self) -> "AssetUploadQueue":
         """
         Wraps around start method, for use as context manager

cognite/extractorutils/uploader/events.py

@@ -12,13 +12,13 @@
 #  See the License for the specific language governing permissions and
 #  limitations under the License.
 
-import threading
 from types import TracebackType
 from typing import Callable, List, Optional, Type
 
 from cognite.client import CogniteClient
 from cognite.client.data_classes import Event
-from cognite.client.exceptions import CogniteAPIError, CogniteDuplicatedError
+from cognite.client.exceptions import CogniteDuplicatedError
+from cognite.extractorutils.threading import CancellationToken
 from cognite.extractorutils.uploader._base import (
     RETRIES,
     RETRY_BACKOFF_FACTOR,
@@ -31,7 +31,7 @@ from cognite.extractorutils.uploader._metrics import (
     EVENTS_UPLOADER_QUEUED,
     EVENTS_UPLOADER_WRITTEN,
 )
-from cognite.extractorutils.util import retry
+from cognite.extractorutils.util import cognite_exceptions, retry
 
 
 class EventUploadQueue(AbstractUploadQueue):
@@ -57,7 +57,7 @@ class EventUploadQueue(AbstractUploadQueue):
         max_upload_interval: Optional[int] = None,
         trigger_log_level: str = "DEBUG",
         thread_name: Optional[str] = None,
-        cancellation_token: threading.Event = threading.Event(),
+        cancellation_token: Optional[CancellationToken] = None,
     ):
         # Super sets post_upload and threshold
         super().__init__(
@@ -96,11 +96,38 @@ class EventUploadQueue(AbstractUploadQueue):
         """
         Trigger an upload of the queue, clears queue afterwards
         """
+
+        @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:
+            try:
+                self.cdf_client.events.create([e for e in 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]
+                to_create = []
+                to_update = []
+                for evt in failed:
+                    if evt.external_id is not None and evt.external_id in duplicated_ids:
+                        to_update.append(evt)
+                    else:
+                        to_create.append(evt)
+                if to_create:
+                    self.cdf_client.events.create(to_create)
+                if to_update:
+                    self.cdf_client.events.update(to_update)
+
         if len(self.upload_queue) == 0:
             return
 
         with self.lock:
-            self._upload_batch()
+            _upload_batch()
 
             self.events_written.inc(self.upload_queue_size)
 
@@ -113,31 +140,6 @@ class EventUploadQueue(AbstractUploadQueue):
             self.upload_queue_size = 0
             self.queue_size.set(self.upload_queue_size)
 
-    @retry(
-        exceptions=(CogniteAPIError, ConnectionError),
-        tries=RETRIES,
-        delay=RETRY_DELAY,
-        max_delay=RETRY_MAX_DELAY,
-        backoff=RETRY_BACKOFF_FACTOR,
-    )
-    def _upload_batch(self) -> None:
-        try:
-            self.cdf_client.events.create([e for e in 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]
-            to_create = []
-            to_update = []
-            for evt in failed:
-                if evt.external_id is not None and evt.external_id in duplicated_ids:
-                    to_update.append(evt)
-                else:
-                    to_create.append(evt)
-            if to_create:
-                self.cdf_client.events.create(to_create)
-            if to_update:
-                self.cdf_client.events.update(to_update)
-
     def __enter__(self) -> "EventUploadQueue":
         """
         Wraps around start method, for use as context manager

cognite/extractorutils/uploader/files.py

@@ -17,13 +17,11 @@ from concurrent.futures import Future, ThreadPoolExecutor
 from io import BytesIO
 from os import PathLike
 from types import TracebackType
-from typing import BinaryIO, Callable, List, Optional, Tuple, Type, Union
-
-from requests import ConnectionError
+from typing import Any, BinaryIO, Callable, Dict, List, Optional, Tuple, Type, Union
 
 from cognite.client import CogniteClient
 from cognite.client.data_classes import FileMetadata
-from cognite.client.exceptions import CogniteAPIError
+from cognite.extractorutils.threading import CancellationToken
 from cognite.extractorutils.uploader._base import (
     RETRIES,
     RETRY_BACKOFF_FACTOR,
@@ -36,7 +34,10 @@ from cognite.extractorutils.uploader._metrics import (
     FILES_UPLOADER_QUEUED,
     FILES_UPLOADER_WRITTEN,
 )
-from cognite.extractorutils.util import retry
+from cognite.extractorutils.util import cognite_exceptions, retry
+
+_QUEUES: int = 0
+_QUEUES_LOCK: threading.RLock = threading.RLock()
 
 
 class IOFileUploadQueue(AbstractUploadQueue):
@@ -50,9 +51,7 @@ class IOFileUploadQueue(AbstractUploadQueue):
         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).
+        max_queue_size: Maximum size of upload queue.
         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,
@@ -66,11 +65,10 @@ class IOFileUploadQueue(AbstractUploadQueue):
         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(),
+        cancellation_token: Optional[CancellationToken] = None,
         max_parallelism: int = 0,
     ):
         # Super sets post_upload and threshold
@@ -78,13 +76,18 @@ class IOFileUploadQueue(AbstractUploadQueue):
             cdf_client,
             post_upload_function,
             max_queue_size,
-            max_upload_interval,
+            None,
             trigger_log_level,
             thread_name,
             cancellation_token,
         )
 
-        self.upload_queue: List[Tuple[FileMetadata, Union[str, Callable[[], BinaryIO]]]] = []
+        if self.threshold <= 0:
+            raise ValueError("Max queue size must be positive for file upload queues")
+
+        self.upload_queue: List[Future] = []
+        self.errors: List[Exception] = []
+
         self.overwrite_existing = overwrite_existing
 
         self.parallelism = self.cdf_client.config.max_workers
@@ -97,85 +100,118 @@ class IOFileUploadQueue(AbstractUploadQueue):
         self.files_written = FILES_UPLOADER_WRITTEN
         self.queue_size = FILES_UPLOADER_QUEUE_SIZE
 
-    def add_io_to_upload_queue(self, file_meta: FileMetadata, read_file: Callable[[], BinaryIO]) -> None:
+        self._update_queue_thread = threading.Thread(target=self._remove_done_from_queue, daemon=True)
+
+        self._full_queue = threading.Condition()
+
+        global _QUEUES, _QUEUES_LOCK
+        with _QUEUES_LOCK:
+            self._pool = ThreadPoolExecutor(
+                max_workers=self.parallelism, thread_name_prefix=f"FileUploadQueue-{_QUEUES}"
+            )
+            _QUEUES += 1
+
+    def _remove_done_from_queue(self) -> None:
+        while not self.cancellation_token.is_cancelled:
+            with self.lock:
+                self.upload_queue = list(filter(lambda f: f.running(), self.upload_queue))
+
+            self.cancellation_token.wait(5)
+
+    def add_io_to_upload_queue(
+        self,
+        file_meta: FileMetadata,
+        read_file: Callable[[], BinaryIO],
+        extra_retries: Optional[
+            Union[Tuple[Type[Exception], ...], Dict[Type[Exception], Callable[[Any], bool]]]
+        ] = None,
+    ) -> 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 immedeately. If the size of the queue is larger than
+        the specified max size, this call will block until it's
 
         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
+            extra_retries: Exception types that might be raised by ``read_file`` that should be retried
         """
+        retries = cognite_exceptions()
+        if isinstance(extra_retries, tuple):
+            retries.update({exc: lambda _e: True for exc in extra_retries or []})
+        elif isinstance(extra_retries, dict):
+            retries.update(extra_retries)
+
+        @retry(
+            exceptions=retries,
+            cancellation_token=self.cancellation_token,
+            tries=RETRIES,
+            delay=RETRY_DELAY,
+            max_delay=RETRY_MAX_DELAY,
+            backoff=RETRY_BACKOFF_FACTOR,
+        )
+        def _upload_single(read_file: Callable[[], BinaryIO], file_meta: FileMetadata) -> None:
+            try:
+                # Upload file
+                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,
+                    )
+
+                if self.post_upload_function:
+                    try:
+                        self.post_upload_function([file_meta])
+                    except Exception as e:
+                        self.logger.error("Error in upload callback: %s", str(e))
+
+            except Exception as e:
+                self.logger.exception("Unexpected error while uploading file")
+                self.errors.append(e)
+
+            finally:
+                with self.lock:
+                    self.files_written.inc()
+                    self.upload_queue_size -= 1
+                    self.queue_size.set(self.upload_queue_size)
+                with self._full_queue:
+                    self._full_queue.notify()
+
+        if self.upload_queue_size >= self.threshold:
+            with self._full_queue:
+                while not self._full_queue.wait(timeout=2) and not self.cancellation_token.is_cancelled:
+                    pass
+
         with self.lock:
-            self.upload_queue.append((file_meta, read_file))
+            self.upload_queue.append(self._pool.submit(_upload_single, read_file, file_meta))
             self.upload_queue_size += 1
             self.files_queued.inc()
             self.queue_size.set(self.upload_queue_size)
 
-            self._check_triggers()
-
-    def upload(self) -> None:
+    def upload(self, fail_on_errors: bool = True, timeout: Optional[float] = None) -> None:
         """
-        Trigger an upload of the queue, clears queue afterwards
+        Wait for all uploads to finish
         """
-        if len(self.upload_queue) == 0:
-            return
-
+        for future in self.upload_queue:
+            future.result(timeout=timeout)
         with self.lock:
-            self._upload_batch()
-
-            self.files_written.inc(self.upload_queue_size)
-
-            try:
-                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()
-            self.logger.info(f"Uploaded {self.upload_queue_size} files")
-            self.upload_queue_size = 0
             self.queue_size.set(self.upload_queue_size)
-
-    @retry(
-        exceptions=(CogniteAPIError, ConnectionError),
-        tries=RETRIES,
-        delay=RETRY_DELAY,
-        max_delay=RETRY_MAX_DELAY,
-        backoff=RETRY_BACKOFF_FACTOR,
-    )
-    def _upload_single(self, index: int, read_file: Callable[[], BinaryIO], file_meta: FileMetadata) -> None:
-        # Upload file
-        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, read_file)
-
-    def _upload_batch(self) -> None:
-        # Concurrently execute file-uploads
-
-        futures: List[Future] = []
-        with ThreadPoolExecutor(self.parallelism) as pool:
-            for i, (file_meta, file_name) in enumerate(self.upload_queue):
-                futures.append(pool.submit(self._upload_single, i, file_name, file_meta))
-        for fut in futures:
-            fut.result(0.0)
+        if fail_on_errors and self.errors:
+            # There might be more errors, but we can only have one as the cause, so pick the first
+            raise RuntimeError(f"{len(self.errors)} upload(s) finished with errors") from self.errors[0]
 
     def __enter__(self) -> "IOFileUploadQueue":
         """
@@ -185,6 +221,8 @@ class IOFileUploadQueue(AbstractUploadQueue):
             self
         """
         self.start()
+        self._pool.__enter__()
+        self._update_queue_thread.start()
         return self
 
     def __exit__(
@@ -198,6 +236,7 @@ class IOFileUploadQueue(AbstractUploadQueue):
             exc_val: Exception value
             exc_tb: Traceback
         """
+        self._pool.__exit__(exc_type, exc_val, exc_tb)
         self.stop()
 
     def __len__(self) -> int:
@@ -218,9 +257,7 @@ class FileUploadQueue(IOFileUploadQueue):
         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).
+        max_queue_size: Maximum size of upload queue.
         trigger_log_level: Log level to log upload triggers to.
         thread_name: Thread name of uploader thread.
     """
@@ -234,14 +271,13 @@ class FileUploadQueue(IOFileUploadQueue):
         trigger_log_level: str = "DEBUG",
         thread_name: Optional[str] = None,
         overwrite_existing: bool = False,
-        cancellation_token: threading.Event = threading.Event(),
+        cancellation_token: Optional[CancellationToken] = None,
     ):
         # Super sets post_upload and threshold
         super().__init__(
             cdf_client,
             post_upload_function,
             max_queue_size,
-            max_upload_interval,
             trigger_log_level,
             thread_name,
             overwrite_existing,
@@ -273,9 +309,7 @@ class BytesUploadQueue(IOFileUploadQueue):
         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).
+        max_queue_size: Maximum size of upload queue.
         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
@@ -286,17 +320,15 @@ class BytesUploadQueue(IOFileUploadQueue):
         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(),
+        cancellation_token: Optional[CancellationToken] = None,
     ) -> None:
         super().__init__(
             cdf_client,
             post_upload_function,
             max_queue_size,
-            max_upload_interval,
             trigger_log_level,
             thread_name,
             overwrite_existing,

cognite/extractorutils/uploader/raw.py

@@ -12,17 +12,15 @@
 #  See the License for the specific language governing permissions and
 #  limitations under the License.
 
-import threading
 from types import TracebackType
 from typing import Any, Callable, Dict, List, Optional, Type
 
 import arrow
 from arrow import Arrow
-from requests import ConnectionError
 
 from cognite.client import CogniteClient
 from cognite.client.data_classes import Row
-from cognite.client.exceptions import CogniteAPIError, CogniteReadTimeout
+from cognite.extractorutils.threading import CancellationToken
 from cognite.extractorutils.uploader._base import (
     RETRIES,
     RETRY_BACKOFF_FACTOR,
@@ -37,7 +35,7 @@ from cognite.extractorutils.uploader._metrics import (
     RAW_UPLOADER_ROWS_QUEUED,
     RAW_UPLOADER_ROWS_WRITTEN,
 )
-from cognite.extractorutils.util import retry
+from cognite.extractorutils.util import cognite_exceptions, retry
 
 
 class RawUploadQueue(AbstractUploadQueue):
@@ -63,7 +61,7 @@ class RawUploadQueue(AbstractUploadQueue):
         max_upload_interval: Optional[int] = None,
         trigger_log_level: str = "DEBUG",
         thread_name: Optional[str] = None,
-        cancellation_token: threading.Event = threading.Event(),
+        cancellation_token: Optional[CancellationToken] = None,
     ):
         # Super sets post_upload and thresholds
         super().__init__(
@@ -112,6 +110,19 @@ class RawUploadQueue(AbstractUploadQueue):
         """
         Trigger an upload of the queue, clears queue afterwards
         """
+
+        @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(database: str, table: str, patch: List[Row]) -> None:
+            # Upload
+            self.cdf_client.raw.rows.insert(db_name=database, table_name=table, row=patch, ensure_parent=True)
+
         if len(self.upload_queue) == 0:
             return
 
@@ -125,7 +136,7 @@ class RawUploadQueue(AbstractUploadQueue):
                     patch: Dict[str, Row] = {r.payload.key: r.payload for r in rows}
                     self.rows_duplicates.labels(_labels).inc(len(rows) - len(patch))
 
-                    self._upload_batch(database=database, table=table, patch=list(patch.values()))
+                    _upload_batch(database=database, table=table, patch=list(patch.values()))
                     self.rows_written.labels(_labels).inc(len(patch))
                     _written: Arrow = arrow.utcnow()
 
@@ -140,17 +151,6 @@ class RawUploadQueue(AbstractUploadQueue):
             self.upload_queue_size = 0
             self.queue_size.set(self.upload_queue_size)
 
-    @retry(
-        exceptions=(CogniteAPIError, ConnectionError, CogniteReadTimeout),
-        tries=RETRIES,
-        delay=RETRY_DELAY,
-        max_delay=RETRY_MAX_DELAY,
-        backoff=RETRY_BACKOFF_FACTOR,
-    )
-    def _upload_batch(self, database: str, table: str, patch: List[Row]) -> None:
-        # Upload
-        self.cdf_client.raw.rows.insert(db_name=database, table_name=table, row=patch, ensure_parent=True)
-
     def __enter__(self) -> "RawUploadQueue":
         """
         Wraps around start method, for use as context manager