cognite-extractor-utils 6.1.0__tar.gz → 6.2.0__tar.gz

{cognite_extractor_utils-6.1.0 → cognite_extractor_utils-6.2.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: cognite-extractor-utils
-Version: 6.1.0
+Version: 6.2.0
 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-6.1.0 → cognite_extractor_utils-6.2.0}/cognite/extractorutils/__init__.py

@@ -16,5 +16,5 @@
 Cognite extractor utils is a Python package that simplifies the development of new extractors.
 """
 
-__version__ = "6.1.0"
+__version__ = "6.2.0"
 from .base import Extractor

{cognite_extractor_utils-6.1.0 → cognite_extractor_utils-6.2.0}/cognite/extractorutils/configtools/elements.py

@@ -27,12 +27,25 @@ import yaml
 from prometheus_client import REGISTRY, start_http_server
 
 from cognite.client import ClientConfig, CogniteClient
-from cognite.client.credentials import CredentialProvider, OAuthClientCertificate, OAuthClientCredentials
+from cognite.client.credentials import (
+    CredentialProvider,
+    OAuthClientCertificate,
+    OAuthClientCredentials,
+)
 from cognite.client.data_classes import Asset, DataSet, ExtractionPipeline
 from cognite.extractorutils.configtools._util import _load_certificate_data
 from cognite.extractorutils.exceptions import InvalidConfigError
-from cognite.extractorutils.metrics import AbstractMetricsPusher, CognitePusher, PrometheusPusher
-from cognite.extractorutils.statestore import AbstractStateStore, LocalStateStore, NoStateStore, RawStateStore
+from cognite.extractorutils.metrics import (
+    AbstractMetricsPusher,
+    CognitePusher,
+    PrometheusPusher,
+)
+from cognite.extractorutils.statestore import (
+    AbstractStateStore,
+    LocalStateStore,
+    NoStateStore,
+    RawStateStore,
+)
 from cognite.extractorutils.util import EitherId
 
 _logger = logging.getLogger(__name__)
@@ -93,6 +106,14 @@ class TimeIntervalConfig(yaml.YAMLObject):
     def __init__(self, expression: str) -> None:
         self._interval, self._expression = TimeIntervalConfig._parse_expression(expression)
 
+    def __eq__(self, other: object) -> bool:
+        if not isinstance(other, TimeIntervalConfig):
+            return NotImplemented
+        return self._interval == other._interval
+
+    def __hash__(self) -> int:
+        return hash(self._interval)
+
     @classmethod
     def _parse_expression(cls, expression: str) -> Tuple[int, str]:
         # First, try to parse pure number and assume seconds (for backwards compatibility)
@@ -170,7 +191,10 @@ class FileSizeConfig(yaml.YAMLObject):
         expression_lower = expression.lower()
         for size in sizes:
             if expression_lower.endswith(size):
-                return int(float(expression_lower.replace(size, "")) * sizes[size]), expression
+                return (
+                    int(float(expression_lower.replace(size, "")) * sizes[size]),
+                    expression,
+                )
         else:
             raise InvalidConfigError(f"Invalid unit for file size: {expression}. Valid units: {sizes.keys()}")
 
@@ -241,7 +265,10 @@ class CogniteConfig:
     host: str = "https://api.cognitedata.com"
 
     def get_cognite_client(
-        self, client_name: str, token_custom_args: Optional[Dict[str, str]] = None, use_experimental_sdk: bool = False
+        self,
+        client_name: str,
+        token_custom_args: Optional[Dict[str, str]] = None,
+        use_experimental_sdk: bool = False,
     ) -> CogniteClient:
         from cognite.client.config import global_config
 
@@ -264,7 +291,8 @@ class CogniteConfig:
             else:
                 raise InvalidConfigError("Either authority-url or tenant is required for certificate authentication")
             (thumprint, key) = _load_certificate_data(
-                self.idp_authentication.certificate.path, self.idp_authentication.certificate.password
+                self.idp_authentication.certificate.path,
+                self.idp_authentication.certificate.password,
             )
             credential_provider = OAuthClientCertificate(
                 authority_url=authority_url,
@@ -479,7 +507,10 @@ class MetricsConfig:
             asset = None
 
             if self.cognite.asset_name is not None and self.cognite.asset_external_id:
-                asset = Asset(name=self.cognite.asset_name, external_id=self.cognite.asset_external_id)
+                asset = Asset(
+                    name=self.cognite.asset_name,
+                    external_id=self.cognite.asset_external_id,
+                )
 
             pusher = CognitePusher(
                 cdf_client=cdf_client,
@@ -586,7 +617,10 @@ class StateStoreConfig:
             )
 
         if self.local:
-            return LocalStateStore(file_path=self.local.path, save_interval=self.local.save_interval.seconds)
+            return LocalStateStore(
+                file_path=self.local.path,
+                save_interval=self.local.save_interval.seconds,
+            )
 
         if default_to_local:
             return LocalStateStore(file_path="states.json")

{cognite_extractor_utils-6.1.0 → cognite_extractor_utils-6.2.0}/cognite/extractorutils/uploader/__init__.py

@@ -66,7 +66,7 @@ instead. If both are used, the condition being met first will trigger the upload
 
 from .assets import AssetUploadQueue
 from .events import EventUploadQueue
-from .files import BytesUploadQueue, FileUploadQueue
+from .files import BytesUploadQueue, FileUploadQueue, IOFileUploadQueue
 from .raw import RawUploadQueue
 from .time_series import (
     DataPoint,

{cognite_extractor_utils-6.1.0 → cognite_extractor_utils-6.2.0}/cognite/extractorutils/uploader/files.py

@@ -13,10 +13,11 @@
 #  limitations under the License.
 
 import threading
-from concurrent.futures import ThreadPoolExecutor
+from concurrent.futures import Future, ThreadPoolExecutor
+from io import BytesIO
 from os import PathLike
 from types import TracebackType
-from typing import Any, Callable, List, Optional, Tuple, Type, Union
+from typing import Any, BinaryIO, Callable, List, Optional, Tuple, Type, Union
 
 from requests import ConnectionError
 
@@ -31,9 +32,6 @@ from cognite.extractorutils.uploader._base import (
     AbstractUploadQueue,
 )
 from cognite.extractorutils.uploader._metrics import (
-    BYTES_UPLOADER_QUEUE_SIZE,
-    BYTES_UPLOADER_QUEUED,
-    BYTES_UPLOADER_WRITTEN,
     FILES_UPLOADER_QUEUE_SIZE,
     FILES_UPLOADER_QUEUED,
     FILES_UPLOADER_WRITTEN,
@@ -41,9 +39,12 @@ from cognite.extractorutils.uploader._metrics import (
 from cognite.extractorutils.util import retry
 
 
-class FileUploadQueue(AbstractUploadQueue):
+class IOFileUploadQueue(AbstractUploadQueue):
     """
