cognite-extractor-utils 7.1.0__py3-none-any.whl → 7.1.2__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.0"
+__version__ = "7.1.2"
 from .base import Extractor

cognite/extractorutils/base.py

@@ -152,17 +152,17 @@ class Extractor(Generic[CustomConfigClass]):
     def _reload_config(self) -> None:
         self.logger.info("Config file has changed")
 
-        if self.reload_config_action == ReloadConfigAction.REPLACE_ATTRIBUTE:
+        if self.reload_config_action is ReloadConfigAction.REPLACE_ATTRIBUTE:
             self.logger.info("Loading in new config file")
             self.config_resolver.accept_new_config()
             self.config = self.config_resolver.config
             Extractor._config_singleton = self.config  # type: ignore
 
-        elif self.reload_config_action == ReloadConfigAction.SHUTDOWN:
+        elif self.reload_config_action is ReloadConfigAction.SHUTDOWN:
             self.logger.info("Shutting down, expecting to be restarted")
             self.cancellation_token.cancel()
 
-        elif self.reload_config_action == ReloadConfigAction.CALLBACK:
+        elif self.reload_config_action is ReloadConfigAction.CALLBACK:
             self.logger.info("Loading in new config file")
             self.config_resolver.accept_new_config()
             self.config = self.config_resolver.config

cognite/extractorutils/configtools/elements.py

@@ -372,7 +372,7 @@ class CogniteConfig:
 
         return cdf_client.data_sets.retrieve(
             id=self.data_set.either_id.internal_id,
-            external_id=self.data_set.either_id.external_id,  # type: ignore
+            external_id=self.data_set.either_id.external_id,
         )
 
     def get_extraction_pipeline(self, cdf_client: CogniteClient) -> Optional[ExtractionPipeline]:
@@ -381,8 +381,8 @@ class CogniteConfig:
 
         either_id = self.extraction_pipeline.either_id
         extraction_pipeline = cdf_client.extraction_pipelines.retrieve(
-            id=either_id.internal_id,  # type: ignore
-            external_id=either_id.external_id,  # type: ignore
+            id=either_id.internal_id,
+            external_id=either_id.external_id,
         )
         if extraction_pipeline is None:
             raise ValueError(f"Extraction pipeline with {either_id.type()} {either_id.content()} not found")

cognite/extractorutils/configtools/loaders.py

@@ -13,6 +13,7 @@
 #  limitations under the License.
 
 import argparse
+import dataclasses
 import json
 import logging
 import os
@@ -21,7 +22,7 @@ import sys
 from enum import Enum
 from hashlib import sha256
 from pathlib import Path
-from typing import Any, Callable, Dict, Generic, Iterable, Optional, TextIO, Type, TypeVar, Union
+from typing import Any, Callable, Dict, Generic, Iterable, Optional, TextIO, Type, TypeVar, Union, cast
 
 import dacite
 import yaml
@@ -349,45 +350,43 @@ class ConfigResolver(Generic[CustomConfigClass]):
         return cls(args.config[0], config_type)
 
     def _inject_cognite(self, local_part: _BaseConfig, remote_part: Dict[str, Any]) -> Dict[str, Any]:
-        if "cognite" not in remote_part:
-            remote_part["cognite"] = {}
-
-        remote_part["cognite"]["idp-authentication"] = {
-            "client_id": local_part.cognite.idp_authentication.client_id,
-            "scopes": local_part.cognite.idp_authentication.scopes,
-            "secret": local_part.cognite.idp_authentication.secret,
-            "tenant": local_part.cognite.idp_authentication.tenant,
-            "token_url": local_part.cognite.idp_authentication.token_url,
-            "resource": local_part.cognite.idp_authentication.resource,
-            "authority": local_part.cognite.idp_authentication.authority,
-        }
+        # We can not dump 'local_part.cognite' directly because e.g. 'data_set' may be set remote only...
+        remote_part.setdefault("cognite", {})
+        remote_part["cognite"]["idp_authentication"] = dataclasses.asdict(local_part.cognite.idp_authentication)
+        remote_part["cognite"]["extraction-pipeline"] = dataclasses.asdict(
+            local_part.cognite.extraction_pipeline  # type: ignore [arg-type]
+        )
+
         if local_part.cognite.host is not None:
             remote_part["cognite"]["host"] = local_part.cognite.host
         remote_part["cognite"]["project"] = local_part.cognite.project
 
