diff --git a/docs/filesystem.md b/docs/filesystem.md index 91a631b6c..fe2040e58 100644 --- a/docs/filesystem.md +++ b/docs/filesystem.md @@ -86,6 +86,49 @@ fs.cp("s3://YOUR_S3_BUCKET/src", "s3://YOUR_S3_BUCKET/dst") fs.rm("s3://YOUR_S3_BUCKET/path/to/", recursive=True) ``` +`cat`, `cat_ranges`, `get`, `put`, `pipe` with a dict of values, `cp`, and `mv` transfer +their objects in parallel, `batch_size` objects at a time: + +```python +fs.get("s3://YOUR_S3_BUCKET/path/to/", "local/dir/", recursive=True, batch_size=16) +``` + +`batch_size` is an argument of these methods, of `rm`, and of `S3FileSystem`. By default, +and with `-1`, it is the filesystem's `max_workers`. The transfers of a batch run in +threads of their own. With several transfers at once, each object sends at most +`max(1, max_workers // n)` range reads, parts, or part copies at once, `n` being the +smaller of `batch_size` and the number of objects, so that a batch keeps about +`max_workers` requests in the request pool and as many parts in memory; a smaller +`max_workers` given to the method is kept. A single object keeps the filesystem's +`max_workers`. `cat`, `cat_ranges`, and a small `pipe` send their single request from the +transfer's thread, so a `batch_size` larger than `max_workers` sends that many requests +at once. + +When a transfer fails, no further transfer starts, the retries of the running ones stop, +and the error is raised once they have ended. A running multipart upload or copy sends +no further part and is aborted instead of being completed, and a running download reads +and writes no further block. A request already sent, such as the PutObject of a small +object, runs to its end; a multipart copy that has been completed still copies its +annotations. With `cache_type="background"`, a block that fsspec's cache prefetches in +its own thread may still be read after the stop. The transfers that ended before the failure stay done, and their number is +not bounded by `batch_size`. `mv` deletes the sources only once every copy has +succeeded. `on_error` of `cat`, `cat_ranges`, and `cp` keeps its fsspec meaning; with +`on_error="omit"` or `"return"` of `cat`, `"return"` of `cat_ranges`, or `"ignore"` of +`cp` for a missing source, those errors do not stop the batch. + +Pairs whose result would depend on the order of the transfers raise `ValueError` before +anything is transferred: two sources of `get`, `put`, `cp`, or `pipe` with the same +destination, including two paths of one existing local file or one file-like object, +and a `cp` destination that is another source. Local paths that do not exist yet are +compared as strings, after `os.path.normcase`, so on a case-insensitive filesystem other +than Windows, such as the default one of macOS, give `get` destinations that differ by +more than case. A `cp` of an object onto itself, such as with +`MetadataDirective="REPLACE"`, is no conflict. `get` with `outfile` writes the objects +into it one after another, in order. + +`rm` deletes up to 1,000 objects per DeleteObjects request, `batch_size` requests at a +time. Its `max_workers` argument is a deprecated alias of `batch_size`. + Writes with `pipe`/`pipe_file` issue a single PutObject request for data up to the block size (5 MiB by default); larger data is uploaded as a parallel multipart upload through the buffered file path. Inside an @@ -128,7 +171,8 @@ cancels the range reads that have not started, and the running ones make no furt attempt. The interrupt or the error is raised once they have ended, and their data is released. A second interrupt stops that wait, and the reads still finish. `rm` and `clear_multipart_uploads` send every request also after one fails; an interrupt cancels -those that have not started and is raised once the running ones have ended. Closing +those that have not started and is raised once the running ones have ended. An interrupt of +a batch operation, such as `get` or `cp`, does the same as a failed transfer. Closing the file waits for its running requests; an interrupt of that wait stops it, and the requests still finish. An interrupt while a request is being submitted can let one more request than `max_workers` run at once diff --git a/pyathena/filesystem/s3.py b/pyathena/filesystem/s3.py index 5368911fa..af67f7fd6 100644 --- a/pyathena/filesystem/s3.py +++ b/pyathena/filesystem/s3.py @@ -8,11 +8,14 @@ import logging import mimetypes import os.path +import threading import time +import warnings from collections.abc import Callable, Iterable, Mapping from copy import deepcopy from dataclasses import dataclass, field from datetime import datetime +from glob import has_magic from io import BytesIO from typing import Any, BinaryIO, cast from urllib.parse import unquote_plus @@ -27,7 +30,7 @@ from fsspec.core import get_compression from fsspec.implementations.local import LocalFileSystem, make_path_posix, trailing_sep from fsspec.spec import AbstractBufferedFile -from fsspec.utils import check_contained, isfilelike, tokenize +from fsspec.utils import check_contained, isfilelike, other_paths, tokenize import pyathena from pyathena.connection import Connection @@ -53,7 +56,7 @@ from pyathena.filesystem.s3_path import S3Path from pyathena.filesystem.s3_path_pairing import S3PathPairing from pyathena.filesystem.s3_writer import _S3WriteSession -from pyathena.util import RetryConfig, override +from pyathena.util import RetryConfig, override, raise_if_stopped, unstoppable_retries _logger = logging.getLogger(__name__) @@ -109,6 +112,109 @@ def compress(cls, value: bytes | bytearray | memoryview, compression: str) -> by return buffer.getvalue() +class _S3LockedCallback(Callback): + """A progress callback whose uses, and those of its branches, are serialized. + + The per-object calls of a batch operation run in parallel threads and + update the callback of the operation and the branches that it creates. + fsspec does not require a callback to be thread-safe, nor its branches to + be independent of it, so every use goes through one lock. + + Args: + callback: The callback of the operation, or a branch of it. + lock: The lock of the operation. + """ + + def __init__(self, callback: Callback, lock: threading.Lock) -> None: + super().__init__() + self._callback = callback + self._lock = lock + + def __enter__(self) -> _S3LockedCallback: + """Enter the callback. + + Returns: + This callback, or a proxy of the callback that the wrapped one + returned from ``__enter__()``, if that is another one. + """ + with self._lock: + entered = self._callback.__enter__() + if entered is self._callback: + return self + return _S3LockedCallback(entered, self._lock) + + def __exit__(self, *exc_args: Any) -> Any: + """Exit the callback. + + Args: + *exc_args: The exception of the block, if any. + + Returns: + What the callback's ``__exit__()`` returns. + """ + with self._lock: + return self._callback.__exit__(*exc_args) + + def close(self) -> None: + """Close the callback.""" + with self._lock: + self._callback.close() + + def branched(self, path_1, path_2, **kwargs) -> Callback: + """Return a branch of the callback for the transfer of one file. + + Args: + path_1: The source of the transfer. + path_2: The destination of the transfer. + **kwargs: Passed to the callback's ``branched()``. + + Returns: + The branch, whose uses take the lock of this callback. + """ + with self._lock: + return _S3LockedCallback(self._callback.branched(path_1, path_2, **kwargs), self._lock) + + def set_size(self, size) -> None: + """Set the size of the callback. + + Args: + size: The total. + """ + with self._lock: + self._callback.set_size(size) + + def absolute_update(self, value) -> None: + """Set the progress of the callback. + + Args: + value: The progress. + """ + with self._lock: + self._callback.absolute_update(value) + + def relative_update(self, inc=1) -> None: + """Advance the progress of the callback. + + Args: + inc: The increment. + """ + with self._lock: + self._callback.relative_update(inc) + + def call(self, hook_name=None, **kwargs) -> Any: + """Run the hooks of the callback. + + Args: + hook_name: The hook to run, or None for all of them. + **kwargs: Passed to the hooks. + + Returns: + What the callback's ``call()`` returns. + """ + with self._lock: + return self._callback.call(hook_name, **kwargs) + + class S3FileSystem(AbstractFileSystem): """A filesystem interface for Amazon S3 that implements the fsspec protocol. @@ -180,6 +286,7 @@ def __init__( allow_bucket_deletion: bool = False, version_aware: bool = False, *args, + batch_size: int | None = None, **kwargs, ) -> None: """Create a filesystem for Amazon S3. @@ -223,12 +330,20 @@ def __init__( version_aware: Whether reads pin the object version observed at open time. *args: Passed to ``fsspec.AbstractFileSystem``. + batch_size: The number of per-object calls of a batch operation + (``cat``, ``get``, ``put``, ``copy``, and the like) that run at + once, and of DeleteObjects requests of ``rm``, unless the call + gives its own. None or -1 is ``max_workers``. **kwargs: Passed to ``fsspec.AbstractFileSystem``; without a ``connection``, also s3fs-compatible client arguments. ``requester_pays=True`` sends requester-pays requests with the operations that accept ``RequestPayer``. + + Raises: + ValueError: If ``batch_size`` is neither positive, -1, nor None. """ super().__init__(*args, **kwargs) + self.batch_size = self._check_batch_size(batch_size) if connection: client = connection.s3_client retry_config = connection.retry_config @@ -1148,7 +1263,7 @@ def rm_file(self, path: str, **kwargs) -> None: self.core.delete_object(s3_path, **self.core.operation_params("delete_object", kwargs)) self.invalidate_cache(path) - def rm(self, path, recursive=False, maxdepth=None, **kwargs) -> None: + def rm(self, path, recursive=False, maxdepth=None, batch_size=None, **kwargs) -> None: """Delete objects with DeleteObjects requests. Expands the paths with ``expand_path`` and deletes the matched objects @@ -1160,15 +1275,31 @@ def rm(self, path, recursive=False, maxdepth=None, **kwargs) -> None: path: S3 path (s3://bucket/key) or list of paths to delete. recursive: Whether to delete all objects below the paths. maxdepth: Maximum depth to expand when ``recursive`` is True. + batch_size: Number of DeleteObjects requests to send at the same + time. By default the ``batch_size`` of the filesystem; None or + -1 there is ``max_workers``. ``max_workers`` is a deprecated + alias. **kwargs: Additional parameters passed to the DeleteObjects API. ``Quiet`` (default True) sets the quiet mode of the requests. Raises: - ValueError: If a path is a bucket. + ValueError: If a path is a bucket, or ``batch_size`` is neither + positive, -1, nor None. + TypeError: If both ``batch_size`` and ``max_workers`` are given. OSError: If S3 could not delete some of the objects. """ + if "max_workers" in kwargs: + if batch_size is not None: + raise TypeError("rm() got both batch_size and its alias max_workers.") + warnings.warn( + "The max_workers argument of rm() is deprecated; use batch_size.", + DeprecationWarning, + stacklevel=2, + ) + batch_size = kwargs.pop("max_workers") + max_workers = self._batch_limit(batch_size) paths = self._delete_paths(path, recursive=recursive, maxdepth=maxdepth) - self._delete_objects(paths, **kwargs) + self._delete_objects(paths, max_workers=max_workers, **kwargs) def _delete_paths( self, path: str | list[str], recursive: bool = False, maxdepth: int | None = None @@ -1246,6 +1377,105 @@ def _map( finally: executor.shutdown(wait=False) + @staticmethod + def _check_batch_size(batch_size: int | None) -> int | None: + """Check a ``batch_size`` argument. + + Args: + batch_size: The argument. + + Returns: + The argument. + + Raises: + ValueError: If it is neither positive, -1, nor None. + """ + if batch_size is not None and batch_size != -1 and batch_size <= 0: + raise ValueError("batch_size must be a positive integer, -1, or None.") + return batch_size + + def _batch_limit(self, batch_size: int | None) -> int: + """Return the number of calls of a batch operation that may run at once. + + Args: + batch_size: The ``batch_size`` of the call. None takes that of the + filesystem; None or -1 there is ``max_workers``, the size of + the request pool. + + Returns: + The number of calls. + + Raises: + ValueError: If ``batch_size`` is neither positive, -1, nor None. + """ + batch_size = self._check_batch_size(batch_size) + if batch_size is None: + batch_size = self.batch_size + if batch_size is None or batch_size == -1: + return self.max_workers + return batch_size + + def _batch_kwargs(self, concurrency: int, kwargs: dict[str, Any]) -> dict[str, Any]: + """Return the arguments of the per-object calls of a batch operation. + + With several calls at once, each object gets ``max_workers // + concurrency`` parallel requests, at least one, so that a batch keeps + about ``max_workers`` requests in the request pool. A ``max_workers`` + of the caller is capped at that; None or 0 takes it. + + Args: + concurrency: The number of calls that run at once. + kwargs: The arguments that the caller gave the operation. + + Returns: + The arguments of each call. + """ + if concurrency <= 1: + return kwargs + window = max(1, self.max_workers // concurrency) + given = kwargs.get("max_workers") + return {**kwargs, "max_workers": min(given, window) if given else window} + + def _run_batch( + self, + fn: Callable[[ItemT], T], + items: list[ItemT], + concurrency: int, + counts: Callable[[BaseException], bool] | None = None, + ) -> S3MapResult[T]: + """Run the per-object calls of a batch operation. + + The calls are composite: each one waits for requests of its object in + the request pool. They run in threads of a pool of their own, at most + ``concurrency`` at once, and never in the request pool. A single + item runs in the calling thread. See :meth:`S3Executor.map() + ` for the failures, + interrupts, and ``counts``. + + Args: + fn: The call to run for each item. + items: The items, in the order of the outcomes. + concurrency: The number of calls that run at once. + counts: Whether a failure stops the batch, by default every one. + + Returns: + The outcomes of the calls. + """ + if len(items) == 1: + try: + return S3MapResult([fn(items[0])], None) + except Exception as e: + if counts is None or counts(e): + raise + return S3MapResult([e], None) + executor = S3ThreadPool(max_workers=concurrency) + try: + return executor.map(fn, items, counts=counts) + finally: + # The map has waited for the calls that it started, unless a + # second interrupt stopped it. + executor.shutdown(wait=False) + def _delete_objects(self, paths: list[str], max_workers: int | None = None, **kwargs) -> None: """Delete objects with DeleteObjects requests grouped by bucket. @@ -1511,7 +1741,7 @@ def touch(self, path: str, truncate: bool = True, **kwargs) -> dict[str, Any]: self.invalidate_cache(path) return object_.to_dict() - def mv(self, path1, path2, recursive=False, maxdepth=None, **kwargs) -> None: + def mv(self, path1, path2, recursive=False, maxdepth=None, batch_size=None, **kwargs) -> None: """Move files from one S3 location to another. Copies the files to the destinations that ``copy()`` gives them, and @@ -1531,12 +1761,18 @@ def mv(self, path1, path2, recursive=False, maxdepth=None, **kwargs) -> None: requires ``s3:GetBucketVersioning``. The result is not cached across moves. Directory buckets do not support versioning and need no lookup. + The files are copied in parallel as ``copy()`` copies them, and are + deleted once every copy has succeeded; nothing is deleted after a + failed copy. + Args: path1: Source S3 path, glob pattern, or list of paths. path2: Destination S3 path, or list of paths when ``path1`` is a list. recursive: Whether to move the directories with their contents. maxdepth: Maximum depth of a recursive move. + batch_size: Number of copies to run at the same time, as in + ``copy()``. **kwargs: Additional S3 copy parameters passed to ``cp_file()``. The lookups and the deletion of the sources receive ``RequestPayer`` and the source's expected bucket owner (see @@ -1548,23 +1784,42 @@ def mv(self, path1, path2, recursive=False, maxdepth=None, **kwargs) -> None: destination is another source, including one left in place, which is checked before anything is copied. A directory with no object at its key, which is not copied, does not conflict. + Also if ``batch_size`` is neither positive, -1, nor None. OSError: If bucket versioning cannot be read. """ if path1 == path2: return + limit = self._batch_limit(batch_size) source_params = self.core.copy_source_params(kwargs) - copied = [ - p1 - for p1, p2 in self._move_pairs( - path1, - path2, - recursive=recursive, - maxdepth=maxdepth, - source_params=source_params, - ) - if self._copy_file(p1, p2, **kwargs) - ] - self._delete_objects(copied, **self.core.operation_params("delete_objects", source_params)) + pairs = self._move_pairs( + path1, + path2, + recursive=recursive, + maxdepth=maxdepth, + source_params=source_params, + ) + if not pairs: + return + concurrency = min(limit, len(pairs)) + call_kwargs = self._batch_kwargs(concurrency, kwargs) + + def copy_pair(pair: tuple[str, str]) -> bool: + """Copy a pair of the move. + + Args: + pair: The source and its destination. + + Returns: + Whether the source was copied, as ``_copy_file()`` returns it. + """ + return self._copy_file(*pair, **call_kwargs) + + # Every copy has ended, and succeeded, before anything is deleted. + copied = self._run_batch(copy_pair, pairs, concurrency).outcomes + self._delete_objects( + [p1 for (p1, _), c in zip(pairs, copied, strict=True) if c], + **self.core.operation_params("delete_objects", source_params), + ) def _move_pairs( self, @@ -1595,6 +1850,38 @@ def _move_pairs( """ pairing = S3PathPairing(path1, path2, recursive=recursive, maxdepth=maxdepth) pairs = self._copy_pairs(pairing) + missing, versioned_buckets = self._conflict_lookups(pairing, pairs, source_params) + return pairing.move_pairs( + pairs, missing=missing, versioning_enabled_buckets=versioned_buckets + ) + + def _conflict_lookups( + self, + pairing: S3PathPairing, + pairs: list[tuple[str, str]], + source_params: Mapping[str, Any] | None = None, + copy: bool = False, + ) -> tuple[set[str], set[str]]: + """Look up what the conflict checks of a move or a copy need. + + See :meth:`S3PathPairing.conflict_candidates`. + + Args: + pairing: The paths of the move or the copy. + pairs: The pairs of the paths. + source_params: The parameters of the requests that read the + sources (see :meth:`S3Core.copy_source_params`), sent with + the lookups of the sources. + copy: Whether the pairs are those of a copy. + + Returns: + The conflict candidates with no object at their key, and the + buckets whose versioning is enabled among those with a ``null`` + version and its unversioned key in the pairs. + + Raises: + OSError: If bucket versioning cannot be read. + """ paths = {p: S3Path.parse(p) for pair in pairs for p in pair} unversioned = {p.name for p in paths.values() if not p.version_id} buckets = { @@ -1609,16 +1896,17 @@ def _move_pairs( missing = { source for source in pairing.conflict_candidates( - pairs, versioning_enabled_buckets=versioned_buckets + pairs, versioning_enabled_buckets=versioned_buckets, copy=copy ) if self._head_object(source, lookup_kwargs=lookup_kwargs) is None } - return pairing.move_pairs( - pairs, missing=missing, versioning_enabled_buckets=versioned_buckets - ) + return missing, versioned_buckets def _copy_pairs( - self, pairing: S3PathPairing, isdir: Callable[[str], bool] | None = None + self, + pairing: S3PathPairing, + isdir: Callable[[str], bool] | None = None, + **kwargs: Any, ) -> list[tuple[str, str]]: """Expand and pair the paths of a copy (see :meth:`S3PathPairing.copy_pairs`). @@ -1628,6 +1916,8 @@ def _copy_pairs( pairing: The paths of the copy. isdir: Whether the destination is a directory, by default ``self.isdir``; ``get()`` passes the local filesystem's. + **kwargs: Passed to ``expand_path()``, as fsspec's ``get()`` passes + its keyword arguments. Returns: The sources and their destinations. @@ -1635,7 +1925,7 @@ def _copy_pairs( if not pairing.expands: return pairing.copy_pairs() sources = self.expand_path( - pairing.path1, recursive=pairing.recursive, maxdepth=pairing.maxdepth + pairing.path1, recursive=pairing.recursive, maxdepth=pairing.maxdepth, **kwargs ) if pairing.skips_directories: sources = [p for p in sources if not (trailing_sep(p) or self.isdir(p))] @@ -1647,13 +1937,30 @@ def _copy_pairs( ) return pairing.copy_pairs(sources, destination_is_dir) - def copy(self, path1, path2, recursive=False, maxdepth=None, on_error=None, **kwargs) -> None: + def copy( + self, + path1, + path2, + recursive=False, + maxdepth=None, + on_error=None, + batch_size=None, + **kwargs, + ) -> None: """Copy files within S3. As fsspec's ``copy()``, except that a source with a version ID copies that version of the object to a destination named after its key, as :meth:`~pyathena.filesystem.s3_path_pairing.S3PathPairing.copy_pairs` - pairs them. + pairs them, and that the files are copied in parallel, ``batch_size`` + at a time. Pairs that would depend on the order of the copies raise + before anything is copied (see + :meth:`~pyathena.filesystem.s3_path_pairing.S3PathPairing.check_copy_pairs`). + + An error that ``on_error`` does not ignore starts no further copy, + stops the running ones before their next request where it can (a + multipart copy sends no further part and is aborted), and is raised + once they have ended. Copies that ended before it stay in place. Args: path1: Source S3 path, glob pattern, or list of them. @@ -1664,57 +1971,163 @@ def copy(self, path1, path2, recursive=False, maxdepth=None, on_error=None, **kw on_error: ``"raise"`` or ``"ignore"`` for a missing source; by default ``"ignore"`` with ``recursive`` and ``"raise"`` otherwise. + batch_size: Number of copies to run at the same time. By default + the ``batch_size`` of the filesystem; None or -1 there is + ``max_workers``. With ``n`` copies at once, the smaller of + ``batch_size`` and the number of pairs, each multipart copy + sends at most ``max(1, max_workers // n)`` parts at once. **kwargs: Additional S3 copy parameters passed to ``cp_file()``. + + Raises: + ValueError: If two sources have the same destination, or a + destination is another source, except for a directory with no + object at its key, which is not copied. Also if + ``batch_size`` is neither positive, -1, nor None. + OSError: If bucket versioning cannot be read for the checks. """ - sources = [path1] if isinstance(path1, (str, os.PathLike)) else path1 - if isinstance(path2, str) and any(S3Path.has_version_id(p) for p in sources): - pairs = self._copy_pairs( - S3PathPairing(path1, path2, recursive=recursive, maxdepth=maxdepth) - ) - if not pairs: - return - path1, path2 = [p1 for p1, _ in pairs], [p2 for _, p2 in pairs] - super().copy( - path1, path2, recursive=recursive, maxdepth=maxdepth, on_error=on_error, **kwargs + if on_error is None: + on_error = "ignore" if recursive else "raise" + limit = self._batch_limit(batch_size) + pairing = S3PathPairing(path1, path2, recursive=recursive, maxdepth=maxdepth) + pairs = self._copy_pairs(pairing) + if not pairs: + return + missing, versioned_buckets = self._conflict_lookups( + pairing, pairs, self.core.copy_source_params(kwargs), copy=True + ) + pairing.check_copy_pairs( + pairs, missing=missing, versioning_enabled_buckets=versioned_buckets + ) + concurrency = min(limit, len(pairs)) + call_kwargs = self._batch_kwargs(concurrency, kwargs) + + def copy_pair(pair: tuple[str, str]) -> None: + """Copy a pair. + + Args: + pair: The source and its destination. + """ + self.cp_file(*pair, **call_kwargs) + + self._run_batch( + copy_pair, + pairs, + concurrency, + counts=None if on_error == "raise" else lambda e: not isinstance(e, FileNotFoundError), ) def get( - self, rpath, lpath, recursive=False, callback=_DEFAULT_CALLBACK, maxdepth=None, **kwargs + self, + rpath, + lpath, + recursive=False, + callback=_DEFAULT_CALLBACK, + maxdepth=None, + batch_size=None, + **kwargs, ) -> None: """Copy files from S3 to the local filesystem. As fsspec's ``get()``, except that a source with a version ID downloads that version of the object to a local path named after its key, as :meth:`~pyathena.filesystem.s3_path_pairing.S3PathPairing.copy_pairs` - pairs them. Those destinations are checked to lie under ``lpath``. + pairs them, and that the files are downloaded in parallel, as with + ``copy()``. The destinations named after the sources are checked to + lie under ``lpath``. With ``outfile``, the files are written to it + one after another, in order. + + A failed download starts no further one, the running ones read and + write no further block, and the error is raised once they have ended. Args: rpath: Source S3 path, glob pattern, or list of them. - lpath: Local destination path, or list of paths when ``rpath`` is - a list. + lpath: Local destination path, or list of paths or file-like + objects when ``rpath`` is a list. recursive: Whether to copy the directories with their contents. - callback: Progress callback. + callback: Progress callback. It and its branches for the files + are used from the threads that download them, one at a time. maxdepth: Maximum depth of a recursive copy. + batch_size: Number of downloads to run at the same time, as in + ``copy()``. **kwargs: Additional parameters passed to ``get_file()``. Raises: - ValueError: If a source with a version ID is paired, and a - destination lies outside ``lpath``. - """ - sources = [rpath] if isinstance(rpath, (str, os.PathLike)) else rpath - if isinstance(lpath, (str, os.PathLike)) and any(S3Path.has_version_id(p) for p in sources): - root = make_path_posix(lpath) - pairs = self._copy_pairs( - S3PathPairing(rpath, root, recursive=recursive, maxdepth=maxdepth), - isdir=LocalFileSystem().isdir, - ) - rpath, lpath = [p1 for p1, _ in pairs], [p2 for _, p2 in pairs] - check_contained(root, lpath) - if not rpath: - return - super().get( - rpath, lpath, recursive=recursive, callback=callback, maxdepth=maxdepth, **kwargs + ValueError: If a destination named after a source lies outside + ``lpath``, two sources have the same destination, or + ``batch_size`` is neither positive, -1, nor None. + """ + limit = self._batch_limit(batch_size) + root = make_path_posix(lpath) if isinstance(lpath, (str, os.PathLike)) else lpath + pairs = self._copy_pairs( + S3PathPairing(rpath, root, recursive=recursive, maxdepth=maxdepth), + isdir=LocalFileSystem().isdir, + **kwargs, ) + if isinstance(root, str): + check_contained(root, [lp for _, lp in pairs]) + if not pairs: + return + # get_file() writes to an outfile instead of the local paths. + outfile = kwargs.get("outfile") is not None + if not outfile: + self._check_local_destinations([lp for _, lp in pairs]) + concurrency = 1 if outfile else min(limit, len(pairs)) + call_kwargs = self._batch_kwargs(concurrency, kwargs) + locked = _S3LockedCallback(callback, threading.Lock()) + + def get_pair(pair: tuple[str, Any]) -> None: + """Download a pair with a branch of the callback. + + Args: + pair: The S3 source and its local destination. + """ + rp, lp = pair + with locked.branched(rp, lp) as child: + self.get_file(rp, lp, callback=child, **call_kwargs) + locked.relative_update(1) + + locked.set_size(len(pairs)) + if outfile: + # One after another, in order, in the calling thread, which may + # own the stream. + for pair in pairs: + get_pair(pair) + return + self._run_batch(get_pair, pairs, concurrency) + + @staticmethod + def _check_local_destinations(destinations: list[Any]) -> None: + """Check that no two downloads of a batch write the same local file. + + Destinations are compared by the object for file-like ones, by the + file for existing paths, such as hard links to one file, and by the + resolved path otherwise. + + Args: + destinations: The local paths or file-like objects. + + Raises: + ValueError: If two destinations write the same file. + """ + seen: set[tuple[Any, ...]] = set() + for destination in destinations: + key: tuple[Any, ...] + if isfilelike(destination): + key = ("object", id(destination)) + else: + path = os.fspath(destination) + try: + stat = os.stat(path) + except OSError: + key = ("path", os.path.normcase(os.path.realpath(path))) + else: + if os.path.isdir(path): + # Directories are created, not written. + continue + key = ("file", stat.st_dev, stat.st_ino) + if key in seen: + raise ValueError("Cannot get several paths to the same destination.") + seen.add(key) def cp_file( self, path1: str, path2: str, recursive=False, maxdepth=None, on_error=None, **kwargs @@ -1789,6 +2202,7 @@ def _copy_file(self, path1: str, path2: str, **kwargs) -> bool: size1 = info1.get("size", 0) try: if size1 <= self.core.MULTIPART_UPLOAD_MAX_PART_SIZE: + raise_if_stopped(f"copy {source.uri} to {destination.uri}") self.core.copy_object(source, destination, **kwargs) else: self._copy_object_with_multipart_upload( @@ -1825,7 +2239,11 @@ def _copy_object_with_multipart_upload( interrupts while the creation is waited for are suppressed; one while the part copies are waited for or the upload is aborted stops the cleanup, and the upload ID is logged. A failed annotation copy is - raised and leaves the destination in place. If HeadObject reports a + raised and leaves the destination in place. The annotations are + copied, retries included, also once an enclosing stop scope, such as + that of a batch operation, has been stopped, since the destination is + committed; until then, a stop sends no further part or completion + and the upload is aborted. If HeadObject reports a size that fits in a single CopyObject request, the reported version is copied with CopyObject instead. @@ -1848,6 +2266,7 @@ def _copy_object_with_multipart_upload( # The size that the caller found, which may come from a cached # listing, was larger than the copied version, which fits in a # single CopyObject request. + raise_if_stopped(f"copy {plan.source.uri} to {plan.destination.uri}") self.core.copy_object(plan.source, plan.destination, **kwargs) return with self._create_executor(max_workers=max_workers) as executor: @@ -1891,15 +2310,61 @@ def _copy_object_with_multipart_upload( f"to s3://{upload.bucket}/{upload.key}." ) raise - for name in plan.annotations: - self.core.copy_object_annotation( - name, - plan.source, - plan.destination, - completed.version_id, - completed.etag, - **kwargs, - ) + # The destination is committed, so its annotations are copied to the + # end also when an enclosing batch has stopped. + with unstoppable_retries(): + for name in plan.annotations: + self.core.copy_object_annotation( + name, + plan.source, + plan.destination, + completed.version_id, + completed.etag, + **kwargs, + ) + + def pipe(self, path, value=None, batch_size=None, **kwargs) -> None: + """Write bytes into one or more paths. + + As fsspec's ``pipe()``, except that the values of a dict are written + in parallel with ``pipe_file()``, as ``copy()`` copies its files. + + Args: + path: S3 path to write ``value`` to, or a dict of the values to + write by path. + value: The bytes to write to a single path; ignored with a dict. + batch_size: Number of writes to run at the same time, as in + ``copy()``. + **kwargs: Additional parameters passed to ``pipe_file()``. + + Raises: + ValueError: If ``path`` is neither a string nor a dict, two paths + of a dict name the same object, or ``batch_size`` is neither + positive, -1, nor None. + """ + limit = self._batch_limit(batch_size) + if isinstance(path, str): + self.pipe_file(self._strip_protocol(path), value, **kwargs) + return + if not isinstance(path, dict): + raise ValueError("path must be str or dict") + items = [(self._strip_protocol(k), v) for k, v in path.items()] + if len({k for k, _ in items}) != len(items): + raise ValueError("Cannot pipe several values to the same destination.") + if not items: + return + concurrency = min(limit, len(items)) + call_kwargs = self._batch_kwargs(concurrency, kwargs) + + def pipe_item(item: tuple[str, Any]) -> None: + """Write a value. + + Args: + item: The path and its value. + """ + self.pipe_file(*item, **call_kwargs) + + self._run_batch(pipe_item, items, concurrency) def pipe_file( self, path: str, value: bytes | bytearray | memoryview, mode: str = "overwrite", **kwargs @@ -1979,9 +2444,126 @@ def pipe_file( # buffered path does. value = bytes(value) + raise_if_stopped(f"put {s3_path.uri}") self.core.put_object(s3_path, value, **request_kwargs) self.invalidate_cache(path) + def cat(self, path, recursive=False, on_error="raise", batch_size=None, **kwargs): + """Read the contents of one or more paths. + + As fsspec's ``cat()``, except that the objects are read in parallel + with ``cat_file()``, as ``copy()`` copies its files. With + ``on_error="raise"``, an error starts no further read and is raised + once the running reads have ended. + + Args: + path: S3 path, glob pattern, or list of them. + recursive: Whether to read the files below the directories. + on_error: ``"raise"`` to raise an error of a read, ``"omit"`` to + leave out the paths whose read failed, or ``"return"`` to + return the error in place of their contents. + batch_size: Number of reads to run at the same time, as in + ``copy()``. + **kwargs: Additional parameters passed to ``expand_path()`` and + ``cat_file()``. + + Returns: + The contents of the path, or a dict of the contents by path if + there are several paths or the path was expanded. + + Raises: + ValueError: If ``batch_size`` is neither positive, -1, nor None. + """ + limit = self._batch_limit(batch_size) + paths = self.expand_path(path, recursive=recursive, **kwargs) + if not (len(paths) > 1 or isinstance(path, list) or paths[0] != self._strip_protocol(path)): + return self.cat_file(paths[0], **kwargs) + outcomes: list[Any] | None = None + contents: dict[str, Any] | None = None + outcome = None + try: + outcomes = self._run_batch( + functools.partial(self.cat_file, **kwargs), + paths, + min(limit, len(paths)), + counts=None if on_error == "raise" else lambda e: False, + ).outcomes + contents = {} + for p, outcome in zip(paths, outcomes, strict=True): + if on_error != "omit" or not isinstance(outcome, Exception): + contents[p] = outcome + return contents + except BaseException: + # Not kept by this frame, which the traceback of an interrupt + # keeps (F5). + contents = outcomes = outcome = None + raise + + def cat_ranges( + self, paths, starts, ends, max_gap=None, on_error="return", batch_size=None, **kwargs + ) -> list[bytes | Exception]: + """Read byte ranges of one or more paths. + + As fsspec's ``cat_ranges()``, except that the ranges are read in + parallel with ``cat_file()``, as ``copy()`` copies its files. Unless + ``on_error`` is ``"return"``, an error starts no further read and is + raised once the running reads have ended. + + Args: + paths: The S3 paths, one per range. + starts: The start offset of each range, or one for all of them. + ends: The end offset of each range, or one for all of them. + max_gap: Not supported; must be None. + on_error: ``"return"`` to return the error of a read in place of + its bytes; otherwise the error is raised. + batch_size: Number of reads to run at the same time, as in + ``copy()``. + **kwargs: Additional parameters passed to ``cat_file()``. + + Returns: + The bytes of each range, or the error of its read. + + Raises: + NotImplementedError: If ``max_gap`` is given. + TypeError: If ``paths`` is not a list. + ValueError: If the numbers of paths, starts, and ends differ, or + if ``batch_size`` is neither positive, -1, nor None. + """ + if max_gap is not None: + raise NotImplementedError + if not isinstance(paths, list): + raise TypeError + if not isinstance(starts, list): + starts = [starts] * len(paths) + if not isinstance(ends, list): + ends = [ends] * len(paths) + if len(starts) != len(paths) or len(ends) != len(paths): + raise ValueError + limit = self._batch_limit(batch_size) + if not paths: + return [] + + def cat_range(item: tuple[str, int | None, int | None]) -> bytes: + """Read a range. + + Args: + item: The path, start, and end. + + Returns: + The bytes of the range. + """ + return self.cat_file(*item, **kwargs) + + return cast( + list[bytes | Exception], + self._run_batch( + cat_range, + list(zip(paths, starts, ends, strict=True)), + min(limit, len(paths)), + counts=None if on_error != "return" else lambda e: False, + ).outcomes, + ) + def cat_file( self, path: str, start: int | None = None, end: int | None = None, **kwargs ) -> bytes: @@ -2055,6 +2637,88 @@ def cat_file( return b"" raise + def put( + self, + lpath, + rpath, + recursive=False, + callback=_DEFAULT_CALLBACK, + maxdepth=None, + batch_size=None, + **kwargs, + ) -> None: + """Copy files from the local filesystem to S3. + + As fsspec's ``put()``, except that the files are uploaded in parallel + with ``put_file()``, as ``copy()`` copies its files. A failed upload + starts no further one, the running ones send no further part and do + not commit, and the error is raised once they have ended. + + Args: + lpath: Local source path, glob pattern, or list of them. + rpath: S3 destination path, or list of paths when ``lpath`` is a + list. + recursive: Whether to copy the directories with their contents. + callback: Progress callback. It and its branches for the files + are used from the threads that upload them, one at a time. + maxdepth: Maximum depth of a recursive copy. + batch_size: Number of uploads to run at the same time, as in + ``copy()``. + **kwargs: Additional parameters passed to ``put_file()``. + + Raises: + ValueError: If two sources have the same destination, or + ``batch_size`` is neither positive, -1, nor None. + """ + limit = self._batch_limit(batch_size) + if isinstance(lpath, list) and isinstance(rpath, list): + pairs = [(lp, self._strip_protocol(rp)) for lp, rp in zip(lpath, rpath, strict=False)] + else: + # As fsspec's put() pairs them. + source_is_str = isinstance(lpath, str) + if source_is_str: + lpath = make_path_posix(lpath) + local = LocalFileSystem() + lpaths = local.expand_path(lpath, recursive=recursive, maxdepth=maxdepth, **kwargs) + if source_is_str and (not recursive or maxdepth is not None): + # A string source without recursion does not copy directories. + lpaths = [p for p in lpaths if not (trailing_sep(p) or local.isdir(p))] + if not lpaths: + return + dest_is_dir = isinstance(rpath, str) and (trailing_sep(rpath) or self.isdir(rpath)) + rpath = ( + self._strip_protocol(rpath) + if isinstance(rpath, str) + else [self._strip_protocol(p) for p in rpath] + ) + exists = source_is_str and ( + (has_magic(lpath) and len(lpaths) == 1) + or (not has_magic(lpath) and dest_is_dir and not trailing_sep(lpath)) + ) + rpaths = other_paths(lpaths, rpath, exists=exists, flatten=not source_is_str) + pairs = list(zip(lpaths, rpaths, strict=True)) + if not pairs: + return + if len({rp for _, rp in pairs}) != len(pairs): + raise ValueError("Cannot put several paths to the same destination.") + concurrency = min(limit, len(pairs)) + call_kwargs = self._batch_kwargs(concurrency, kwargs) + locked = _S3LockedCallback(callback, threading.Lock()) + + def put_pair(pair: tuple[str, str]) -> None: + """Upload a pair with a branch of the callback. + + Args: + pair: The local source and its S3 destination. + """ + lp, rp = pair + with locked.branched(lp, rp) as child: + self.put_file(lp, rp, callback=child, **call_kwargs) + locked.relative_update(1) + + locked.set_size(len(pairs)) + self._run_batch(put_pair, pairs, concurrency) + def put_file( self, lpath: str, @@ -2171,6 +2835,8 @@ def get_file(self, rpath: str, lpath=None, callback=_DEFAULT_CALLBACK, outfile=N outfile = stack.enter_context(open(lpath, "wb")) callback.set_size(remote.size) while data := remote.read(remote.blocksize): + # A block read when the stop came is not written. + raise_if_stopped(f"write the contents of {rpath}") outfile.write(data) callback.relative_update(len(data)) @@ -3681,7 +4347,9 @@ def _fetch_range(self, start: int, end: int) -> bytes: :meth:`S3Executor.map() `); when one fails, those that have not started are cancelled, the running ones make no further attempt, and its error is raised once they have - ended. + ended. Once an enclosing stop scope is stopped, such as that of a + batch operation after another object failed, no further request is + sent (see :func:`~pyathena.util.raise_if_stopped`). Args: start: The offset of the first byte to read. @@ -3689,6 +4357,10 @@ def _fetch_range(self, start: int, end: int) -> bytes: Returns: The bytes read, empty if the clamped range is empty. + + Raises: + concurrent.futures.CancelledError: If an enclosing stop scope has + been stopped. """ end = min(end, self.size) if start >= end: @@ -3705,12 +4377,34 @@ def _fetch_range(self, start: int, end: int) -> bytes: cast( list[bytes], self._executor.map( - functools.partial(self.fs.core.get_object, path, **request_kwargs), + functools.partial(self._get_range, path, **request_kwargs), ranges, ).outcomes, ) ) - return self.fs.core.get_object(path, ranges[0], **request_kwargs) + return self._get_range(path, ranges[0], **request_kwargs) + + def _get_range(self, path: S3Path, range_: tuple[int, int], **kwargs: Any) -> bytes: + """Read a byte range with GetObject unless the read has been stopped. + + A range that waited behind other requests checks the stop scopes of + its context when it starts, since the retry stop does not cover the + first attempt. + + Args: + path: The object, with its version ID, if any. + range_: The range, with an exclusive end. + **kwargs: The GetObject parameters. + + Returns: + The bytes of the range. + + Raises: + concurrent.futures.CancelledError: If an enclosing stop scope has + been stopped. + """ + raise_if_stopped(f"read {path.uri}") + return self.fs.core.get_object(path, range_, **kwargs) @staticmethod def _get_ranges( diff --git a/pyathena/filesystem/s3_path_pairing.py b/pyathena/filesystem/s3_path_pairing.py index 5b5a80f98..68e4a6a0a 100644 --- a/pyathena/filesystem/s3_path_pairing.py +++ b/pyathena/filesystem/s3_path_pairing.py @@ -145,13 +145,16 @@ def conflict_candidates( pairs: Sequence[tuple[str, str]], *, versioning_enabled_buckets: Collection[str] = (), + copy: bool = False, ) -> list[str]: """Return the sources of a move whose conflicts depend on an object at their key. A source with another source below it may be a directory. If no object exists at its key, it is not copied, so it writes no destination and is left out of the conflict checks of :meth:`move_pairs`. A path with - a version always names an object. + a version always names an object. For a copy, any source may be a + directory, such as a prefix given in a list, which the copy skips + (see :meth:`check_copy_pairs`). Args: pairs: The sources and destinations of the move, as @@ -159,6 +162,9 @@ def conflict_candidates( versioning_enabled_buckets: Buckets whose versioning is enabled. Their ``null`` versions are distinct from their unversioned keys. Use the same collection for :meth:`move_pairs`. + copy: Whether the pairs are those of a copy, whose candidates + include every source without a version whose destination + conflicts. Returns: The sources, in ``bucket/key`` form (their @@ -170,7 +176,7 @@ def conflict_candidates( Raises: TypeError: If ``versioning_enabled_buckets`` is a string. """ - return self._moves(pairs, versioning_enabled_buckets)[2] + return self._moves(pairs, versioning_enabled_buckets, copy)[2] def move_pairs( self, @@ -204,9 +210,72 @@ def move_pairs( TypeError: If ``missing`` is a string instead of a collection. Also if ``versioning_enabled_buckets`` is a string. """ + named = self._check(pairs, missing, versioning_enabled_buckets, "move") + return [(p1, p2) for p1, p2, _, source, dest in named if source != dest] + + def check_copy_pairs( + self, + pairs: Sequence[tuple[str, str]], + missing: Collection[str] | None = None, + *, + versioning_enabled_buckets: Collection[str] = (), + ) -> None: + """Check the pairs of a copy whose objects are copied in parallel. + + A copy has the conflicts of a move (see :meth:`move_pairs`), since a + parallel copy has no order in which one pair would overwrite another + or read a source after another pair has written it. A pair whose + destination is its own source is no conflict: it rewrites the object + in place, such as with ``MetadataDirective="REPLACE"``. + + Args: + pairs: The sources and destinations of the copy, as + :meth:`copy_pairs` of this pairing returns them. + missing: The :meth:`conflict_candidates` without an object at their + key, in any form that names them; needed when there are + candidates. + versioning_enabled_buckets: Buckets whose versioning is enabled. + Use the same collection for :meth:`conflict_candidates`. + + Raises: + ValueError: If two sources have the same destination, or a + destination is another source, except for a directory with no + object at its key, which is not copied. Also if ``missing`` is + needed and None. + TypeError: If ``missing`` is a string instead of a collection. + Also if ``versioning_enabled_buckets`` is a string. + """ + self._check(pairs, missing, versioning_enabled_buckets, "copy") + + def _check( + self, + pairs: Sequence[tuple[str, str]], + missing: Collection[str] | None, + versioning_enabled_buckets: Collection[str], + operation: str, + ) -> list[tuple[str, str, bool, str, str]]: + """Raise for the conflicts of a move or a copy. + + Args: + pairs: The sources and destinations. + missing: The conflict candidates without an object at their key. + versioning_enabled_buckets: Buckets whose versioning is enabled. + operation: ``"move"`` or ``"copy"``, for the error messages. + + Returns: + The pairs as :meth:`_moves` names them. + + Raises: + ValueError: If the pairs conflict, or ``missing`` is needed and + None. + TypeError: If ``missing`` or ``versioning_enabled_buckets`` is a + string. + """ if isinstance(missing, str): raise TypeError("missing is a collection of paths, not a path.") - named, sources, candidates = self._moves(pairs, versioning_enabled_buckets) + named, sources, candidates = self._moves( + pairs, versioning_enabled_buckets, operation == "copy" + ) if missing is None and candidates: raise ValueError("missing is needed to check the pairs.") # A directory without an object at its key writes no destination. @@ -222,10 +291,11 @@ def move_pairs( ) for dest in writers: if writers[dest] > 1: - raise ValueError("Cannot move several paths to the same destination.") + raise ValueError(f"Cannot {operation} several paths to the same destination.") if dest in sources: - raise ValueError("Cannot move a path onto another path that is moved.") - return [(p1, p2) for p1, p2, _, source, dest in named if source != dest] + done = "moved" if operation == "move" else "copied" + raise ValueError(f"Cannot {operation} a path onto another path that is {done}.") + return named @classmethod def delete_paths(cls, path: str | list[str]) -> tuple[list[str], list[str]]: @@ -284,12 +354,14 @@ def _target(path: S3Path, versioning_enabled_buckets: Collection[str]) -> str: def _moves( pairs: Sequence[tuple[str, str]], versioning_enabled_buckets: Collection[str] = (), + copy: bool = False, ) -> tuple[list[tuple[str, str, bool, str, str]], set[str], list[str]]: """Compare the paths of a move by what they name. Args: pairs: The sources and destinations of the move. versioning_enabled_buckets: Buckets whose versioning is enabled. + copy: Whether every conflicting source is a candidate. Returns: Each pair with whether its source has a version and the targets @@ -331,7 +403,7 @@ def _moves( for _, _, versioned, source, dest in named if source != dest and (counts[dest] > 1 or dest in sources) - and source in directories + and (copy or source in directories) and not versioned ] return named, sources, candidates diff --git a/pyathena/filesystem/s3_writer.py b/pyathena/filesystem/s3_writer.py index 8abc70fb3..7727bc825 100644 --- a/pyathena/filesystem/s3_writer.py +++ b/pyathena/filesystem/s3_writer.py @@ -25,7 +25,7 @@ S3PutObject, ) from pyathena.filesystem.s3_path import S3Path -from pyathena.util import stop_retries_on +from pyathena.util import raise_if_stopped, stop_retries_on T = TypeVar("T") @@ -485,7 +485,10 @@ def start(self) -> None: Raises: S3MultipartStateError: If a request has already been sent. + concurrent.futures.CancelledError: If an enclosing stop scope has + been stopped (see :func:`~pyathena.util.raise_if_stopped`). """ + raise_if_stopped(f"create a multipart upload to {self._path.uri}") creation: Future[S3MultipartUpload] = Future() creation.add_done_callback(self._on_created) with self._lock: @@ -625,11 +628,11 @@ def put(self, body: bytes | None) -> S3PutObject: Raises: S3MultipartStateError: If the write started a multipart upload, failed, or has been discarded. + concurrent.futures.CancelledError: If an enclosing stop scope has + been stopped before the request (see + :func:`~pyathena.util.raise_if_stopped`); the write fails then. """ - return self._commit( - _S3WriteState.IDLE, - lambda: self._core.put_object(self._path, body, **self._params("put_object")), - ) + return self._commit(_S3WriteState.IDLE, lambda: self._put_object(body)) def complete(self) -> S3CompleteMultipartUpload: """Commit the write by completing its multipart upload. @@ -648,6 +651,9 @@ def complete(self) -> S3CompleteMultipartUpload: raised yet, or the failure of the completion. S3MultipartStateError: If no upload is active, or the write failed earlier or has been discarded. + concurrent.futures.CancelledError: If an enclosing stop scope has + been stopped before the request (see + :func:`~pyathena.util.raise_if_stopped`); the write fails then. """ return self._commit(_S3WriteState.ACTIVE, self._complete_parts) @@ -831,7 +837,13 @@ def _submit(self, fn: Callable[..., S3MultipartUploadPart], **kwargs: Any) -> No fn: The writer method that sends the request. **kwargs: The arguments of the method, whose ``part_number`` must be the next one. + + Raises: + concurrent.futures.CancelledError: If an enclosing stop scope has + been stopped, before the request is tracked, so that nothing + waits for a request that never runs. """ + raise_if_stopped(f"upload part {kwargs['part_number']} to {self._path.uri}") request: Future[S3MultipartUploadPart] = Future() # Added first, so that a cancellation of the tracked request always # stops tracking it. @@ -853,7 +865,7 @@ def _submit(self, fn: Callable[..., S3MultipartUploadPart], **kwargs: Any) -> No self._next_part_number += 1 try: with stop_retries_on(self._stop): - self._executor.submit_to(request, fn, **kwargs) + self._executor.submit_to(request, self._send_part, fn, **kwargs) except BaseException as e: # Recorded before the cancellation, which would record its own. with self._update(): @@ -863,6 +875,45 @@ def _submit(self, fn: Callable[..., S3MultipartUploadPart], **kwargs: Any) -> No request.cancel() raise + def _send_part( + self, fn: Callable[..., S3MultipartUploadPart], **kwargs: Any + ) -> S3MultipartUploadPart: + """Send a part request unless the write has been stopped meanwhile. + + A part that waited behind other requests of the executor checks the + stop scopes of its context, which :meth:`_submit` copied, when it + starts: the retry stop does not cover the first attempt. + + Args: + fn: The writer method that sends the request. + **kwargs: The arguments of the method. + + Returns: + The uploaded part. + + Raises: + concurrent.futures.CancelledError: If an enclosing stop scope has + been stopped. + """ + raise_if_stopped(f"upload part {kwargs['part_number']} to {self._path.uri}") + return fn(**kwargs) + + def _put_object(self, body: bytes | None) -> S3PutObject: + """Send the PutObject request of :meth:`put` unless the write has been stopped. + + Args: + body: The data of the object, or None for an empty object. + + Returns: + The result of the request. + + Raises: + concurrent.futures.CancelledError: If an enclosing stop scope has + been stopped. + """ + raise_if_stopped(f"put {self._path.uri}") + return self._core.put_object(self._path, body, **self._params("put_object")) + def _commit(self, state: _S3WriteState, send: Callable[[], T]) -> T: """Reserve the commit, send its request, and record the outcome. @@ -927,6 +978,10 @@ def _complete_parts(self) -> S3CompleteMultipartUpload: # invariant that no part is completed without its data. missing = sorted(set(range(1, self._next_part_number)) - set(self._parts)) raise RuntimeError(f"The parts {missing} of {self._path.uri} have not been uploaded.") + # After the wait for the parts, which a stop may end, and inside the + # commit, so that a stopped completion leaves the write failed for + # the caller to abort. + raise_if_stopped(f"complete the multipart upload to {self._path.uri}") return self._writer.complete(parts, **self._params("complete_multipart_upload")) def _on_created(self, creation: Future[S3MultipartUpload]) -> None: diff --git a/pyathena/util.py b/pyathena/util.py index 9e6343739..b70e78e77 100644 --- a/pyathena/util.py +++ b/pyathena/util.py @@ -7,7 +7,7 @@ import threading import time from collections.abc import Callable, Iterable, Iterator -from concurrent.futures import Future, wait +from concurrent.futures import CancelledError, Future, wait from contextlib import contextmanager from contextvars import ContextVar from re import Pattern @@ -328,6 +328,28 @@ def unstoppable_retries() -> Iterator[None]: _RETRY_STOP_EVENTS.reset(token) +def raise_if_stopped(operation: str) -> None: + """Raise if an event of the enclosing stop scopes is set. + + A writer calls this before a request that adds to or commits a write, + such as a part or the completion of a multipart upload, so that a write + whose owner has stopped, such as a batch operation after the failure of + another object, sends nothing more. The retries that + :func:`stop_retries_on` stops do not cover the first attempt of a new + request. The scopes that :func:`unstoppable_retries` clears do not stop + the writer either. + + Args: + operation: What the caller was about to do, for the error message. + + Raises: + concurrent.futures.CancelledError: If an event of the enclosing + :func:`stop_retries_on` scopes is set. + """ + if any(event.is_set() for event in _RETRY_STOP_EVENTS.get()): + raise CancelledError(f"Stopped before it could {operation}.") + + def wait_with_retry_stop(future: Future[_T], *, stop: threading.Event) -> _T: """Wait for work whose retries an event stops, relaying the stop of this context. diff --git a/tests/pyathena/filesystem/test_s3.py b/tests/pyathena/filesystem/test_s3.py index b4ba251f5..06a291ee1 100644 --- a/tests/pyathena/filesystem/test_s3.py +++ b/tests/pyathena/filesystem/test_s3.py @@ -18,9 +18,10 @@ import urllib.parse import urllib.request import uuid +import warnings import weakref from base64 import b64encode -from concurrent.futures import Future, ThreadPoolExecutor, wait +from concurrent.futures import CancelledError, Future, ThreadPoolExecutor, wait from datetime import UTC, datetime from itertools import chain from pathlib import Path @@ -37,6 +38,7 @@ from fsspec.compression import compr from fsspec.dircache import DirCache from fsspec.implementations.dirfs import DirFileSystem +from fsspec.implementations.local import LocalFileSystem from fsspec.implementations.memory import MemoryFileSystem import pyathena @@ -48,7 +50,7 @@ from pyathena.filesystem.s3_object import S3MultipartUpload, S3Object, S3ObjectType, S3StorageClass from pyathena.filesystem.s3_path import S3Path from pyathena.filesystem.s3_writer import S3MultipartStateError, _S3WriteSession -from pyathena.util import RetryConfig +from pyathena.util import RetryConfig, raise_if_stopped, stop_retries_on from tests import ENV from tests.pyathena.conftest import connect from tests.pyathena.util import ( @@ -105,6 +107,7 @@ def _make_fs(): fs.s3_additional_kwargs = {} fs._intrans = False fs.version_aware = False + fs.batch_size = None return fs @staticmethod @@ -2341,6 +2344,547 @@ def test_mv_nothing_within_maxdepth(self): fs.mv("s3://bucket/src", "s3://bucket/dst/", recursive=True, maxdepth=1) fs._call.assert_not_called() + # The per-object method that each batch operation calls. + BATCH_OPERATIONS = { + "cat": "cat_file", + "cat_ranges": "cat_file", + "pipe": "pipe_file", + "put": "put_file", + "get": "get_file", + "copy": "cp_file", + "mv": "_copy_file", + } + + @staticmethod + def _batch(fs, operation, count, **kwargs): + # Run a batch operation over count objects, with lists of sources and + # destinations, which are paired as given without lookups. + sources = [f"s3://bucket/src/{i}" for i in range(count)] + targets = [f"s3://bucket/dst/{i}" for i in range(count)] + local = [f"local/{i}" for i in range(count)] + if operation == "cat": + return fs.cat(sources, **kwargs) + if operation == "cat_ranges": + return fs.cat_ranges(sources, 0, 1, **kwargs) + if operation == "pipe": + return fs.pipe(dict.fromkeys(targets, b"x"), **kwargs) + if operation == "put": + return fs.put(local, targets, **kwargs) + if operation == "get": + return fs.get(sources, local, **kwargs) + return getattr(fs, operation)(sources, targets, **kwargs) + + @pytest.mark.parametrize("operation", list(BATCH_OPERATIONS)) + def test_batch_operation_runs_calls_concurrently(self, operation): + # GH-1322: the per-object calls of a batch operation run batch_size at + # a time. Run one at a time, the first would wait for the others + # until the barrier breaks. + fs = self._make_fs() + fs._delete_objects = mock.MagicMock() + barrier = threading.Barrier(3, timeout=EVENT_TIMEOUT) + + def call(*args, **kwargs): + barrier.wait() + # The bytes of a read, which also count as a copied file. + return b"x" + + setattr(fs, self.BATCH_OPERATIONS[operation], call) + self._batch(fs, operation, 3, batch_size=3) + + @pytest.mark.parametrize( + ("fs_batch_size", "batch_size", "expected"), + [(None, None, 4), (None, -1, 4), (None, 3, 3), (2, None, 2), (2, 3, 3), (-1, None, 4)], + ) + def test_batch_size(self, fs_batch_size, batch_size, expected): + # GH-1322 (O1): at most batch_size calls run at once, by default that + # of the filesystem; None or -1 is the size of the request pool + # (max_workers, 4). + fs = self._make_fs() + fs.batch_size = fs_batch_size + barrier = threading.Barrier(expected, timeout=EVENT_TIMEOUT) + lock = threading.Lock() + running = peak = 0 + + def cat_file(path, **kwargs): + nonlocal running, peak + with lock: + running += 1 + peak = max(peak, running) + try: + barrier.wait() + return b"" + finally: + with lock: + running -= 1 + + fs.cat_file = cat_file + # A multiple of each expected value, so that every round fills the barrier. + self._batch(fs, "cat", 12, batch_size=batch_size) + assert peak == expected + + @pytest.mark.parametrize("batch_size", [0, -2]) + def test_batch_size_invalid(self, batch_size): + # GH-1322: a batch_size that is neither positive, -1, nor None raises, + # per call and per filesystem. + fs = self._make_fs() + fs.cat_file = mock.MagicMock() + with pytest.raises(ValueError, match="batch_size"): + self._batch(fs, "cat", 2, batch_size=batch_size) + fs.cat_file.assert_not_called() + with pytest.raises(ValueError, match="batch_size"): + S3FileSystem(batch_size=batch_size, skip_instance_cache=True) + + @pytest.mark.parametrize( + ("count", "kwargs", "expected"), + [ + # A single object keeps the window of the filesystem. + (1, {}, None), + (1, {"max_workers": 3}, 3), + # Two at once: max_workers (4) // 2. + (2, {}, 2), + (2, {"max_workers": None}, 2), + (2, {"max_workers": 0}, 2), + (2, {"max_workers": 1}, 1), + # Four at once: a larger max_workers is capped at 4 // 4. + (4, {"max_workers": 3}, 1), + ], + ) + def test_batch_window(self, count, kwargs, expected): + # GH-1322 (O2): inside a batch, each object gets max_workers // n + # parallel requests, n being the calls that run at once, so that the + # batch keeps about max_workers requests in the request pool. + fs = self._make_fs() + windows = [] + lock = threading.Lock() + + def put_file(lpath, rpath, **kwargs): + with lock: + windows.append(kwargs.get("max_workers")) + + fs.put_file = put_file + self._batch(fs, "put", count, **kwargs) + assert windows == [expected] * count + + def test_batch_failure_stops_the_batch(self): + # GH-1322 (B2): a failed call starts no further call and stops the + # retries of the running one, whose backoff wait would take at least + # EVENT_TIMEOUT seconds; its error is raised once that call has ended. + fs = self._make_fs() + operation = ThrottledOperation(operation_name="GetObject") + throttled = retried(operation, RetryConfig(attempt=2, multiplier=EVENT_TIMEOUT)) + started, ended = set(), [] + + def cat_file(path, **kwargs): + started.add(path) + if path == "bucket/src/0": + assert operation.started.wait(EVENT_TIMEOUT) + raise PermissionError("Access Denied") + try: + return throttled() + finally: + ended.append(path) + + fs.cat_file = cat_file + begin = time.monotonic() + with pytest.raises(PermissionError, match="Access Denied"): + self._batch(fs, "cat", 4, batch_size=2) + assert time.monotonic() - begin < EVENT_TIMEOUT + assert operation.calls == 1 + assert started == {"bucket/src/0", "bucket/src/1"} + assert ended == ["bucket/src/1"] + + def test_batch_interrupt_waits_for_running_calls(self): + # GH-1322 (B3): an interrupt of the waiting thread starts no further + # call and is raised once the running calls have ended. + fs = self._make_fs() + lock = threading.Lock() + started, ended = [], [] + + def get_file(rpath, lpath, **kwargs): + with lock: + started.append(rpath) + # Runs until the interrupt stops the batch. + with pytest.raises(CancelledError): + poll_until(lambda: raise_if_stopped("end") or False, "the batch did not stop") + with lock: + ended.append(rpath) + + fs.get_file = get_file + with sigint_when(lambda: len(started) == 2), pytest.raises(KeyboardInterrupt): + self._batch(fs, "get", 4, batch_size=2) + assert sorted(started) == sorted(ended) == ["s3://bucket/src/0", "s3://bucket/src/1"] + + @pytest.mark.parametrize("on_error", ["omit", "return"]) + def test_cat_on_error(self, on_error): + # GH-1322 (B5): as in fsspec, a failed read does not stop the others + # with "omit", which leaves out its path, or "return", which returns + # its error in place of the contents. + fs = self._make_fs() + error = FileNotFoundError("bucket/src/0") + + def cat_file(path, **kwargs): + if path == "bucket/src/0": + raise error + return path.encode() + + fs.cat_file = cat_file + expected = {"bucket/src/1": b"bucket/src/1", "bucket/src/2": b"bucket/src/2"} + if on_error == "return": + expected = {"bucket/src/0": error, **expected} + assert self._batch(fs, "cat", 3, on_error=on_error, batch_size=2) == expected + + @pytest.mark.parametrize("count", [1, 2]) + def test_cat_on_error_stops_on_interrupt(self, count): + # GH-1322 (B5): a failure that is not an Exception, such as an + # interrupt raised in a call, stops the batch also with on_error + # "return", as fsspec catches Exception only. A single path runs in + # the calling thread. + fs = self._make_fs() + fs.cat_file = mock.MagicMock(side_effect=[KeyboardInterrupt, b"x"]) + with pytest.raises(KeyboardInterrupt): + self._batch(fs, "cat", count, on_error="return", batch_size=1) + assert fs.cat_file.call_count == 1 + + def test_cat_interrupted_while_assembling_keeps_no_contents(self): + # GH-1322 (F5): an interrupt while cat() builds its dict, after the + # reads have ended, keeps no contents in the traceback. + fs = self._make_fs() + + class Contents: + pass + + class Path(str): + # Interrupts the insertion of the first path into the dict. + def __hash__(self): + raise KeyboardInterrupt + + refs = [] + + def cat_file(path, **kwargs): + contents = Contents() + refs.append(weakref.ref(contents)) + return contents + + fs.cat_file = cat_file + fs.expand_path = mock.MagicMock(return_value=[Path("bucket/a"), Path("bucket/b")]) + gc.disable() + try: + with pytest.raises(KeyboardInterrupt) as info: + fs.cat(["s3://bucket/a", "s3://bucket/b"], batch_size=2) + assert info.value.__traceback__ is not None + assert len(refs) == 2 + assert all(ref() is None for ref in refs) + finally: + gc.enable() + + def test_copy_on_error_ignore(self): + # GH-1322 (B5): as in fsspec, on_error="ignore" skips a missing source + # without stopping the other copies, and raises the other errors. + fs = self._make_fs() + copied = [] + + def cp_file(path1, path2, **kwargs): + if path1 == "s3://bucket/src/0": + raise FileNotFoundError(path1) + if path1 == "s3://bucket/src/2": + raise PermissionError(path1) + copied.append(path1) + + fs.cp_file = cp_file + self._batch(fs, "copy", 2, on_error="ignore", batch_size=2) + assert copied == ["s3://bucket/src/1"] + with pytest.raises(PermissionError): + self._batch(fs, "copy", 3, on_error="ignore", batch_size=1) + + def test_mv_deletes_nothing_after_a_failed_copy(self): + # GH-1322 (O3): mv() deletes once every copy has succeeded, also + # when another copy ends after the failure. + fs = self._make_fs() + fs._delete_objects = mock.MagicMock() + both = threading.Barrier(2, timeout=EVENT_TIMEOUT) + failed = threading.Event() + succeeded = [] + + def copy_file(path1, path2, **kwargs): + both.wait() + if path1 == "s3://bucket/src/0": + failed.set() + raise PermissionError(path1) + # Succeeds once the failure has stopped the batch. + assert failed.wait(EVENT_TIMEOUT) + with pytest.raises(CancelledError): + poll_until(lambda: raise_if_stopped("end") or False, "the batch did not stop") + succeeded.append(path1) + return True + + fs._copy_file = copy_file + with pytest.raises(PermissionError): + self._batch(fs, "mv", 2, batch_size=2) + # The other copy succeeded after the failure, and nothing was deleted. + assert succeeded == ["s3://bucket/src/1"] + fs._delete_objects.assert_not_called() + + @pytest.mark.parametrize( + "case", ["copy-destination", "copy-source", "pipe", "put", "get-stream"] + ) + def test_batch_conflicting_pairs(self, case): + # GH-1322: pairs whose outcome depended on the order of the sequential + # loop raise before anything is transferred, as in mv(). + fs = self._make_fs() + methods = set(self.BATCH_OPERATIONS.values()) + for method in methods: + setattr(fs, method, mock.MagicMock()) + stream = io.BytesIO() + operations = { + # Two sources with one destination. + "copy-destination": lambda: fs.copy( + ["s3://bucket/a", "s3://bucket/b"], ["s3://bucket/c"] * 2 + ), + # A destination that is another source. + "copy-source": lambda: fs.copy( + ["s3://bucket/a", "s3://bucket/b"], ["s3://bucket/b", "s3://bucket/c"] + ), + # Two keys of one object. + "pipe": lambda: fs.pipe({"s3://bucket/a": b"1", "bucket/a": b"2"}), + "put": lambda: fs.put(["x", "y"], ["s3://bucket/a", "bucket/a"]), + # One stream for two downloads. + "get-stream": lambda: fs.get(["s3://bucket/a", "s3://bucket/b"], [stream] * 2), + } + with pytest.raises(ValueError, match="Cannot"): + operations[case]() + assert not any(getattr(fs, method).called for method in methods) + + def test_get_hard_links_conflict(self, tmp_path): + # GH-1322: two local paths of one file conflict. + fs = self._make_fs() + fs.get_file = mock.MagicMock() + (tmp_path / "a").write_bytes(b"") + os.link(tmp_path / "a", tmp_path / "b") + with pytest.raises(ValueError, match="same destination"): + fs.get(["s3://bucket/a", "s3://bucket/b"], [str(tmp_path / "a"), str(tmp_path / "b")]) + fs.get_file.assert_not_called() + + def test_copy_skipped_directory_does_not_conflict(self): + # GH-1322: a listed source without an object at its key, a prefix + # that cp_file() skips, writes no destination, so its destination + # being another source is no conflict. + fs = self._make_fs() + fs.cp_file = mock.MagicMock() + fs._head_object = mock.MagicMock( + side_effect=lambda path, **kwargs: ( + None if path.endswith("dir") else self._file_object("x") + ) + ) + fs.copy(["s3://bucket/dir", "s3://bucket/x"], ["s3://bucket/x", "s3://bucket/y"]) + assert fs.cp_file.call_count == 2 + + def test_get_put_pass_kwargs_to_expansion(self, tmp_path): + # GH-1322: as in fsspec, get() and put() pass their keyword arguments + # to the expansion of the sources, such as prefix or assume_literal. + fs = self._make_fs() + fs.put_file = mock.MagicMock() + (tmp_path / "a").write_bytes(b"") + with mock.patch.object( + LocalFileSystem, "expand_path", autospec=True, return_value=[str(tmp_path / "a")] + ) as expand_path: + fs.put(str(tmp_path / "a"), "s3://bucket/dst/", assume_literal=True) + assert expand_path.call_args.kwargs["assume_literal"] is True + fs.expand_path = mock.MagicMock(return_value=["bucket/root/a"]) + fs.get_file = mock.MagicMock() + fs.isdir = mock.MagicMock(return_value=False) + fs.get("s3://bucket/root/", str(tmp_path / "out") + "/", recursive=True, prefix="a") + assert fs.expand_path.call_args.kwargs["prefix"] == "a" + + @pytest.mark.parametrize("on_error", ["return", "raise"]) + def test_cat_ranges_on_error(self, on_error): + # GH-1322 (B5): with on_error="return", a failed range returns its + # error in place; otherwise it is raised. + fs = self._make_fs() + error = PermissionError("bucket/src/0") + fs.cat_file = lambda path, start, end, **kwargs: ( + (_ for _ in ()).throw(error) if path == "s3://bucket/src/0" else b"x" + ) + if on_error == "return": + assert self._batch(fs, "cat_ranges", 2, on_error=on_error) == [error, b"x"] + else: + with pytest.raises(PermissionError): + self._batch(fs, "cat_ranges", 2, on_error=on_error) + + def test_range_read_stopped_sends_no_request(self): + # GH-1322 (B2): a range read whose scope has stopped, such as one that + # waited behind other requests of a stopped batch, sends no GetObject. + file = mock.MagicMock() + stop = threading.Event() + stop.set() + with stop_retries_on(stop), pytest.raises(CancelledError): + S3File._get_range(file, S3Path("bucket", "key"), (0, 1)) + file.fs.core.get_object.assert_not_called() + + def test_copy_in_place_does_not_conflict(self): + # GH-1322: a pair whose destination is its own source rewrites the + # object in place, such as with MetadataDirective="REPLACE". + fs = self._make_fs() + fs.cp_file = mock.MagicMock() + fs.copy(["s3://bucket/a", "s3://bucket/b"], ["s3://bucket/a", "s3://bucket/c"]) + assert fs.cp_file.call_count == 2 + + def test_get_outfile_runs_in_order(self): + # GH-1322: downloads into one output stream run one after another, in + # the order of the sources, in the calling thread, which may own the + # stream; get_file() ignores the local paths then, so they do not + # conflict. + fs = self._make_fs() + threads = set() + + def get_file(rpath, lpath, outfile, **kwargs): + threads.add(threading.current_thread()) + outfile.write(rpath.encode() + b";") + + fs.get_file = get_file + sources = [f"s3://bucket/src/{i}" for i in range(3)] + outfile = io.BytesIO() + fs.get(sources, ["unused"] * 3, outfile=outfile) + assert outfile.getvalue() == b"".join(p.encode() + b";" for p in sources) + assert threads == {threading.current_thread()} + + def test_batch_callback_is_used_one_thread_at_a_time(self): + # GH-1322: the callback of a batch and its branches are used from the + # threads of the batch one at a time; fsspec callbacks need not be + # thread-safe, nor their branches independent of them. + fs = self._make_fs() + busy = threading.Lock() + overlaps = [] + + exited = [] + + class Child(Callback): + # A branch that needs its context: fsspec enters it with "with", + # and passes on the callback that __enter__() returns. + entered = False + + def __enter__(self): + SharedCallback._use(self) + entered = Child() + entered.entered = True + return entered + + def __exit__(self, *exc_args): + SharedCallback._use(self) + exited.append(self.entered) + + def set_size(self, size): + assert self.entered + SharedCallback._use(self) + + def relative_update(self, inc=1): + assert self.entered + SharedCallback._use(self) + + def close(self): + SharedCallback._use(self) + + class SharedCallback(Callback): + def _use(self): + if not busy.acquire(blocking=False): + overlaps.append(threading.current_thread().name) + return + try: + time.sleep(0.01) + finally: + busy.release() + + def branched(self, path_1, path_2, **kwargs): + self._use() + # A branch that shares the lock of its parent. + return Child() + + def set_size(self, size): + self._use() + + def relative_update(self, inc=1): + self._use() + + def put_file(lpath, rpath, callback, **kwargs): + callback.set_size(5) + for _ in range(5): + callback.relative_update(1) + + fs.put_file = put_file + self._batch(fs, "put", 4, batch_size=4, callback=SharedCallback()) + assert overlaps == [] + # The branches that branched() returned are exited. + assert exited == [False] * 4 + + @pytest.mark.parametrize("kwargs", [{"batch_size": 2}, {"max_workers": 2}]) + def test_rm_batch_size(self, kwargs): + # GH-1322 (O4): rm(batch_size=) counts the DeleteObjects requests at + # once; max_workers is its deprecated alias. + fs = self._make_fs() + fs._delete_objects = mock.MagicMock() + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + fs.rm(["s3://bucket/a?versionId=v1"], **kwargs) + assert fs._delete_objects.call_args.kwargs["max_workers"] == 2 + deprecations = [w for w in caught if issubclass(w.category, DeprecationWarning)] + assert len(deprecations) == ("max_workers" in kwargs) + with pytest.raises(TypeError, match="max_workers"): + fs.rm(["s3://bucket/a?versionId=v1"], batch_size=2, max_workers=2) + + def test_get_file_stopped_writes_no_further_block(self): + # GH-1322 (B2): a download whose batch has stopped writes no further + # block, including the one read when the stop came. + fs = self._make_fs() + stop = threading.Event() + remote = mock.MagicMock(size=8, blocksize=4) + remote.__enter__.return_value = remote + + blocks = [b"abcd", b"efgh", b""] + + def read(size): + # The batch stops while the first block is being read. + stop.set() + return blocks.pop(0) + + remote.read.side_effect = read + fs.open = mock.MagicMock(return_value=remote) + outfile = io.BytesIO() + with stop_retries_on(stop), pytest.raises(CancelledError): + fs.get_file("s3://bucket/key?versionId=v1", outfile=outfile) + assert outfile.getvalue() == b"" + + def test_multipart_copy_annotations_retry_after_stop(self): + # GH-1322: the annotations of a committed multipart copy are copied + # to the end also once an enclosing batch has stopped, retries + # included, so that the object keeps all of them. + fs = self._make_fs() + stop = threading.Event() + plan = S3MultipartCopyPlan( + source=S3Path("bucket", "src"), + destination=S3Path("bucket", "dst"), + size=2, + ranges=((0, 1), (1, 2)), + annotations=("a",), + ) + fs.core.plan_multipart_copy = mock.MagicMock(return_value=plan) + fs.core.create_multipart_upload = mock.MagicMock( + return_value=S3MultipartUpload({"Bucket": "bucket", "Key": "dst", "UploadId": "id"}) + ) + fs.core.upload_part_copy = lambda part_number, **kwargs: SimpleNamespace( + etag='"e"', part_number=part_number + ) + + def complete(*args, **kwargs): + # The batch stops once the object is committed. + stop.set() + return SimpleNamespace(version_id="v1", etag='"e"') + + fs.core.complete_multipart_upload = complete + operation = ThrottledOperation(failures=1, operation_name="PutObjectAnnotation") + fs.core.copy_object_annotation = retried(operation, RetryConfig(attempt=2, multiplier=0)) + with stop_retries_on(stop): + fs._copy_object_with_multipart_upload(plan.source, plan.destination, max_workers=2) + assert operation.calls == 2 + def test_cp_file_multipart_parameters(self): # GH-967: block_size and max_workers control a multipart copy and are # not sent to S3. @@ -6635,9 +7179,11 @@ def test_mv_source_lookup_and_delete_params(self): fs = self._stubbed_fs() with Stubber(fs._client) as stubber: stub_move_with_source_params(stubber) + # One copy at a time, in the order of the stubbed responses. fs.mv( ["s3://bucket/a", "s3://bucket/a/x"], ["s3://bucket/c", "s3://bucket/c"], + batch_size=1, **MOVE_SOURCE_PARAMS, ) stubber.assert_no_pending_responses() diff --git a/tests/pyathena/filesystem/test_s3_writer.py b/tests/pyathena/filesystem/test_s3_writer.py index 5a38092a9..8839a8cba 100644 --- a/tests/pyathena/filesystem/test_s3_writer.py +++ b/tests/pyathena/filesystem/test_s3_writer.py @@ -30,7 +30,8 @@ S3MultipartWriter, _S3WriteSession, ) -from pyathena.util import RetryConfig +from pyathena.util import RetryConfig, stop_retries_on +from tests.pyathena.util import poll_until def _make_writer(block_size=4, **kwargs): @@ -257,8 +258,9 @@ def _make_session(executor, *, max_workers=1, part_errors=None, gate=None): gate: An event that each part request waits for, if any. Returns: - The session and its recorded requests: ``completed`` (the part - numbers of each completion) and ``aborts``. + The session and its recorded requests: ``created``, ``parts`` (the + numbers of the part requests sent), ``completed`` (the part numbers + of each completion), and ``aborts``. """ client = boto3.client( "s3", region_name="us-east-1", aws_access_key_id="dummy", aws_secret_access_key="dummy" @@ -267,9 +269,14 @@ def _make_session(executor, *, max_workers=1, part_errors=None, gate=None): core.MULTIPART_UPLOAD_MIN_PART_SIZE = 4 core.MULTIPART_UPLOAD_MAX_PART_SIZE = 8 errors = dict(part_errors or {}) - requests = SimpleNamespace(completed=[], aborts=0) + requests = SimpleNamespace(created=0, parts=[], completed=[], aborts=0) + + def create(path, **params): + requests.created += 1 + return S3MultipartUpload({"Bucket": "bucket", "Key": "key", "UploadId": "uploadid"}) def upload_part(upload, part_number, body, **params): + requests.parts.append(part_number) if gate is not None: assert gate.wait(5) if part_number in errors: @@ -282,9 +289,7 @@ def complete(upload, parts, **params): def abort(upload, **params): requests.aborts += 1 - core.create_multipart_upload = lambda path, **params: S3MultipartUpload( # type: ignore[method-assign] - {"Bucket": "bucket", "Key": "key", "UploadId": "uploadid"} - ) + core.create_multipart_upload = create # type: ignore[method-assign] core.upload_part = upload_part # type: ignore[method-assign] core.complete_multipart_upload = complete # type: ignore[method-assign] core.abort_multipart_upload = abort # type: ignore[method-assign] @@ -301,6 +306,45 @@ def _start(session): class TestS3WriteSession: + @pytest.mark.parametrize("stopped", ["start", "submit", "queued", "completion"]) + def test_stopped_write_sends_nothing_more(self, stopped): + # GH-1322 (B2): once an enclosing stop scope is stopped, such as that + # of a batch operation after another object failed, the write sends + # no creation, part, or completion: not when it is asked to, not for + # a part that waited behind another request, and not after the wait + # for the parts. The stop of the retries does not cover these first + # attempts. The write fails, and the abort aborts the upload. + stop = threading.Event() + gate = threading.Event() + with S3ThreadPool(max_workers=1) as executor, stop_retries_on(stop): + session, requests = _make_session(executor, max_workers=2, gate=gate) + if stopped == "start": + stop.set() + with pytest.raises(CancelledError): + session.start() + else: + _start(session) + if stopped == "submit": + stop.set() + with pytest.raises(CancelledError): + session.submit_part(1, b"abcd") + else: + session.submit_part(1, b"abcd") + poll_until(lambda: requests.parts == [1], "part 1 did not start") + if stopped == "queued": + # Waits behind part 1 in the only thread. + session.submit_part(2, b"efgh") + stop.set() + gate.set() + else: + threading.Timer(0.2, lambda: (stop.set(), gate.set())).start() + with pytest.raises(CancelledError): + session.complete() + session.abort() + expected_parts = [1] if stopped in ("queued", "completion") else [] + assert (requests.parts, requests.completed) == (expected_parts, []) + assert (requests.created, requests.aborts) == ((0, 0) if stopped == "start" else (1, 1)) + def test_failed_part_is_released(self): # GH-1211: the session drops the request of a failed part, whose body # the executor frees, and keeps no reference to its error once it is