cognite-extractor-utils 7.2.3__tar.gz → 7.4.0__tar.gz

{cognite_extractor_utils-7.2.3 → cognite_extractor_utils-7.4.0}/PKG-INFO

@@ -1,15 +1,14 @@
 Metadata-Version: 2.1
 Name: cognite-extractor-utils
-Version: 7.2.3
+Version: 7.4.0
 Summary: Utilities for easier development of extractors for CDF
 Home-page: https://github.com/cognitedata/python-extractor-utils
 License: Apache-2.0
 Author: Mathias Lohne
 Author-email: mathias.lohne@cognite.com
-Requires-Python: >=3.8.0,<4.0.0
+Requires-Python: >=3.9.0,<4.0.0
 Classifier: License :: OSI Approved :: Apache Software License
 Classifier: Programming Language :: Python :: 3
-Classifier: Programming Language :: Python :: 3.8
 Classifier: Programming Language :: Python :: 3.9
 Classifier: Programming Language :: Python :: 3.10
 Classifier: Programming Language :: Python :: 3.11
@@ -18,9 +17,10 @@ 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.43.3,<8.0.0)
+Requires-Dist: cognite-sdk (>=7.54.17,<8.0.0)
 Requires-Dist: dacite (>=1.6.0,<2.0.0)
 Requires-Dist: decorator (>=5.1.1,<6.0.0)
+Requires-Dist: httpx (>=0.27.0,<0.28.0)
 Requires-Dist: more-itertools (>=10.0.0,<11.0.0)
 Requires-Dist: orjson (>=3.10.3,<4.0.0)
 Requires-Dist: prometheus-client (>0.7.0,<=1.0.0)

{cognite_extractor_utils-7.2.3 → cognite_extractor_utils-7.4.0}/cognite/extractorutils/__init__.py

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

cognite_extractor_utils-7.4.0/cognite/extractorutils/unstable/__init__.py

@@ -0,0 +1,8 @@
+"""
+The unstable package contains experimental functions and classes currently
+deemed unstable. The contents of this package is subject to change without
+notice, even in minor or patch releases.
+
+Whenever you import anything from the unstable package, you should make sure to
+run a type checker such as mypy to help catch these changes.
+"""

{cognite_extractor_utils-7.2.3 → cognite_extractor_utils-7.4.0}/cognite/extractorutils/uploader/files.py

@@ -18,12 +18,16 @@ from io import BytesIO, RawIOBase
 from math import ceil
 from os import PathLike
 from types import TracebackType
-from typing import Any, BinaryIO, Callable, Dict, List, Optional, Tuple, Type, Union
+from typing import Any, BinaryIO, Callable, Dict, Iterator, List, Optional, Tuple, Type, Union
 
+from httpx import URL, Client, Headers, Request, StreamConsumed, SyncByteStream
 from requests.utils import super_len
 
 from cognite.client import CogniteClient
 from cognite.client.data_classes import FileMetadata