-        # Ignoring None type, extraction pipelines is required at this point
-        remote_part["cognite"]["extraction-pipeline"] = {}
-        remote_part["cognite"]["extraction-pipeline"]["id"] = local_part.cognite.extraction_pipeline.id  # type: ignore
-        remote_part["cognite"]["extraction-pipeline"][
-            "external_id"
-        ] = local_part.cognite.extraction_pipeline.external_id  # type: ignore
-
         return remote_part
 
+    def _use_cached_cognite_client(self, tmp_config: _BaseConfig) -> bool:
+        # Ideally we'd check tmp_config == self._config, but due to 'is_remote & _inject_...', this is not
+        # reliable to avoid new unneeded instantiations of CogniteClient:
+        return (
+            self.cognite_client is not None
+            and self._config is not None
+            and tmp_config.cognite.host == self._config.cognite.host
+            and tmp_config.cognite.project == self._config.cognite.project
+            and tmp_config.cognite.idp_authentication == self._config.cognite.idp_authentication
+        )
+
     def _resolve_config(self) -> None:
         self._reload_file()
 
         if self.is_remote:
             _logger.debug("Loading remote config file")
             tmp_config: _BaseConfig = load_yaml(self._config_text, _BaseConfig)  # type: ignore
-            if self.cognite_client is None or self._config is None or tmp_config.cognite != self._config.cognite:
-                # Credentials towards CDF may have changed, instantiate (and store) a new client:
-                client = tmp_config.cognite.get_cognite_client("config_resolver")
-                self.cognite_client = client
-            else:
+            if self._use_cached_cognite_client(tmp_config):
                 # Use existing client to avoid invoking a token refresh, if possible. Reason: this is run every 5 min
                 # by default ('ConfigReloader' thread) which for certain OAuth providers like Auth0, incurs a cost:
-                client = self.cognite_client
+                client = cast(CogniteClient, self.cognite_client)
+            else:
+                # Credentials towards CDF may have changed, instantiate (and store) a new client:
+                client = self.cognite_client = tmp_config.cognite.get_cognite_client("config_resolver")
 
             response = client.extraction_pipelines.config.retrieve(
                 tmp_config.cognite.get_extraction_pipeline(client).external_id  # type: ignore  # ignoring extpipe None

cognite/extractorutils/metrics.py

@@ -384,7 +384,7 @@ class CognitePusher(AbstractMetricsPusher):
                 data_set_id = dataset.id
 
         for metric in REGISTRY.collect():
-            if type(metric) == Metric and metric.type in ["gauge", "counter"]:
+            if type(metric) is Metric and metric.type in ["gauge", "counter"]:
                 external_id = self.external_id_prefix + metric.name
 
                 time_series.append(
@@ -393,8 +393,8 @@ class CognitePusher(AbstractMetricsPusher):
                         name=metric.name,
                         legacy_name=external_id,
                         description=metric.documentation,
-                        asset_id=asset_id,  # type: ignore  # this is optional. Type hint in SDK is wrong
-                        data_set_id=data_set_id,  # type: ignore  # this is optional. Type hint in SDK is wrong
+                        asset_id=asset_id,
+                        data_set_id=data_set_id,
                     )
                 )
 

cognite/extractorutils/statestore.py

@@ -389,8 +389,7 @@ class RawStateStore(AbstractStateStore):
             if self._initialized and not force:
                 return
 
-            # ignore type since list _is_ optional, sdk types are wrong
-            rows = self._cdf_client.raw.rows.list(db_name=self.database, table_name=self.table, limit=None)  # type: ignore
+            rows = self._cdf_client.raw.rows.list(db_name=self.database, table_name=self.table, limit=None)
 
             with self.lock:
                 self._local_state.clear()

