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

cognite/extractorutils/uploader/files.py

@@ -18,8 +18,9 @@ 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
@@ -140,6 +141,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
@@ -205,6 +222,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(
@@ -266,44 +285,32 @@ class IOFileUploadQueue(AbstractUploadQueue):
                         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
+
+                    res = self.cdf_client.files._post(
+                        url_path="/files/initmultipartupload",
+                        json=file_meta.dump(camel_case=True),
+                        params={"overwrite": self.overwrite_existing, "parts": chunks.chunk_count},
+                    )
+                    returned_file_metadata = res.json()
+                    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()
+
+                    self.cdf_client.files._post(
+                        url_path="/files/completemultipartupload", json={"id": file_meta.id, "uploadId": upload_id}
+                    )
+
                 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,
+                    file_meta, url = self.cdf_client.files.create(
+                        file_metadata=file_meta, overwrite=self.overwrite_existing
                     )
+                    resp = self._httpx_client.send(self._get_file_upload_request(url, file, size))
+                    resp.raise_for_status()
 
             if self.post_upload_function:
                 try:
@@ -338,6 +345,25 @@ class IOFileUploadQueue(AbstractUploadQueue):
             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 upload(self, fail_on_errors: bool = True, timeout: Optional[float] = None) -> None:
         """
         Wait for all uploads to finish

{cognite_extractor_utils-7.2.3.dist-info → cognite_extractor_utils-7.3.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: cognite-extractor-utils
-Version: 7.2.3
+Version: 7.3.0
 Summary: Utilities for easier development of extractors for CDF
 Home-page: https://github.com/cognitedata/python-extractor-utils
 License: Apache-2.0
@@ -21,6 +21,7 @@ Requires-Dist: azure-keyvault-secrets (>=4.7.0,<5.0.0)
 Requires-Dist: cognite-sdk (>=7.43.3,<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.dist-info → cognite_extractor_utils-7.3.0.dist-info}/RECORD

@@ -1,4 +1,4 @@
-cognite/extractorutils/__init__.py,sha256=No8QzFZvAX7DY63jX70dTCPGuk7k0Rd0k99nU_sqeTo,739
+cognite/extractorutils/__init__.py,sha256=xQex_pCl54rp85pvtXkNBPlRtgkgs5dfw7G_CBOPtpM,739
 cognite/extractorutils/_inner_util.py,sha256=gmz6aqS7jDNsg8z4RHgJjMFohDLOMiaU4gMWBhg3xcE,1558
 cognite/extractorutils/base.py,sha256=q6NU2bPec3WOasVnnIFoh-aUJudVZWZ2R6emz3IRj8Q,16391
 cognite/extractorutils/configtools/__init__.py,sha256=L-daaqInIsmHcjb2forJeY0fW8tz1mlteOUo7IsWnrU,3059
@@ -19,13 +19,13 @@ cognite/extractorutils/uploader/_metrics.py,sha256=J2LJXb19L_SLSJ_voNIQHYLp0pjxU
 cognite/extractorutils/uploader/assets.py,sha256=2E90N1kxsaA6Ah4h0_r_dTVhDYY_68ItRWrHYkkltJw,5628
 cognite/extractorutils/uploader/data_modeling.py,sha256=w35Ix5mu0Cgfn4ywnDyif4VVjo04LVTlkMEevk6ztUs,3639
 cognite/extractorutils/uploader/events.py,sha256=NZP2tMoU_rh_rb-EZiUBsOT5KdNABHN4c9Oddk0OsdE,5680
-cognite/extractorutils/uploader/files.py,sha256=31kPS4fwz8ZSXWss-CKmYTM6ZLVx9LtsDe7LHT7Wy98,18329
+cognite/extractorutils/uploader/files.py,sha256=5dVdG_3QtZ5Gr_GwYcE-FPQkKzC2I5odr6JXG5-arCc,18658
 cognite/extractorutils/uploader/raw.py,sha256=wFjF90PFTjmByOWx_Y4_YfDJ2w2jl0EQJ2Tjx2MP2PM,6738
 cognite/extractorutils/uploader/time_series.py,sha256=HBtQdsQoIOaL-EG5lMsaY-ORwVb0kGiXG86VjE5-_Bg,26815
 cognite/extractorutils/uploader_extractor.py,sha256=E-mpVvbPg_Tk90U4S9JybV0duptJ2SXE88HB6npE3zI,7732
 cognite/extractorutils/uploader_types.py,sha256=wxfrsiKPTzG5lmoYtQsxt8Xyj-s5HnaLl8WDzJNrazg,1020
 cognite/extractorutils/util.py,sha256=UA6mUZ1caHd6vtA45gZXrk6cxo5cSB2PZ32bMwfEU0M,17229
-cognite_extractor_utils-7.2.3.dist-info/LICENSE,sha256=psuoW8kuDP96RQsdhzwOqi6fyWv0ct8CR6Jr7He_P_k,10173
-cognite_extractor_utils-7.2.3.dist-info/METADATA,sha256=u-YGjhv1yLXR0t8OYCx07D7vp1_8Vg_yFwMke0HThZ8,5486
-cognite_extractor_utils-7.2.3.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
-cognite_extractor_utils-7.2.3.dist-info/RECORD,,
+cognite_extractor_utils-7.3.0.dist-info/LICENSE,sha256=psuoW8kuDP96RQsdhzwOqi6fyWv0ct8CR6Jr7He_P_k,10173
+cognite_extractor_utils-7.3.0.dist-info/METADATA,sha256=d5YNb6IXGvnpihhUECrRu3yhspo7ywz9Fs1ejlgpzE4,5526
+cognite_extractor_utils-7.3.0.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
+cognite_extractor_utils-7.3.0.dist-info/RECORD,,