cognite-extractor-utils 7.5.4__py3-none-any.whl → 7.5.6__py3-none-any.whl

cognite/extractorutils/uploader/files.py

@@ -18,7 +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, Iterator, List, Optional, Tuple, Type, Union
+from typing import (
+    Any,
+    BinaryIO,
+    Callable,
+    Iterator,
+    List,
+    Optional,
+    Type,
+    Union,
+)
 from urllib.parse import ParseResult, urlparse
 
 from httpx import URL, Client, Headers, Request, StreamConsumed, SyncByteStream
@@ -27,7 +36,9 @@ from requests.utils import super_len
 from cognite.client import CogniteClient
 from cognite.client.data_classes import FileMetadata, FileMetadataUpdate
 from cognite.client.data_classes.data_modeling import NodeId
-from cognite.client.data_classes.data_modeling.extractor_extensions.v1 import CogniteExtractorFileApply
+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 (
@@ -42,6 +53,7 @@ from cognite.extractorutils.uploader._metrics import (
     FILES_UPLOADER_QUEUED,
     FILES_UPLOADER_WRITTEN,
 )
+from cognite.extractorutils.uploader.upload_failure_handler import FileFailureManager
 from cognite.extractorutils.util import cognite_exceptions, retry
 
 _QUEUES: int = 0
@@ -54,6 +66,7 @@ _MAX_FILE_CHUNK_SIZE = 4 * 1024 * 1024 * 1000
 
 _CDF_ALPHA_VERSION_HEADER = {"cdf-version": "alpha"}
 
+
 FileMetadataOrCogniteExtractorFile = Union[FileMetadata, CogniteExtractorFileApply]
 
 
@@ -97,7 +110,10 @@ class ChunkedStream(RawIOBase, BinaryIO):
         return super().__enter__()
 
     def __exit__(
-        self, exc_type: Optional[Type[BaseException]], exc_val: Optional[BaseException], exc_tb: Optional[TracebackType]
+        self,
+        exc_type: Type[BaseException] | None,
+        exc_val: BaseException | None,
+        exc_tb: TracebackType | None,
     ) -> None:
         return super().__exit__(exc_type, exc_val, exc_tb)
 
@@ -181,18 +197,22 @@ class IOFileUploadQueue(AbstractUploadQueue):
         thread_name: Thread name of uploader thread.
         max_parallelism: Maximum number of parallel uploads. If nothing is given, the parallelism will be capped by the
             max_workers of the cognite client.
+        ssl_verify: Either a string (path to a CA bundle) or a bool (false to turn off completely, true to use standard
+            CA bundle)
     """
 
     def __init__(
         self,
         cdf_client: CogniteClient,
-        post_upload_function: Optional[Callable[[List[FileMetadataOrCogniteExtractorFile]], None]] = None,
-        max_queue_size: Optional[int] = None,
+        post_upload_function: Callable[[list[FileMetadataOrCogniteExtractorFile]], None] | None = None,
+        max_queue_size: int | None = None,
         trigger_log_level: str = "DEBUG",
-        thread_name: Optional[str] = None,
+        thread_name: str | None = None,
         overwrite_existing: bool = False,
         cancellation_token: Optional[CancellationToken] = None,
         max_parallelism: Optional[int] = None,
+        failure_logging_path: None | str = None,
+        ssl_verify: bool | str = True,
     ):
         # Super sets post_upload and threshold
         super().__init__(
@@ -208,6 +228,9 @@ class IOFileUploadQueue(AbstractUploadQueue):
         if self.threshold <= 0:
             raise ValueError("Max queue size must be positive for file upload queues")
 
+        self.failure_logging_path = failure_logging_path or None
+        self.initialize_failure_logging()
+
         self.upload_queue: List[Future] = []
         self.errors: List[Exception] = []
 
@@ -230,15 +253,40 @@ class IOFileUploadQueue(AbstractUploadQueue):
 
         self._full_queue = threading.Condition()
 
-        self._httpx_client = Client(follow_redirects=True, timeout=cdf_client.config.file_transfer_timeout)
+        self._httpx_client = Client(
+            follow_redirects=True,
+            timeout=cdf_client.config.file_transfer_timeout,
+            verify=ssl_verify,
+        )
 
         global _QUEUES, _QUEUES_LOCK
         with _QUEUES_LOCK:
             self._pool = ThreadPoolExecutor(
-                max_workers=self.parallelism, thread_name_prefix=f"FileUploadQueue-{_QUEUES}"
+                max_workers=self.parallelism,
+                thread_name_prefix=f"FileUploadQueue-{_QUEUES}",
             )
             _QUEUES += 1
 
+    def initialize_failure_logging(self) -> None:
+        self._file_failure_manager: FileFailureManager | None = (
+            FileFailureManager(path_to_file=self.failure_logging_path)
+            if self.failure_logging_path is not None
+            else None
+        )
+
+    def get_failure_logger(self) -> FileFailureManager | None:
+        return self._file_failure_manager
+
+    def add_entry_failure_logger(self, file_name: str, error: Exception) -> None:
+        if self._file_failure_manager is not None:
+            error_reason = str(error)
+            self._file_failure_manager.add(file_name=file_name, error_reason=error_reason)
+
+    def flush_failure_logger(self) -> None:
+        if self._file_failure_manager is not None:
+            self.logger.info("Flushing failure logs")
+            self._file_failure_manager.write_to_file()
+
     def _remove_done_from_queue(self) -> None:
         while not self.cancellation_token.is_cancelled:
             with self.lock:
@@ -251,7 +299,7 @@ class IOFileUploadQueue(AbstractUploadQueue):
         node = instance_result.nodes[0]
         return node.as_id()
 
-    def _upload_empty(
+    def _upload_only_metadata(
         self, file_meta: FileMetadataOrCogniteExtractorFile
     ) -> tuple[FileMetadataOrCogniteExtractorFile, str]:
         if isinstance(file_meta, CogniteExtractorFileApply):
@@ -281,11 +329,51 @@ class IOFileUploadQueue(AbstractUploadQueue):
 
         return file_meta_response, url
 
+    def _upload_empty_file(
+        self,
+        file_meta: FileMetadataOrCogniteExtractorFile,
+    ) -> None:
+        file_meta_response, url = self._upload_only_metadata(file_meta)
+
+        self._upload_only_file_reference(file_meta, url)
+
     def _upload_bytes(self, size: int, file: BinaryIO, file_meta: FileMetadataOrCogniteExtractorFile) -> None:
-        file_meta, url = self._upload_empty(file_meta)
+        file_meta, url = self._upload_only_metadata(file_meta)
         resp = self._httpx_client.send(self._get_file_upload_request(url, file, size, file_meta.mime_type))
         resp.raise_for_status()
 
+    def _prepare_request_data_for_empty_file(self, url_str: str) -> Request:
+        FILE_SIZE = 0  # this path is only entered for an empty file
+        EMPTY_CONTENT = ""
+
+        url = URL(url_str)
+        base_url = URL(self.cdf_client.config.base_url)
+
+        if url.host == base_url.host:
+            upload_url = url
+        else:
+            parsed_url: ParseResult = urlparse(url_str)
+            parsed_base_url: ParseResult = urlparse(self.cdf_client.config.base_url)
+            replaced_upload_url = parsed_url._replace(netloc=parsed_base_url.netloc).geturl()
+            upload_url = URL(replaced_upload_url)
+
+        headers = Headers(self._httpx_client.headers)
+        headers.update(
+            {
+                "Accept": "*/*",
+                "Content-Length": str(FILE_SIZE),
+                "Host": upload_url.netloc.decode("ascii"),
+                "x-cdp-app": self.cdf_client._config.client_name,
+            }
+        )
+
+        return Request(method="PUT", url=upload_url, headers=headers, content=EMPTY_CONTENT)
+
+    def _upload_only_file_reference(self, file_meta: FileMetadataOrCogniteExtractorFile, url_str: str) -> None:
+        request_data = self._prepare_request_data_for_empty_file(url_str)
+        resp = self._httpx_client.send(request_data)
+        resp.raise_for_status()
+
     def _upload_multipart(self, size: int, file: BinaryIO, file_meta: FileMetadataOrCogniteExtractorFile) -> None:
         chunks = ChunkedStream(file, self.max_file_chunk_size, size)
         self.logger.debug(
@@ -329,7 +417,10 @@ class IOFileUploadQueue(AbstractUploadQueue):
             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},
+                params={
+                    "overwrite": self.overwrite_existing,
+                    "parts": chunks.chunk_count,
+                },
             )
             res.raise_for_status()
             return res.json()
@@ -338,9 +429,7 @@ class IOFileUploadQueue(AbstractUploadQueue):
         self,
         file_meta: FileMetadataOrCogniteExtractorFile,
         read_file: Callable[[], BinaryIO],
-        extra_retries: Optional[
-            Union[Tuple[Type[Exception], ...], Dict[Type[Exception], Callable[[Any], bool]]]
-        ] = None,
+        extra_retries: tuple[Type[Exception], ...] | dict[Type[Exception], Callable[[Any], bool]] | None = None,
     ) -> None:
         """
         Add file to upload queue. The file will start uploading immedeately. If the size of the queue is larger than
@@ -366,12 +455,14 @@ class IOFileUploadQueue(AbstractUploadQueue):
             max_delay=RETRY_MAX_DELAY,
             backoff=RETRY_BACKOFF_FACTOR,
         )