+from cognite.client.data_classes.data_modeling import NodeId
+from cognite.client.data_classes.data_modeling.extractor_extensions.v1 import CogniteExtractorFileApply
+from cognite.client.utils._identifier import IdentifierSequence
 from cognite.extractorutils.threading import CancellationToken
 from cognite.extractorutils.uploader._base import (
     RETRIES,
@@ -47,6 +51,10 @@ _MAX_SINGLE_CHUNK_FILE_SIZE = 5 * 1024 * 1024 * 1024
 # 4000 MiB
 _MAX_FILE_CHUNK_SIZE = 4 * 1024 * 1024 * 1000
 
+_CDF_ALPHA_VERSION_HEADER = {"cdf-version": "alpha"}
+
+FileMetadataOrCogniteExtractorFile = Union[FileMetadata, CogniteExtractorFileApply]
+
 
 class ChunkedStream(RawIOBase, BinaryIO):
     """
@@ -140,6 +148,22 @@ class ChunkedStream(RawIOBase, BinaryIO):
         return True
 
 
+class IOByteStream(SyncByteStream):
+    CHUNK_SIZE = 65_536
+
+    def __init__(self, stream: BinaryIO) -> None:
+        self._stream = stream
+        self._is_stream_consumed = False
+
+    def __iter__(self) -> Iterator[bytes]:
+        if self._is_stream_consumed:
+            raise StreamConsumed()
+        chunk = self._stream.read(self.CHUNK_SIZE)
+        while chunk:
+            yield chunk
+            chunk = self._stream.read(self.CHUNK_SIZE)
+
+
 class IOFileUploadQueue(AbstractUploadQueue):
     """
     Upload queue for files using BinaryIO
@@ -161,7 +185,7 @@ class IOFileUploadQueue(AbstractUploadQueue):
     def __init__(
         self,
         cdf_client: CogniteClient,
-        post_upload_function: Optional[Callable[[List[FileMetadata]], None]] = None,
+        post_upload_function: Optional[Callable[[List[FileMetadataOrCogniteExtractorFile]], None]] = None,
         max_queue_size: Optional[int] = None,
         trigger_log_level: str = "DEBUG",
         thread_name: Optional[str] = None,
@@ -205,6 +229,8 @@ class IOFileUploadQueue(AbstractUploadQueue):
 
         self._full_queue = threading.Condition()
 
+        self._httpx_client = Client(follow_redirects=True)
+
         global _QUEUES, _QUEUES_LOCK
         with _QUEUES_LOCK:
             self._pool = ThreadPoolExecutor(
@@ -219,9 +245,81 @@ class IOFileUploadQueue(AbstractUploadQueue):
 
             self.cancellation_token.wait(5)
 
+    def _apply_cognite_file(self, file_apply: CogniteExtractorFileApply) -> NodeId:
+        instance_result = self.cdf_client.data_modeling.instances.apply(file_apply)
+        node = instance_result.nodes[0]
+        return node.as_id()
+
+    def _upload_empty(
+        self, meta_or_apply: FileMetadataOrCogniteExtractorFile
+    ) -> tuple[FileMetadataOrCogniteExtractorFile, str]:
+        if isinstance(meta_or_apply, CogniteExtractorFileApply):
+            node_id = self._apply_cognite_file(meta_or_apply)
+            meta_or_apply, url = self._create_cdm(instance_id=node_id)
+        else:
+            meta_or_apply, url = self.cdf_client.files.create(
+                file_metadata=meta_or_apply, overwrite=self.overwrite_existing
+            )
+        return meta_or_apply, url
+
+    def _upload_bytes(self, size: int, file: BinaryIO, meta_or_apply: FileMetadataOrCogniteExtractorFile) -> None:
+        meta_or_apply, url = self._upload_empty(meta_or_apply)
+        resp = self._httpx_client.send(self._get_file_upload_request(url, file, size))
+        resp.raise_for_status()
+
+    def _upload_multipart(self, size: int, file: BinaryIO, meta_or_apply: FileMetadataOrCogniteExtractorFile) -> None:
+        chunks = ChunkedStream(file, self.max_file_chunk_size, size)
+        self.logger.debug(
+            f"File {meta_or_apply.external_id} is larger than 5GiB ({size})"
+            f", uploading in {chunks.chunk_count} chunks"
+        )
+
+        returned_file_metadata = self._create_multi_part(meta_or_apply, chunks)
+        upload_urls = returned_file_metadata["uploadUrls"]
+        upload_id = returned_file_metadata["uploadId"]
+        file_meta = FileMetadata.load(returned_file_metadata)
+
+        for url in upload_urls:
+            chunks.next_chunk()
+            resp = self._httpx_client.send(self._get_file_upload_request(url, chunks, len(chunks)))
+            resp.raise_for_status()
+
+        completed_headers = (
+            _CDF_ALPHA_VERSION_HEADER if isinstance(meta_or_apply, CogniteExtractorFileApply) is not None else None
+        )
+
+        res = self.cdf_client.files._post(
+            url_path="/files/completemultipartupload",
+            json={"id": file_meta.id, "uploadId": upload_id},
+            headers=completed_headers,
+        )
+        res.raise_for_status()
+
+    def _create_multi_part(self, meta_or_apply: FileMetadataOrCogniteExtractorFile, chunks: ChunkedStream) -> dict:
+        if isinstance(meta_or_apply, CogniteExtractorFileApply):
+            node_id = self._apply_cognite_file(meta_or_apply)
+            identifiers = IdentifierSequence.load(instance_ids=node_id).as_singleton()
+            self.cdf_client.files._warn_alpha()
+            res = self.cdf_client.files._post(
+                url_path="/files/multiuploadlink",
+                json={"items": identifiers.as_dicts()},
+                params={"parts": chunks.chunk_count},
+                headers=_CDF_ALPHA_VERSION_HEADER,
+            )
+            res.raise_for_status()
+            return res.json()["items"][0]
+        else:
+            res = self.cdf_client.files._post(
+                url_path="/files/initmultipartupload",
+                json=meta_or_apply.dump(camel_case=True),
+                params={"overwrite": self.overwrite_existing, "parts": chunks.chunk_count},
+            )
+            res.raise_for_status()
+            return res.json()
+
     def add_io_to_upload_queue(
         self,
-        file_meta: FileMetadata,
+        meta_or_apply: FileMetadataOrCogniteExtractorFile,
         read_file: Callable[[], BinaryIO],
         extra_retries: Optional[
             Union[Tuple[Type[Exception], ...], Dict[Type[Exception], Callable[[Any], bool]]]
@@ -229,7 +327,7 @@ class IOFileUploadQueue(AbstractUploadQueue):
     ) -> None:
         """
         Add file to upload queue. The file will start uploading immedeately. If the size of the queue is larger than
-        the specified max size, this call will block until it's
+        the specified max size, this call will block until it's completed the upload.
 
         Args:
             file_meta: File metadata-object
@@ -239,7 +337,7 @@ class IOFileUploadQueue(AbstractUploadQueue):
         """
         retries = cognite_exceptions()
         if isinstance(extra_retries, tuple):
-            retries.update({exc: lambda _e: True for exc in extra_retries or []})
+            retries.update({exc: lambda _: True for exc in extra_retries or []})
         elif isinstance(extra_retries, dict):
             retries.update(extra_retries)
 
@@ -251,72 +349,36 @@ class IOFileUploadQueue(AbstractUploadQueue):
             max_delay=RETRY_MAX_DELAY,
             backoff=RETRY_BACKOFF_FACTOR,
         )
-        def upload_file(read_file: Callable[[], BinaryIO], file_meta: FileMetadata) -> None:
+        def upload_file(read_file: Callable[[], BinaryIO], meta_or_apply: FileMetadataOrCogniteExtractorFile) -> 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
-                    )
+                    meta_or_apply, _ = self._upload_empty(meta_or_apply)
                 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
+                    self._upload_multipart(size, file, meta_or_apply)
+
                 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,
-                    )
+                    self._upload_bytes(size, file, meta_or_apply)
+
+                if isinstance(meta_or_apply, CogniteExtractorFileApply):
+                    meta_or_apply.is_uploaded = True
 
             if self.post_upload_function:
                 try:
-                    self.post_upload_function([file_meta])
+                    self.post_upload_function([meta_or_apply])
                 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:
+        def wrapped_upload(
+            read_file: Callable[[], BinaryIO], meta_or_apply: FileMetadataOrCogniteExtractorFile
+        ) -> None:
             try:
-                upload_file(read_file, file_meta)
+                upload_file(read_file, meta_or_apply)
 
             except Exception as e:
-                self.logger.exception(f"Unexpected error while uploading file: {file_meta.external_id}")
+                self.logger.exception(f"Unexpected error while uploading file: {meta_or_apply.external_id}")
                 self.errors.append(e)
 
             finally:
@@ -333,11 +395,42 @@ class IOFileUploadQueue(AbstractUploadQueue):
                     pass
 
         with self.lock:
-            self.upload_queue.append(self._pool.submit(wrapped_upload, read_file, file_meta))
+            self.upload_queue.append(self._pool.submit(wrapped_upload, read_file, meta_or_apply))
             self.upload_queue_size += 1
             self.files_queued.inc()
             self.queue_size.set(self.upload_queue_size)
 
+    def _get_file_upload_request(self, url_str: str, stream: BinaryIO, size: int) -> Request:
+        url = URL(url_str)
+        headers = Headers(self._httpx_client.headers)
+        headers.update(
+            {
+                "Accept": "*/*",
+                "Content-Length": str(size),
+                "Host": url.netloc.decode("ascii"),
+                "x-cdp-app": self.cdf_client._config.client_name,
+            }
+        )
+
+        return Request(
+            method="PUT",
+            url=url,
+            stream=IOByteStream(stream),
+            headers=headers,
+        )
+
+    def _create_cdm(self, instance_id: NodeId) -> tuple[FileMetadata, str]:
+        self.cdf_client.files._warn_alpha()
+        identifiers = IdentifierSequence.load(instance_ids=instance_id).as_singleton()
+        res = self.cdf_client.files._post(
+            url_path="/files/uploadlink",
+            json={"items": identifiers.as_dicts()},
+            headers=_CDF_ALPHA_VERSION_HEADER,
+        )
+        res.raise_for_status()
+        resp_json = res.json()["items"][0]
+        return FileMetadata.load(resp_json), resp_json["uploadUrl"]
+
     def upload(self, fail_on_errors: bool = True, timeout: Optional[float] = None) -> None:
         """
         Wait for all uploads to finish
@@ -402,7 +495,7 @@ class FileUploadQueue(IOFileUploadQueue):
     def __init__(
         self,
         cdf_client: CogniteClient,
-        post_upload_function: Optional[Callable[[List[FileMetadata]], None]] = None,
+        post_upload_function: Optional[Callable[[List[FileMetadataOrCogniteExtractorFile]], None]] = None,
         max_queue_size: Optional[int] = None,
         max_upload_interval: Optional[int] = None,
         trigger_log_level: str = "DEBUG",
@@ -421,7 +514,9 @@ class FileUploadQueue(IOFileUploadQueue):
             cancellation_token,
         )
 
-    def add_to_upload_queue(self, file_meta: FileMetadata, file_name: Union[str, PathLike]) -> None:
+    def add_to_upload_queue(
+        self, meta_or_apply: FileMetadataOrCogniteExtractorFile, file_name: Union[str, PathLike]
+    ) -> None:
         """
         Add file to upload queue. The queue will be uploaded if the queue size is larger than the threshold
         specified in the __init__.
@@ -435,7 +530,7 @@ class FileUploadQueue(IOFileUploadQueue):
         def load_file_from_path() -> BinaryIO:
             return open(file_name, "rb")
 
-        self.add_io_to_upload_queue(file_meta, load_file_from_path)
+        self.add_io_to_upload_queue(meta_or_apply, load_file_from_path)
 
 
 class BytesUploadQueue(IOFileUploadQueue):
@@ -455,7 +550,7 @@ class BytesUploadQueue(IOFileUploadQueue):
     def __init__(
         self,
         cdf_client: CogniteClient,
-        post_upload_function: Optional[Callable[[List[FileMetadata]], None]] = None,
+        post_upload_function: Optional[Callable[[List[FileMetadataOrCogniteExtractorFile]], None]] = None,
         max_queue_size: Optional[int] = None,
         trigger_log_level: str = "DEBUG",
         thread_name: Optional[str] = None,
@@ -472,7 +567,7 @@ class BytesUploadQueue(IOFileUploadQueue):
             cancellation_token,
         )
 
-    def add_to_upload_queue(self, content: bytes, metadata: FileMetadata) -> None:
+    def add_to_upload_queue(self, content: bytes, meta_or_apply: FileMetadataOrCogniteExtractorFile) -> None:
         """
         Add object to upload queue. The queue will be uploaded if the queue size is larger than the threshold
         specified in the __init__.
@@ -484,4 +579,4 @@ class BytesUploadQueue(IOFileUploadQueue):
         def get_byte_io() -> BinaryIO:
             return BytesIO(content)
 
-        self.add_io_to_upload_queue(metadata, get_byte_io)
+        self.add_io_to_upload_queue(meta_or_apply, get_byte_io)

{cognite_extractor_utils-7.2.3 → cognite_extractor_utils-7.4.0}/cognite/extractorutils/util.py

@@ -19,6 +19,7 @@ extractors.
 
 import logging
 import random
+from datetime import datetime, timezone
 from functools import partial, wraps
 from threading import Thread
 from time import time
@@ -501,3 +502,11 @@ def cognite_exceptions(
         return True
 
     return {CogniteException: handle_cognite_errors}
+
+
+def datetime_to_timestamp(dt: datetime) -> int:
+    return int(dt.timestamp() * 1000)
+
+
+def timestamp_to_datetime(ts: int) -> datetime:
+    return datetime.fromtimestamp(ts / 1000, tz=timezone.utc)

{cognite_extractor_utils-7.2.3 → cognite_extractor_utils-7.4.0}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "cognite-extractor-utils"
-version = "7.2.3"
+version = "7.4.0"
 description = "Utilities for easier development of extractors for CDF"
 authors = ["Mathias Lohne <mathias.lohne@cognite.com>"]
 license = "Apache-2.0"
@@ -38,6 +38,11 @@ target-version = "py310"
 [tool.ruff.isort]
 known-third-party = ["alembic"]
 
+[tool.pytest.ini_options]
+markers = [
+    "unstable: tests for the unstable package (deselect with '-m \"not unstable\"')",
+]
+
 [tool.mypy]
 pretty = true
 check_untyped_defs = true
@@ -52,8 +57,8 @@ warn_unused_ignores = true
 exclude = "tests/*"
 
 [tool.poetry.dependencies]
-python = "^3.8.0"
-cognite-sdk = "^7.43.3"
+python = "^3.9.0"
+cognite-sdk = "^7.54.17"
 prometheus-client = ">0.7.0, <=1.0.0"
 arrow = "^1.0.0"
 pyyaml = ">=5.3.0, <7"
@@ -66,13 +71,14 @@ python-dotenv = "^1.0.0"
 azure-identity = "^1.14.0"
 azure-keyvault-secrets = "^4.7.0"
 orjson = "^3.10.3"
+httpx = "^0.27.0"
 
 [tool.poetry.extras]
 experimental = ["cognite-sdk-experimental"]
 
 [tool.poetry.group.dev.dependencies]
-mypy = "1.10.1"
-ruff = "^0.5.0"
+mypy = "1.11.2"
+ruff = "^0.6.0"
 pytest = "^8.0.0"
 pytest-cov = "^5.0.0"
 sphinx = "^7.0.0"
@@ -85,7 +91,7 @@ parameterized = "*"
 requests = "^2.31.0"
 types-requests = "^2.31.0.20240125"
 httpx = "^0.27.0"
-faker = "^26.0.0"
+faker = "^28.0.0"
 
 [build-system]
 requires = ["poetry-core>=1.0.0"]