cognite/extractorutils/uploader/time_series.py

@@ -522,8 +522,8 @@ class SequenceUploadQueue(AbstractUploadQueue):
 
             try:
                 self.cdf_client.sequences.data.insert(
-                    id=either_id.internal_id,  # type: ignore
-                    external_id=either_id.external_id,  # type: ignore
+                    id=either_id.internal_id,
+                    external_id=either_id.external_id,
                     rows=upload_this,
                     column_external_ids=None,
                 )
@@ -534,8 +534,8 @@ class SequenceUploadQueue(AbstractUploadQueue):
 
                     # Retry
                     self.cdf_client.sequences.data.insert(
-                        id=either_id.internal_id,  # type: ignore
-                        external_id=either_id.external_id,  # type: ignore
+                        id=either_id.internal_id,
+                        external_id=either_id.external_id,
                         rows=upload_this,
                         column_external_ids=None,
                     )
@@ -553,7 +553,6 @@ class SequenceUploadQueue(AbstractUploadQueue):
                 self._resolve_dataset_ids()
 
             for either_id, upload_this in self.upload_queue.items():
-                _labels = str(either_id.content())
                 _upload_single(either_id, upload_this)
                 self.points_written.inc()
 
@@ -582,11 +581,11 @@ class SequenceUploadQueue(AbstractUploadQueue):
         try:
             seq = self.cdf_client.sequences.create(
                 Sequence(
-                    id=either_id.internal_id,  # type: ignore  # these are optional, the SDK types are wrong
-                    external_id=either_id.external_id,  # type: ignore
-                    name=self.sequence_names.get(either_id, None),  # type: ignore
-                    description=self.sequence_descriptions.get(either_id, None),  # type: ignore
-                    metadata=self.sequence_metadata.get(either_id, None),  # type: ignore
+                    id=either_id.internal_id,
+                    external_id=either_id.external_id,
+                    name=self.sequence_names.get(either_id, None),
+                    description=self.sequence_descriptions.get(either_id, None),
+                    metadata=self.sequence_metadata.get(either_id, None),
                     asset_id=self.asset_ids.get(self.sequence_asset_external_ids.get(either_id, None), None),  # type: ignore
                     data_set_id=self.dataset_ids.get(self.sequence_dataset_external_ids.get(either_id, None), None),  # type: ignore
                     columns=column_def,  # type: ignore  # We already  checked for None, mypy is wrong
@@ -595,9 +594,9 @@ class SequenceUploadQueue(AbstractUploadQueue):
 
         except CogniteDuplicatedError:
             self.logger.info("Sequnce already exist: {}".format(either_id))
-            seq = self.cdf_client.sequences.retrieve(
-                id=either_id.internal_id,  # type: ignore  # these are optional, the SDK types are wrong
-                external_id=either_id.external_id,  # type: ignore
+            seq = self.cdf_client.sequences.retrieve(  # type: ignore [assignment]
+                id=either_id.internal_id,
+                external_id=either_id.external_id,
             )
 
         # Update definition of cached sequence

cognite/extractorutils/uploader_types.py

@@ -1,12 +1,12 @@
+import sys
 from typing import Iterable, List, Optional, Union
 
 from cognite.client.data_classes import Event as _Event
 from cognite.client.data_classes import Row as _Row
 
-try:
-    from typing import TypeAlias  # type: ignore
-except ImportError:
-    # Backport for python < 3.10
+if sys.version_info >= (3, 10):
+    from typing import TypeAlias
+else:
     from typing_extensions import TypeAlias
 
 

cognite/extractorutils/util.py

@@ -16,6 +16,7 @@
 The ``util`` package contains miscellaneous functions and classes that can some times be useful while developing
 extractors.
 """
+
 import logging
 import random
 from functools import partial, wraps
@@ -27,7 +28,7 @@ from decorator import decorator
 
 from cognite.client import CogniteClient
 from cognite.client.data_classes import Asset, ExtractionPipelineRun, TimeSeries
-from cognite.client.exceptions import CogniteAPIError, CogniteException, CogniteNotFoundError
+from cognite.client.exceptions import CogniteAPIError, CogniteException, CogniteFileUploadError, CogniteNotFoundError
 from cognite.extractorutils.threading import CancellationToken
 
 
@@ -424,7 +425,7 @@ def requests_exceptions(
     """
     status_codes = status_codes or [408, 425, 429, 500, 502, 503, 504]
     # types ignored, since they are not installed as we don't depend on the package
-    from requests.exceptions import HTTPError, RequestException  # type: ignore
+    from requests.exceptions import HTTPError, RequestException
 
     def handle_http_errors(exception: RequestException) -> bool:
         if isinstance(exception, HTTPError):
@@ -458,7 +459,7 @@ def httpx_exceptions(
     """
     status_codes = status_codes or [408, 425, 429, 500, 502, 503, 504]
     # types ignored, since they are not installed as we don't depend on the package
-    from httpx import HTTPError, HTTPStatusError  # type: ignore
+    from httpx import HTTPError, HTTPStatusError
 
     def handle_http_errors(exception: HTTPError) -> bool:
         if isinstance(exception, HTTPStatusError):
@@ -492,7 +493,7 @@ def cognite_exceptions(
     status_codes = status_codes or [408, 425, 429, 500, 502, 503, 504]
 
     def handle_cognite_errors(exception: CogniteException) -> bool:
-        if isinstance(exception, CogniteAPIError):
+        if isinstance(exception, (CogniteAPIError, CogniteFileUploadError)):
             return exception.code in status_codes
         return True
 

{cognite_extractor_utils-7.1.0.dist-info → cognite_extractor_utils-7.1.2.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: cognite-extractor-utils
-Version: 7.1.0
+Version: 7.1.2
 Summary: Utilities for easier development of extractors for CDF
 Home-page: https://github.com/cognitedata/python-extractor-utils
 License: Apache-2.0
@@ -18,7 +18,7 @@ Provides-Extra: experimental
 Requires-Dist: arrow (>=1.0.0,<2.0.0)
 Requires-Dist: azure-identity (>=1.14.0,<2.0.0)
 Requires-Dist: azure-keyvault-secrets (>=4.7.0,<5.0.0)
-Requires-Dist: cognite-sdk (>=7.28.1,<8.0.0)
+Requires-Dist: cognite-sdk (>=7.28.1,<7.35.0)
 Requires-Dist: dacite (>=1.6.0,<2.0.0)
 Requires-Dist: decorator (>=5.1.1,<6.0.0)
 Requires-Dist: more-itertools (>=10.0.0,<11.0.0)

{cognite_extractor_utils-7.1.0.dist-info → cognite_extractor_utils-7.1.2.dist-info}/RECORD

@@ -1,14 +1,14 @@
-cognite/extractorutils/__init__.py,sha256=tzmeAmZqXbBOYIeP3jFuPUAC_ZDhZ2EYpsRFUawHyVA,739
+cognite/extractorutils/__init__.py,sha256=Jr9_i-43717ZejlWzgVBNpwiLv8OXGDt6xZDnG1Svzo,739
 cognite/extractorutils/_inner_util.py,sha256=gmz6aqS7jDNsg8z4RHgJjMFohDLOMiaU4gMWBhg3xcE,1558
-cognite/extractorutils/base.py,sha256=o5oFzGBaszv7-HWMVG2ATmZwZfzi8enyrYIKwUagyCQ,16148
+cognite/extractorutils/base.py,sha256=0t9HUANPLKSbAyfDtb0X7stwuarW1NkamRF3vyXyQzc,16148
 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=G9-2J-tzao4z5PfOpWifJCMgP7X7ABHK257NtrWsweI,21382
-cognite/extractorutils/configtools/loaders.py,sha256=MvgFxgjExqLIS7z5DVziXq-O51wfcc13B53H0ptfvps,16372
+cognite/extractorutils/configtools/elements.py,sha256=P9oXw07my3Q6UNwef9Ln6-LupChwaSrsF69BHuee9sY,21334
+cognite/extractorutils/configtools/loaders.py,sha256=VmKNfGqwdHycwZB91i-BHarjW-2Mw_Shv31R6uozm88,16317
 cognite/extractorutils/exceptions.py,sha256=XiwyNPSN0YxFYaPw7tfA63B94PL48xDK3EfdGdhgQgc,1084
-cognite/extractorutils/metrics.py,sha256=IzYevOH19N_Nhi2XfhcuGgsGfmF7SexfcHcTQyZmH1o,15635
+cognite/extractorutils/metrics.py,sha256=01ZMRbDisXPxrfCSyTSEkXMsslzmZwEqw18fuu9okdc,15509
 cognite/extractorutils/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-cognite/extractorutils/statestore.py,sha256=98_UchuEw6k69aqVqFljjUtc10-MZ8-RLTHc-9Jdyf8,18733
+cognite/extractorutils/statestore.py,sha256=CiE2E6zOCWCLXmEzxaoM9V_-laH3T7U4-ll_qE6rgRE,18645
 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
@@ -17,11 +17,11 @@ cognite/extractorutils/uploader/assets.py,sha256=2E90N1kxsaA6Ah4h0_r_dTVhDYY_68I
 cognite/extractorutils/uploader/events.py,sha256=NZP2tMoU_rh_rb-EZiUBsOT5KdNABHN4c9Oddk0OsdE,5680
 cognite/extractorutils/uploader/files.py,sha256=-yskmzcS9FcAsT2wmu3G4pd9cHJeiNqxmrERoRC72Dg,18417
 cognite/extractorutils/uploader/raw.py,sha256=wFjF90PFTjmByOWx_Y4_YfDJ2w2jl0EQJ2Tjx2MP2PM,6738
-cognite/extractorutils/uploader/time_series.py,sha256=VptUq129MY0t8yw4rxeL0kOhz2dMibz4XdyvfhfYGj8,26840
+cognite/extractorutils/uploader/time_series.py,sha256=7lLiIcNo1FS6e9zHI3WXbFuqyJFzICHdqIDFYSBVlRM,26548
 cognite/extractorutils/uploader_extractor.py,sha256=E-mpVvbPg_Tk90U4S9JybV0duptJ2SXE88HB6npE3zI,7732
-cognite/extractorutils/uploader_types.py,sha256=5MKT14DUnTFVD5Nx4Zvnfp2SfaICuTKLWAFAYaZishk,1045
-cognite/extractorutils/util.py,sha256=1Atq4-N2Pyi6GrxMa1PEfMCTqhPiLb2YoUXTzd60Roc,17179
-cognite_extractor_utils-7.1.0.dist-info/LICENSE,sha256=psuoW8kuDP96RQsdhzwOqi6fyWv0ct8CR6Jr7He_P_k,10173
-cognite_extractor_utils-7.1.0.dist-info/METADATA,sha256=uBij3gu7rMwr28ENpJ0TRoOshUDRujnfhj22DC-1us4,5446
-cognite_extractor_utils-7.1.0.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
-cognite_extractor_utils-7.1.0.dist-info/RECORD,,
+cognite/extractorutils/uploader_types.py,sha256=wxfrsiKPTzG5lmoYtQsxt8Xyj-s5HnaLl8WDzJNrazg,1020
+cognite/extractorutils/util.py,sha256=p7AGEgeIU0bNjuFJcFR3V5ZYr6QDj_ZC3zGxRJTf4yk,17198
+cognite_extractor_utils-7.1.2.dist-info/LICENSE,sha256=psuoW8kuDP96RQsdhzwOqi6fyWv0ct8CR6Jr7He_P_k,10173
+cognite_extractor_utils-7.1.2.dist-info/METADATA,sha256=qMnjci_XOaGrKVmI6cqMdWN2nwZjPvmR2qPljGE2pP0,5447
+cognite_extractor_utils-7.1.2.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
+cognite_extractor_utils-7.1.2.dist-info/RECORD,,