-    Upload queue for files
+    Upload queue for files using BinaryIO
+
+    Note that if the upload fails, the stream needs to be restarted, so
+    the enqueued callback needs to produce a new IO object for each call.
 
     Args:
         cdf_client: Cognite Data Fusion client to use
@@ -54,6 +55,10 @@ class FileUploadQueue(AbstractUploadQueue):
             methods).
         trigger_log_level: Log level to log upload triggers to.
         thread_name: Thread name of uploader thread.
+        max_parallelism: Maximum number of parallel uploads. If this is greater than 0,
+            the largest of this and client.config.max_workers is used to limit the number
+            of parallel uploads. This may be important if the IO objects being processed
+            also load data from an external system.
     """
 
     def __init__(
@@ -66,6 +71,7 @@ class FileUploadQueue(AbstractUploadQueue):
         thread_name: Optional[str] = None,
         overwrite_existing: bool = False,
         cancellation_token: threading.Event = threading.Event(),
+        max_parallelism: int = 0,
     ):
         # Super sets post_upload and threshold
         super().__init__(
@@ -78,14 +84,20 @@ class FileUploadQueue(AbstractUploadQueue):
             cancellation_token,
         )
 
-        self.upload_queue: List[Tuple[FileMetadata, Union[str, PathLike]]] = []
+        self.upload_queue: List[Tuple[FileMetadata, Union[str, Callable[[], BinaryIO]]]] = []
         self.overwrite_existing = overwrite_existing
 
+        self.parallelism = self.cdf_client.config.max_workers
+        if max_parallelism > 0 and max_parallelism < self.parallelism:
+            self.parallelism = max_parallelism
+        if self.parallelism <= 0:
+            self.parallelism = 4
+
         self.files_queued = FILES_UPLOADER_QUEUED
         self.files_written = FILES_UPLOADER_WRITTEN
         self.queue_size = FILES_UPLOADER_QUEUE_SIZE
 
-    def add_to_upload_queue(self, file_meta: FileMetadata, file_name: Union[str, PathLike]) -> None:
+    def add_io_to_upload_queue(self, file_meta: FileMetadata, read_file: Callable[[], BinaryIO]) -> None:
         """
         Add file to upload queue. The queue will be uploaded if the queue size is larger than the threshold
         specified in the __init__.
