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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: cognite-extractor-utils
-Version: 7.1.3
+Version: 7.1.5
 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-7.1.3 → cognite_extractor_utils-7.1.5}/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.3"
+__version__ = "7.1.5"
 from .base import Extractor

{cognite_extractor_utils-7.1.3 → cognite_extractor_utils-7.1.5}/cognite/extractorutils/base.py

@@ -189,9 +189,17 @@ class Extractor(Generic[CustomConfigClass]):
 
         state_store_config = recursive_find_state_store(self.config.__dict__)
         if state_store_config:
-            self.state_store = state_store_config.create_state_store(self.cognite_client, self.use_default_state_store)
+            self.state_store = state_store_config.create_state_store(
+                cdf_client=self.cognite_client,
+                default_to_local=self.use_default_state_store,
+                cancellation_token=self.cancellation_token,
+            )
         else:
-            self.state_store = LocalStateStore("states.json") if self.use_default_state_store else NoStateStore()
+            self.state_store = (
+                LocalStateStore("states.json", cancellation_token=self.cancellation_token)
+                if self.use_default_state_store
+                else NoStateStore()
+            )
 
         try:
             self.state_store.initialize()

{cognite_extractor_utils-7.1.3 → cognite_extractor_utils-7.1.5}/cognite/extractorutils/configtools/elements.py

@@ -623,7 +623,10 @@ class StateStoreConfig:
     local: Optional[LocalStateStoreConfig] = None
 
     def create_state_store(
-        self, cdf_client: Optional[CogniteClient] = None, default_to_local: bool = True
+        self,
+        cdf_client: Optional[CogniteClient] = None,
+        default_to_local: bool = True,
+        cancellation_token: Optional[CancellationToken] = None,
     ) -> AbstractStateStore:
         """
         Create a state store object based on the config.
@@ -648,15 +651,17 @@ class StateStoreConfig:
                 database=self.raw.database,
                 table=self.raw.table,
                 save_interval=self.raw.upload_interval.seconds,
+                cancellation_token=cancellation_token,
             )
 
         if self.local:
             return LocalStateStore(
                 file_path=self.local.path,
                 save_interval=self.local.save_interval.seconds,
+                cancellation_token=cancellation_token,
             )
 
         if default_to_local:
-            return LocalStateStore(file_path="states.json")
+            return LocalStateStore(file_path="states.json", cancellation_token=cancellation_token)
         else:
             return NoStateStore()

{cognite_extractor_utils-7.1.3 → cognite_extractor_utils-7.1.5}/cognite/extractorutils/statestore.py

@@ -132,7 +132,7 @@ class AbstractStateStore(ABC):
 
         self.logger = logging.getLogger(__name__)
 
-        self.thread = threading.Thread(target=self._run, daemon=True, name=thread_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()
 

{cognite_extractor_utils-7.1.3 → cognite_extractor_utils-7.1.5}/cognite/extractorutils/uploader/_base.py

@@ -58,7 +58,7 @@ class AbstractUploadQueue(ABC):
         self.trigger_log_level = _resolve_log_level(trigger_log_level)
         self.logger = logging.getLogger(__name__)
 
-        self.thread = threading.Thread(target=self._run, daemon=True, name=thread_name)
+        self.thread = threading.Thread(target=self._run, daemon=cancellation_token is None, name=thread_name)
         self.lock = threading.RLock()
         self.cancellation_token: CancellationToken = (
             cancellation_token.create_child_token() if cancellation_token else CancellationToken()

{cognite_extractor_utils-7.1.3 → cognite_extractor_utils-7.1.5}/cognite/extractorutils/uploader/files.py

@@ -251,70 +251,72 @@ class IOFileUploadQueue(AbstractUploadQueue):
             max_delay=RETRY_MAX_DELAY,
             backoff=RETRY_BACKOFF_FACTOR,
         )
-        def _upload_single(read_file: Callable[[], BinaryIO], file_meta: FileMetadata) -> None:
+        def upload_file(read_file: Callable[[], BinaryIO], file_meta: FileMetadata) -> None:
+            with read_file() as file:
+                size = super_len(file)
+                if size == 0:
+                    # upload just the file metadata witout data
+                    file_meta, _url = self.cdf_client.files.create(
+                        file_metadata=file_meta, overwrite=self.overwrite_existing
+                    )
+                elif size >= self.max_single_chunk_file_size:
+                    # The minimum chunk size is 4000MiB.
+                    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"
+                    )
+                    with self.cdf_client.files.multipart_upload_session(
+                        file_meta.name if file_meta.name is not None else "",
+                        parts=chunks.chunk_count,
+                        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,
+                    ) as session:
+                        while chunks.next_chunk():
+                            session.upload_part(chunks.current_chunk, chunks)
+                        file_meta = session.file_metadata
+                else:
+                    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))
+
+        def wrapped_upload(read_file: Callable[[], BinaryIO], file_meta: FileMetadata) -> None:
             try:
-                # Upload file
-                with read_file() as file:
-                    size = super_len(file)
-                    if size == 0:
-                        # upload just the file metadata witout data
-                        file_meta, _url = self.cdf_client.files.create(
-                            file_metadata=file_meta, overwrite=self.overwrite_existing
-                        )
-                    elif size >= self.max_single_chunk_file_size:
-                        # The minimum chunk size is 4000MiB.
-                        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"
-                        )
-                        with self.cdf_client.files.multipart_upload_session(
-                            file_meta.name if file_meta.name is not None else "",
-                            parts=chunks.chunk_count,
-                            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,
-                        ) as session:
-                            while chunks.next_chunk():
-                                session.upload_part(chunks.current_chunk, chunks)
-                            file_meta = session.file_metadata
-                    else:
-                        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))
+                upload_file(read_file, file_meta)
 
             except Exception as e:
-                self.logger.exception("Unexpected error while uploading file")
+                self.logger.exception(f"Unexpected error while uploading file: {file_meta.external_id}")
                 self.errors.append(e)
 
             finally:
@@ -331,7 +333,7 @@ class IOFileUploadQueue(AbstractUploadQueue):
                     pass
 
         with self.lock:
-            self.upload_queue.append(self._pool.submit(_upload_single, read_file, file_meta))
+            self.upload_queue.append(self._pool.submit(wrapped_upload, read_file, file_meta))
             self.upload_queue_size += 1
             self.files_queued.inc()
             self.queue_size.set(self.upload_queue_size)

{cognite_extractor_utils-7.1.3 → cognite_extractor_utils-7.1.5}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "cognite-extractor-utils"
-version = "7.1.3"
+version = "7.1.5"
 description = "Utilities for easier development of extractors for CDF"
 authors = ["Mathias Lohne <mathias.lohne@cognite.com>"]
 license = "Apache-2.0"