-        def upload_file(read_file: Callable[[], BinaryIO], file_meta: FileMetadataOrCogniteExtractorFile) -> None:
+        def upload_file(
+            read_file: Callable[[], BinaryIO],
+            file_meta: FileMetadataOrCogniteExtractorFile,
+        ) -> None:
             with read_file() as file:
                 size = super_len(file)
                 if size == 0:
-                    # upload just the file metadata witout data
-                    file_meta, _ = self._upload_empty(file_meta)
+                    self._upload_empty_file(file_meta)
                 elif size >= self.max_single_chunk_file_size:
                     # The minimum chunk size is 4000MiB.
                     self._upload_multipart(size, file, file_meta)
@@ -388,12 +479,18 @@ class IOFileUploadQueue(AbstractUploadQueue):
                 except Exception as e:
                     self.logger.error("Error in upload callback: %s", str(e))
 
-        def wrapped_upload(read_file: Callable[[], BinaryIO], file_meta: FileMetadataOrCogniteExtractorFile) -> None:
+        def wrapped_upload(
+            read_file: Callable[[], BinaryIO],
+            file_meta: FileMetadataOrCogniteExtractorFile,
+        ) -> None:
             try:
                 upload_file(read_file, file_meta)
 
             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: {file_meta.external_id} {file_meta.name}"