@@ -96,7 +108,7 @@ class FileUploadQueue(AbstractUploadQueue):
                 If none, the file object will still be created, but no data is uploaded
         """
         with self.lock:
-            self.upload_queue.append((file_meta, file_name))
+            self.upload_queue.append((file_meta, read_file))
             self.upload_queue_size += 1
             self.files_queued.inc()
             self.queue_size.set(self.upload_queue_size)
@@ -131,21 +143,41 @@ class FileUploadQueue(AbstractUploadQueue):
         max_delay=RETRY_MAX_DELAY,
         backoff=RETRY_BACKOFF_FACTOR,
     )
-    def _upload_single(self, index: int, file_name: Union[str, PathLike], file_meta: FileMetadata) -> None:
+    def _upload_single(self, index: int, read_file: Callable[[], BinaryIO], file_meta: FileMetadata) -> None:
         # Upload file
-        file_meta = self.cdf_client.files.upload(str(file_name), overwrite=self.overwrite_existing, **file_meta.dump())  # type: ignore
+        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,
+            )
 
         # Update meta-object in queue
-        self.upload_queue[index] = (file_meta, file_name)
+        self.upload_queue[index] = (file_meta, read_file)
 
     def _upload_batch(self) -> None:
         # Concurrently execute file-uploads
 
-        with ThreadPoolExecutor(self.cdf_client.config.max_workers) as pool:
+        futures: List[Future] = []
+        with ThreadPoolExecutor(self.parallelism) as pool:
             for i, (file_meta, file_name) in enumerate(self.upload_queue):
-                pool.submit(self._upload_single, i, file_name, file_meta)
+                futures.append(pool.submit(self._upload_single, i, file_name, file_meta))
+        for fut in futures:
+            fut.result(0.0)
 
-    def __enter__(self) -> "FileUploadQueue":
+    def __enter__(self) -> "IOFileUploadQueue":
         """
         Wraps around start method, for use as context manager
 
@@ -178,9 +210,9 @@ class FileUploadQueue(AbstractUploadQueue):
         return self.upload_queue_size
 
 
-class BytesUploadQueue(AbstractUploadQueue):
+class FileUploadQueue(IOFileUploadQueue):
     """
-    Upload queue for bytes
+    Upload queue for files
 
     Args:
         cdf_client: Cognite Data Fusion client to use
@@ -191,20 +223,20 @@ class BytesUploadQueue(AbstractUploadQueue):
             methods).
         trigger_log_level: Log level to log upload triggers to.
         thread_name: Thread name of uploader thread.
-        overwrite_existing: If 'overwrite' is set to true, fields for the files found for externalIds can be overwritten
     """
 
     def __init__(
         self,
         cdf_client: CogniteClient,
-        post_upload_function: Optional[Callable[[List[Any]], None]] = None,
+        post_upload_function: Optional[Callable[[List[Event]], None]] = None,
         max_queue_size: Optional[int] = None,
         max_upload_interval: Optional[int] = None,
         trigger_log_level: str = "DEBUG",
         thread_name: Optional[str] = None,
         overwrite_existing: bool = False,
         cancellation_token: threading.Event = threading.Event(),
-    ) -> None:
+    ):
+        # Super sets post_upload and threshold
         super().__init__(
             cdf_client,
             post_upload_function,
@@ -212,105 +244,75 @@ class BytesUploadQueue(AbstractUploadQueue):
             max_upload_interval,
             trigger_log_level,
             thread_name,
+            overwrite_existing,
             cancellation_token,
         )
-        self.upload_queue: List[Tuple[bytes, FileMetadata]] = []
-        self.overwrite_existing = overwrite_existing
-        self.upload_queue_size = 0
-
-        self.bytes_queued = BYTES_UPLOADER_QUEUED
-        self.queue_size = BYTES_UPLOADER_QUEUE_SIZE
-        self.bytes_written = BYTES_UPLOADER_WRITTEN
 
-    def add_to_upload_queue(self, content: bytes, metadata: FileMetadata) -> None:
+    def add_to_upload_queue(self, file_meta: FileMetadata, file_name: Union[str, PathLike]) -> None:
         """
-        Add object to upload queue. The queue will be uploaded if the queue size is larger than the threshold
+        Add file to upload queue. The queue will be uploaded if the queue size is larger than the threshold
         specified in the __init__.
-        Args:
-            content: bytes object to upload
-            metadata: metadata for the given bytes object
-        """
-        with self.lock:
-            self.upload_queue.append((content, metadata))
-            self.upload_queue_size += 1
-            self.bytes_queued.inc()
-            self.queue_size.set(self.upload_queue_size)
 
-    def upload(self) -> None:
-        """
-        Trigger an upload of the queue, clears queue afterwards
+        Args:
+            file_meta: File metadata-object
+            file_name: Path to file to be uploaded.
+                If none, the file object will still be created, but no data is uploaded
         """
-        if len(self.upload_queue) == 0:
-            return
 
-        with self.lock:
-            # Upload frames in batches
-            self._upload_batch()
+        def load_file_from_path() -> BinaryIO:
+            return open(file_name, "rb")
 
