2610 lines
106 KiB
Python
2610 lines
106 KiB
Python
from __future__ import annotations
|
|
|
|
import base64
|
|
import datetime
|
|
import logging
|
|
import math
|
|
import os
|
|
import pathlib
|
|
import platform
|
|
import re
|
|
import shutil
|
|
import sys
|
|
import xml.etree.ElementTree as ET
|
|
from abc import ABC, abstractmethod
|
|
from collections import deque
|
|
from collections.abc import Iterator
|
|
from concurrent.futures import ThreadPoolExecutor, as_completed
|
|
from dataclasses import dataclass
|
|
from datetime import timedelta
|
|
from io import BytesIO
|
|
from queue import Empty, Full, Queue
|
|
from tempfile import mkstemp
|
|
from threading import Event, Thread
|
|
from types import TracebackType
|
|
from typing import (TYPE_CHECKING, AnyStr, BinaryIO, Callable, Generator,
|
|
Iterable, Optional, Type, Union)
|
|
from urllib import parse
|
|
|
|
import requests
|
|
import requests.adapters
|
|
from requests import RequestException
|
|
|
|
from .._base_client import _BaseClient, _RawResponse, _StreamingResponse
|
|
from .._property import _cached_property
|
|
from ..config import Config
|
|
from ..errors import AlreadyExists, InternalError, NotFound, PermissionDenied
|
|
from ..errors.mapper import _error_mapper
|
|
from ..retries import retried
|
|
from ..service import files
|
|
from ..service._internal import _escape_multi_segment_path_parameter
|
|
from ..service.files import DownloadResponse
|
|
from .files_utils import (CreateDownloadUrlResponse, _ConcatenatedInputStream,
|
|
_PresignedUrlDistributor)
|
|
|
|
if TYPE_CHECKING:
|
|
from _typeshed import Self
|
|
|
|
_LOG = logging.getLogger(__name__)
|
|
|
|
|
|
class _DbfsIO(BinaryIO):
|
|
MAX_CHUNK_SIZE = 1024 * 1024
|
|
|
|
_status: files.FileInfo = None
|
|
_created: files.CreateResponse = None
|
|
_offset = 0
|
|
_closed = False
|
|
|
|
def __init__(
|
|
self,
|
|
api: files.DbfsAPI,
|
|
path: str,
|
|
*,
|
|
read: bool = False,
|
|
write: bool = False,
|
|
overwrite: bool = False,
|
|
):
|
|
self._api = api
|
|
self._path = path
|
|
if write and read:
|
|
raise IOError(f"can open either for reading or writing")
|
|
if read:
|
|
self._status = api.get_status(path)
|
|
elif write:
|
|
self._created = api.create(path, overwrite=overwrite)
|
|
else:
|
|
raise IOError(f"need to open either for reading or writing")
|
|
|
|
def __enter__(self) -> Self:
|
|
return self
|
|
|
|
@property
|
|
def name(self) -> str:
|
|
return self._path
|
|
|
|
def writable(self) -> bool:
|
|
"""
|
|
Return whether object was opened for writing.
|
|
|
|
If False, write() will raise OSError.
|
|
"""
|
|
return self._created is not None
|
|
|
|
def write(self, buffer: bytes) -> int:
|
|
"""Write bytes to file.
|
|
|
|
:return: Return the number of bytes written.
|
|
"""
|
|
if not self.writable():
|
|
raise IOError("file not open for writing")
|
|
if type(buffer) is not bytes:
|
|
# Python doesn't strictly enforce types. Even if they're specified.
|
|
raise TypeError(f"a bytes-like object is required, not {type(buffer)}")
|
|
total = 0
|
|
while total < len(buffer):
|
|
chunk = buffer[total:]
|
|
if len(chunk) > self.MAX_CHUNK_SIZE:
|
|
chunk = chunk[: self.MAX_CHUNK_SIZE]
|
|
encoded = base64.b64encode(chunk).decode()
|
|
self._api.add_block(self._created.handle, encoded)
|
|
total += len(chunk)
|
|
return total
|
|
|
|
def close(self) -> None:
|
|
"""Disable all I/O operations."""
|
|
if self.writable():
|
|
self._api.close(self._created.handle)
|
|
self._closed = True
|
|
|
|
@property
|
|
def closed(self) -> bool:
|
|
return self._closed
|
|
|
|
def __exit__(
|
|
self,
|
|
__t: Type[BaseException] | None,
|
|
__value: BaseException | None,
|
|
__traceback: TracebackType | None,
|
|
):
|
|
self.close()
|
|
|
|
def readable(self) -> bool:
|
|
return self._status is not None
|
|
|
|
def read(self, size: int = ...) -> bytes:
|
|
"""Read at most size bytes, returned as a bytes object.
|
|
|
|
:param size: If the size argument is negative, read until EOF is reached.
|
|
Return an empty bytes object at EOF.
|
|
:return: bytes
|
|
"""
|
|
if not self.readable():
|
|
raise IOError("file not open for reading")
|
|
|
|
# call __iter__() and read until EOF is reached
|
|
if size is ... or size < 0:
|
|
buffer = b""
|
|
for chunk in self:
|
|
buffer += chunk
|
|
return buffer
|
|
|
|
response = self._api.read(self._path, length=size, offset=self._offset)
|
|
# The guard against offset >= size happens above, so this can only happen
|
|
# if the file is modified or truncated while reading. If this happens,
|
|
# the read contents will likely be corrupted, so we return an error.
|
|
if response.bytes_read == 0:
|
|
# as per Python interface convention, return an empty bytes object at EOF,
|
|
# and not the EOFError as in other SDKs
|
|
return b""
|
|
|
|
raw = base64.b64decode(response.data)
|
|
self._offset += response.bytes_read
|
|
return raw
|
|
|
|
def __iter__(self) -> Iterator[bytes]:
|
|
while self._offset < self._status.file_size:
|
|
yield self.__next__()
|
|
|
|
def __next__(self) -> bytes:
|
|
# TODO: verify semantics
|
|
return self.read(self.MAX_CHUNK_SIZE)
|
|
|
|
def fileno(self) -> int:
|
|
return 0
|
|
|
|
def flush(self) -> None:
|
|
pass
|
|
|
|
def isatty(self) -> bool:
|
|
return False
|
|
|
|
def readline(self, __limit: int = ...) -> AnyStr:
|
|
raise NotImplementedError
|
|
|
|
def readlines(self, __hint: int = ...) -> list[AnyStr]:
|
|
raise NotImplementedError
|
|
|
|
def seek(self, __offset: int, __whence: int = ...) -> int:
|
|
raise NotImplementedError
|
|
|
|
def seekable(self) -> bool:
|
|
return False
|
|
|
|
def tell(self) -> int:
|
|
return self._offset
|
|
|
|
def truncate(self, __size: int | None = ...) -> int:
|
|
raise NotImplementedError
|
|
|
|
def writelines(self, __lines: Iterable[AnyStr]) -> None:
|
|
raise NotImplementedError
|
|
|
|
def __repr__(self) -> str:
|
|
return f"<_DbfsIO {self._path} {'read' if self.readable() else 'write'}=True>"
|
|
|
|
|
|
class _VolumesIO(BinaryIO):
|
|
|
|
def __init__(
|
|
self,
|
|
api: files.FilesAPI,
|
|
path: str,
|
|
*,
|
|
read: bool,
|
|
write: bool,
|
|
overwrite: bool,
|
|
):
|
|
self._buffer = []
|
|
self._api = api
|
|
self._path = path
|
|
self._read = read
|
|
self._write = write
|
|
self._overwrite = overwrite
|
|
self._closed = False
|
|
self._read_handle = None
|
|
self._offset = 0
|
|
|
|
def __enter__(self):
|
|
if self._read:
|
|
self.__open_read()
|
|
return self
|
|
|
|
def close(self):
|
|
if self._closed:
|
|
return
|
|
if self._write:
|
|
to_write = b"".join(self._buffer)
|
|
self._api.upload(
|
|
self._path,
|
|
contents=BytesIO(to_write),
|
|
overwrite=self._overwrite,
|
|
)
|
|
elif self._read:
|
|
self._read_handle.close()
|
|
self._closed = True
|
|
|
|
def fileno(self) -> int:
|
|
return 0
|
|
|
|
def flush(self):
|
|
raise NotImplementedError()
|
|
|
|
def isatty(self) -> bool:
|
|
return False
|
|
|
|
def __check_closed(self):
|
|
if self._closed:
|
|
raise ValueError("I/O operation on closed file")
|
|
|
|
def __open_read(self):
|
|
if self._read_handle is None:
|
|
self._read_handle = self._api.download(self._path).contents
|
|
|
|
def read(self, __n=...):
|
|
self.__check_closed()
|
|
self.__open_read()
|
|
return self._read_handle.read(__n)
|
|
|
|
def readable(self):
|
|
return self._read
|
|
|
|
def readline(self, __limit=...):
|
|
raise NotImplementedError()
|
|
|
|
def readlines(self, __hint=...):
|
|
raise NotImplementedError()
|
|
|
|
def seek(self, __offset, __whence=...):
|
|
raise NotImplementedError()
|
|
|
|
def seekable(self):
|
|
return False
|
|
|
|
def tell(self):
|
|
if self._read_handle is not None:
|
|
return self._read_handle.tell()
|
|
return self._offset
|
|
|
|
def truncate(self, __size=...):
|
|
raise NotImplementedError()
|
|
|
|
def writable(self):
|
|
return self._write
|
|
|
|
def write(self, __s):
|
|
self.__check_closed()
|
|
self._buffer.append(__s)
|
|
|
|
def writelines(self, __lines):
|
|
raise NotImplementedError()
|
|
|
|
def __next__(self):
|
|
self.__check_closed()
|
|
return self._read_handle.__next__()
|
|
|
|
def __iter__(self):
|
|
self.__check_closed()
|
|
return self._read_handle.__iter__()
|
|
|
|
def __exit__(self, __t, __value, __traceback):
|
|
self.close()
|
|
|
|
def __repr__(self) -> str:
|
|
return f"<_VolumesIO {self._path} {'read' if self.readable() else 'write'}=True>"
|
|
|
|
|
|
class _Path(ABC):
|
|
|
|
@abstractmethod
|
|
def __init__(self): ...
|
|
|
|
@property
|
|
def is_local(self) -> bool:
|
|
return self._is_local()
|
|
|
|
@abstractmethod
|
|
def _is_local(self) -> bool: ...
|
|
|
|
@property
|
|
def is_dbfs(self) -> bool:
|
|
return self._is_dbfs()
|
|
|
|
@abstractmethod
|
|
def _is_dbfs(self) -> bool: ...
|
|
|
|
@abstractmethod
|
|
def child(self, path: str) -> str: ...
|
|
|
|
@_cached_property
|
|
def is_dir(self) -> bool:
|
|
return self._is_dir()
|
|
|
|
@abstractmethod
|
|
def _is_dir(self) -> bool: ...
|
|
|
|
@abstractmethod
|
|
def exists(self) -> bool: ...
|
|
|
|
@abstractmethod
|
|
def open(self, *, read=False, write=False, overwrite=False): ...
|
|
|
|
def list(self, *, recursive=False) -> Generator[files.FileInfo, None, None]: ...
|
|
|
|
@abstractmethod
|
|
def mkdir(self): ...
|
|
|
|
@abstractmethod
|
|
def delete(self, *, recursive=False): ...
|
|
|
|
@property
|
|
def name(self) -> str:
|
|
return self._path.name
|
|
|
|
@property
|
|
def as_string(self) -> str:
|
|
return str(self._path)
|
|
|
|
|
|
class _LocalPath(_Path):
|
|
|
|
def __init__(self, path: str):
|
|
if platform.system() == "Windows":
|
|
self._path = pathlib.Path(str(path).replace("file:///", "").replace("file:", ""))
|
|
else:
|
|
self._path = pathlib.Path(str(path).replace("file:", ""))
|
|
|
|
def _is_local(self) -> bool:
|
|
return True
|
|
|
|
def _is_dbfs(self) -> bool:
|
|
return False
|
|
|
|
def child(self, path: str) -> Self:
|
|
return _LocalPath(str(self._path / path))
|
|
|
|
def _is_dir(self) -> bool:
|
|
return self._path.is_dir()
|
|
|
|
def mkdir(self):
|
|
self._path.mkdir(mode=0o755, parents=True, exist_ok=True)
|
|
|
|
def exists(self) -> bool:
|
|
return self._path.exists()
|
|
|
|
def open(self, *, read=False, write=False, overwrite=False):
|
|
# make local fs follow the similar semantics as DBFS
|
|
self._path.parent.mkdir(mode=0o755, parents=True, exist_ok=True)
|
|
return self._path.open(mode="wb" if overwrite else "rb" if read else "xb")
|
|
|
|
def list(self, recursive=False) -> Generator[files.FileInfo, None, None]:
|
|
if not self.is_dir:
|
|
st = self._path.stat()
|
|
yield files.FileInfo(
|
|
path="file:" + str(self._path.absolute()),
|
|
is_dir=False,
|
|
file_size=st.st_size,
|
|
modification_time=int(st.st_mtime_ns / 1e6),
|
|
)
|
|
return
|
|
queue = deque([self._path])
|
|
while queue:
|
|
path = queue.popleft()
|
|
for leaf in path.iterdir():
|
|
if leaf.is_dir():
|
|
if recursive:
|
|
queue.append(leaf)
|
|
continue
|
|
info = leaf.stat()
|
|
yield files.FileInfo(
|
|
path="file:" + str(leaf.absolute()),
|
|
is_dir=False,
|
|
file_size=info.st_size,
|
|
modification_time=int(info.st_mtime_ns / 1e6),
|
|
)
|
|
|
|
def delete(self, *, recursive=False):
|
|
if self.is_dir:
|
|
if recursive:
|
|
for leaf in self.list(recursive=True):
|
|
_LocalPath(leaf.path).delete()
|
|
self._path.rmdir()
|
|
else:
|
|
kw = {}
|
|
if sys.version_info[:2] > (3, 7):
|
|
kw["missing_ok"] = True
|
|
self._path.unlink(**kw)
|
|
|
|
def __repr__(self) -> str:
|
|
return f"<_LocalPath {self._path}>"
|
|
|
|
|
|
class _VolumesPath(_Path):
|
|
|
|
def __init__(self, api: files.FilesAPI, src: Union[str, pathlib.Path]):
|
|
self._path = pathlib.PurePosixPath(str(src).replace("dbfs:", "").replace("file:", ""))
|
|
self._api = api
|
|
|
|
def _is_local(self) -> bool:
|
|
return False
|
|
|
|
def _is_dbfs(self) -> bool:
|
|
return False
|
|
|
|
def child(self, path: str) -> Self:
|
|
return _VolumesPath(self._api, str(self._path / path))
|
|
|
|
def _is_dir(self) -> bool:
|
|
try:
|
|
self._api.get_directory_metadata(self.as_string)
|
|
return True
|
|
except NotFound:
|
|
return False
|
|
|
|
def mkdir(self):
|
|
self._api.create_directory(self.as_string)
|
|
|
|
def exists(self) -> bool:
|
|
try:
|
|
self._api.get_metadata(self.as_string)
|
|
return True
|
|
except NotFound:
|
|
return self.is_dir
|
|
|
|
def open(self, *, read=False, write=False, overwrite=False) -> BinaryIO:
|
|
return _VolumesIO(
|
|
self._api,
|
|
self.as_string,
|
|
read=read,
|
|
write=write,
|
|
overwrite=overwrite,
|
|
)
|
|
|
|
def list(self, *, recursive=False) -> Generator[files.FileInfo, None, None]:
|
|
if not self.is_dir:
|
|
meta = self._api.get_metadata(self.as_string)
|
|
yield files.FileInfo(
|
|
path=self.as_string,
|
|
is_dir=False,
|
|
file_size=meta.content_length,
|
|
modification_time=meta.last_modified,
|
|
)
|
|
return
|
|
queue = deque([self])
|
|
while queue:
|
|
next_path = queue.popleft()
|
|
for file in self._api.list_directory_contents(next_path.as_string):
|
|
if recursive and file.is_directory:
|
|
queue.append(self.child(file.name))
|
|
if not recursive or not file.is_directory:
|
|
yield files.FileInfo(
|
|
path=file.path,
|
|
is_dir=file.is_directory,
|
|
file_size=file.file_size,
|
|
modification_time=file.last_modified,
|
|
)
|
|
|
|
def delete(self, *, recursive=False):
|
|
if self.is_dir:
|
|
for entry in self.list(recursive=False):
|
|
_VolumesPath(self._api, entry.path).delete(recursive=True)
|
|
self._api.delete_directory(self.as_string)
|
|
else:
|
|
self._api.delete(self.as_string)
|
|
|
|
def __repr__(self) -> str:
|
|
return f"<_VolumesPath {self._path}>"
|
|
|
|
|
|
class _DbfsPath(_Path):
|
|
|
|
def __init__(self, api: files.DbfsAPI, src: str):
|
|
self._path = pathlib.PurePosixPath(str(src).replace("dbfs:", "").replace("file:", ""))
|
|
self._api = api
|
|
|
|
def _is_local(self) -> bool:
|
|
return False
|
|
|
|
def _is_dbfs(self) -> bool:
|
|
return True
|
|
|
|
def child(self, path: str) -> Self:
|
|
child = self._path / path
|
|
return _DbfsPath(self._api, str(child))
|
|
|
|
def _is_dir(self) -> bool:
|
|
try:
|
|
remote = self._api.get_status(self.as_string)
|
|
return remote.is_dir
|
|
except NotFound:
|
|
return False
|
|
|
|
def mkdir(self):
|
|
self._api.mkdirs(self.as_string)
|
|
|
|
def exists(self) -> bool:
|
|
try:
|
|
self._api.get_status(self.as_string)
|
|
return True
|
|
except NotFound:
|
|
return False
|
|
|
|
def open(self, *, read=False, write=False, overwrite=False) -> BinaryIO:
|
|
return _DbfsIO(
|
|
self._api,
|
|
self.as_string,
|
|
read=read,
|
|
write=write,
|
|
overwrite=overwrite,
|
|
)
|
|
|
|
def list(self, *, recursive=False) -> Generator[files.FileInfo, None, None]:
|
|
if not self.is_dir:
|
|
meta = self._api.get_status(self.as_string)
|
|
yield files.FileInfo(
|
|
path=self.as_string,
|
|
is_dir=False,
|
|
file_size=meta.file_size,
|
|
modification_time=meta.modification_time,
|
|
)
|
|
return
|
|
queue = deque([self])
|
|
while queue:
|
|
next_path = queue.popleft()
|
|
for file in self._api.list(next_path.as_string):
|
|
if recursive and file.is_dir:
|
|
queue.append(self.child(file.path))
|
|
if not recursive or not file.is_dir:
|
|
yield file
|
|
|
|
def delete(self, *, recursive=False):
|
|
self._api.delete(self.as_string, recursive=recursive)
|
|
|
|
def __repr__(self) -> str:
|
|
return f"<_DbfsPath {self._path}>"
|
|
|
|
|
|
class _RetryableException(Exception):
|
|
"""Base class for retryable exceptions in DBFS operations."""
|
|
|
|
def __init__(self, message: str, http_status_code: int):
|
|
super().__init__()
|
|
self.message = message
|
|
self.http_status_code = http_status_code
|
|
|
|
def __str__(self) -> str:
|
|
return f"{self.message} (HTTP Status: {self.http_status_code})"
|
|
|
|
@staticmethod
|
|
def make_error(response: requests.Response) -> "_RetryableException":
|
|
"""Map the response to a retryable exception."""
|
|
|
|
return _RetryableException(
|
|
message=response.text,
|
|
http_status_code=response.status_code,
|
|
)
|
|
|
|
|
|
class DbfsExt(files.DbfsAPI):
|
|
__doc__ = files.DbfsAPI.__doc__
|
|
|
|
def __init__(self, api_client):
|
|
super().__init__(api_client)
|
|
self._files_api = files.FilesAPI(api_client)
|
|
self._dbfs_api = files.DbfsAPI(api_client)
|
|
|
|
def open(
|
|
self,
|
|
path: str,
|
|
*,
|
|
read: bool = False,
|
|
write: bool = False,
|
|
overwrite: bool = False,
|
|
) -> BinaryIO:
|
|
return self._path(path).open(read=read, write=write, overwrite=overwrite)
|
|
|
|
def upload(self, path: str, src: BinaryIO, *, overwrite: bool = False):
|
|
"""Upload file to DBFS"""
|
|
with self.open(path, write=True, overwrite=overwrite) as dst:
|
|
shutil.copyfileobj(src, dst, length=_DbfsIO.MAX_CHUNK_SIZE)
|
|
|
|
def download(self, path: str) -> BinaryIO:
|
|
"""Download file from DBFS"""
|
|
return self.open(path, read=True)
|
|
|
|
def list(self, path: str, *, recursive=False) -> Iterator[files.FileInfo]:
|
|
"""List directory contents or file details.
|
|
|
|
List the contents of a directory, or details of the file. If the file or directory does not exist,
|
|
this call throws an exception with `RESOURCE_DOES_NOT_EXIST`.
|
|
|
|
When calling list on a large directory, the list operation will time out after approximately 60
|
|
seconds.
|
|
|
|
:param path: the DBFS or UC Volume path to list
|
|
:param recursive: traverse deep into directory tree
|
|
:returns iterator of metadata for every file
|
|
"""
|
|
p = self._path(path)
|
|
yield from p.list(recursive=recursive)
|
|
|
|
def mkdirs(self, path: str):
|
|
"""Create directory on DBFS"""
|
|
p = self._path(path)
|
|
p.mkdir()
|
|
|
|
def exists(self, path: str) -> bool:
|
|
"""If file exists on DBFS"""
|
|
p = self._path(path)
|
|
return p.exists()
|
|
|
|
__ALLOWED_SCHEMES = [None, "file", "dbfs"]
|
|
|
|
def _path(self, src):
|
|
src = parse.urlparse(str(src))
|
|
if src.scheme and src.scheme not in self.__ALLOWED_SCHEMES:
|
|
raise ValueError(
|
|
f'unsupported scheme "{src.scheme}". DBUtils in the SDK only supports local, root DBFS, and '
|
|
"UC Volumes paths, not external locations or DBFS mount points."
|
|
)
|
|
if src.scheme == "file":
|
|
return _LocalPath(src.geturl())
|
|
if src.path.startswith("/Volumes"):
|
|
return _VolumesPath(self._files_api, src.geturl())
|
|
return _DbfsPath(self._dbfs_api, src.geturl())
|
|
|
|
def copy(self, src: str, dst: str, *, recursive=False, overwrite=False):
|
|
"""Copy files between DBFS and local filesystems"""
|
|
src = self._path(src)
|
|
dst = self._path(dst)
|
|
if src.is_local and dst.is_local:
|
|
raise IOError("both destinations are on local FS")
|
|
if dst.exists() and dst.is_dir:
|
|
# if target is a folder, make file with the same name there
|
|
dst = dst.child(src.name)
|
|
if src.is_dir:
|
|
queue = [self._path(x.path) for x in src.list(recursive=recursive) if not x.is_dir]
|
|
else:
|
|
queue = [src]
|
|
for child in queue:
|
|
child_dst = dst.child(os.path.relpath(child.as_string, src.as_string))
|
|
with child.open(read=True) as reader:
|
|
with child_dst.open(write=True, overwrite=overwrite) as writer:
|
|
shutil.copyfileobj(reader, writer, length=_DbfsIO.MAX_CHUNK_SIZE)
|
|
|
|
def move_(self, src: str, dst: str, *, recursive=False, overwrite=False):
|
|
"""Move files between local and DBFS systems"""
|
|
source = self._path(src)
|
|
target = self._path(dst)
|
|
if source.is_dbfs and target.is_dbfs:
|
|
# Moves a file from one location to another location within DBFS.
|
|
# this operation is recursive by default.
|
|
return self.move(source.as_string, target.as_string)
|
|
if source.is_local and target.is_local:
|
|
raise IOError("both destinations are on local FS")
|
|
if source.is_dir and not recursive:
|
|
src_type = "local" if source.is_local else "DBFS" if source.is_dbfs else "UC Volume"
|
|
dst_type = "local" if target.is_local else "DBFS" if target.is_dbfs else "UC Volume"
|
|
raise IOError(f"moving a directory from {src_type} to {dst_type} requires recursive flag")
|
|
# do cross-fs moving
|
|
self.copy(src, dst, recursive=recursive, overwrite=overwrite)
|
|
self.delete(src, recursive=recursive)
|
|
|
|
def delete(self, path: str, *, recursive=False):
|
|
"""Delete file or directory on DBFS"""
|
|
p = self._path(path)
|
|
if p.is_dir and not recursive:
|
|
raise IOError("deleting directories requires recursive flag")
|
|
p.delete(recursive=recursive)
|
|
|
|
|
|
class FallbackToUploadUsingFilesApi(Exception):
|
|
"""Custom exception that signals to fallback to FilesAPI for upload"""
|
|
|
|
def __init__(self, buffer, message):
|
|
super().__init__(message)
|
|
self.buffer = buffer
|
|
|
|
|
|
class FallbackToDownloadUsingFilesApi(Exception):
|
|
"""Custom exception that signals to fallback to FilesAPI for download"""
|
|
|
|
def __init__(self, message):
|
|
super().__init__(message)
|
|
|
|
|
|
@dataclass
|
|
class UploadStreamResult:
|
|
"""Result of an upload from stream operation. Currently empty, but can be extended in the future."""
|
|
|
|
|
|
@dataclass
|
|
class UploadFileResult:
|
|
"""Result of an upload from file operation. Currently empty, but can be extended in the future."""
|
|
|
|
|
|
@dataclass
|
|
class DownloadFileResult:
|
|
"""Result of a download to file operation. Currently empty, but can be extended in the future."""
|
|
|
|
|
|
class FilesExt(files.FilesAPI):
|
|
__doc__ = files.FilesAPI.__doc__
|
|
|
|
# note that these error codes are retryable only for idempotent operations
|
|
_RETRYABLE_STATUS_CODES: list[int] = [408, 429, 502, 503, 504]
|
|
|
|
@dataclass(frozen=True)
|
|
class _UploadContext:
|
|
target_path: str
|
|
"""The absolute remote path of the target file, e.g. /Volumes/path/to/your/file."""
|
|
overwrite: Optional[bool]
|
|
"""If true, an existing file will be overwritten. When unspecified, default behavior of the cloud storage provider is performed."""
|
|
part_size: int
|
|
"""The size of each part in bytes for multipart upload."""
|
|
batch_size: int
|
|
"""The number of urls to request in a single batch."""
|
|
content_length: Optional[int] = None
|
|
"""The total size of the content being uploaded, if known."""
|
|
source_file_path: Optional[str] = None
|
|
"""The local path of the file being uploaded, if applicable."""
|
|
use_parallel: Optional[bool] = None
|
|
"""If true, the upload will be performed using multiple threads."""
|
|
parallelism: Optional[int] = None
|
|
"""The number of threads to use for parallel upload, if applicable."""
|
|
|
|
def __init__(self, api_client, config: Config):
|
|
super().__init__(api_client)
|
|
self._config = config.copy()
|
|
self._multipart_upload_read_ahead_bytes = 1
|
|
|
|
def download(
|
|
self,
|
|
file_path: str,
|
|
) -> DownloadResponse:
|
|
"""Download a file.
|
|
|
|
Downloads a file as a stream into memory.
|
|
|
|
Use this when you want to process the downloaded file in memory or pipe it into another system. Supports files of any size in SDK v0.72.0+. Earlier versions have a 5 GB file size limit.
|
|
|
|
If the download is successful, the function returns the downloaded file result. If the download is unsuccessful, the function raises an exception.
|
|
|
|
:param file_path: str
|
|
The remote path of the file, e.g. /Volumes/path/to/your/file
|
|
|
|
:returns: :class:`DownloadResponse`
|
|
"""
|
|
if self._config.disable_experimental_files_api_client:
|
|
_LOG.info("Disable experimental files API client, will use the original download method.")
|
|
return super().download(file_path)
|
|
|
|
initial_response: DownloadResponse = self._open_download_stream(
|
|
file_path=file_path, start_byte_offset=0, if_unmodified_since_timestamp=None
|
|
)
|
|
|
|
wrapped_response = self._wrap_stream(file_path, initial_response)
|
|
initial_response.contents._response = wrapped_response
|
|
return initial_response
|
|
|
|
def download_to(
|
|
self,
|
|
file_path: str,
|
|
destination: str,
|
|
*,
|
|
overwrite: bool = True,
|
|
use_parallel: bool = False,
|
|
parallelism: Optional[int] = None,
|
|
) -> DownloadFileResult:
|
|
"""Downloads a file directly to a local file path.
|
|
|
|
Use this when you want to write the file straight to disk instead of holding it in memory. Supports files of any size in SDK v0.72.0+. Earlier versions have a 5 GB file size limit.
|
|
|
|
Supports parallel download (use_parallel=True), which may improve performance for large files. This is available on all operating systems except Windows.
|
|
|
|
:param file_path: str
|
|
The remote path of the file, e.g. /Volumes/path/to/your/file
|
|
:param destination: str
|
|
The local path where the file will be saved.
|
|
:param overwrite: bool
|
|
If true, an existing file will be overwritten. When not specified, defaults to True.
|
|
:param use_parallel: bool
|
|
If true, the download will be performed using multiple threads.
|
|
:param parallelism: int
|
|
The number of parallel threads to use for downloading. If not specified, defaults to the number of CPU cores.
|
|
|
|
:returns: :class:`DownloadFileResult`
|
|
"""
|
|
if self._config.disable_experimental_files_api_client:
|
|
raise NotImplementedError(
|
|
"Experimental files API features are disabled, download_to is not supported. Please use download instead."
|
|
)
|
|
|
|
# The existence of the target file is checked before starting the download. This is a best-effort check
|
|
# to avoid overwriting an existing file. However, there is nothing preventing a file from being created
|
|
# at the destination path after this check and before the file is written, and no way to prevent other
|
|
# actor from writing to the destination path concurrently.
|
|
if not overwrite and os.path.exists(destination):
|
|
raise FileExistsError(destination)
|
|
if use_parallel:
|
|
# Parallel download is not supported for Windows due to the limit of only one open file handle
|
|
# for writing. If parallel download is requested on Windows, fall back to sequential download with
|
|
# a warning.
|
|
if platform.system() == "Windows":
|
|
_LOG.warning("Parallel download is not supported on Windows. Falling back to sequential download.")
|
|
self._sequential_download_to_file(destination, remote_path=file_path)
|
|
return DownloadFileResult()
|
|
if parallelism is None:
|
|
parallelism = self._config.files_ext_parallel_download_default_parallelism
|
|
if parallelism < 1 or parallelism > 64:
|
|
raise ValueError("parallelism must be between 1 and 64")
|
|
self._parallel_download_with_fallback(file_path, destination, parallelism=parallelism)
|
|
else:
|
|
self._sequential_download_to_file(destination, remote_path=file_path)
|
|
return DownloadFileResult()
|
|
|
|
def _parallel_download_with_fallback(self, remote_path: str, destination: str, parallelism: int) -> None:
|
|
"""Download a file in parallel to a local path. There would be no responses returned if the download is successful.
|
|
This method first tries to use the Presigned URL for parallel download. If it fails due to permission issues,
|
|
it falls back to using Files API.
|
|
|
|
:param remote_path: str
|
|
The remote path of the file, e.g. /Volumes/path/to/your/file
|
|
:param destination: str
|
|
The local path where the file will be saved.
|
|
:param parallelism: int
|
|
The number of parallel threads to use for downloading.
|
|
|
|
:returns: None
|
|
"""
|
|
try:
|
|
self._parallel_download_presigned_url(remote_path, destination, parallelism)
|
|
except FallbackToDownloadUsingFilesApi as e:
|
|
_LOG.info("Falling back to Files API download due to permission issues with Presigned URL: %s", e)
|
|
self._parallel_download_files_api(remote_path, destination, parallelism)
|
|
|
|
def _sequential_download_to_file(
|
|
self, destination: str, remote_path: str, last_modified: Optional[str] = None
|
|
) -> None:
|
|
with open(destination, "wb") as f:
|
|
response = self._open_download_stream(
|
|
file_path=remote_path,
|
|
start_byte_offset=0,
|
|
if_unmodified_since_timestamp=last_modified,
|
|
)
|
|
wrapped_response = self._wrap_stream(remote_path, response, 0)
|
|
response.contents._response = wrapped_response
|
|
shutil.copyfileobj(response.contents, f)
|
|
|
|
def _do_parallel_download(
|
|
self, remote_path: str, destination: str, parallelism: int, download_chunk: Callable
|
|
) -> None:
|
|
|
|
file_info = self.get_metadata(remote_path)
|
|
file_size = file_info.content_length
|
|
last_modified = file_info.last_modified
|
|
# If the file is smaller than the threshold, do not use parallel download.
|
|
if file_size <= self._config.files_ext_parallel_download_min_file_size:
|
|
self._sequential_download_to_file(destination, remote_path, last_modified)
|
|
return
|
|
part_size = self._config.files_ext_parallel_download_default_part_size
|
|
part_count = int(math.ceil(file_size / part_size))
|
|
|
|
fd, temp_file = mkstemp()
|
|
# We are preallocate the file size to the same as the remote file to avoid seeking beyond the file size.
|
|
os.truncate(temp_file, file_size)
|
|
os.close(fd)
|
|
try:
|
|
aborted = Event()
|
|
|
|
def wrapped_download_chunk(start: int, end: int, last_modified: Optional[str], temp_file: str) -> None:
|
|
if aborted.is_set():
|
|
return
|
|
additional_headers = {
|
|
"Range": f"bytes={start}-{end}",
|
|
"If-Unmodified-Since": last_modified,
|
|
}
|
|
try:
|
|
contents = download_chunk(additional_headers)
|
|
with open(temp_file, "r+b") as f:
|
|
f.seek(start)
|
|
shutil.copyfileobj(contents, f)
|
|
except Exception as e:
|
|
aborted.set()
|
|
raise e
|
|
|
|
with ThreadPoolExecutor(max_workers=parallelism) as executor:
|
|
futures = []
|
|
# Start the threads to download parts of the file.
|
|
for i in range(part_count):
|
|
start = i * part_size
|
|
end = min(start + part_size - 1, file_size - 1)
|
|
futures.append(executor.submit(wrapped_download_chunk, start, end, last_modified, temp_file))
|
|
|
|
# Wait for all threads to complete and check for exceptions.
|
|
for future in as_completed(futures):
|
|
exception = future.exception()
|
|
if exception:
|
|
raise exception
|
|
# Finally, move the temp file to the destination.
|
|
shutil.move(temp_file, destination)
|
|
finally:
|
|
if os.path.exists(temp_file):
|
|
os.remove(temp_file)
|
|
|
|
def _parallel_download_presigned_url(self, remote_path: str, destination: str, parallelism: int) -> None:
|
|
"""Download a file in parallel to a local path. There would be no responses returned if the download is successful.
|
|
|
|
:param remote_path: str
|
|
The remote path of the file, e.g. /Volumes/path/to/your/file
|
|
:param destination: str
|
|
The local path where the file will be saved.
|
|
:param parallelism: int
|
|
The number of parallel threads to use for downloading.
|
|
|
|
:returns: None
|
|
"""
|
|
|
|
cloud_session = self._create_cloud_provider_session()
|
|
url_distributor = _PresignedUrlDistributor(lambda: self._create_download_url(remote_path))
|
|
# An event to indicate if any download chunk has succeeded. If any chunk succeeds, we do not fall back to Files API.
|
|
any_success = Event()
|
|
|
|
def download_chunk(additional_headers: dict[str, str]) -> BinaryIO:
|
|
retry_count = 0
|
|
while retry_count < self._config.files_ext_parallel_download_max_retries:
|
|
url_and_header, version = url_distributor.get_url()
|
|
|
|
headers = {**url_and_header.headers, **additional_headers}
|
|
|
|
def get_content() -> requests.Response:
|
|
return cloud_session.get(url_and_header.url, headers=headers)
|
|
|
|
raw_resp = self._retry_cloud_idempotent_operation(get_content)
|
|
|
|
if FilesExt._is_url_expired_response(raw_resp):
|
|
_LOG.info("Presigned URL expired, fetching a new one.")
|
|
url_distributor.invalidate_url(version)
|
|
retry_count += 1
|
|
continue
|
|
elif raw_resp.status_code == 403 and not any_success.is_set():
|
|
raise FallbackToDownloadUsingFilesApi("Received 403 Forbidden from presigned URL")
|
|
elif not any_success.is_set():
|
|
# For other errors, we raise a retryable exception to trigger retry logic.
|
|
raise FallbackToDownloadUsingFilesApi(f"Received {raw_resp.status_code} from presigned URL")
|
|
|
|
raw_resp.raise_for_status()
|
|
any_success.set()
|
|
return BytesIO(raw_resp.content)
|
|
raise ValueError("Exceeded maximum retries for downloading with presigned URL: URL expired too many times")
|
|
|
|
self._do_parallel_download(remote_path, destination, parallelism, download_chunk)
|
|
|
|
def _parallel_download_files_api(self, remote_path: str, destination: str, parallelism: int) -> None:
|
|
"""Download a file in parallel to a local path using FilesAPI. There would be no responses returned if the download is successful.
|
|
|
|
:param remote_path: str
|
|
The remote path of the file, e.g. /Volumes/path/to/your/file
|
|
:param destination: str
|
|
The local path where the file will be saved.
|
|
:param parallelism: int
|
|
The number of parallel threads to use for downloading.
|
|
|
|
:returns: None
|
|
"""
|
|
|
|
def download_chunk(additional_headers: dict[str, str]) -> BinaryIO:
|
|
raw_response: dict = self._api.do(
|
|
method="GET",
|
|
path=f"/api/2.0/fs/files{remote_path}",
|
|
headers=additional_headers,
|
|
raw=True,
|
|
)
|
|
return raw_response["contents"]
|
|
|
|
self._do_parallel_download(remote_path, destination, parallelism, download_chunk)
|
|
|
|
def _get_optimized_performance_parameters_for_upload(
|
|
self, content_length: Optional[int], part_size_overwrite: Optional[int]
|
|
) -> (int, int):
|
|
"""Get optimized part size and batch size for upload based on content length and provided part size.
|
|
|
|
Returns tuple of (part_size, batch_size).
|
|
"""
|
|
chosen_part_size = None
|
|
|
|
# 1. decide on the part size
|
|
if part_size_overwrite is not None: # If a part size is provided, we use it directly after validation.
|
|
if part_size_overwrite > self._config.files_ext_multipart_upload_max_part_size:
|
|
raise ValueError(
|
|
f"Part size {part_size_overwrite} exceeds maximum allowed size {self._config.files_ext_multipart_upload_max_part_size} bytes."
|
|
)
|
|
chosen_part_size = part_size_overwrite
|
|
_LOG.debug(f"Using provided part size: {chosen_part_size} bytes")
|
|
else: # If no part size is provided, we will optimize based on the content length.
|
|
if content_length is not None:
|
|
# Choosing the smallest part size that allows for a maximum of 100 parts.
|
|
for part_size in self._config.files_ext_multipart_upload_part_size_options:
|
|
part_num = (content_length + part_size - 1) // part_size
|
|
if part_num <= 100:
|
|
chosen_part_size = part_size
|
|
_LOG.debug(
|
|
f"Optimized part size for upload: {chosen_part_size} bytes for content length {content_length} bytes"
|
|
)
|
|
break
|
|
if chosen_part_size is None: # If no part size was chosen, we default to the maximum allowed part size.
|
|
chosen_part_size = self._config.files_ext_multipart_upload_max_part_size
|
|
|
|
# Use defaults if not determined yet
|
|
if chosen_part_size is None:
|
|
chosen_part_size = self._config.files_ext_multipart_upload_default_part_size
|
|
|
|
# 2. decide on the batch size
|
|
if content_length is not None and chosen_part_size is not None:
|
|
part_num = (content_length + chosen_part_size - 1) // chosen_part_size
|
|
chosen_batch_size = int(
|
|
math.ceil(math.sqrt(part_num))
|
|
) # Using the square root of the number of parts as a heuristic for batch size.
|
|
else:
|
|
chosen_batch_size = self._config.files_ext_multipart_upload_batch_url_count
|
|
|
|
return chosen_part_size, chosen_batch_size
|
|
|
|
def upload(
|
|
self,
|
|
file_path: str,
|
|
contents: BinaryIO,
|
|
*,
|
|
overwrite: Optional[bool] = None,
|
|
part_size: Optional[int] = None,
|
|
use_parallel: bool = True,
|
|
parallelism: Optional[int] = None,
|
|
) -> UploadStreamResult:
|
|
"""
|
|
Uploads a file from memory or a stream interface.
|
|
|
|
Use this when you want to upload data already in memory or piped from another system. Supports files of any size in SDK v0.72.0+. Earlier versions have a 5 GB file size limit.
|
|
|
|
Limitations: If the storage account is on Azure and has firewall enabled, the maximum file size is 5GB.
|
|
|
|
:param file_path: str
|
|
The absolute remote path of the target file, e.g. /Volumes/path/to/your/file
|
|
:param contents: BinaryIO
|
|
The contents of the file to upload. This must be a BinaryIO stream.
|
|
:param overwrite: bool (optional)
|
|
If true, an existing file will be overwritten. When not specified, defaults to True.
|
|
:param part_size: int (optional)
|
|
If set, multipart upload will use the value as its size per uploading part. If not set, an appropriate value will be automatically used.
|
|
:param use_parallel: bool (optional)
|
|
If true, the upload will be performed using multiple threads. Note that this will consume more memory
|
|
because multiple parts will be buffered in memory before being uploaded. The amount of memory used is proportional
|
|
to `parallelism * part_size`.
|
|
If false, the upload will be performed in a single thread.
|
|
Default is True.
|
|
:param parallelism: int (optional)
|
|
The number of threads to use for parallel uploads. This is only used if `use_parallel` is True.
|
|
|
|
:returns: :class:`UploadStreamResult`
|
|
"""
|
|
|
|
if self._config.disable_experimental_files_api_client:
|
|
_LOG.info("Disable experimental files API client, will use the original upload method.")
|
|
super().upload(file_path=file_path, contents=contents, overwrite=overwrite)
|
|
return UploadStreamResult()
|
|
|
|
_LOG.debug(f"Uploading file from BinaryIO stream")
|
|
if parallelism is not None and not use_parallel:
|
|
raise ValueError("parallelism can only be set if use_parallel is True")
|
|
if parallelism is None and use_parallel:
|
|
parallelism = self._config.files_ext_multipart_upload_default_parallelism
|
|
|
|
# Determine content length if the stream is seekable
|
|
content_length = None
|
|
if contents.seekable():
|
|
_LOG.debug(f"Uploading using seekable mode")
|
|
# If the stream is seekable, we can read its size.
|
|
contents.seek(0, os.SEEK_END)
|
|
content_length = contents.tell()
|
|
contents.seek(0)
|
|
|
|
# Get optimized part size and batch size based on content length and provided part size
|
|
optimized_part_size, optimized_batch_size = self._get_optimized_performance_parameters_for_upload(
|
|
content_length, part_size
|
|
)
|
|
|
|
# Create context with all final parameters
|
|
ctx = self._UploadContext(
|
|
target_path=file_path,
|
|
overwrite=overwrite,
|
|
part_size=optimized_part_size,
|
|
batch_size=optimized_batch_size,
|
|
content_length=content_length,
|
|
use_parallel=use_parallel,
|
|
parallelism=parallelism,
|
|
)
|
|
|
|
_LOG.debug(
|
|
f"Upload context: part_size={ctx.part_size}, batch_size={ctx.batch_size}, content_length={ctx.content_length}"
|
|
)
|
|
|
|
if ctx.use_parallel and (
|
|
ctx.content_length is None or ctx.content_length >= self._config.files_ext_multipart_upload_min_stream_size
|
|
):
|
|
self._parallel_upload_from_stream(ctx, contents)
|
|
return UploadStreamResult()
|
|
elif ctx.content_length is not None:
|
|
self._upload_single_thread_with_known_size(ctx, contents)
|
|
return UploadStreamResult()
|
|
else:
|
|
_LOG.debug(f"Uploading using non-seekable mode")
|
|
# If the stream is not seekable, we cannot determine its size.
|
|
# We will use a multipart upload.
|
|
_LOG.debug(f"Using multipart upload for non-seekable input stream of unknown size for file {file_path}")
|
|
self._single_thread_multipart_upload(ctx, contents)
|
|
return UploadStreamResult()
|
|
|
|
def upload_from(
|
|
self,
|
|
file_path: str,
|
|
source_path: str,
|
|
*,
|
|
overwrite: Optional[bool] = None,
|
|
part_size: Optional[int] = None,
|
|
use_parallel: bool = True,
|
|
parallelism: Optional[int] = None,
|
|
) -> UploadFileResult:
|
|
"""
|
|
Uploads a file from a local file path.
|
|
|
|
Use this when your data already exists on disk and you want to upload it directly without manually opening it yourself. Supports files of any size in SDK v0.72.0+. Earlier versions have a 5 GB file size limit.
|
|
|
|
:param file_path: str
|
|
The absolute remote path of the target file.
|
|
:param source_path: str
|
|
The local path of the file to upload. This must be a path to a local file.
|
|
:param part_size: int (optional)
|
|
If set, multipart upload will use the value as its size per uploading part. If not set, an appropriate default value will be automatically used.
|
|
:param overwrite: bool (optional)
|
|
If true, an existing file will be overwritten. When not specified, defaults True.
|
|
:param use_parallel: bool (optional)
|
|
If true, the upload will be performed using multiple threads. Default is True.
|
|
:param parallelism: int (optional)
|
|
The number of threads to use for parallel uploads. This is only used if `use_parallel` is True.
|
|
If not specified, the default parallelism will be set to config.multipart_upload_default_parallelism
|
|
|
|
:returns: :class:`UploadFileResult`
|
|
"""
|
|
if self._config.disable_experimental_files_api_client:
|
|
raise NotImplementedError(
|
|
"Experimental files API features are disabled, upload_from is not supported. Please use upload instead."
|
|
)
|
|
|
|
_LOG.debug(f"Uploading file from local path: {source_path}")
|
|
|
|
if parallelism is not None and not use_parallel:
|
|
raise ValueError("parallelism can only be set if use_parallel is True")
|
|
if parallelism is None and use_parallel:
|
|
parallelism = self._config.files_ext_multipart_upload_default_parallelism
|
|
# Get the file size
|
|
file_size = os.path.getsize(source_path)
|
|
|
|
# Get optimized part size and batch size based on content length and provided part size
|
|
optimized_part_size, optimized_batch_size = self._get_optimized_performance_parameters_for_upload(
|
|
file_size, part_size
|
|
)
|
|
|
|
# Create context with all final parameters
|
|
ctx = self._UploadContext(
|
|
target_path=file_path,
|
|
overwrite=overwrite,
|
|
part_size=optimized_part_size,
|
|
batch_size=optimized_batch_size,
|
|
content_length=file_size,
|
|
source_file_path=source_path,
|
|
use_parallel=use_parallel,
|
|
parallelism=parallelism,
|
|
)
|
|
if ctx.use_parallel and ctx.content_length >= self._config.files_ext_multipart_upload_min_stream_size:
|
|
self._parallel_upload_from_file(ctx)
|
|
return UploadFileResult()
|
|
else:
|
|
with open(source_path, "rb") as f:
|
|
self._upload_single_thread_with_known_size(ctx, f)
|
|
return UploadFileResult()
|
|
|
|
def _upload_single_thread_with_known_size(self, ctx: _UploadContext, contents: BinaryIO) -> None:
|
|
"""Upload a file with a known size."""
|
|
if ctx.content_length < self._config.files_ext_multipart_upload_min_stream_size:
|
|
_LOG.debug(f"Using single-shot upload for input stream of size {ctx.content_length} bytes")
|
|
return self._single_thread_single_shot_upload(ctx, contents)
|
|
else:
|
|
_LOG.debug(f"Using multipart upload for input stream of size {ctx.content_length} bytes")
|
|
return self._single_thread_multipart_upload(ctx, contents)
|
|
|
|
def _single_thread_single_shot_upload(self, ctx: _UploadContext, contents: BinaryIO) -> None:
|
|
"""Upload a file with a known size."""
|
|
_LOG.debug(f"Using single-shot upload for input stream")
|
|
return super().upload(file_path=ctx.target_path, contents=contents, overwrite=ctx.overwrite)
|
|
|
|
def _initiate_multipart_upload(self, ctx: _UploadContext) -> dict:
|
|
"""Initiate a multipart upload and return the response."""
|
|
query = {"action": "initiate-upload"}
|
|
if ctx.overwrite is not None:
|
|
query["overwrite"] = ctx.overwrite
|
|
|
|
# Method _api.do() takes care of retrying and will raise an exception in case of failure.
|
|
initiate_upload_response = self._api.do(
|
|
"POST", f"/api/2.0/fs/files{_escape_multi_segment_path_parameter(ctx.target_path)}", query=query
|
|
)
|
|
return initiate_upload_response
|
|
|
|
def _single_thread_multipart_upload(self, ctx: _UploadContext, contents: BinaryIO) -> None:
|
|
|
|
# Upload empty and small files with one-shot upload.
|
|
pre_read_buffer = contents.read(self._config.files_ext_multipart_upload_min_stream_size)
|
|
if len(pre_read_buffer) < self._config.files_ext_multipart_upload_min_stream_size:
|
|
_LOG.debug(
|
|
f"Using one-shot upload for input stream of size {len(pre_read_buffer)} below {self._config.files_ext_multipart_upload_min_stream_size} bytes"
|
|
)
|
|
return self._single_thread_single_shot_upload(ctx, BytesIO(pre_read_buffer))
|
|
|
|
# Initiate the multipart upload.
|
|
initiate_upload_response = self._initiate_multipart_upload(ctx)
|
|
|
|
if initiate_upload_response.get("multipart_upload"):
|
|
cloud_provider_session = self._create_cloud_provider_session()
|
|
session_token = initiate_upload_response["multipart_upload"].get("session_token")
|
|
if not session_token:
|
|
raise ValueError(f"Unexpected server response: {initiate_upload_response}")
|
|
|
|
try:
|
|
self._perform_multipart_upload(ctx, contents, session_token, pre_read_buffer, cloud_provider_session)
|
|
except FallbackToUploadUsingFilesApi as e:
|
|
try:
|
|
self._abort_multipart_upload(ctx, session_token, cloud_provider_session)
|
|
except BaseException as ex:
|
|
# Ignore abort exceptions as it is a best-effort.
|
|
_LOG.warning(f"Failed to abort upload: {ex}")
|
|
|
|
_LOG.info(f"Falling back to single-shot upload with Files API: {e}")
|
|
# Concatenate the buffered part and the rest of the stream.
|
|
full_stream = _ConcatenatedInputStream(BytesIO(e.buffer), contents)
|
|
return self._single_thread_single_shot_upload(ctx, full_stream)
|
|
|
|
except Exception as e:
|
|
_LOG.info(f"Aborting multipart upload on error: {e}")
|
|
try:
|
|
self._abort_multipart_upload(ctx, session_token, cloud_provider_session)
|
|
except BaseException as ex:
|
|
# Ignore abort exceptions as it is a best-effort.
|
|
_LOG.warning(f"Failed to abort upload: {ex}")
|
|
finally:
|
|
# Rethrow the original exception
|
|
raise e from None
|
|
|
|
elif initiate_upload_response.get("resumable_upload"):
|
|
cloud_provider_session = self._create_cloud_provider_session()
|
|
session_token = initiate_upload_response["resumable_upload"]["session_token"]
|
|
|
|
try:
|
|
self._perform_resumable_upload(ctx, contents, session_token, pre_read_buffer, cloud_provider_session)
|
|
except FallbackToUploadUsingFilesApi as e:
|
|
_LOG.info(f"Falling back to single-shot upload with Files API: {e}")
|
|
# Concatenate the buffered part and the rest of the stream.
|
|
full_stream = _ConcatenatedInputStream(BytesIO(e.buffer), contents)
|
|
return self._single_thread_single_shot_upload(ctx, full_stream)
|
|
else:
|
|
raise ValueError(f"Unexpected server response: {initiate_upload_response}")
|
|
|
|
def _parallel_upload_from_stream(self, ctx: _UploadContext, contents: BinaryIO) -> None:
|
|
"""
|
|
Upload a stream using multipart upload with multiple threads.
|
|
This method is not implemented in this example, but it would typically
|
|
involve creating multiple threads to upload different parts of the stream concurrently.
|
|
"""
|
|
initiate_upload_response = self._initiate_multipart_upload(ctx)
|
|
|
|
if initiate_upload_response.get("resumable_upload"):
|
|
_LOG.warning("GCP does not support parallel resumable uploads, falling back to single-threaded upload")
|
|
return self._single_thread_multipart_upload(ctx, contents)
|
|
elif initiate_upload_response.get("multipart_upload"):
|
|
session_token = initiate_upload_response["multipart_upload"].get("session_token")
|
|
cloud_provider_session = self._create_cloud_provider_session()
|
|
if not session_token:
|
|
raise ValueError(f"Unexpected server response: {initiate_upload_response}")
|
|
try:
|
|
self._parallel_multipart_upload_from_stream(ctx, session_token, contents, cloud_provider_session)
|
|
except FallbackToUploadUsingFilesApi as e:
|
|
try:
|
|
self._abort_multipart_upload(ctx, session_token, cloud_provider_session)
|
|
except Exception as abort_ex:
|
|
_LOG.warning(f"Failed to abort upload: {abort_ex}")
|
|
_LOG.info(f"Falling back to single-shot upload with Files API: {e}")
|
|
# Concatenate the buffered part and the rest of the stream.
|
|
full_stream = _ConcatenatedInputStream(BytesIO(e.buffer), contents)
|
|
return self._single_thread_single_shot_upload(ctx, full_stream)
|
|
except Exception as e:
|
|
_LOG.info(f"Aborting multipart upload on error: {e}")
|
|
try:
|
|
self._abort_multipart_upload(ctx, session_token, cloud_provider_session)
|
|
except Exception as abort_ex:
|
|
_LOG.warning(f"Failed to abort upload: {abort_ex}")
|
|
finally:
|
|
# Rethrow the original exception.
|
|
raise e from None
|
|
else:
|
|
raise ValueError(f"Unexpected server response: {initiate_upload_response}")
|
|
|
|
def _parallel_upload_from_file(
|
|
self,
|
|
ctx: _UploadContext,
|
|
) -> None:
|
|
"""
|
|
Upload a file using multipart upload with multiple threads.
|
|
This method is not implemented in this example, but it would typically
|
|
involve creating multiple threads to upload different parts of the file concurrently.
|
|
"""
|
|
|
|
initiate_upload_response = self._initiate_multipart_upload(ctx)
|
|
|
|
if initiate_upload_response.get("multipart_upload"):
|
|
cloud_provider_session = self._create_cloud_provider_session()
|
|
session_token = initiate_upload_response["multipart_upload"].get("session_token")
|
|
if not session_token:
|
|
raise ValueError(f"Unexpected server response: {initiate_upload_response}")
|
|
try:
|
|
self._parallel_multipart_upload_from_file(ctx, session_token)
|
|
except FallbackToUploadUsingFilesApi as e:
|
|
try:
|
|
self._abort_multipart_upload(ctx, session_token, cloud_provider_session)
|
|
except Exception as abort_ex:
|
|
_LOG.warning(f"Failed to abort upload: {abort_ex}")
|
|
|
|
_LOG.info(f"Falling back to single-shot upload with Files API: {e}")
|
|
# Concatenate the buffered part and the rest of the stream.
|
|
with open(ctx.source_file_path, "rb") as f:
|
|
return self._single_thread_single_shot_upload(ctx, f)
|
|
|
|
except Exception as e:
|
|
_LOG.info(f"Aborting multipart upload on error: {e}")
|
|
try:
|
|
self._abort_multipart_upload(ctx, session_token, cloud_provider_session)
|
|
except Exception as abort_ex:
|
|
_LOG.warning(f"Failed to abort upload: {abort_ex}")
|
|
finally:
|
|
# Rethrow the original exception.
|
|
raise e from None
|
|
|
|
elif initiate_upload_response.get("resumable_upload"):
|
|
_LOG.warning("GCP does not support parallel resumable uploads, falling back to single-threaded upload")
|
|
with open(ctx.source_file_path, "rb") as f:
|
|
return self._upload_single_thread_with_known_size(ctx, f)
|
|
else:
|
|
raise ValueError(f"Unexpected server response: {initiate_upload_response}")
|
|
|
|
@dataclass
|
|
class _MultipartUploadPart:
|
|
ctx: FilesExt._UploadContext
|
|
part_index: int
|
|
part_offset: int
|
|
part_size: int
|
|
session_token: str
|
|
|
|
def _parallel_multipart_upload_from_file(
|
|
self,
|
|
ctx: _UploadContext,
|
|
session_token: str,
|
|
) -> None:
|
|
# Calculate the number of parts.
|
|
file_size = os.path.getsize(ctx.source_file_path)
|
|
part_size = ctx.part_size
|
|
num_parts = (file_size + part_size - 1) // part_size
|
|
_LOG.debug(f"Uploading file of size {file_size} bytes in {num_parts} parts using {ctx.parallelism} threads")
|
|
cloud_provider_session = self._create_cloud_provider_session()
|
|
|
|
# Upload one part to verify the upload can proceed.
|
|
with open(ctx.source_file_path, "rb") as f:
|
|
f.seek(0)
|
|
first_part_size = min(part_size, file_size)
|
|
first_part_buffer = f.read(first_part_size)
|
|
try:
|
|
etag = self._do_upload_one_part(
|
|
ctx,
|
|
cloud_provider_session,
|
|
1,
|
|
0,
|
|
first_part_size,
|
|
session_token,
|
|
BytesIO(first_part_buffer),
|
|
is_first_part=True,
|
|
)
|
|
except FallbackToUploadUsingFilesApi as e:
|
|
raise FallbackToUploadUsingFilesApi(None, "Falling back to single-shot upload with Files API") from e
|
|
if num_parts == 1:
|
|
self._complete_multipart_upload(ctx, {1: etag}, session_token)
|
|
return
|
|
|
|
# Create queues and worker threads.
|
|
task_queue = Queue()
|
|
etags_result_queue = Queue()
|
|
etags_result_queue.put_nowait((1, etag))
|
|
exception_queue = Queue()
|
|
aborted = Event()
|
|
workers = [
|
|
Thread(
|
|
target=self._upload_file_consumer,
|
|
args=(cloud_provider_session, task_queue, etags_result_queue, exception_queue, aborted),
|
|
)
|
|
for _ in range(ctx.parallelism)
|
|
]
|
|
_LOG.debug(f"Starting {len(workers)} worker threads for parallel upload")
|
|
|
|
# Enqueue all parts. Since the task queue is populated before starting the workers, we don't need to signal completion.
|
|
for part_index in range(2, num_parts + 1):
|
|
part_offset = (part_index - 1) * part_size
|
|
part_size = min(part_size, file_size - part_offset)
|
|
part = self._MultipartUploadPart(ctx, part_index, part_offset, part_size, session_token)
|
|
task_queue.put(part)
|
|
|
|
# Start the worker threads for parallel upload.
|
|
for worker in workers:
|
|
worker.start()
|
|
|
|
# Wait for all tasks to be processed.
|
|
for worker in workers:
|
|
worker.join()
|
|
|
|
# Check for exceptions: if any worker encountered an exception, raise the first one.
|
|
if not exception_queue.empty():
|
|
first_exception = exception_queue.get()
|
|
raise first_exception
|
|
|
|
# Collect results from the etags queue.
|
|
etags: dict = {}
|
|
while not etags_result_queue.empty():
|
|
part_number, etag = etags_result_queue.get()
|
|
etags[part_number] = etag
|
|
|
|
self._complete_multipart_upload(ctx, etags, session_token)
|
|
|
|
def _parallel_multipart_upload_from_stream(
|
|
self,
|
|
ctx: _UploadContext,
|
|
session_token: str,
|
|
content: BinaryIO,
|
|
cloud_provider_session: requests.Session,
|
|
) -> None:
|
|
|
|
task_queue = Queue(maxsize=ctx.parallelism) # Limit queue size to control memory usage
|
|
etags_result_queue = Queue()
|
|
exception_queue = Queue()
|
|
all_produced = Event()
|
|
aborted = Event()
|
|
|
|
# Do the first part read ahead
|
|
pre_read_buffer = content.read(ctx.part_size)
|
|
if not pre_read_buffer:
|
|
raise FallbackToUploadUsingFilesApi(
|
|
b"", "Falling back to single-shot upload with Files API due to empty input stream"
|
|
)
|
|
try:
|
|
etag = self._do_upload_one_part(
|
|
ctx,
|
|
cloud_provider_session,
|
|
1,
|
|
0,
|
|
len(pre_read_buffer),
|
|
session_token,
|
|
BytesIO(pre_read_buffer),
|
|
is_first_part=True,
|
|
)
|
|
etags_result_queue.put((1, etag))
|
|
except FallbackToUploadUsingFilesApi as e:
|
|
raise FallbackToUploadUsingFilesApi(
|
|
pre_read_buffer, "Falling back to single-shot upload with Files API"
|
|
) from e
|
|
|
|
if len(pre_read_buffer) < ctx.part_size:
|
|
self._complete_multipart_upload(ctx, {1: etag}, session_token)
|
|
return
|
|
|
|
def producer() -> None:
|
|
part_index = 2
|
|
part_size = ctx.part_size
|
|
while not aborted.is_set():
|
|
part_content = content.read(part_size)
|
|
if not part_content:
|
|
break
|
|
part_offset = (part_index - 1) * part_size
|
|
part = self._MultipartUploadPart(ctx, part_index, part_offset, len(part_content), session_token)
|
|
while not aborted.is_set():
|
|
try:
|
|
task_queue.put((part, part_content), timeout=0.1)
|
|
break
|
|
except Full:
|
|
continue
|
|
part_index += 1
|
|
all_produced.set()
|
|
|
|
producer_thread = Thread(target=producer)
|
|
consumers = [
|
|
Thread(
|
|
target=self._upload_stream_consumer,
|
|
args=(task_queue, etags_result_queue, exception_queue, all_produced, aborted),
|
|
)
|
|
for _ in range(ctx.parallelism)
|
|
]
|
|
_LOG.debug(f"Starting {len(consumers)} worker threads for parallel upload")
|
|
# Start producer and consumer threads
|
|
producer_thread.start()
|
|
for consumer in consumers:
|
|
consumer.start()
|
|
|
|
# Wait for producer to finish
|
|
_LOG.debug(f"threads started, waiting for producer to finish")
|
|
producer_thread.join()
|
|
# Wait for all tasks to be processed
|
|
_LOG.debug(f"producer finished, waiting for consumers to finish")
|
|
# task_queue.join()
|
|
for consumer in consumers:
|
|
consumer.join()
|
|
|
|
# Check for exceptions: if any worker encountered an exception, raise the first one.
|
|
if not exception_queue.empty():
|
|
first_exception = exception_queue.get()
|
|
raise first_exception
|
|
|
|
# Collect results from the etags queue
|
|
etags: dict = {}
|
|
while not etags_result_queue.empty():
|
|
part_number, etag = etags_result_queue.get()
|
|
etags[part_number] = etag
|
|
|
|
self._complete_multipart_upload(ctx, etags, session_token)
|
|
|
|
def _complete_multipart_upload(self, ctx, etags, session_token):
|
|
query = {"action": "complete-upload", "upload_type": "multipart", "session_token": session_token}
|
|
headers = {"Content-Type": "application/json"}
|
|
body: dict = {}
|
|
parts = []
|
|
for part_number, etag in sorted(etags.items()):
|
|
part = {"part_number": part_number, "etag": etag}
|
|
parts.append(part)
|
|
body["parts"] = parts
|
|
self._api.do(
|
|
"POST",
|
|
f"/api/2.0/fs/files{_escape_multi_segment_path_parameter(ctx.target_path)}",
|
|
query=query,
|
|
headers=headers,
|
|
body=body,
|
|
)
|
|
|
|
def _upload_file_consumer(
|
|
self,
|
|
cloud_provider_session: requests.Session,
|
|
task_queue: Queue[FilesExt._MultipartUploadPart],
|
|
etags_queue: Queue[tuple[int, str]],
|
|
exception_queue: Queue[Exception],
|
|
aborted: Event,
|
|
) -> None:
|
|
while not aborted.is_set():
|
|
try:
|
|
part = task_queue.get(block=False)
|
|
except Empty:
|
|
# The task_queue was populated before the workers were started, so we can exit if it's empty.
|
|
break
|
|
|
|
try:
|
|
with open(part.ctx.source_file_path, "rb") as f:
|
|
f.seek(part.part_offset, os.SEEK_SET)
|
|
part_content = BytesIO(f.read(part.part_size))
|
|
etag = self._do_upload_one_part(
|
|
part.ctx,
|
|
cloud_provider_session,
|
|
part.part_index,
|
|
part.part_offset,
|
|
part.part_size,
|
|
part.session_token,
|
|
part_content,
|
|
)
|
|
etags_queue.put((part.part_index, etag))
|
|
except Exception as e:
|
|
aborted.set()
|
|
exception_queue.put(e)
|
|
finally:
|
|
task_queue.task_done()
|
|
|
|
def _upload_stream_consumer(
|
|
self,
|
|
task_queue: Queue[tuple[FilesExt._MultipartUploadPart, bytes]],
|
|
etags_queue: Queue[tuple[int, str]],
|
|
exception_queue: Queue[Exception],
|
|
all_produced: Event,
|
|
aborted: Event,
|
|
) -> None:
|
|
cloud_provider_session = self._create_cloud_provider_session()
|
|
while not aborted.is_set():
|
|
try:
|
|
part, content = task_queue.get(block=False, timeout=0.1)
|
|
except Empty:
|
|
if all_produced.is_set():
|
|
break # No more parts will be produced and the queue is empty
|
|
else:
|
|
continue
|
|
try:
|
|
etag = self._do_upload_one_part(
|
|
part.ctx,
|
|
cloud_provider_session,
|
|
part.part_index,
|
|
part.part_offset,
|
|
part.part_size,
|
|
part.session_token,
|
|
BytesIO(content),
|
|
)
|
|
etags_queue.put((part.part_index, etag))
|
|
except Exception as e:
|
|
aborted.set()
|
|
exception_queue.put(e)
|
|
finally:
|
|
task_queue.task_done()
|
|
|
|
def _do_upload_one_part(
|
|
self,
|
|
ctx: _UploadContext,
|
|
cloud_provider_session: requests.Session,
|
|
part_index: int,
|
|
part_offset: int,
|
|
part_size: int,
|
|
session_token: str,
|
|
part_content: BinaryIO,
|
|
is_first_part: bool = False,
|
|
) -> str:
|
|
retry_count = 0
|
|
|
|
# Try to upload the part, retrying if the upload URL expires.
|
|
while True:
|
|
body: dict = {
|
|
"path": ctx.target_path,
|
|
"session_token": session_token,
|
|
"start_part_number": part_index,
|
|
"count": 1,
|
|
"expire_time": self._get_upload_url_expire_time(),
|
|
}
|
|
|
|
headers = {"Content-Type": "application/json"}
|
|
|
|
# Requesting URLs for the same set of parts is an idempotent operation and is safe to retry.
|
|
try:
|
|
# The _api.do() method handles retries and will raise an exception in case of failure.
|
|
upload_part_urls_response = self._api.do(
|
|
"POST", "/api/2.0/fs/create-upload-part-urls", headers=headers, body=body
|
|
)
|
|
except Exception as e:
|
|
if is_first_part:
|
|
raise FallbackToUploadUsingFilesApi(
|
|
None,
|
|
f"Failed to obtain upload URL for part {part_index}: {e}, falling back to single shot upload",
|
|
)
|
|
else:
|
|
raise e
|
|
|
|
upload_part_urls = upload_part_urls_response.get("upload_part_urls", [])
|
|
if len(upload_part_urls) == 0:
|
|
raise ValueError(f"Unexpected server response: {upload_part_urls_response}")
|
|
upload_part_url = upload_part_urls[0]
|
|
url = upload_part_url["url"]
|
|
required_headers = upload_part_url.get("headers", [])
|
|
assert part_index == upload_part_url["part_number"]
|
|
|
|
headers: dict = {"Content-Type": "application/octet-stream"}
|
|
for h in required_headers:
|
|
headers[h["name"]] = h["value"]
|
|
|
|
_LOG.debug(f"Uploading part {part_index}: [{part_offset}, {part_offset + part_size - 1}]")
|
|
|
|
def rewind() -> None:
|
|
part_content.seek(0, os.SEEK_SET)
|
|
|
|
def perform_upload() -> requests.Response:
|
|
return cloud_provider_session.request(
|
|
"PUT",
|
|
url,
|
|
headers=headers,
|
|
data=part_content,
|
|
timeout=self._config.files_ext_network_transfer_inactivity_timeout_seconds,
|
|
)
|
|
|
|
upload_response = self._retry_cloud_idempotent_operation(perform_upload, rewind)
|
|
|
|
if upload_response.status_code in (200, 201):
|
|
etag = upload_response.headers.get("ETag", "")
|
|
return etag
|
|
elif FilesExt._is_url_expired_response(upload_response):
|
|
if retry_count < self._config.files_ext_multipart_upload_max_retries:
|
|
retry_count += 1
|
|
_LOG.debug("Upload URL expired, retrying...")
|
|
continue
|
|
else:
|
|
raise ValueError(f"Unsuccessful chunk upload: upload URL expired after {retry_count} retries")
|
|
elif upload_response.status_code == 403 and is_first_part:
|
|
raise FallbackToUploadUsingFilesApi(None, f"Direct upload forbidden: {upload_response.content}")
|
|
elif is_first_part:
|
|
message = f"Unsuccessful chunk upload. Response status: {upload_response.status_code}, body: {upload_response.content}"
|
|
raise FallbackToUploadUsingFilesApi(None, message)
|
|
else:
|
|
message = f"Unsuccessful chunk upload. Response status: {upload_response.status_code}, body: {upload_response.content}"
|
|
_LOG.warning(message)
|
|
mapped_error = _error_mapper(upload_response, {})
|
|
raise mapped_error or ValueError(message)
|
|
|
|
def _perform_multipart_upload(
|
|
self,
|
|
ctx: _UploadContext,
|
|
input_stream: BinaryIO,
|
|
session_token: str,
|
|
pre_read_buffer: bytes,
|
|
cloud_provider_session: requests.Session,
|
|
) -> None:
|
|
"""
|
|
Performs multipart upload using presigned URLs on AWS and Azure:
|
|
https://docs.aws.amazon.com/AmazonS3/latest/userguide/mpuoverview.html
|
|
"""
|
|
current_part_number = 1
|
|
etags: dict = {}
|
|
|
|
# Why are we buffering the current chunk?
|
|
# AWS and Azure don't support traditional "Transfer-encoding: chunked", so we must
|
|
# provide each chunk size up front. In case of a non-seekable input stream we need
|
|
# to buffer a chunk before uploading to know its size. This also allows us to rewind
|
|
# the stream before retrying on request failure.
|
|
# AWS signed chunked upload: https://docs.aws.amazon.com/AmazonS3/latest/API/sigv4-streaming.html
|
|
# https://learn.microsoft.com/en-us/azure/storage/blobs/storage-blobs-tune-upload-download-python#buffering-during-uploads
|
|
|
|
chunk_offset = 0
|
|
|
|
# This buffer is expected to contain at least multipart_upload_chunk_size bytes.
|
|
# Note that initially buffer can be bigger (from pre_read_buffer).
|
|
buffer = pre_read_buffer
|
|
|
|
retry_count = 0
|
|
eof = False
|
|
while not eof:
|
|
# If needed, buffer the next chunk.
|
|
buffer = FilesExt._fill_buffer(buffer, ctx.part_size, input_stream)
|
|
if len(buffer) == 0:
|
|
# End of stream, no need to request the next block of upload URLs.
|
|
break
|
|
|
|
_LOG.debug(
|
|
f"Multipart upload: requesting next {ctx.batch_size} upload URLs starting from part {current_part_number}"
|
|
)
|
|
|
|
body: dict = {
|
|
"path": ctx.target_path,
|
|
"session_token": session_token,
|
|
"start_part_number": current_part_number,
|
|
"count": ctx.batch_size,
|
|
"expire_time": self._get_upload_url_expire_time(),
|
|
}
|
|
|
|
headers = {"Content-Type": "application/json"}
|
|
|
|
# Requesting URLs for the same set of parts is an idempotent operation, safe to retry.
|
|
try:
|
|
# Method _api.do() takes care of retrying and will raise an exception in case of failure.
|
|
upload_part_urls_response = self._api.do(
|
|
"POST", "/api/2.0/fs/create-upload-part-urls", headers=headers, body=body
|
|
)
|
|
except Exception as e:
|
|
if chunk_offset == 0:
|
|
raise FallbackToUploadUsingFilesApi(
|
|
buffer, f"Failed to obtain upload URLs: {e}, falling back to single shot upload"
|
|
) from e
|
|
else:
|
|
raise e
|
|
|
|
upload_part_urls = upload_part_urls_response.get("upload_part_urls", [])
|
|
if len(upload_part_urls) == 0:
|
|
raise ValueError(f"Unexpected server response: {upload_part_urls_response}")
|
|
|
|
for upload_part_url in upload_part_urls:
|
|
buffer = FilesExt._fill_buffer(buffer, ctx.part_size, input_stream)
|
|
actual_buffer_length = len(buffer)
|
|
if actual_buffer_length == 0:
|
|
eof = True
|
|
break
|
|
|
|
url = upload_part_url["url"]
|
|
required_headers = upload_part_url.get("headers", [])
|
|
assert current_part_number == upload_part_url["part_number"]
|
|
|
|
headers: dict = {"Content-Type": "application/octet-stream"}
|
|
for h in required_headers:
|
|
headers[h["name"]] = h["value"]
|
|
|
|
actual_chunk_length = min(actual_buffer_length, ctx.part_size)
|
|
_LOG.debug(
|
|
f"Uploading part {current_part_number}: [{chunk_offset}, {chunk_offset + actual_chunk_length - 1}]"
|
|
)
|
|
|
|
chunk = BytesIO(buffer[:actual_chunk_length])
|
|
|
|
def rewind():
|
|
chunk.seek(0, os.SEEK_SET)
|
|
|
|
def perform():
|
|
return cloud_provider_session.request(
|
|
"PUT",
|
|
url,
|
|
headers=headers,
|
|
data=chunk,
|
|
timeout=self._config.files_ext_network_transfer_inactivity_timeout_seconds,
|
|
)
|
|
|
|
upload_response = self._retry_cloud_idempotent_operation(perform, rewind)
|
|
|
|
if upload_response.status_code in (200, 201):
|
|
# Chunk upload successful
|
|
|
|
chunk_offset += actual_chunk_length
|
|
|
|
etag = upload_response.headers.get("ETag", "")
|
|
etags[current_part_number] = etag
|
|
|
|
# Discard uploaded bytes
|
|
buffer = buffer[actual_chunk_length:]
|
|
|
|
# Reset retry count when progressing along the stream
|
|
retry_count = 0
|
|
|
|
elif FilesExt._is_url_expired_response(upload_response):
|
|
if retry_count < self._config.files_ext_multipart_upload_max_retries:
|
|
retry_count += 1
|
|
_LOG.debug("Upload URL expired")
|
|
# Preserve the buffer so we'll upload the current part again using next upload URL
|
|
else:
|
|
# don't confuse user with unrelated "Permission denied" error.
|
|
raise ValueError(f"Unsuccessful chunk upload: upload URL expired")
|
|
|
|
elif upload_response.status_code == 403 and chunk_offset == 0:
|
|
# We got 403 failure when uploading the very first chunk (we can't tell if it is Azure for sure yet).
|
|
# This might happen due to Azure firewall enabled for the customer bucket.
|
|
# Let's fallback to using Files API which might be allowlisted to upload, passing
|
|
# currently buffered (but not yet uploaded) part of the stream.
|
|
raise FallbackToUploadUsingFilesApi(buffer, f"Direct upload forbidden: {upload_response.content}")
|
|
elif chunk_offset == 0:
|
|
# We got an upload failure when uploading the very first chunk.
|
|
# Let's fallback to using Files API which might be more reliable in this case,
|
|
# passing currently buffered (but not yet uploaded) part of the stream.
|
|
raise FallbackToUploadUsingFilesApi(
|
|
buffer,
|
|
f"Unsuccessful chunk upload: {upload_response.status_code}, falling back to single shot upload",
|
|
)
|
|
else:
|
|
message = f"Unsuccessful chunk upload. Response status: {upload_response.status_code}, body: {upload_response.content}"
|
|
_LOG.warning(message)
|
|
mapped_error = _error_mapper(upload_response, {})
|
|
raise mapped_error or ValueError(message)
|
|
|
|
current_part_number += 1
|
|
|
|
_LOG.debug(f"Completing multipart upload after uploading {len(etags)} parts of up to {ctx.part_size} bytes")
|
|
self._complete_multipart_upload(ctx, etags, session_token)
|
|
|
|
@staticmethod
|
|
def _fill_buffer(buffer: bytes, desired_min_size: int, input_stream: BinaryIO) -> bytes:
|
|
"""
|
|
Tries to fill given buffer to contain at least `desired_min_size` bytes by reading from input stream.
|
|
"""
|
|
bytes_to_read = max(0, desired_min_size - len(buffer))
|
|
if bytes_to_read > 0:
|
|
next_buf = input_stream.read(bytes_to_read)
|
|
new_buffer = buffer + next_buf
|
|
return new_buffer
|
|
else:
|
|
# we have already buffered enough data
|
|
return buffer
|
|
|
|
@staticmethod
|
|
def _is_url_expired_response(response: requests.Response) -> bool:
|
|
"""
|
|
Checks if response matches one of the known "URL expired" responses from the cloud storage providers.
|
|
"""
|
|
if response.status_code != 403:
|
|
return False
|
|
|
|
try:
|
|
xml_root = ET.fromstring(response.content)
|
|
if xml_root.tag != "Error":
|
|
return False
|
|
|
|
code = xml_root.find("Code")
|
|
if code is None:
|
|
return False
|
|
|
|
if code.text == "AuthenticationFailed":
|
|
# Azure
|
|
details = xml_root.find("AuthenticationErrorDetail")
|
|
if details is not None and "Signature not valid in the specified time frame" in details.text:
|
|
return True
|
|
|
|
if code.text == "AccessDenied":
|
|
# AWS
|
|
message = xml_root.find("Message")
|
|
if message is not None and message.text == "Request has expired":
|
|
return True
|
|
|
|
except ET.ParseError:
|
|
pass
|
|
|
|
return False
|
|
|
|
def _is_presigned_urls_disabled_error(self, e: PermissionDenied) -> bool:
|
|
error_infos = e.get_error_info()
|
|
for error_info in error_infos:
|
|
if error_info.reason == "FILES_API_API_IS_NOT_ENABLED":
|
|
return True
|
|
return False
|
|
|
|
def _is_presigned_urls_network_zone_error(self, e: InternalError) -> bool:
|
|
error_infos = e.get_error_info()
|
|
for error_info in error_infos:
|
|
if error_info.reason == "FILES_API_REQUESTER_NETWORK_ZONE_UNKNOWN":
|
|
return True
|
|
return False
|
|
|
|
def _perform_resumable_upload(
|
|
self,
|
|
ctx: _UploadContext,
|
|
input_stream: BinaryIO,
|
|
session_token: str,
|
|
pre_read_buffer: bytes,
|
|
cloud_provider_session: requests.Session,
|
|
) -> None:
|
|
"""
|
|
Performs resumable upload on GCP: https://cloud.google.com/storage/docs/performing-resumable-uploads
|
|
"""
|
|
|
|
# Session URI we're using expires after a week
|
|
|
|
# Why are we buffering the current chunk?
|
|
# When using resumable upload API we're uploading data in chunks. During chunk upload
|
|
# server responds with the "received offset" confirming how much data it stored so far,
|
|
# so we should continue uploading from that offset. (Note this is not a failure but an
|
|
# expected behaviour as per the docs.) But, input stream might be consumed beyond that
|
|
# offset, since server might have read more data than it confirmed received, or some data
|
|
# might have been pre-cached by e.g. OS or a proxy. So, to continue upload, we must rewind
|
|
# the input stream back to the byte next to "received offset". This is not possible
|
|
# for non-seekable input stream, so we must buffer the whole last chunk and seek inside
|
|
# the buffer. By always uploading from the buffer we fully support non-seekable streams.
|
|
|
|
# Why are we doing read-ahead?
|
|
# It's not possible to upload an empty chunk as "Content-Range" header format does not
|
|
# support this. So if current chunk happens to finish exactly at the end of the stream,
|
|
# we need to know that and mark the chunk as last (by passing real file size in the
|
|
# "Content-Range" header) when uploading it. To detect if we're at the end of the stream
|
|
# we're reading "ahead" an extra bytes but not uploading them immediately. If
|
|
# nothing has been read ahead, it means we're at the end of the stream.
|
|
# On the contrary, in multipart upload we can decide to complete upload *after*
|
|
# last chunk has been sent.
|
|
|
|
body: dict = {"path": ctx.target_path, "session_token": session_token}
|
|
|
|
headers = {"Content-Type": "application/json"}
|
|
|
|
try:
|
|
# Method _api.do() takes care of retrying and will raise an exception in case of failure.
|
|
resumable_upload_url_response = self._api.do(
|
|
"POST", "/api/2.0/fs/create-resumable-upload-url", headers=headers, body=body
|
|
)
|
|
except Exception as e:
|
|
raise FallbackToUploadUsingFilesApi(
|
|
pre_read_buffer, f"Failed to obtain resumable upload URL: {e}, falling back to single shot upload"
|
|
) from e
|
|
|
|
resumable_upload_url_node = resumable_upload_url_response.get("resumable_upload_url")
|
|
if not resumable_upload_url_node:
|
|
raise ValueError(f"Unexpected server response: {resumable_upload_url_response}")
|
|
|
|
resumable_upload_url = resumable_upload_url_node.get("url")
|
|
if not resumable_upload_url:
|
|
raise ValueError(f"Unexpected server response: {resumable_upload_url_response}")
|
|
|
|
required_headers = resumable_upload_url_node.get("headers", [])
|
|
base_headers: dict = {}
|
|
for h in required_headers:
|
|
base_headers[h["name"]] = h["value"]
|
|
|
|
try:
|
|
# We will buffer this many bytes: one chunk + read-ahead block.
|
|
# Note buffer may contain more data initially (from pre_read_buffer).
|
|
min_buffer_size = ctx.part_size + self._multipart_upload_read_ahead_bytes
|
|
|
|
buffer = pre_read_buffer
|
|
|
|
# How many bytes in the buffer were confirmed to be received by the server.
|
|
# All the remaining bytes in the buffer must be uploaded.
|
|
uploaded_bytes_count = 0
|
|
|
|
chunk_offset = 0
|
|
|
|
retry_count = 0
|
|
while True:
|
|
# If needed, fill the buffer to contain at least min_buffer_size bytes
|
|
# (unless end of stream), discarding already uploaded bytes.
|
|
bytes_to_read = max(0, min_buffer_size - (len(buffer) - uploaded_bytes_count))
|
|
next_buf = input_stream.read(bytes_to_read)
|
|
buffer = buffer[uploaded_bytes_count:] + next_buf
|
|
|
|
if len(next_buf) < bytes_to_read:
|
|
# This is the last chunk in the stream.
|
|
# Let's upload all the remaining bytes in one go.
|
|
actual_chunk_length = len(buffer)
|
|
file_size = chunk_offset + actual_chunk_length
|
|
else:
|
|
# More chunks expected, let's upload current chunk (excluding read-ahead block).
|
|
actual_chunk_length = ctx.part_size
|
|
file_size = "*"
|
|
|
|
headers: dict = {"Content-Type": "application/octet-stream", **base_headers}
|
|
|
|
chunk_last_byte_offset = chunk_offset + actual_chunk_length - 1
|
|
content_range_header = f"bytes {chunk_offset}-{chunk_last_byte_offset}/{file_size}"
|
|
_LOG.debug(f"Uploading chunk: {content_range_header}")
|
|
headers["Content-Range"] = content_range_header
|
|
|
|
def retrieve_upload_status() -> Optional[requests.Response]:
|
|
def perform():
|
|
return cloud_provider_session.request(
|
|
"PUT",
|
|
resumable_upload_url,
|
|
headers={"Content-Range": "bytes */*"},
|
|
data=b"",
|
|
timeout=self._config.files_ext_network_transfer_inactivity_timeout_seconds,
|
|
)
|
|
|
|
try:
|
|
return self._retry_cloud_idempotent_operation(perform)
|
|
except RequestException:
|
|
_LOG.warning("Failed to retrieve upload status")
|
|
return None
|
|
|
|
try:
|
|
upload_response = cloud_provider_session.request(
|
|
"PUT",
|
|
resumable_upload_url,
|
|
headers=headers,
|
|
data=BytesIO(buffer[:actual_chunk_length]),
|
|
timeout=self._config.files_ext_network_transfer_inactivity_timeout_seconds,
|
|
)
|
|
|
|
# https://cloud.google.com/storage/docs/performing-resumable-uploads#resume-upload
|
|
# If an upload request is terminated before receiving a response, or if you receive
|
|
# a 503 or 500 response, then you need to resume the interrupted upload from where it left off.
|
|
|
|
# Let's follow that for all potentially retryable status codes.
|
|
if upload_response.status_code in self._RETRYABLE_STATUS_CODES:
|
|
if retry_count < self._config.files_ext_multipart_upload_max_retries:
|
|
retry_count += 1
|
|
# let original upload_response be handled as an error
|
|
upload_response = retrieve_upload_status() or upload_response
|
|
else:
|
|
# we received non-retryable response, reset retry count
|
|
retry_count = 0
|
|
|
|
except RequestException as e:
|
|
# Let's do the same for retryable network errors.
|
|
if (
|
|
_BaseClient._is_retryable(e)
|
|
and retry_count < self._config.files_ext_multipart_upload_max_retries
|
|
):
|
|
retry_count += 1
|
|
upload_response = retrieve_upload_status()
|
|
if not upload_response:
|
|
# rethrow original exception
|
|
raise e from None
|
|
else:
|
|
# rethrow original exception
|
|
raise e from None
|
|
|
|
if upload_response.status_code in (200, 201):
|
|
if file_size == "*":
|
|
raise ValueError(
|
|
f"Received unexpected status {upload_response.status_code} before reaching end of stream"
|
|
)
|
|
|
|
# upload complete
|
|
break
|
|
|
|
elif upload_response.status_code == 308:
|
|
# chunk accepted (or check-status succeeded), let's determine received offset to resume from there
|
|
range_string = upload_response.headers.get("Range")
|
|
confirmed_offset = self._extract_range_offset(range_string)
|
|
_LOG.debug(f"Received confirmed offset: {confirmed_offset}")
|
|
|
|
if confirmed_offset:
|
|
if confirmed_offset < chunk_offset - 1 or confirmed_offset > chunk_last_byte_offset:
|
|
raise ValueError(
|
|
f"Unexpected received offset: {confirmed_offset} is outside of expected range, chunk offset: {chunk_offset}, chunk last byte offset: {chunk_last_byte_offset}"
|
|
)
|
|
else:
|
|
if chunk_offset > 0:
|
|
raise ValueError(
|
|
f"Unexpected received offset: {confirmed_offset} is outside of expected range, chunk offset: {chunk_offset}, chunk last byte offset: {chunk_last_byte_offset}"
|
|
)
|
|
|
|
# We have just uploaded a part of chunk starting from offset "chunk_offset" and ending
|
|
# at offset "confirmed_offset" (inclusive), so the next chunk will start at
|
|
# offset "confirmed_offset + 1"
|
|
if confirmed_offset:
|
|
next_chunk_offset = confirmed_offset + 1
|
|
else:
|
|
next_chunk_offset = chunk_offset
|
|
uploaded_bytes_count = next_chunk_offset - chunk_offset
|
|
chunk_offset = next_chunk_offset
|
|
|
|
elif upload_response.status_code == 412 and not ctx.overwrite:
|
|
# Assuming this is only possible reason
|
|
# Full message in this case: "At least one of the pre-conditions you specified did not hold."
|
|
raise AlreadyExists("The file being created already exists.")
|
|
|
|
else:
|
|
message = f"Unsuccessful chunk upload. Response status: {upload_response.status_code}, body: {upload_response.content}"
|
|
_LOG.warning(message)
|
|
mapped_error = _error_mapper(upload_response, {})
|
|
raise mapped_error or ValueError(message)
|
|
|
|
except Exception as e:
|
|
_LOG.info(f"Aborting resumable upload on error: {e}")
|
|
try:
|
|
self._abort_resumable_upload(resumable_upload_url, base_headers, cloud_provider_session)
|
|
except BaseException as ex:
|
|
_LOG.warning(f"Failed to abort upload: {ex}")
|
|
# ignore, abort is a best-effort
|
|
finally:
|
|
# rethrow original exception
|
|
raise e from None
|
|
|
|
@staticmethod
|
|
def _extract_range_offset(range_string: Optional[str]) -> Optional[int]:
|
|
"""Parses the response range header to extract the last byte."""
|
|
if not range_string:
|
|
return None # server did not yet confirm any bytes
|
|
|
|
if match := re.match("bytes=0-(\\d+)", range_string):
|
|
return int(match.group(1))
|
|
else:
|
|
raise ValueError(f"Cannot parse response header: Range: {range_string}")
|
|
|
|
def _get_rfc339_timestamp_with_future_offset(self, base_time: datetime.datetime, offset: timedelta) -> str:
|
|
"""Generates an offset timestamp in an RFC3339 format suitable for URL generation"""
|
|
offset_timestamp = base_time + offset
|
|
# From Google Protobuf doc:
|
|
# In JSON format, the Timestamp type is encoded as a string in the
|
|
# * [RFC 3339](https://www.ietf.org/rfc/rfc3339.txt) format. That is, the
|
|
# * format is "{year}-{month}-{day}T{hour}:{min}:{sec}[.{frac_sec}]Z"
|
|
return offset_timestamp.strftime("%Y-%m-%dT%H:%M:%SZ")
|
|
|
|
def _get_upload_url_expire_time(self) -> str:
|
|
"""Generates expiration time in the required format."""
|
|
current_time = datetime.datetime.now(datetime.timezone.utc)
|
|
return self._get_rfc339_timestamp_with_future_offset(
|
|
current_time, self._config.files_ext_multipart_upload_url_expiration_duration
|
|
)
|
|
|
|
def _get_download_url_expire_time(self) -> str:
|
|
"""Generates expiration time in the required format."""
|
|
current_time = datetime.datetime.now(datetime.timezone.utc)
|
|
return self._get_rfc339_timestamp_with_future_offset(
|
|
current_time, self._config.files_ext_presigned_download_url_expiration_duration
|
|
)
|
|
|
|
def _abort_multipart_upload(
|
|
self, ctx: _UploadContext, session_token: str, cloud_provider_session: requests.Session
|
|
) -> None:
|
|
"""Aborts ongoing multipart upload session to clean up incomplete file."""
|
|
body: dict = {
|
|
"path": ctx.target_path,
|
|
"session_token": session_token,
|
|
"expire_time": self._get_upload_url_expire_time(),
|
|
}
|
|
|
|
headers = {"Content-Type": "application/json"}
|
|
|
|
# Method _api.do() takes care of retrying and will raise an exception in case of failure.
|
|
abort_url_response = self._api.do("POST", "/api/2.0/fs/create-abort-upload-url", headers=headers, body=body)
|
|
|
|
abort_upload_url_node = abort_url_response["abort_upload_url"]
|
|
abort_url = abort_upload_url_node["url"]
|
|
required_headers = abort_upload_url_node.get("headers", [])
|
|
|
|
headers: dict = {"Content-Type": "application/octet-stream"}
|
|
for h in required_headers:
|
|
headers[h["name"]] = h["value"]
|
|
|
|
def perform() -> requests.Response:
|
|
return cloud_provider_session.request(
|
|
"DELETE",
|
|
abort_url,
|
|
headers=headers,
|
|
data=b"",
|
|
timeout=self._config.files_ext_network_transfer_inactivity_timeout_seconds,
|
|
)
|
|
|
|
abort_response = self._retry_cloud_idempotent_operation(perform)
|
|
|
|
if abort_response.status_code not in (200, 201):
|
|
raise ValueError(abort_response)
|
|
|
|
def _abort_resumable_upload(
|
|
self, resumable_upload_url: str, headers: dict[str, str], cloud_provider_session: requests.Session
|
|
) -> None:
|
|
"""Aborts ongoing resumable upload session to clean up incomplete file."""
|
|
|
|
def perform() -> requests.Response:
|
|
return cloud_provider_session.request(
|
|
"DELETE",
|
|
resumable_upload_url,
|
|
headers=headers,
|
|
data=b"",
|
|
timeout=self._config.files_ext_network_transfer_inactivity_timeout_seconds,
|
|
)
|
|
|
|
abort_response = self._retry_cloud_idempotent_operation(perform)
|
|
|
|
if abort_response.status_code not in (200, 201):
|
|
raise ValueError(abort_response)
|
|
|
|
def _create_cloud_provider_session(self) -> requests.Session:
|
|
"""Creates a separate session which does not inherit auth headers from BaseClient session."""
|
|
session = requests.Session()
|
|
|
|
# following session config in _BaseClient
|
|
http_adapter = requests.adapters.HTTPAdapter(
|
|
self._config.max_connection_pools or 20, self._config.max_connections_per_pool or 20, pool_block=True
|
|
)
|
|
session.mount("https://", http_adapter)
|
|
# presigned URL for storage proxy can use plain HTTP
|
|
session.mount("http://", http_adapter)
|
|
return session
|
|
|
|
def _retry_cloud_idempotent_operation(
|
|
self, operation: Callable[[], requests.Response], before_retry: Optional[Callable] = None
|
|
) -> requests.Response:
|
|
"""Perform given idempotent operation with necessary retries for requests to non Databricks APIs.
|
|
For cloud APIs, we will retry on network errors and on server response codes.
|
|
Since operation is idempotent it's safe to retry it for response codes where server state might have changed.
|
|
"""
|
|
|
|
def delegate() -> requests.Response:
|
|
response = operation()
|
|
if response.status_code in self._RETRYABLE_STATUS_CODES:
|
|
raise _RetryableException.make_error(response)
|
|
else:
|
|
return response
|
|
|
|
def extended_is_retryable(e: BaseException) -> Optional[str]:
|
|
retry_reason_from_base = _BaseClient._is_retryable(e)
|
|
if retry_reason_from_base is not None:
|
|
return retry_reason_from_base
|
|
|
|
if isinstance(e, _RetryableException):
|
|
# this is a retriable exception, but not a network error
|
|
return f"retryable exception (status_code:{e.http_status_code})"
|
|
return None
|
|
|
|
# following _BaseClient timeout
|
|
retry_timeout_seconds = self._config.retry_timeout_seconds or 300
|
|
|
|
return retried(
|
|
timeout=timedelta(seconds=retry_timeout_seconds),
|
|
max_attempts=self._config.experimental_files_ext_cloud_api_max_retries,
|
|
# also retry on network errors (connection error, connection timeout)
|
|
# where we believe request didn't reach the server
|
|
is_retryable=extended_is_retryable,
|
|
before_retry=before_retry,
|
|
clock=self._config.clock,
|
|
)(delegate)()
|
|
|
|
def _open_download_stream(
|
|
self,
|
|
file_path: str,
|
|
start_byte_offset: int,
|
|
if_unmodified_since_timestamp: Optional[str] = None,
|
|
) -> DownloadResponse:
|
|
"""Opens a download stream from given offset, performing necessary retries."""
|
|
headers = {
|
|
"Accept": "application/octet-stream",
|
|
}
|
|
|
|
if start_byte_offset and not if_unmodified_since_timestamp:
|
|
raise Exception("if_unmodified_since_timestamp is required if start_byte_offset is specified")
|
|
|
|
if start_byte_offset > 0:
|
|
headers["Range"] = f"bytes={start_byte_offset}-"
|
|
|
|
if if_unmodified_since_timestamp:
|
|
headers["If-Unmodified-Since"] = if_unmodified_since_timestamp
|
|
|
|
response_headers = [
|
|
"content-length",
|
|
"content-type",
|
|
"last-modified",
|
|
]
|
|
|
|
result = self._init_download_response_mode_csp_with_fallback(file_path, headers, response_headers)
|
|
|
|
if not isinstance(result.contents, _StreamingResponse):
|
|
raise Exception(
|
|
"Internal error: response contents is of unexpected type: " + type(result.contents).__name__
|
|
)
|
|
|
|
return result
|
|
|
|
def _init_download_response_files_api(
|
|
self, file_path: str, headers: dict[str, str], response_headers: list[str]
|
|
) -> DownloadResponse:
|
|
"""
|
|
Initiates a download response using the Files API.
|
|
"""
|
|
|
|
# Method _api.do() takes care of retrying and will raise an exception in case of failure.
|
|
res = self._api.do(
|
|
"GET",
|
|
f"/api/2.0/fs/files{_escape_multi_segment_path_parameter(file_path)}",
|
|
headers=headers,
|
|
response_headers=response_headers,
|
|
raw=True,
|
|
)
|
|
return DownloadResponse.from_dict(res)
|
|
|
|
def _create_download_url(self, file_path: str) -> CreateDownloadUrlResponse:
|
|
"""
|
|
Creates a presigned download URL using the CSP presigned URL API.
|
|
|
|
Wrapped in similar retry logic to the internal API.do call:
|
|
1. Call _.api.do to obtain the presigned URL
|
|
2. Return the presigned URL
|
|
"""
|
|
|
|
# Method _api.do() takes care of retrying and will raise an exception in case of failure.
|
|
try:
|
|
raw_response = self._api.do(
|
|
"POST",
|
|
f"/api/2.0/fs/create-download-url",
|
|
query={
|
|
"path": file_path,
|
|
"expire_time": self._get_download_url_expire_time(),
|
|
},
|
|
)
|
|
|
|
return CreateDownloadUrlResponse.from_dict(raw_response)
|
|
except Exception as e:
|
|
raise FallbackToDownloadUsingFilesApi(f"Failed to create download URL: {e}") from e
|
|
|
|
def _init_download_response_presigned_api(self, file_path: str, added_headers: dict[str, str]) -> DownloadResponse:
|
|
"""
|
|
Initiates a download response using the CSP presigned URL API.
|
|
|
|
Wrapped in similar retry logic to the internal API.do call:
|
|
1. Call _.api.do to obtain the presigned URL
|
|
2. Attempt to establish a streaming connection via the presigned URL
|
|
3. Construct a StreamingResponse from the presigned URL
|
|
"""
|
|
|
|
url_and_headers = self._create_download_url(file_path)
|
|
cloud_provider_session = self._create_cloud_provider_session()
|
|
|
|
header_overlap = added_headers.keys() & url_and_headers.headers.keys()
|
|
if header_overlap:
|
|
raise ValueError(
|
|
f"Provided headers overlap with required headers from the CSP API bundle: {header_overlap}"
|
|
)
|
|
|
|
merged_headers = {**added_headers, **url_and_headers.headers}
|
|
|
|
def perform() -> requests.Response:
|
|
return cloud_provider_session.request(
|
|
"GET",
|
|
url_and_headers.url,
|
|
headers=merged_headers,
|
|
timeout=self._config.files_ext_network_transfer_inactivity_timeout_seconds,
|
|
stream=True,
|
|
)
|
|
|
|
csp_response: _RawResponse = self._retry_cloud_idempotent_operation(perform)
|
|
|
|
# Mapping the error if the response is not successful.
|
|
if csp_response.status_code in (200, 201, 206):
|
|
resp = DownloadResponse(
|
|
content_length=int(csp_response.headers.get("content-length")),
|
|
content_type=csp_response.headers.get("content-type"),
|
|
last_modified=csp_response.headers.get("last-modified"),
|
|
contents=_StreamingResponse(csp_response, self._config.files_ext_client_download_streaming_chunk_size),
|
|
)
|
|
return resp
|
|
else:
|
|
message = (
|
|
f"Unsuccessful download. Response status: {csp_response.status_code}, body: {csp_response.content}"
|
|
)
|
|
raise FallbackToDownloadUsingFilesApi(message)
|
|
|
|
def _init_download_response_mode_csp_with_fallback(
|
|
self, file_path: str, headers: dict[str, str], response_headers: list[str]
|
|
) -> DownloadResponse:
|
|
"""
|
|
Initiates a download response using the CSP presigned URL API or the Files API, depending on the configuration.
|
|
If the CSP presigned download API is enabled, it will attempt to use that first.
|
|
If the CSP API call fails, it will fall back to the Files API.
|
|
If the CSP presigned download API is disabled, it will use the Files API directly.
|
|
"""
|
|
|
|
try:
|
|
_LOG.debug(f"Attempting download of {file_path} via CSP APIs")
|
|
return self._init_download_response_presigned_api(file_path, headers)
|
|
except FallbackToDownloadUsingFilesApi as e:
|
|
_LOG.info(f"Falling back to download via Files API: {e}")
|
|
_LOG.debug(f"Attempt via CSP APIs for {file_path} failed. Falling back to download via Files API")
|
|
ret = self._init_download_response_files_api(file_path, headers, response_headers)
|
|
return ret
|
|
|
|
def _wrap_stream(
|
|
self,
|
|
file_path: str,
|
|
download_response: DownloadResponse,
|
|
start_byte_offset: int = 0,
|
|
) -> "_ResilientResponse":
|
|
underlying_response = _ResilientIterator._extract_raw_response(download_response)
|
|
return _ResilientResponse(
|
|
self,
|
|
file_path,
|
|
download_response.last_modified,
|
|
offset=start_byte_offset,
|
|
underlying_response=underlying_response,
|
|
)
|
|
|
|
|
|
class _ResilientResponse(_RawResponse):
|
|
|
|
def __init__(
|
|
self,
|
|
api: FilesExt,
|
|
file_path: str,
|
|
file_last_modified: str,
|
|
offset: int,
|
|
underlying_response: _RawResponse,
|
|
) -> None:
|
|
self.api = api
|
|
self.file_path = file_path
|
|
self.underlying_response = underlying_response
|
|
self.offset = offset
|
|
self.file_last_modified = file_last_modified
|
|
|
|
def iter_content(self, chunk_size: int = 1, decode_unicode: bool = False) -> Iterator[bytes]:
|
|
if decode_unicode:
|
|
raise ValueError("Decode unicode is not supported")
|
|
|
|
iterator = self.underlying_response.iter_content(chunk_size=chunk_size, decode_unicode=False)
|
|
self.iterator = _ResilientIterator(
|
|
iterator, self.file_path, self.file_last_modified, self.offset, self.api, chunk_size
|
|
)
|
|
return self.iterator
|
|
|
|
def close(self) -> None:
|
|
self.iterator.close()
|
|
|
|
|
|
class _ResilientIterator(Iterator):
|
|
# This class tracks current offset (returned to the client code)
|
|
# and recovers from failures by requesting download from the current offset.
|
|
|
|
@staticmethod
|
|
def _extract_raw_response(
|
|
download_response: DownloadResponse,
|
|
) -> _RawResponse:
|
|
streaming_response: _StreamingResponse = download_response.contents
|
|
return streaming_response._response
|
|
|
|
def __init__(
|
|
self,
|
|
underlying_iterator: Iterator[bytes],
|
|
file_path: str,
|
|
file_last_modified: str,
|
|
offset: int,
|
|
api: FilesExt,
|
|
chunk_size: int,
|
|
) -> None:
|
|
self._underlying_iterator = underlying_iterator
|
|
self._api = api
|
|
self._file_path = file_path
|
|
|
|
# Absolute current offset (0-based), i.e. number of bytes from the beginning of the file
|
|
# that were so far returned to the caller code.
|
|
self._offset = offset
|
|
self._file_last_modified = file_last_modified
|
|
self._chunk_size = chunk_size
|
|
|
|
self._total_recovers_count: int = 0
|
|
self._recovers_without_progressing_count: int = 0
|
|
self._closed: bool = False
|
|
|
|
def _should_recover(self) -> bool:
|
|
if self._total_recovers_count == self._api._config.files_ext_client_download_max_total_recovers:
|
|
_LOG.debug("Total recovers limit exceeded")
|
|
return False
|
|
if (
|
|
self._api._config.files_ext_client_download_max_total_recovers_without_progressing is not None
|
|
and self._recovers_without_progressing_count
|
|
>= self._api._config.files_ext_client_download_max_total_recovers_without_progressing
|
|
):
|
|
_LOG.debug("No progression recovers limit exceeded")
|
|
return False
|
|
return True
|
|
|
|
def _recover(self) -> bool:
|
|
if not self._should_recover():
|
|
return False # recover suppressed, rethrow original exception
|
|
|
|
self._total_recovers_count += 1
|
|
self._recovers_without_progressing_count += 1
|
|
|
|
try:
|
|
self._underlying_iterator.close()
|
|
|
|
_LOG.debug(f"Trying to recover from offset {self._offset}")
|
|
|
|
# following call includes all the required network retries
|
|
downloadResponse = self._api._open_download_stream(self._file_path, self._offset, self._file_last_modified)
|
|
underlying_response = _ResilientIterator._extract_raw_response(downloadResponse)
|
|
self._underlying_iterator = underlying_response.iter_content(
|
|
chunk_size=self._chunk_size, decode_unicode=False
|
|
)
|
|
_LOG.debug("Recover succeeded")
|
|
return True
|
|
except:
|
|
return False # recover failed, rethrow original exception
|
|
|
|
def __next__(self) -> bytes:
|
|
if self._closed:
|
|
# following _BaseClient
|
|
raise ValueError("I/O operation on closed file")
|
|
|
|
while True:
|
|
try:
|
|
returned_bytes = next(self._underlying_iterator)
|
|
self._offset += len(returned_bytes)
|
|
self._recovers_without_progressing_count = 0
|
|
return returned_bytes
|
|
|
|
except StopIteration:
|
|
raise
|
|
|
|
# https://requests.readthedocs.io/en/latest/user/quickstart/#errors-and-exceptions
|
|
except RequestException:
|
|
if not self._recover():
|
|
raise
|
|
|
|
def close(self) -> None:
|
|
self._underlying_iterator.close()
|
|
self._closed = True
|