cognite-extractor-utils 7.1.3__py3-none-any.whl → 7.1.5__py3-none-any.whl

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/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/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/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/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/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.dist-info → cognite_extractor_utils-7.1.5.dist-info}/METADATA

@@ -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.dist-info → cognite_extractor_utils-7.1.5.dist-info}/RECORD

@@ -1,27 +1,27 @@
-cognite/extractorutils/__init__.py,sha256=-IiYwZmQ1ixTVlM0nkB38zxXs4AsB04AKutQnXLoKTM,739
+cognite/extractorutils/__init__.py,sha256=vRCQ7k-wwiC8W2oUIVpZoFbuyYyNOOarwLEFnZPau_A,739
 cognite/extractorutils/_inner_util.py,sha256=gmz6aqS7jDNsg8z4RHgJjMFohDLOMiaU4gMWBhg3xcE,1558
-cognite/extractorutils/base.py,sha256=0t9HUANPLKSbAyfDtb0X7stwuarW1NkamRF3vyXyQzc,16148
+cognite/extractorutils/base.py,sha256=q6NU2bPec3WOasVnnIFoh-aUJudVZWZ2R6emz3IRj8Q,16391
 cognite/extractorutils/configtools/__init__.py,sha256=L-daaqInIsmHcjb2forJeY0fW8tz1mlteOUo7IsWnrU,3059
 cognite/extractorutils/configtools/_util.py,sha256=SZycZm_py9v9WZbDiDQbgS6_PiLtu-TtwuuH7tG2YCI,4739
-cognite/extractorutils/configtools/elements.py,sha256=P9oXw07my3Q6UNwef9Ln6-LupChwaSrsF69BHuee9sY,21334
+cognite/extractorutils/configtools/elements.py,sha256=FbkKg1YQOpCs4Bwjs4RacRswttM08TJHnkA9HNweYFs,21564
 cognite/extractorutils/configtools/loaders.py,sha256=VmKNfGqwdHycwZB91i-BHarjW-2Mw_Shv31R6uozm88,16317
 cognite/extractorutils/exceptions.py,sha256=XiwyNPSN0YxFYaPw7tfA63B94PL48xDK3EfdGdhgQgc,1084
 cognite/extractorutils/metrics.py,sha256=01ZMRbDisXPxrfCSyTSEkXMsslzmZwEqw18fuu9okdc,15509
 cognite/extractorutils/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-cognite/extractorutils/statestore.py,sha256=CiE2E6zOCWCLXmEzxaoM9V_-laH3T7U4-ll_qE6rgRE,18645
+cognite/extractorutils/statestore.py,sha256=-KC71NyXaSvzS-j1kUJM9YE70qCmWS_iVzBVhDC0i74,18667
 cognite/extractorutils/threading.py,sha256=2Hke5cFvP-wA45Crvh58JahoKXB64P3tr7R4y_BhBqM,3605
 cognite/extractorutils/uploader/__init__.py,sha256=W22u6QHA4cR0j78LN5LTL5YGbfC-uTApagTyP5ab7uQ,3110
-cognite/extractorutils/uploader/_base.py,sha256=-aFfoMSBGd9YUUMHL3ZQpLIuNMA7TNklWCEjPA18ER8,5282
+cognite/extractorutils/uploader/_base.py,sha256=wktbV8dpb8zBOsNaECZkBNoJSpOz437NlNMER3-a3xQ,5304
 cognite/extractorutils/uploader/_metrics.py,sha256=J2LJXb19L_SLSJ_voNIQHYLp0pjxUKevpH1q_xKX6Hk,3247
 cognite/extractorutils/uploader/assets.py,sha256=2E90N1kxsaA6Ah4h0_r_dTVhDYY_68ItRWrHYkkltJw,5628
 cognite/extractorutils/uploader/events.py,sha256=NZP2tMoU_rh_rb-EZiUBsOT5KdNABHN4c9Oddk0OsdE,5680
-cognite/extractorutils/uploader/files.py,sha256=-yskmzcS9FcAsT2wmu3G4pd9cHJeiNqxmrERoRC72Dg,18417
+cognite/extractorutils/uploader/files.py,sha256=31kPS4fwz8ZSXWss-CKmYTM6ZLVx9LtsDe7LHT7Wy98,18329
 cognite/extractorutils/uploader/raw.py,sha256=wFjF90PFTjmByOWx_Y4_YfDJ2w2jl0EQJ2Tjx2MP2PM,6738
 cognite/extractorutils/uploader/time_series.py,sha256=WAtEQy7k5IjG-sw1oWwCujIM6PjHZYl4LKa4wy2tBPw,26817
 cognite/extractorutils/uploader_extractor.py,sha256=E-mpVvbPg_Tk90U4S9JybV0duptJ2SXE88HB6npE3zI,7732
 cognite/extractorutils/uploader_types.py,sha256=wxfrsiKPTzG5lmoYtQsxt8Xyj-s5HnaLl8WDzJNrazg,1020
 cognite/extractorutils/util.py,sha256=p7AGEgeIU0bNjuFJcFR3V5ZYr6QDj_ZC3zGxRJTf4yk,17198
-cognite_extractor_utils-7.1.3.dist-info/LICENSE,sha256=psuoW8kuDP96RQsdhzwOqi6fyWv0ct8CR6Jr7He_P_k,10173
-cognite_extractor_utils-7.1.3.dist-info/METADATA,sha256=yHvMX5Sjas6K_XAFBdxNt-YFxe9q1QzI7k7PyDM6dx0,5446
-cognite_extractor_utils-7.1.3.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
-cognite_extractor_utils-7.1.3.dist-info/RECORD,,
+cognite_extractor_utils-7.1.5.dist-info/LICENSE,sha256=psuoW8kuDP96RQsdhzwOqi6fyWv0ct8CR6Jr7He_P_k,10173
+cognite_extractor_utils-7.1.5.dist-info/METADATA,sha256=YUu6fg06uVoH0m9HCeWuobsiWhFZpe8C_qJrzGl8Ey0,5446
+cognite_extractor_utils-7.1.5.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
+cognite_extractor_utils-7.1.5.dist-info/RECORD,,