+                )
+                self.add_entry_failure_logger(file_name=str(file_meta.name), error=e)
                 self.errors.append(e)
 
             finally:
@@ -416,7 +513,7 @@ class IOFileUploadQueue(AbstractUploadQueue):
             self.queue_size.set(self.upload_queue_size)
 
     def _get_file_upload_request(
-        self, url_str: str, stream: BinaryIO, size: int, mime_type: Optional[str] = None
+        self, url_str: str, stream: BinaryIO, size: int, mime_type: str | None = None
     ) -> Request:
         url = URL(url_str)
         base_url = URL(self.cdf_client.config.base_url)
@@ -460,7 +557,7 @@ class IOFileUploadQueue(AbstractUploadQueue):
         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:
+    def upload(self, fail_on_errors: bool = True, timeout: float | None = None) -> None:
         """
         Wait for all uploads to finish
         """
@@ -470,6 +567,7 @@ class IOFileUploadQueue(AbstractUploadQueue):
             self.queue_size.set(self.upload_queue_size)
         if fail_on_errors and self.errors:
             # There might be more errors, but we can only have one as the cause, so pick the first
+            self.flush_failure_logger()
             raise RuntimeError(f"{len(self.errors)} upload(s) finished with errors") from self.errors[0]
 
     def __enter__(self) -> "IOFileUploadQueue":