-            # Log stats
-            self.bytes_written.inc(self.upload_queue_size)
+        self.add_io_to_upload_queue(file_meta, load_file_from_path)
 
-            try:
-                self._post_upload(self.upload_queue)
-            except Exception as e:
-                self.logger.error("Error in upload callback: %s", str(e))
 
-            # Clear queue
-            self.upload_queue.clear()
-            self.upload_queue_size = 0
-            self.logger.info(f"Uploaded {self.upload_queue_size} files")
-            self.queue_size.set(self.upload_queue_size)
+class BytesUploadQueue(IOFileUploadQueue):
+    """
+    Upload queue for bytes
 
-    def _upload_batch(self) -> None:
-        # Concurrently execute bytes-uploads
-        with ThreadPoolExecutor(self.cdf_client.config.max_workers) as pool:
-            for i, (frame, metadata) in enumerate(self.upload_queue):
-                pool.submit(self._upload_single, i, frame, metadata)
+    Args:
+        cdf_client: Cognite Data Fusion client to use
+        post_upload_function: A function that will be called after each upload. The function will be given one argument:
+            A list of the events that were uploaded.
+        max_queue_size: Maximum size of upload queue. Defaults to no max size.
+        max_upload_interval: Automatically trigger an upload each m seconds when run as a thread (use start/stop
+            methods).
+        trigger_log_level: Log level to log upload triggers to.
+        thread_name: Thread name of uploader thread.
+        overwrite_existing: If 'overwrite' is set to true, fields for the files found for externalIds can be overwritten
+    """
 
-    @retry(
-        exceptions=(CogniteAPIError, ConnectionError),
-        tries=RETRIES,
-        delay=RETRY_DELAY,
-        max_delay=RETRY_MAX_DELAY,
-        backoff=RETRY_BACKOFF_FACTOR,
-    )
-    def _upload_single(self, index: int, content: bytes, metadata: FileMetadata) -> None:
-        # Upload object
-        file_meta_data: FileMetadata = self.cdf_client.files.upload_bytes(
-            content, overwrite=self.overwrite_existing, **metadata.dump()
+    def __init__(
+        self,
+        cdf_client: CogniteClient,
+        post_upload_function: Optional[Callable[[List[Any]], None]] = None,
+        max_queue_size: Optional[int] = None,
+        max_upload_interval: Optional[int] = None,
+        trigger_log_level: str = "DEBUG",
+        thread_name: Optional[str] = None,
+        overwrite_existing: bool = False,
+        cancellation_token: threading.Event = threading.Event(),
+    ) -> None:
+        super().__init__(
+            cdf_client,
+            post_upload_function,
+            max_queue_size,
+            max_upload_interval,
+            trigger_log_level,
+            thread_name,
+            overwrite_existing,
+            cancellation_token,
         )
 
-        # Update meta-object in queue
-        self.upload_queue[index] = (content, file_meta_data)
-
-    def __enter__(self) -> "BytesUploadQueue":
-        """
-        Wraps around start method, for use as context manager
-
-        Returns:
-            self
-        """
-        self.start()
-        return self
-
-    def __exit__(
-        self, exc_type: Optional[Type[BaseException]], exc_val: Optional[BaseException], exc_tb: Optional[TracebackType]
-    ) -> None:
+    def add_to_upload_queue(self, content: bytes, metadata: FileMetadata) -> None:
         """
-        Wraps around stop method, for use as context manager
-
+        Add object to upload queue. The queue will be uploaded if the queue size is larger than the threshold
+        specified in the __init__.
         Args:
-            exc_type: Exception type
-            exc_val: Exception value
-            exc_tb: Traceback
+            content: bytes object to upload
+            metadata: metadata for the given bytes object
         """
-        self.stop()
 
-    def __len__(self) -> int:
-        """
-        The size of the upload queue
+        def get_byte_io() -> BinaryIO:
+            return BytesIO(content)
 
-        Returns:
-            Number of events in queue
-        """
-        return self.upload_queue_size
+        self.add_io_to_upload_queue(metadata, get_byte_io)

{cognite_extractor_utils-6.1.0 → cognite_extractor_utils-6.2.0}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "cognite-extractor-utils"
-version = "6.1.0"
+version = "6.2.0"
 description = "Utilities for easier development of extractors for CDF"
 authors = ["Mathias Lohne <mathias.lohne@cognite.com>"]
 license = "Apache-2.0"
@@ -67,7 +67,7 @@ jq = [{version = "^1.3.0", platform = "darwin"}, {version = "^1.3.0", platform =
 experimental = ["cognite-sdk-experimental"]
 
 [tool.poetry.group.dev.dependencies]
-mypy = "1.7.1"
+mypy = "1.8.0"
 ruff = "^0.1.0"
 pytest = "^7.0.0"
 pytest-cov = "^4.0.0"