cognite-extractor-utils 7.0.0__tar.gz → 7.0.2__tar.gz

{cognite_extractor_utils-7.0.0 → cognite_extractor_utils-7.0.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: cognite-extractor-utils
-Version: 7.0.0
+Version: 7.0.2
 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.0.0 → cognite_extractor_utils-7.0.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.0.0"
+__version__ = "7.0.2"
 from .base import Extractor

{cognite_extractor_utils-7.0.0 → cognite_extractor_utils-7.0.2}/cognite/extractorutils/configtools/__init__.py

@@ -85,10 +85,15 @@ Get a state store object as configured:
 However, all of these things will be automatically done for you if you are using the base Extractor class.
 """
 
+from cognite.extractorutils.exceptions import InvalidConfigError
+
 from .elements import (
+    AuthenticatorConfig,
     BaseConfig,
+    CertificateConfig,
     CogniteConfig,
     ConfigType,
+    ConnectionConfig,
     EitherIdConfig,
     FileSizeConfig,
     LocalStateStoreConfig,
@@ -99,4 +104,4 @@ from .elements import (
     StateStoreConfig,
     TimeIntervalConfig,
 )
-from .loaders import ConfigResolver, load_yaml
+from .loaders import ConfigResolver, KeyVaultAuthenticationMethod, KeyVaultLoader, load_yaml, load_yaml_dict

{cognite_extractor_utils-7.0.0 → cognite_extractor_utils-7.0.2}/cognite/extractorutils/configtools/elements.py

@@ -53,6 +53,10 @@ _logger = logging.getLogger(__name__)
 
 @dataclass
 class CertificateConfig:
+    """
+    Configuration parameters for certificates
+    """
+
     path: str
     password: Optional[str]
     authority_url: Optional[str] = None
@@ -94,6 +98,11 @@ class ConnectionConfig:
 
 @dataclass
 class EitherIdConfig:
+    """
+    Configuration parameter representing an ID in CDF, which can either be an external or internal ID.
+    An EitherId can only hold one ID type, not both.
+    """
+
     id: Optional[int]
     external_id: Optional[str]
 
@@ -103,6 +112,10 @@ class EitherIdConfig:
 
 
 class TimeIntervalConfig(yaml.YAMLObject):
+    """
+    Configuration parameter for setting a time interval
+    """
+
     def __init__(self, expression: str) -> None:
         self._interval, self._expression = TimeIntervalConfig._parse_expression(expression)
 
@@ -167,6 +180,10 @@ class TimeIntervalConfig(yaml.YAMLObject):
 
 
 class FileSizeConfig(yaml.YAMLObject):
+    """
+    Configuration parameter for setting a file size
+    """
+
     def __init__(self, expression: str) -> None:
         self._bytes, self._expression = FileSizeConfig._parse_expression(expression)
 
@@ -569,23 +586,39 @@ class BaseConfig(_BaseConfig):
 
 @dataclass
 class RawDestinationConfig:
+    """
+    Configuration parameters for using Raw
+    """
+
     database: str
     table: str
 
 
 @dataclass
 class RawStateStoreConfig(RawDestinationConfig):
+    """
+    Configuration of a state store based on CDF RAW
+    """
+
     upload_interval: TimeIntervalConfig = TimeIntervalConfig("30s")
 
 
 @dataclass
 class LocalStateStoreConfig:
+    """
+    Configuration of a state store using a local JSON file
+    """
+
     path: str
     save_interval: TimeIntervalConfig = TimeIntervalConfig("30s")
 
 
 @dataclass
 class StateStoreConfig:
+    """
+    Configuration of the State Store, containing ``LocalStateStoreConfig`` or ``RawStateStoreConfig``
+    """
+
     raw: Optional[RawStateStoreConfig] = None
     local: Optional[LocalStateStoreConfig] = None
 

{cognite_extractor_utils-7.0.0 → cognite_extractor_utils-7.0.2}/cognite/extractorutils/configtools/loaders.py

@@ -47,6 +47,10 @@ class KeyVaultAuthenticationMethod(Enum):
 
 
 class KeyVaultLoader:
+    """
+    Class responsible for configuring keyvault for clients using Azure
+    """
+
     def __init__(self, config: Optional[dict]):
         self.config = config
 

{cognite_extractor_utils-7.0.0 → cognite_extractor_utils-7.0.2}/cognite/extractorutils/exceptions.py

@@ -15,7 +15,7 @@
 
 class InvalidConfigError(Exception):
     """
-    Exception thrown from ``load_yaml`` if config file is invalid. This can be due to
+    Exception thrown from ``load_yaml`` and ``load_yaml_dict`` if config file is invalid. This can be due to
 
       * Missing fields
       * Incompatible types

{cognite_extractor_utils-7.0.0 → cognite_extractor_utils-7.0.2}/cognite/extractorutils/threading.py

@@ -1,6 +1,7 @@
 import logging
 import signal
 from threading import Condition
+from time import time
 from typing import Any, Optional
 
 
@@ -59,11 +60,21 @@ class CancellationToken:
         self.cancel()
 
     def wait(self, timeout: Optional[float] = None) -> bool:
+        endtime = None
+        if timeout is not None:
+            endtime = time() + timeout
+
         while not self.is_cancelled:
             with self._cv:
-                did_not_time_out = self._cv.wait(timeout)
-                if not did_not_time_out:
-                    return False
+                if endtime is not None:
+                    remaining_time = endtime - time()
+                    if remaining_time <= 0.0:
+                        return True
+                    timed_out = not self._cv.wait(remaining_time)
+                    if timed_out:
+                        return False
+                else:
+                    self._cv.wait()
         return True
 
     def create_child_token(self) -> "CancellationToken":

{cognite_extractor_utils-7.0.0 → cognite_extractor_utils-7.0.2}/cognite/extractorutils/uploader/files.py

@@ -19,6 +19,8 @@ from os import PathLike
 from types import TracebackType
 from typing import Any, BinaryIO, Callable, Dict, List, Optional, Tuple, Type, Union
 
+from requests.utils import super_len
+
 from cognite.client import CogniteClient
 from cognite.client.data_classes import FileMetadata
 from cognite.extractorutils.threading import CancellationToken
@@ -154,23 +156,34 @@ class IOFileUploadQueue(AbstractUploadQueue):
             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,
-                    )
+                    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 > pow(5, 9):
+                        # File bigger than 5Gb
+                        self.logger.warning(f"File {file_meta.source} has more than 5Gb")
+                        file_meta = FileMetadata()
+                    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:

{cognite_extractor_utils-7.0.0 → cognite_extractor_utils-7.0.2}/cognite/extractorutils/util.py

@@ -79,7 +79,7 @@ class EitherId:
 
     Args:
         id: Internal ID
-        externalId or external_id: external ID
+        external_id: external ID. It can be `external_id` or `externalId`
 
     Raises:
         TypeError: If none of both of id types are set.
@@ -175,24 +175,25 @@ def add_extraction_pipeline(
     added_message: str = "",
 ) -> Callable[[Callable[..., _T1]], Callable[..., _T1]]:
     """
-    This is to be used as a decorator for extractor functions to add extraction pipeline information
+    This is to be used as a decorator for extractor functions to add extraction pipeline information.
 
     Args:
-        extraction_pipeline_ext_id:
-        cognite_client:
-        heartbeat_waiting_time:
-        added_message:
+        extraction_pipeline_ext_id: External ID of the extraction pipeline
+        cognite_client: Client to use when communicating with CDF
+        heartbeat_waiting_time: Target interval between heartbeats, in seconds
 
     Usage:
         If you have a function named "extract_data(*args, **kwargs)" and want to connect it to an extraction
         pipeline, you can use this decorator function as:
-        @add_extraction_pipeline(
-            extraction_pipeline_ext_id=<INSERT EXTERNAL ID>,
-            cognite_client=<INSERT COGNITE CLIENT OBJECT>,
-            logger=<INSERT LOGGER>,
-        )
-        def extract_data(*args, **kwargs):
-            <INSERT FUNCTION BODY>
+
+        .. code-block:: python
+
+            @add_extraction_pipeline(
+                extraction_pipeline_ext_id=<INSERT EXTERNAL ID>,
+                cognite_client=<INSERT COGNITE CLIENT OBJECT>,
+            )
+            def extract_data(*args, **kwargs):
+                <INSERT FUNCTION BODY>
     """
 
     # TODO 1. Consider refactoring this decorator to share methods with the Extractor context manager in .base.py

{cognite_extractor_utils-7.0.0 → cognite_extractor_utils-7.0.2}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "cognite-extractor-utils"
-version = "7.0.0"
+version = "7.0.2"
 description = "Utilities for easier development of extractors for CDF"
 authors = ["Mathias Lohne <mathias.lohne@cognite.com>"]
 license = "Apache-2.0"
@@ -81,7 +81,7 @@ pytest-order = "^1.0.1"
 parameterized = "*"
 requests = "^2.31.0"
 types-requests = "^2.31.0.20240125"
-httpx = "^0.26.0"
+httpx = "^0.27.0"
 
 [build-system]
 requires = ["poetry-core>=1.0.0"]