@@ -485,7 +583,10 @@ class IOFileUploadQueue(AbstractUploadQueue):
         return self
 
     def __exit__(
-        self, exc_type: Optional[Type[BaseException]], exc_val: Optional[BaseException], exc_tb: Optional[TracebackType]
+        self,
+        exc_type: Optional[Type[BaseException]],
+        exc_val: Optional[BaseException],
+        exc_tb: Optional[TracebackType],
     ) -> None:
         """
         Wraps around stop method, for use as context manager
@@ -524,27 +625,31 @@ class FileUploadQueue(IOFileUploadQueue):
     def __init__(
         self,
         cdf_client: CogniteClient,
-        post_upload_function: Optional[Callable[[List[FileMetadataOrCogniteExtractorFile]], None]] = None,
-        max_queue_size: Optional[int] = None,
-        max_upload_interval: Optional[int] = None,
+        post_upload_function: Callable[[list[FileMetadataOrCogniteExtractorFile]], None] | None = None,
+        max_queue_size: int | None = None,
+        max_upload_interval: int | None = None,
         trigger_log_level: str = "DEBUG",
-        thread_name: Optional[str] = None,
+        thread_name: str | None = None,
         overwrite_existing: bool = False,
-        cancellation_token: Optional[CancellationToken] = None,
+        cancellation_token: CancellationToken | None = None,
+        ssl_verify: bool | str = True,
     ):
         # Super sets post_upload and threshold
         super().__init__(
-            cdf_client,
-            post_upload_function,
-            max_queue_size,
-            trigger_log_level,
-            thread_name,
-            overwrite_existing,
-            cancellation_token,
+            cdf_client=cdf_client,
+            post_upload_function=post_upload_function,
+            max_queue_size=max_queue_size,
+            trigger_log_level=trigger_log_level,
+            thread_name=thread_name,
+            overwrite_existing=overwrite_existing,
+            cancellation_token=cancellation_token,
+            ssl_verify=ssl_verify,
         )
 
     def add_to_upload_queue(
-        self, file_meta: FileMetadataOrCogniteExtractorFile, file_name: Union[str, PathLike]
+        self,
+        file_meta: 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
@@ -579,21 +684,23 @@ class BytesUploadQueue(IOFileUploadQueue):
     def __init__(
         self,
         cdf_client: CogniteClient,
-        post_upload_function: Optional[Callable[[List[FileMetadataOrCogniteExtractorFile]], None]] = None,
-        max_queue_size: Optional[int] = None,
+        post_upload_function: Callable[[list[FileMetadataOrCogniteExtractorFile]], None] | None = None,
+        max_queue_size: int | None = None,
         trigger_log_level: str = "DEBUG",
-        thread_name: Optional[str] = None,
+        thread_name: str | None = None,
         overwrite_existing: bool = False,
-        cancellation_token: Optional[CancellationToken] = None,
+        cancellation_token: CancellationToken | None = None,
+        ssl_verify: bool | str = True,
     ) -> None:
         super().__init__(
-            cdf_client,
-            post_upload_function,
-            max_queue_size,
-            trigger_log_level,
-            thread_name,
-            overwrite_existing,
-            cancellation_token,
+            cdf_client=cdf_client,
+            post_upload_function=post_upload_function,
+            max_queue_size=max_queue_size,
+            trigger_log_level=trigger_log_level,
+            thread_name=thread_name,
+            overwrite_existing=overwrite_existing,
+            cancellation_token=cancellation_token,
+            ssl_verify=ssl_verify,
         )
 
     def add_to_upload_queue(self, content: bytes, file_meta: FileMetadataOrCogniteExtractorFile) -> None:

cognite/extractorutils/uploader/raw.py

@@ -13,7 +13,7 @@
 #  limitations under the License.
 
 from types import TracebackType
-from typing import Any, Callable, Dict, List, Optional, Type
+from typing import Any, Callable, Type
 
 import arrow
 from arrow import Arrow
@@ -56,12 +56,12 @@ class RawUploadQueue(AbstractUploadQueue):
     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,
+        post_upload_function: Callable[[list[Any]], None] | None = None,
+        max_queue_size: int | None = None,
+        max_upload_interval: int | None = None,
         trigger_log_level: str = "DEBUG",
-        thread_name: Optional[str] = None,
-        cancellation_token: Optional[CancellationToken] = None,
+        thread_name: str | None = None,
+        cancellation_token: CancellationToken | None = None,
     ):
         # Super sets post_upload and thresholds
         super().__init__(
@@ -73,7 +73,7 @@ class RawUploadQueue(AbstractUploadQueue):
             thread_name,
             cancellation_token,
         )
-        self.upload_queue: Dict[str, Dict[str, List[TimestampedObject]]] = {}
+        self.upload_queue: dict[str, dict[str, list[TimestampedObject]]] = {}
 
         # It is a hack since Prometheus client registers metrics on object creation, so object has to be created once
         self.rows_queued = RAW_UPLOADER_ROWS_QUEUED
@@ -119,7 +119,7 @@ class RawUploadQueue(AbstractUploadQueue):
             max_delay=RETRY_MAX_DELAY,
             backoff=RETRY_BACKOFF_FACTOR,
         )
-        def _upload_batch(database: str, table: str, patch: List[Row]) -> None:
+        def _upload_batch(database: str, table: str, patch: list[Row]) -> None:
             # Upload
             self.cdf_client.raw.rows.insert(db_name=database, table_name=table, row=patch, ensure_parent=True)
 
@@ -133,7 +133,7 @@ class RawUploadQueue(AbstractUploadQueue):
 
                     # Deduplicate
                     # In case of duplicate keys, the first key is preserved, and the last value is preserved.
-                    patch: Dict[str, Row] = {r.payload.key: r.payload for r in rows}
+                    patch: dict[str, Row] = {r.payload.key: r.payload for r in rows}
                     self.rows_duplicates.labels(_labels).inc(len(rows) - len(patch))
 
                     _upload_batch(database=database, table=table, patch=list(patch.values()))
@@ -162,7 +162,7 @@ class RawUploadQueue(AbstractUploadQueue):
         return self
 
     def __exit__(
-        self, exc_type: Optional[Type[BaseException]], exc_val: Optional[BaseException], exc_tb: Optional[TracebackType]
+        self, exc_type: Type[BaseException] | None, exc_val: BaseException | None, exc_tb: TracebackType | None
     ) -> None:
         """
         Wraps around stop method, for use as context manager