cognite-extractor-utils 7.1.0__tar.gz → 7.1.2__tar.gz

{cognite_extractor_utils-7.1.0 → cognite_extractor_utils-7.1.2}/PKG-INFO

@@ -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 → cognite_extractor_utils-7.1.2}/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_extractor_utils-7.1.0 → cognite_extractor_utils-7.1.2}/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_extractor_utils-7.1.0 → cognite_extractor_utils-7.1.2}/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_extractor_utils-7.1.0 → cognite_extractor_utils-7.1.2}/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_extractor_utils-7.1.0 → cognite_extractor_utils-7.1.2}/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_extractor_utils-7.1.0 → cognite_extractor_utils-7.1.2}/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_extractor_utils-7.1.0 → cognite_extractor_utils-7.1.2}/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_extractor_utils-7.1.0 → cognite_extractor_utils-7.1.2}/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_extractor_utils-7.1.0 → cognite_extractor_utils-7.1.2}/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 → cognite_extractor_utils-7.1.2}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "cognite-extractor-utils"
-version = "7.1.0"
+version = "7.1.2"
 description = "Utilities for easier development of extractors for CDF"
 authors = ["Mathias Lohne <mathias.lohne@cognite.com>"]
 license = "Apache-2.0"
@@ -47,11 +47,13 @@ follow_imports = "normal"
 namespace_packages = true
 explicit_package_bases = true
 show_error_codes = true
+warn_redundant_casts = true
+warn_unused_ignores = true
 exclude = "tests/*"
 
 [tool.poetry.dependencies]
 python = "^3.8.0"
-cognite-sdk = "^7.28.1"
+cognite-sdk = ">=7.28.1, <7.35.0"
 prometheus-client = ">0.7.0, <=1.0.0"
 arrow = "^1.0.0"
 pyyaml = ">=5.3.0, <7"
@@ -71,7 +73,7 @@ experimental = ["cognite-sdk-experimental"]
 mypy = "1.9.0"
 ruff = "^0.3.0"
 pytest = "^8.0.0"
-pytest-cov = "^4.0.0"
+pytest-cov = "^5.0.0"
 sphinx = "^7.0.0"
 sphinx-rtd-theme = "^2.0.0"
 pre-commit = "^3.3.0"