diff --git a/docs/aio.md b/docs/aio.md index d37bdfe5d..ca8f3a90c 100644 --- a/docs/aio.md +++ b/docs/aio.md @@ -283,7 +283,7 @@ its coroutines run the blocking boto3 calls in threads that PyAthena owns, sized | **Blocking calls** | Caller's thread | Thread pool of the filesystem | | **Parallelism** | Request pool of the filesystem | `asyncio.gather` on a thread pool; requests in the request pool | | **File handles** | `S3File` with thread pool | `AioS3File` with thread pool; `AioS3StreamedFile` from `open_async()` | -| **Bulk delete** | Request pool, `max_workers` requests at once | Request pool, `max_workers` requests at once | +| **Bulk delete** | Request pool, `batch_size` requests at once | Request pool, `batch_size` requests at once | | **Multipart copy** | Request pool, `max_workers` requests per copy | Request pool, `max_workers` requests per copy | | **Best for** | Synchronous applications | Async frameworks (FastAPI, aiohttp, etc.) | @@ -330,14 +330,26 @@ async with AioS3FileSystem(asynchronous=True) as fs: Cancelling a coroutine of the filesystem also stops the retries of the S3 requests that its call makes in these threads, including the part uploads and range reads that the call runs in parallel: no further attempt is made, and the wait before the next attempt -ends. A request that has already been sent still runs to its end in its thread. The -requests that run to their end after a cancellation keep their retries: the commit of a -file, the CompleteMultipartUpload request of a multipart copy, and the aborts of -multipart uploads. -The range reads of a file from `open_async()`, the DeleteObjects requests of `_rm()`, and the annotation -copies of a multipart copy are waited for before the cancellation is raised, also when -the shutdown of the event loop at the end of `asyncio.run()` cancels the coroutine; a -second cancellation stops that wait. +ends. A write, such as one of `_put_file()` or `_pipe_file()`, sends no further part and +aborts its multipart upload instead of completing it, and `_get_file()` reads and writes +no further block. In a transaction, a cancelled `_put_file()` or `_pipe_file()` is +discarded and left out of the commit. A request that has already been sent still runs to its end in its +thread. The requests that run to their end after a cancellation keep their retries: the +commit of a file, the CompleteMultipartUpload request of a multipart copy, and the +aborts of multipart uploads. +The call, the range reads of a file from `open_async()`, the DeleteObjects requests of +`_rm()`, and the annotation copies of a multipart copy are waited for before the +cancellation is raised, also when the shutdown of the event loop at the end of +`asyncio.run()` cancels the coroutine; a second cancellation stops that wait. + +`_cat()`, `_cat_ranges()`, `_pipe()`, `_put()`, `_get()`, `_copy()`, and `_mv()` run the +coroutines of their objects as tasks, `batch_size` of them at a time as in fsspec, except +`_mv()`, which runs all its copies at once. Cancelling one of them cancels those tasks +and waits for them before the cancellation is raised, also when the shutdown of the +event loop cancels them together with the coroutine. +A recursive `_put()` sends no request for the local directories. `_rm()` sends +`batch_size` DeleteObjects requests at a time, `max_workers` by default and with `-1`; its +`max_workers` argument is a deprecated alias of `batch_size`. ### Executor strategy diff --git a/pyathena/filesystem/s3.py b/pyathena/filesystem/s3.py index af67f7fd6..74d0ebbe5 100644 --- a/pyathena/filesystem/s3.py +++ b/pyathena/filesystem/s3.py @@ -1127,6 +1127,31 @@ def find( return {f.name: f for f in files} return [f.name for f in files] + def glob(self, path, maxdepth=None, **kwargs): + """Find the paths that match a glob pattern. + + As fsspec's ``glob()``, except that only the keys that start with the + part of the last path component before the first wildcard are listed, + as with fsspec's asynchronous ``_glob()``. + + Args: + path: S3 glob pattern. + maxdepth: Maximum depth of the paths below the literal part of the + pattern, at least 1. + **kwargs: Additional arguments passed to ``find``, such as + ``detail``. + + Returns: + The matching paths, or a dictionary of them and their objects with + ``detail=True``. + """ + stripped = self._strip_protocol(path) + wildcards = [i for i in (stripped.find(c) for c in "*?[") if i >= 0] + stem = stripped[: min(wildcards)].rpartition("/")[2] if wildcards else "" + if stem and "prefix" not in kwargs: + kwargs["prefix"] = stem + return super().glob(path, maxdepth=maxdepth, **kwargs) + def expand_path(self, path, recursive=False, maxdepth=None, **kwargs) -> list[str]: """Expand glob patterns and directories into the paths they match. @@ -1288,18 +1313,35 @@ def rm(self, path, recursive=False, maxdepth=None, batch_size=None, **kwargs) -> TypeError: If both ``batch_size`` and ``max_workers`` are given. OSError: If S3 could not delete some of the objects. """ + max_workers = self._delete_batch_limit(batch_size, kwargs) + paths = self._delete_paths(path, recursive=recursive, maxdepth=maxdepth) + self._delete_objects(paths, max_workers=max_workers, **kwargs) + + def _delete_batch_limit(self, batch_size: int | None, kwargs: dict[str, Any]) -> int: + """Return the number of DeleteObjects requests of ``rm()`` to send at once. + + Args: + batch_size: The ``batch_size`` argument of ``rm()``. + kwargs: The other arguments of ``rm()``, from which its deprecated + alias ``max_workers`` is removed. + + Returns: + The number of requests. + + Raises: + ValueError: If the number is neither positive, -1, nor None. + TypeError: If both ``batch_size`` and ``max_workers`` are given. + """ 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, + stacklevel=3, ) 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, max_workers=max_workers, **kwargs) + return self._batch_limit(batch_size) def _delete_paths( self, path: str | list[str], recursive: bool = False, maxdepth: int | None = None @@ -2671,32 +2713,7 @@ def put( ``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)) + pairs = self._put_pairs(lpath, rpath, recursive=recursive, maxdepth=maxdepth, **kwargs) if not pairs: return if len({rp for _, rp in pairs}) != len(pairs): @@ -2719,6 +2736,46 @@ def put_pair(pair: tuple[str, str]) -> None: locked.set_size(len(pairs)) self._run_batch(put_pair, pairs, concurrency) + def _put_pairs( + self, lpath, rpath, recursive=False, maxdepth=None, **kwargs + ) -> list[tuple[str, str]]: + """Pair the local sources of ``put()`` with their S3 destinations as fsspec does. + + 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 include the directories with their contents. + maxdepth: Maximum depth of a recursive copy. + **kwargs: Additional arguments passed to the ``expand_path()`` of + the local filesystem. + + Returns: + The local paths, including the directories of a recursive copy, + with their destinations without the protocol. + """ + if isinstance(lpath, list) and isinstance(rpath, list): + return [(lp, self._strip_protocol(rp)) for lp, rp in zip(lpath, rpath, strict=False)] + 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)) + # A list of destinations or one path, also path-like, as in fsspec. + rpath = self._strip_protocol(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) + return list(zip(lpaths, rpaths, strict=True)) + def put_file( self, lpath: str, diff --git a/pyathena/filesystem/s3_async.py b/pyathena/filesystem/s3_async.py index 34add38cb..6c9880b1d 100644 --- a/pyathena/filesystem/s3_async.py +++ b/pyathena/filesystem/s3_async.py @@ -11,15 +11,17 @@ import asyncio import contextlib +import contextvars import functools import logging import mimetypes import os +import threading import warnings -from collections.abc import Callable, Iterable -from concurrent.futures import Future, wait +from collections.abc import Awaitable, Callable, Iterable +from concurrent.futures import CancelledError, Future, wait from io import BytesIO -from typing import TYPE_CHECKING, Any, cast +from typing import TYPE_CHECKING, Any, ClassVar, cast from fsspec.asyn import AbstractAsyncStreamedFile, AsyncFileSystem, sync from fsspec.callbacks import _DEFAULT_CALLBACK @@ -42,6 +44,7 @@ from pyathena.filesystem.s3_path import S3Path from pyathena.filesystem.s3_path_pairing import S3PathPairing from pyathena.filesystem.s3_writer import _S3WriteSession, _S3WriteState +from pyathena.util import retries_stopped if TYPE_CHECKING: from datetime import datetime @@ -51,6 +54,194 @@ _logger = logging.getLogger(__name__) +class _S3BatchTasks: + """The tasks of the per-object calls of a batch coroutine of fsspec. + + fsspec runs the per-object coroutines of its batch coroutines, such as + ``_cat()`` and ``_put()``, as tasks, and a cancellation of the batch + coroutine leaves them running. Each per-object coroutine of + ``AioS3FileSystem`` adds its task with :meth:`track` to the batch that + runs in its context, which cancels them and waits for them when the + batch coroutine fails or is cancelled. + """ + + _current: ClassVar[contextvars.ContextVar[_S3BatchTasks | None]] = contextvars.ContextVar( + "_S3BatchTasks", default=None + ) + + def __init__(self) -> None: + """Initialize the batch with no task.""" + self._tasks: set[asyncio.Task[Any]] = set() + self._owner: asyncio.Task[Any] | None = None + + @classmethod + def track(cls) -> None: + """Add the current task to the batch that runs in its context, if any. + + A task that the batch coroutine itself runs in is not added. + """ + batch = cls._current.get() + task = asyncio.current_task() + if batch is not None and task is not None and task is not batch._owner: + batch._tasks.add(task) + task.add_done_callback(batch._tasks.discard) + + async def run(self, batch: Awaitable[T]) -> T: + """Await a batch coroutine, and on its failure or cancellation wait for its tasks. + + A task that is already being cancelled, as the shutdown of the event + loop with ``asyncio.run()`` cancels every task, is not cancelled + again, so that the wait for its own calls is not stopped. A repeated + cancellation stops the wait for the tasks. + + Args: + batch: The batch coroutine, created and not yet started. + + Returns: + The return value of the batch coroutine. + """ + self._owner = asyncio.current_task() + token = self._current.set(self) + try: + return await batch + except BaseException: + tasks = [task for task in self._tasks if not task.done()] + for task in tasks: + if not task.cancelling(): + task.cancel() + if tasks: + await asyncio.wait(tasks) + raise + finally: + self._current.reset(token) + + +class _S3TransactionWrite: + """A write of a file of an ``AioS3FileSystem`` transaction, which a cancellation discards. + + The write runs with :meth:`run` in a thread of the filesystem, through + :meth:`S3Executor.run() `, + whose cancellation stops the stop scope of the write, and opens its file + with :meth:`open`, which adds the file to the transaction unless the + write has been stopped. A write that ends stopped drops its file from + the transaction and discards it, so that the wait of the cancelled + coroutine for the write includes the discard. The coroutine calls + :meth:`cancel` once the cancellation has been raised, which drops the + file from the transaction also when the write had ended before it was + stopped, and then discards it. + + Args: + fs: The filesystem, which is in a transaction. + write: The write, which opens its file with the callable that it + receives, which takes the arguments of ``open()``. + """ + + def __init__(self, fs: AioS3FileSystem, write: Callable[[Callable[..., S3File]], None]) -> None: + """Initialize the write. + + Args: + fs: The filesystem, which is in a transaction. + write: The write, which opens its file with the callable that it + receives, which takes the arguments of ``open()``. + """ + self._fs = fs + self._transaction = fs.transaction + self._write = write + self._lock = threading.Lock() + self._files: list[S3File] = [] + self._cancelled = False + self._ended = False + + def run(self) -> None: + """Run the write, and drop and discard its file if it ends stopped. + + Raises: + concurrent.futures.CancelledError: If the write was stopped, also + when the write itself succeeded, since its file is discarded. + """ + try: + self._write(self.open) + finally: + with self._lock: + self._ended = True + stopped = self._cancelled or retries_stopped() + if stopped: + self._cancelled = True + # Also here, for a write that a scope other than the + # cancellation of the coroutine has stopped. + self._detach_all() + if stopped: + self.discard() + if stopped: + raise CancelledError("Stopped before the transaction could commit the write.") + + def open(self, path: str, mode: str, **kwargs) -> S3File: + """Open the file of the write, which joins the transaction unless the write is stopped. + + Args: + path: S3 path (s3://bucket/key) of the file. + mode: The binary write mode of the file. + **kwargs: The other arguments of ``open()``, except + ``compression``. With ``autocommit=True``, as with + ``open()``, the file commits when it is closed and does not + join the transaction. + + Returns: + The file. + """ + autocommit = kwargs.pop("autocommit", False) + file = self._fs._open(self._fs._strip_protocol(path), mode, autocommit=autocommit, **kwargs) + with self._lock: + self._files.append(file) + # Under the lock, so that a cancellation either finds the file + # in the transaction or keeps it out. + if not autocommit and not self._cancelled and not retries_stopped(): + self._transaction.files.append(file) + return file + + def cancel(self) -> bool: + """Drop the file of the write from the transaction. + + Returns: + Whether the write ended before it was stopped, in which case + the caller discards the file with :meth:`discard`; otherwise the + thread of the write discards it when the write ends, or has. + """ + with self._lock: + ended_unstopped = self._ended and not self._cancelled + self._cancelled = True + self._detach_all() + return ended_unstopped + + def discard(self) -> None: + """Discard the file of the write. + + A failed discard is logged instead of raised, so that the + cancellation is raised. + """ + with self._lock: + files = list(self._files) + for file in files: + try: + file.discard() + except Exception: + _logger.exception(f"Failed to discard the cancelled write to s3://{file.path}.") + + def _detach_all(self) -> None: + """Drop the files of the write from the transaction; called under the lock.""" + for file in self._files: + while True: + try: + self._transaction.files.remove(file) + except ValueError: + break + except (IndexError, RuntimeError): + # Another write changed the transaction during the search, + # which compares the files with their __eq__(). + continue + break + + class AioS3FileSystem(AsyncFileSystem): """An async filesystem interface for Amazon S3 using fsspec's AsyncFileSystem. @@ -203,6 +394,7 @@ async def _ls(self, path: str, detail: bool = False, **kwargs) -> list[S3Object] async def _cat_file( self, path: str, start: int | None = None, end: int | None = None, **kwargs ) -> bytes: + _S3BatchTasks.track() return await self._thread_pool.run( self._sync_fs.cat_file, path, start=start, end=end, **kwargs ) @@ -216,15 +408,46 @@ async def _rm_file(self, path: str, **kwargs) -> None: async def _pipe_file( self, path: str, value: bytes | bytearray | memoryview, mode: str = "overwrite", **kwargs ) -> None: + _S3BatchTasks.track() if self._intrans: # The transaction belongs to this filesystem, not to the internal # S3FileSystem, so write through open() to defer the commit to it. - await self._thread_pool.run(self._pipe_file_in_transaction, path, value, mode, **kwargs) + await self._write_in_transaction( + functools.partial(self._pipe_file_in_transaction, path, value, mode, **kwargs) + ) return await self._thread_pool.run(self._sync_fs.pipe_file, path, value, mode=mode, **kwargs) + async def _write_in_transaction(self, write: Callable[[Callable[..., S3File]], None]) -> None: + """Run a write of a file of this filesystem's transaction in a thread of the filesystem. + + A cancellation keeps the file out of the transaction, which would + otherwise commit it, and discards it: a write whose last part was + running when it was cancelled has nothing left that its stop would + skip. A write that the cancellation stops discards its file before + it ends, which the cancellation waits for; the file of a write that + ended before is discarded in a thread of the filesystem without + being waited for. See :class:`_S3TransactionWrite`. + + Args: + write: The write, which opens its file with the callable that it + receives, which takes the arguments of ``open()``. + """ + transaction_write = _S3TransactionWrite(self, write) + try: + await self._thread_pool.run(transaction_write.run) + except asyncio.CancelledError: + if transaction_write.cancel(): + self._thread_pool.submit(transaction_write.discard) + raise + def _pipe_file_in_transaction( - self, path: str, value: bytes | bytearray | memoryview, mode: str, **kwargs + self, + path: str, + value: bytes | bytearray | memoryview, + mode: str, + open_file: Callable[..., S3File], + **kwargs, ) -> None: """Write bytes into the path as a file of this filesystem's transaction. @@ -235,6 +458,7 @@ def _pipe_file_in_transaction( opened in ``xb`` mode: raise FileExistsError when the object already exists, including one created before the transaction is committed, which is not replaced. + open_file: Opens the file, with the arguments of ``open()``. **kwargs: Additional parameters passed to ``open()``, except ``compression``, with which the value is compressed before it is written, as in :meth:`S3FileSystem.pipe_file`. @@ -252,7 +476,7 @@ def _pipe_file_in_transaction( block_size = kwargs.get("block_size") or self._sync_fs.default_block_size # The size in bytes; the length of a memoryview counts its items. self.core.check_multipart_upload_size(memoryview(value).nbytes, block_size) - self.open(path, "xb" if mode == "create" else "wb", **kwargs)._write_and_close(value) + open_file(path, "xb" if mode == "create" else "wb", **kwargs)._write_and_close(value) async def _put_file( self, @@ -262,10 +486,13 @@ async def _put_file( mode: str = "overwrite", **kwargs, ) -> None: + _S3BatchTasks.track() if self._intrans: # See _pipe_file. - await self._thread_pool.run( - self._put_file_in_transaction, lpath, rpath, callback, mode, **kwargs + await self._write_in_transaction( + functools.partial( + self._put_file_in_transaction, lpath, rpath, callback, mode, **kwargs + ) ) return await self._thread_pool.run( @@ -273,7 +500,13 @@ async def _put_file( ) def _put_file_in_transaction( - self, lpath: str, rpath: str, callback, mode: str, **kwargs + self, + lpath: str, + rpath: str, + callback, + mode: str, + open_file: Callable[..., S3File], + **kwargs, ) -> None: """Upload a local file as a file of this filesystem's transaction. @@ -288,6 +521,7 @@ def _put_file_in_transaction( opened in ``xb`` mode: raise FileExistsError when the object already exists, including one created before the transaction is committed, which is not replaced. + open_file: Opens the file, with the arguments of ``open()``. **kwargs: Additional S3 parameters (e.g., ContentType, StorageClass). The ``block_size``, ``max_workers``, and ``s3_additional_kwargs`` parameters of ``open()`` are also accepted. @@ -317,7 +551,7 @@ def _put_file_in_transaction( # See S3FileSystem.put_file. with open(lpath, "rb") as local: - self.open( + open_file( rpath, "xb" if mode == "create" else "wb", block_size=block_size, @@ -327,6 +561,7 @@ def _put_file_in_transaction( self.invalidate_cache(rpath) async def _get_file(self, rpath: str, lpath: str, callback=_DEFAULT_CALLBACK, **kwargs) -> None: + _S3BatchTasks.track() await self._thread_pool.run( self._sync_fs.get_file, rpath, lpath, callback=callback, **kwargs ) @@ -350,34 +585,45 @@ async def _rm( """Delete objects with DeleteObjects requests. See :meth:`S3FileSystem.rm`. The DeleteObjects requests run in - parallel, at most ``max_workers`` at once, in the request pool of the + parallel, at most ``batch_size`` at once, in the request pool of the internal ``S3FileSystem``. A cancellation cancels the requests that have not started and waits for the others. Args: path: S3 path (s3://bucket/key) or list of paths to delete. recursive: Whether to delete all objects below the paths. - batch_size: Accepted for fsspec compatibility; not used. + batch_size: Number of DeleteObjects requests to send at the same + time. None or -1 is ``max_workers``; the ``batch_size`` of + the filesystem, which is that of fsspec's other batch + coroutines, does not apply. ``max_workers`` is a deprecated + alias. maxdepth: Maximum depth to expand when ``recursive`` is True. **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. """ + max_workers = self._sync_fs._delete_batch_limit(batch_size, kwargs) paths = await self._thread_pool.run( self._sync_fs._delete_paths, path, recursive=recursive, maxdepth=maxdepth ) - await self._delete_objects(paths, **kwargs) + await self._delete_objects(paths, max_workers=max_workers, **kwargs) - async def _delete_objects(self, paths: list[str], **kwargs) -> None: + async def _delete_objects( + self, paths: list[str], max_workers: int | None = None, **kwargs + ) -> None: """Delete objects with parallel DeleteObjects requests. See :meth:`S3FileSystem._delete_objects`. Args: paths: Paths of the objects to delete. Bucket paths are skipped. + max_workers: Maximum number of parallel requests. Defaults to the + ``max_workers`` of the internal ``S3FileSystem``. **kwargs: Additional parameters passed to the DeleteObjects API. ``Quiet`` (default True) sets the quiet mode of the requests. @@ -394,6 +640,7 @@ async def _delete_objects(self, paths: list[str], **kwargs) -> None: results = await self._map_async( functools.partial(self._sync_fs._delete_batch, **params), batches, + max_workers=max_workers, stop_on_failure=False, ) self._sync_fs._raise_delete_errors(results.outcomes) @@ -431,8 +678,9 @@ async def _map_async( async def _mv(self, path1, path2, recursive=False, maxdepth=None, **kwargs) -> None: """Move files from one S3 location to another. - See :meth:`S3FileSystem.mv`. The copies and the deletions run in - parallel with ``asyncio.gather``. + See :meth:`S3FileSystem.mv`. The copies run in parallel with + ``asyncio.gather``, whose cancellation cancels them and waits for + them, and the deletions in the request pool. Args: path1: Source S3 path, glob pattern, or list of paths. @@ -506,7 +754,10 @@ async def _copy( """Copy files within S3. See :meth:`S3FileSystem.copy`. The copies run as in fsspec's - ``_copy()``. + ``_copy()``, ``batch_size`` of them at once as tasks. A cancellation + cancels the copies and waits for them; a copy that is already being + cancelled, such as by the shutdown of the event loop, is not + cancelled again. Args: path1: Source S3 path, glob pattern, or list of them. @@ -527,14 +778,16 @@ async def _copy( if not pairs: return path1, path2 = [p1 for p1, _ in pairs], [p2 for _, p2 in pairs] - await super()._copy( - path1, - path2, - recursive=recursive, - on_error=on_error, - maxdepth=maxdepth, - batch_size=batch_size, - **kwargs, + await _S3BatchTasks().run( + super()._copy( + path1, + path2, + recursive=recursive, + on_error=on_error, + maxdepth=maxdepth, + batch_size=batch_size, + **kwargs, + ) ) async def _get( @@ -543,7 +796,8 @@ async def _get( """Copy files from S3 to the local filesystem. See :meth:`S3FileSystem.get`. The downloads run as in fsspec's - ``_get()``. + ``_get()``, and a cancellation cancels them and waits for them, as + in :meth:`_copy`. Args: rpath: Source S3 path, glob pattern, or list of them. @@ -570,11 +824,145 @@ async def _get( check_contained(root, lpath) if not rpath: return - await super()._get( - rpath, lpath, recursive=recursive, callback=callback, maxdepth=maxdepth, **kwargs + await _S3BatchTasks().run( + super()._get( + rpath, lpath, recursive=recursive, callback=callback, maxdepth=maxdepth, **kwargs + ) + ) + + async def _put( + self, + lpath, + rpath, + recursive=False, + callback=_DEFAULT_CALLBACK, + batch_size=None, + maxdepth=None, + **kwargs, + ) -> list[Any]: + """Copy files from the local filesystem to S3. + + See :meth:`S3FileSystem.put`. The uploads run as in fsspec's + ``_put()``, and a cancellation cancels them and waits for them, as + in :meth:`_copy`. Unlike fsspec's ``_put()``, which creates a + directory for each local directory, it makes no other request: S3 + has no directories, and creating one checks or creates the bucket. + + 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. + batch_size: Number of uploads to run at the same time. + maxdepth: Maximum depth of a recursive copy. + **kwargs: Additional parameters passed to ``_put_file()``. + + Returns: + The return values of the ``_put_file()`` calls, as in fsspec. + """ + pairs = await self._thread_pool.run( + self._sync_fs._put_pairs, lpath, rpath, recursive=recursive, maxdepth=maxdepth, **kwargs ) + files = [(lp, rp) for lp, rp in pairs if not os.path.isdir(lp)] + return await _S3BatchTasks().run( + super()._put( + [lp for lp, _ in files], + [rp for _, rp in files], + callback=callback, + batch_size=batch_size, + **kwargs, + ) + ) + + async def _pipe(self, path, value=None, batch_size=None, **kwargs) -> list[Any]: + """Write bytes into one or more paths. - async def _cp_file(self, path1: str, path2: str, **kwargs) -> None: + See :meth:`S3FileSystem.pipe`. The writes run as in fsspec's + ``_pipe()``, and a cancellation cancels them and waits for them, as + in :meth:`_copy`. + + Args: + path: S3 path, or a dictionary of paths and their bytes. + value: The bytes to write when ``path`` is a path. + batch_size: Number of writes to run at the same time. + **kwargs: Additional parameters passed to ``_pipe_file()``. + + Returns: + The return values of the ``_pipe_file()`` calls, as in fsspec. + """ + return await _S3BatchTasks().run( + super()._pipe(path, value, batch_size=batch_size, **kwargs) + ) + + async def _cat(self, path, recursive=False, on_error="raise", batch_size=None, **kwargs): + """Read the contents of one or more paths. + + See :meth:`S3FileSystem.cat`. The reads run as in fsspec's + ``_cat()``, and a cancellation cancels them and waits for them, as + in :meth:`_copy`. + + Args: + path: S3 path, glob pattern, or list of them. + recursive: Whether to read the files below the directories. + on_error: ``"raise"``, ``"omit"``, or ``"return"``. + batch_size: Number of reads to run at the same time. + **kwargs: Additional parameters passed to ``_cat_file()``. + + Returns: + The bytes of a single path, or a dictionary of the paths and + their bytes. + """ + return await _S3BatchTasks().run( + super()._cat( + path, recursive=recursive, on_error=on_error, batch_size=batch_size, **kwargs + ) + ) + + async def _cat_ranges( + self, + paths, + starts, + ends, + max_gap=None, + batch_size=None, + on_error="return", + **kwargs, + ): + """Read byte ranges of one or more paths. + + See :meth:`S3FileSystem.cat_ranges`. The reads run as in fsspec's + ``_cat_ranges()``, and a cancellation cancels them and waits for + them, as in :meth:`_copy`. + + Args: + paths: List of S3 paths. + starts: Start offset of each range, or one for all of them. + ends: End offset of each range, or one for all of them. + max_gap: Not supported; must be None. + batch_size: Number of reads to run at the same time. + on_error: ``"return"`` to return the error of a read in place of + its bytes; otherwise the first error is raised. + **kwargs: Additional parameters passed to ``_cat_file()``. + + Returns: + The bytes or the error of each range, in order. + """ + return await _S3BatchTasks().run( + super()._cat_ranges( + paths, + starts, + ends, + max_gap=max_gap, + batch_size=batch_size, + on_error=on_error, + **kwargs, + ) + ) + + async def _cp_file( + self, path1: str, path2: str, recursive=False, maxdepth=None, on_error=None, **kwargs + ) -> None: """Copy an S3 object, using async parallel multipart upload for large files. A directory ``path1``, which recursive ``copy()`` passes along with @@ -583,11 +971,16 @@ async def _cp_file(self, path1: str, path2: str, **kwargs) -> None: Args: path1: Source S3 path (s3://bucket/key). path2: Destination S3 path (s3://bucket/key). - **kwargs: Additional S3 copy parameters. + recursive: Unused parameter for fsspec compatibility. + maxdepth: Unused parameter for fsspec compatibility. + on_error: Unused parameter for fsspec compatibility. + **kwargs: Additional S3 copy parameters, as for + :meth:`S3FileSystem.cp_file`. Raises: ValueError: If trying to copy to a versioned file or copy buckets. """ + _S3BatchTasks.track() await self._copy_file(path1, path2, **kwargs) async def _copy_file(self, path1: str, path2: str, **kwargs) -> bool: @@ -795,16 +1188,15 @@ async def _find( async def _expand_path(self, path, recursive=False, maxdepth=None, **kwargs) -> list[str]: """Expand glob patterns and directories into the paths they match. - See :meth:`S3FileSystem.expand_path`. The other paths are expanded by - fsspec's ``_expand_path()``, whose glob lists only the keys under the - literal part of the pattern. + Runs :meth:`S3FileSystem.expand_path`, as :meth:`_find` runs + ``find()``. Args: path: S3 path, glob pattern, or list of them. recursive: Whether to include the paths below the directories. maxdepth: Maximum depth of the expansion, at least 1. - **kwargs: Additional arguments passed to fsspec's - ``_expand_path()``, such as ``assume_literal``. + **kwargs: Additional arguments passed to ``glob`` and ``find``, + and ``assume_literal``. Returns: The sorted matching paths. @@ -813,23 +1205,9 @@ async def _expand_path(self, path, recursive=False, maxdepth=None, **kwargs) -> ValueError: If ``maxdepth`` is less than 1. FileNotFoundError: If nothing matches. """ - if maxdepth is not None and maxdepth < 1: - raise ValueError("maxdepth must be at least 1") - versions, others = self._sync_fs._split_version_paths(path) - out = {p for p in versions if not recursive or await self._exists(p)} - if others: - try: - out.update( - await super()._expand_path( - others, recursive=recursive, maxdepth=maxdepth, **kwargs - ) - ) - except FileNotFoundError: - if not out: - raise - if not out: - raise FileNotFoundError(path) - return sorted(out) + return await self._thread_pool.run( + self._sync_fs.expand_path, path, recursive=recursive, maxdepth=maxdepth, **kwargs + ) def _open( self, diff --git a/pyathena/filesystem/s3_executor.py b/pyathena/filesystem/s3_executor.py index 832d5ffa0..95e8760f1 100644 --- a/pyathena/filesystem/s3_executor.py +++ b/pyathena/filesystem/s3_executor.py @@ -498,10 +498,14 @@ def shutdown(self, wait: bool = True) -> None: async def run(self, fn: Callable[..., T], *args: Any, **kwargs: Any) -> T: """Run a blocking call in the executor and wait for it. - Cancelling the waiting task cancels a call that has not started, and - stops the retries of the AWS API calls that a running call makes, as - :func:`~pyathena.aio.util.run_in_thread` does; the call itself keeps - running until it returns or raises. + Cancelling the waiting task cancels a call that has not started. A + running call makes no further attempt of its AWS API calls, as with + :func:`~pyathena.aio.util.run_in_thread`, and a write in it sends no + further part and does not commit (see + :func:`~pyathena.util.raise_if_stopped`); the task waits for the call + to end before it re-raises the cancellation, also when the event loop + shuts down with ``asyncio.run()``. A repeated cancellation stops that + wait, and the call then ends in the executor. Args: fn: The blocking callable. @@ -531,12 +535,24 @@ def call() -> T: with lock: returned = True + # Created before it is scheduled, so that a cancellation while it is + # being scheduled can still cancel it or wait for it. + future: Future[T] = Future() + waiter = self.waiter(future) try: - return await asyncio.wrap_future(self.submit(call)) + self.submit_to(future, call) + # Unlike awaiting the waiter, a cancellation of this task leaves + # the waiter to the wait below. + await asyncio.wait([waiter]) + return waiter.result() except asyncio.CancelledError: with lock: if not returned: stop.set() + future.cancel() + # The states are also checked every second, as in map_async(). + while not future.done(): + await asyncio.wait([waiter], timeout=1.0) raise async def run_shielded(self, fn: Callable[..., T], *args: Any, **kwargs: Any) -> T: diff --git a/pyathena/util.py b/pyathena/util.py index b70e78e77..0097e41ea 100644 --- a/pyathena/util.py +++ b/pyathena/util.py @@ -328,6 +328,16 @@ def unstoppable_retries() -> Iterator[None]: _RETRY_STOP_EVENTS.reset(token) +def retries_stopped() -> bool: + """Return whether an event of the enclosing stop scopes is set. + + Returns: + Whether an event of the enclosing :func:`stop_retries_on` scopes, + apart from those that :func:`unstoppable_retries` clears, is set. + """ + return any(event.is_set() for event in _RETRY_STOP_EVENTS.get()) + + def raise_if_stopped(operation: str) -> None: """Raise if an event of the enclosing stop scopes is set. @@ -346,7 +356,7 @@ def raise_if_stopped(operation: str) -> None: 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()): + if retries_stopped(): raise CancelledError(f"Stopped before it could {operation}.") diff --git a/tests/pyathena/filesystem/test_s3.py b/tests/pyathena/filesystem/test_s3.py index 06a291ee1..4210067bb 100644 --- a/tests/pyathena/filesystem/test_s3.py +++ b/tests/pyathena/filesystem/test_s3.py @@ -1056,6 +1056,17 @@ def test_glob_matches_fsspec(self, pattern, maxdepth): expected = [p.lstrip("/") for p in memory.glob(f"/bucket/{pattern}", maxdepth=maxdepth)] assert sorted(fs.glob(f"s3://bucket/{pattern}", maxdepth=maxdepth)) == expected + def test_glob_lists_stem_prefix(self): + # GH-1337: glob() lists only the keys that start with the stem before + # the first wildcard, as fsspec's asynchronous glob does, which + # AioS3FileSystem used before it expanded its paths with this + # filesystem. It listed every key of the directory. + fs = self._make_fs() + self._serve_keys(fs, self.FIND_KEYS) + + assert fs.glob("s3://bucket/dir/s*") == ["bucket/dir/sub"] + assert [c.kwargs["Prefix"] for c in fs._call.call_args_list] == ["dir/s"] + def test_find_withdirs_omits_bucket(self): fs = self._make_fs() self._serve_keys(fs, self.FIND_KEYS) @@ -1975,14 +1986,18 @@ def test_mv_versions_onto_their_key(self): @staticmethod def _moves(fs): - # The CopyObject requests, as (CopySource, Key), and the objects of - # the DeleteObjects requests that the filesystem sent. + # The CopyObject requests, as (CopySource, Key) sorted by Key, since + # the copies run in parallel, and the objects of the DeleteObjects + # requests that the filesystem sent. calls = fs._call.call_args_list - copies = [ - (c.kwargs["CopySource"], c.kwargs["Key"]) - for c in calls - if c.args[0] is fs._client.copy_object - ] + copies = sorted( + ( + (c.kwargs["CopySource"], c.kwargs["Key"]) + for c in calls + if c.args[0] is fs._client.copy_object + ), + key=lambda copy: copy[1], + ) deletes = [ c.kwargs["Delete"]["Objects"] for c in calls if c.args[0] is fs._client.delete_objects ] diff --git a/tests/pyathena/filesystem/test_s3_async.py b/tests/pyathena/filesystem/test_s3_async.py index 95cfaca98..60bf10e1e 100644 --- a/tests/pyathena/filesystem/test_s3_async.py +++ b/tests/pyathena/filesystem/test_s3_async.py @@ -1,5 +1,6 @@ import array import asyncio +import collections import contextlib import gc import gzip @@ -11,8 +12,9 @@ import uuid import warnings import weakref -from concurrent.futures import Future, ThreadPoolExecutor, wait +from concurrent.futures import CancelledError, Future, ThreadPoolExecutor, wait from datetime import UTC, datetime +from pathlib import PurePosixPath from types import SimpleNamespace from unittest import mock @@ -21,7 +23,12 @@ from fsspec.callbacks import Callback from pyathena.filesystem.s3 import S3FileSystem -from pyathena.filesystem.s3_async import AioS3File, AioS3FileSystem, AioS3StreamedFile +from pyathena.filesystem.s3_async import ( + AioS3File, + AioS3FileSystem, + AioS3StreamedFile, + _S3TransactionWrite, +) from pyathena.filesystem.s3_core import S3Core, S3MultipartCopyPlan from pyathena.filesystem.s3_executor import S3ThreadPool, _S3PoolScope from pyathena.filesystem.s3_object import ( @@ -34,7 +41,7 @@ ) 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 ( @@ -49,6 +56,7 @@ interrupt_each_opcode, mock_connection, overlapping_calls, + poll_until, retried, run_bounded, run_in_forked_child, @@ -1665,7 +1673,7 @@ def test_transaction_pipe_put_file_exceeding_max_parts(self, tmp_path, kwargs): fs._sync_fs.core.MULTIPART_UPLOAD_MAX_PARTS = 3 fs._sync_fs.default_block_size = 4 fs._sync_fs._call = fs._sync_fs._core.call = mock.MagicMock() - fs.open = mock.MagicMock() + fs._open = mock.MagicMock() local = tmp_path / "local" local.write_bytes(b"a" * 13) @@ -1674,7 +1682,7 @@ def test_transaction_pipe_put_file_exceeding_max_parts(self, tmp_path, kwargs): fs.pipe_file("s3://bucket/k1", b"a" * 13, **kwargs) with pytest.raises(ValueError, match="in 3 parts with a block size of 4 bytes"): fs.put_file(str(local), "s3://bucket/k2", **kwargs) - fs.open.assert_not_called() + fs._open.assert_not_called() fs._sync_fs._call.assert_not_called() @pytest.mark.asyncio @@ -1780,13 +1788,563 @@ async def test_rm_maxdepth(self): await fs._rm("s3://bucket/dir", recursive=True, maxdepth=1, batch_size=10) stubber.assert_no_pending_responses() + @pytest.mark.asyncio + async def test_put_file_cancelled_stops_upload(self, tmp_path): + # GH-1337: a cancelled _put_file() sends no further part, does not + # complete the upload, aborts it, and raises once the abort has ended. + # The upload used to run on in its thread and commit. + fs = AioS3FileSystem( + key="dummy", secret="dummy", region_name="us-east-1", skip_instance_cache=True + ) + block_size = S3Core.MULTIPART_UPLOAD_MIN_PART_SIZE + local = tmp_path / "local" + local.write_bytes(bytes(3 * block_size + 1)) + started = threading.Event() + release = threading.Event() + requests = [] + + def call(method, **request): + name = method.__name__ + requests.append(name) + if name == "create_multipart_upload": + return {"Bucket": "bucket", "Key": "key", "UploadId": "uploadid"} + if name == "upload_part": + started.set() + release.wait(EVENT_TIMEOUT) + return {"ETag": '"e"'} + return {} + + fs._sync_fs._call = fs._sync_fs._core.call = call + task = asyncio.ensure_future( + fs._put_file(str(local), "s3://bucket/key", block_size=block_size, max_workers=1) + ) + try: + assert await asyncio.to_thread(started.wait, EVENT_TIMEOUT) + task.cancel() + done, _ = await asyncio.wait([task], timeout=0.1) + assert not done + release.set() + with pytest.raises(asyncio.CancelledError): + await await_bounded(task) + finally: + release.set() + assert requests == ["create_multipart_upload", "upload_part", "abort_multipart_upload"] + + @staticmethod + def _held_until_stopped(state, release, result=None): + # Return a blocking call that records in state that it started, waits + # until a stop scope of its thread stops, records that, and then + # waits for release. A call that is never stopped ends after + # EVENT_TIMEOUT without being recorded as stopped. + lock = threading.Lock() + + def call(*args, **kwargs): + with lock: + state["started"] += 1 + deadline = time.monotonic() + EVENT_TIMEOUT + while time.monotonic() < deadline: + try: + raise_if_stopped("continue") + except CancelledError: + with lock: + state["stopped"] += 1 + release.wait(EVENT_TIMEOUT) + break + time.sleep(0.01) + with lock: + state["ended"] += 1 + return result + + return call + + @pytest.mark.parametrize( + ("method", "cancellation"), + [ + (method, cancellation) + for method in ["_cat", "_cat_ranges", "_pipe", "_put", "_get", "_copy", "_mv"] + for cancellation in ["batch", "calls_first", "shutdown"] + # asyncio.gather() cancels the copies of _mv() again. + if (method, cancellation) != ("_mv", "calls_first") + ], + ) + def test_batch_cancelled_waits_for_calls(self, tmp_path, method, cancellation): + # GH-1337: a cancelled batch coroutine cancels its per-object calls, + # which make no further request, and raises once they have ended. + # fsspec's batch coroutines left them running after the + # cancellation. A call whose task is already being cancelled, as the + # shutdown of the event loop at the end of asyncio.run() cancels + # every task, possibly before the batch handles its own + # cancellation (calls_first), is not cancelled again, which would + # stop its wait. + fs = AioS3FileSystem(connection=mock_connection(), skip_instance_cache=True) + sync_fs = fs._sync_fs + state = {"started": 0, "ended": 0, "stopped": 0} + release = threading.Event() + sources = ["s3://bucket/a", "s3://bucket/b"] + destinations = ["s3://bucket/c", "s3://bucket/d"] + local = [str(tmp_path / "a"), str(tmp_path / "b")] + held = self._held_until_stopped(state, release) + sync_fs.cat_file = self._held_until_stopped(state, release, b"") + sync_fs.pipe_file = sync_fs.put_file = sync_fs.get_file = held + sync_fs.info = self._held_until_stopped( + state, release, {"name": "bucket/a", "type": "file", "size": 0} + ) + sync_fs._move_pairs = mock.MagicMock( + return_value=list(zip(sources, destinations, strict=True)) + ) + batches = { + "_cat": lambda: fs._cat(sources), + "_cat_ranges": lambda: fs._cat_ranges(sources, 0, 1), + "_pipe": lambda: fs._pipe(dict.fromkeys(sources, b"a")), + "_put": lambda: fs._put(local, sources), + "_get": lambda: fs._get(sources, local), + "_copy": lambda: fs._copy(sources, destinations), + "_mv": lambda: fs._mv(sources, destinations), + } + # The calls that had ended when the batch finished, recorded in the + # event loop. + ended_when_done = [] + + def release_once_stopped(): + # The shutdown: the stopped calls are held until the end of the + # batch has been recorded, which must not happen before they end, + # or for 0.5 s. + poll_until(lambda: state["stopped"] == 2, "the calls were not stopped") + deadline = time.monotonic() + 0.5 + while not ended_when_done and time.monotonic() < deadline: + time.sleep(0.01) + release.set() + + async def main(): + task = asyncio.ensure_future(batches[method]()) + task.add_done_callback(lambda _: ended_when_done.append(state["ended"])) + await wait_until(lambda: state["started"] == 2) + if cancellation == "shutdown": + # asyncio.run() cancels the batch and its calls once main() + # has returned. + threading.Thread(target=release_once_stopped, daemon=True).start() + return + if cancellation == "calls_first": + current = asyncio.current_task() + for call in asyncio.all_tasks() - {current, task}: + call.cancel() + await wait_until(lambda: state["stopped"] == 2) + task.cancel() + await wait_until(lambda: state["stopped"] == 2) + done, _ = await asyncio.wait([task], timeout=0.1) + assert not done + release.set() + with pytest.raises(asyncio.CancelledError): + await await_bounded(task) + + try: + run_bounded(main) + finally: + release.set() + assert state == {"started": 2, "ended": 2, "stopped": 2} + assert ended_when_done == [2] + + @pytest.mark.parametrize("method", ["_pipe", "_put", "_get"]) + def test_batch_failure_waits_for_calls(self, tmp_path, method): + # GH-1337: when a per-object call of _pipe(), _put(), or _get() fails, + # fsspec cancels the others, which make no further request, and the + # failure is raised once they have ended. They used to run on, and + # commit, after the failure was raised. + fs = AioS3FileSystem(connection=mock_connection(), skip_instance_cache=True) + sync_fs = fs._sync_fs + state = {"started": 0, "ended": 0, "stopped": 0} + release = threading.Event() + held = self._held_until_stopped(state, release) + + def call(source, *args, **kwargs): + if source.endswith("/a"): + # Once the other call has started. + poll_until(lambda: state["started"] == 1, "the other call did not start") + raise ValueError("failed") + return held(source, *args, **kwargs) + + sync_fs.pipe_file = sync_fs.put_file = sync_fs.get_file = call + sources = ["s3://bucket/b", "s3://bucket/a"] + local = [str(tmp_path / "b"), str(tmp_path / "a")] + batches = { + "_pipe": lambda: fs._pipe(dict.fromkeys(sources, b"a")), + "_put": lambda: fs._put(local, sources), + "_get": lambda: fs._get(sources, local), + } + + async def main(): + task = asyncio.ensure_future(batches[method]()) + await wait_until(lambda: state["stopped"] == 1) + done, _ = await asyncio.wait([task], timeout=0.1) + assert not done + release.set() + with pytest.raises(ValueError, match="failed"): + await await_bounded(task) + assert state == {"started": 1, "ended": 1, "stopped": 1} + + try: + run_bounded(main) + finally: + release.set() + + @pytest.mark.asyncio + async def test_transaction_put_file_cancelled_is_discarded(self, tmp_path): + # GH-1337: a _put_file() in a transaction that is cancelled while its + # last part runs is discarded, and the transaction does not commit + # it. The part ended normally, so the write had nothing left to stop, + # and the file stayed in the transaction. + fs = AioS3FileSystem( + key="dummy", secret="dummy", region_name="us-east-1", skip_instance_cache=True + ) + block_size = S3Core.MULTIPART_UPLOAD_MIN_PART_SIZE + local = tmp_path / "local" + local.write_bytes(bytes(2 * block_size)) + started = threading.Event() + release = threading.Event() + requests = [] + + def call(method, **request): + name = method.__name__ + requests.append(name) + if name == "create_multipart_upload": + return {"Bucket": "bucket", "Key": "key", "UploadId": "uploadid"} + if name == "upload_part": + if request["PartNumber"] == 2: + started.set() + release.wait(EVENT_TIMEOUT) + return {"ETag": '"e"'} + return {} + + fs._sync_fs._call = fs._sync_fs._core.call = call + try: + with fs.transaction: + task = asyncio.ensure_future( + fs._put_file( + str(local), "s3://bucket/key", block_size=block_size, max_workers=1 + ) + ) + assert await asyncio.to_thread(started.wait, EVENT_TIMEOUT) + task.cancel() + await asyncio.sleep(0.1) + release.set() + with pytest.raises(asyncio.CancelledError): + await await_bounded(task) + finally: + release.set() + assert requests == [ + "create_multipart_upload", + "upload_part", + "upload_part", + "abort_multipart_upload", + ] + + @pytest.mark.asyncio + @pytest.mark.parametrize("stage", ["before_open", "last_part", "write_ended", "returned"]) + async def test_transaction_put_file_cancelled_late(self, tmp_path, stage): + # GH-1337: the transaction does not commit the file of a _put_file() + # whose task is cancelled, and a cancellation adds no wait for a + # discard behind other calls of the filesystem. The cases: cancelled + # twice before the write opens its file (before_open), while it runs + # its last part (last_part), or just after the write has ended, the + # second time (write_ended); and cancelled once after the write has + # returned (returned), whose file the coroutine discards. + queued = stage in ("write_ended", "returned") + fs = AioS3FileSystem( + key="dummy", + secret="dummy", + region_name="us-east-1", + max_workers=1 if queued else None, + skip_instance_cache=True, + ) + block_size = S3Core.MULTIPART_UPLOAD_MIN_PART_SIZE + local = tmp_path / "local" + multipart = stage in ("last_part", "returned") + local.write_bytes(bytes(2 * block_size if multipart else 1)) + state = {"started": 0, "ended": 0, "stopped": 0} + release = threading.Event() + held = self._held_until_stopped(state, release) + requests = [] + + def call(method, **request): + requests.append(method.__name__) + if method.__name__ == "create_multipart_upload": + return {"Bucket": "bucket", "Key": "key", "UploadId": "uploadid"} + if ( + stage == "last_part" + and method.__name__ == "upload_part" + and request["PartNumber"] == 2 + ): + held() + return {"ETag": '"e"'} + + class HeldCallback(Callback): + def set_size(self, size): + if stage == "before_open": + held() + + def relative_update(self, inc=1): + if stage == "write_ended": + held() + elif stage == "returned" and state["started"] == 0: + # Held without a stop, until the test releases it. + state["started"] = 1 + release.wait(EVENT_TIMEOUT) + + fs._sync_fs._call = fs._sync_fs._core.call = call + blocker_started = threading.Event() + blocker_release = threading.Event() + blocker_ended = threading.Event() + + def blocker(*args, **kwargs): + blocker_started.set() + blocker_release.wait(EVENT_TIMEOUT) + blocker_ended.set() + return {} + + fs._sync_fs.info = blocker + try: + with fs.transaction: + task = asyncio.ensure_future( + fs._put_file( + str(local), + "s3://bucket/key", + callback=HeldCallback(), + block_size=block_size, + max_workers=1, + ) + ) + await wait_until(lambda: state["started"] == 1) + if queued: + # Takes the only thread once the write has ended. + info = asyncio.ensure_future(fs._info("s3://bucket/other")) + await asyncio.sleep(0.1) + if stage != "returned": + task.cancel() + await wait_until(lambda: state["stopped"] == 1) + if queued: + + def end_write_then_cancel(): + # In one callback of the event loop, so that the task + # is cancelled before it sees the end of the write, + # once the only thread has taken the next call. + release.set() + assert blocker_started.wait(EVENT_TIMEOUT) + task.cancel() + + asyncio.get_running_loop().call_soon(end_write_then_cancel) + else: + task.cancel() + # While the write is held, or the next call holds the thread. + with pytest.raises(asyncio.CancelledError): + await await_bounded(task, timeout=2) + assert not blocker_ended.is_set() + if not queued: + # The write ends while the transaction is still open. + release.set() + await asyncio.to_thread(shutdown_bounded, fs._thread_pool) + blocker_release.set() + if queued: + await info + await asyncio.to_thread(shutdown_bounded, fs._thread_pool) + finally: + release.set() + blocker_release.set() + assert "complete_multipart_upload" not in requests + assert "put_object" not in requests + if multipart: + assert requests[-1] == "abort_multipart_upload" + + @pytest.mark.asyncio + async def test_transaction_write_cancelled_before_it_joins(self, tmp_path): + # GH-1337: a _put_file() in a transaction that is cancelled twice + # before it opens its file never adds the file to the transaction, + # whose exit at any moment therefore commits nothing. The file used + # to join the transaction in open(), and to be dropped from it only + # afterwards. + fs = AioS3FileSystem( + key="dummy", secret="dummy", region_name="us-east-1", skip_instance_cache=True + ) + local = tmp_path / "local" + local.write_bytes(b"a") + state = {"started": 0, "ended": 0, "stopped": 0} + release = threading.Event() + held = self._held_until_stopped(state, release) + requests = [] + joined = threading.Event() + proceed = threading.Event() + + class HeldFiles(collections.deque): + def append(self, file): + # Holds the write once its file has joined the transaction. + super().append(file) + joined.set() + proceed.wait(EVENT_TIMEOUT) + + class HeldCallback(Callback): + def set_size(self, size): + held() + + def call(method, **request): + requests.append(method.__name__) + return {} + + fs._sync_fs._call = fs._sync_fs._core.call = call + write_ended = threading.Event() + run = _S3TransactionWrite.run + + def run_and_signal(transaction_write): + try: + run(transaction_write) + finally: + write_ended.set() + + try: + with ( + mock.patch.object(_S3TransactionWrite, "run", run_and_signal), + fs.transaction, + ): + fs.transaction.files = HeldFiles() + task = asyncio.ensure_future( + fs._put_file(str(local), "s3://bucket/key", callback=HeldCallback()) + ) + await wait_until(lambda: state["started"] == 1) + task.cancel() + await wait_until(lambda: state["stopped"] == 1) + task.cancel() + with pytest.raises(asyncio.CancelledError): + await await_bounded(task, timeout=2) + release.set() + # Exits once the file has joined, or once the write has ended. + await wait_until(lambda: joined.is_set() or write_ended.is_set()) + finally: + proceed.set() + release.set() + await asyncio.to_thread(shutdown_bounded, fs._thread_pool) + assert not joined.is_set() + assert requests == [] + + @pytest.mark.asyncio + async def test_transaction_write_stopped_by_enclosing_scope(self): + # GH-1337: a transaction write that an enclosing stop scope, not a + # cancellation, has stopped raises instead of returning, since its + # file is discarded and left out of the transaction. The + # concurrent.futures.CancelledError of the stop reaches the coroutine + # as asyncio.CancelledError, which asyncio.wrap_future() converts it to. + fs = AioS3FileSystem( + key="dummy", secret="dummy", region_name="us-east-1", skip_instance_cache=True + ) + requests = [] + + def call(method, **request): + requests.append(method.__name__) + return {} + + fs._sync_fs._call = fs._sync_fs._core.call = call + stop = threading.Event() + stop.set() + with fs.transaction: + with stop_retries_on(stop), pytest.raises(asyncio.CancelledError): + await fs._pipe_file("s3://bucket/key", b"a") + assert not fs.transaction.files + assert requests == [] + + @pytest.mark.asyncio + async def test_put_pipe_return_values(self, tmp_path): + # GH-1337: _put() and _pipe() return the values of their per-object + # calls, as fsspec's do, and _put() accepts a path-like destination. + fs = AioS3FileSystem( + key="dummy", secret="dummy", region_name="us-east-1", skip_instance_cache=True + ) + store, _ = self._serve_keys(fs, set()) + local = tmp_path / "local" + local.write_bytes(b"a") + + assert await fs._pipe("s3://bucket/a", b"a") == [None] + assert await fs._put(str(local), PurePosixPath("bucket/b")) == [None] + assert store == {"a", "b"} + + @pytest.mark.asyncio + async def test_put_recursive_creates_no_directories(self, tmp_path): + # GH-1337: a recursive _put() sends no request for the local + # directories. fsspec's _put() created a directory for each, which + # looked up the bucket, created it if it was missing and allowed, and + # otherwise raised PermissionError before any upload. + fs = AioS3FileSystem( + key="dummy", secret="dummy", region_name="us-east-1", skip_instance_cache=True + ) + store, requests = self._serve_keys(fs, set()) + (tmp_path / "src" / "sub").mkdir(parents=True) + (tmp_path / "src" / "a").write_bytes(b"a") + (tmp_path / "src" / "sub" / "b").write_bytes(b"b") + + await fs._put(f"{tmp_path}/src/", "s3://bucket/dst/", recursive=True) + assert store == {"dst/a", "dst/sub/b"} + assert [name for name, _ in requests] == ["put_object", "put_object"] + + @pytest.mark.asyncio + async def test_expand_path_kwargs(self): + # GH-1337: the keyword arguments of _expand_path() reach the listing, + # as with S3FileSystem.expand_path(). fsspec's _expand_path() takes + # none, so they raised TypeError. + fs = AioS3FileSystem( + key="dummy", secret="dummy", region_name="us-east-1", skip_instance_cache=True + ) + self._serve_keys(fs, {"dir/a1", "dir/b1"}) + + assert await fs._expand_path("s3://bucket/dir", recursive=True, prefix="a") == [ + "bucket/dir", + "bucket/dir/a1", + ] + + @pytest.mark.asyncio + @pytest.mark.parametrize("argument", ["batch_size", "max_workers"]) + async def test_rm_batch_size(self, argument): + # GH-1337: batch_size of _rm() is the number of DeleteObjects requests + # sent at once, and max_workers its deprecated alias, as for + # S3FileSystem.rm(). batch_size was ignored, and max_workers was sent + # to DeleteObjects, which rejected it. + fs = AioS3FileSystem(connection=mock_connection(), skip_instance_cache=True) + overlap, state = overlapping_calls(2, timeout=0.2) + requests = [] + + def call(method, **request): + requests.append(request) + overlap() + return {} + + fs._sync_fs._call = fs._sync_fs._core.call = call + with contextlib.ExitStack() as stack: + if argument == "max_workers": + stack.enter_context(pytest.warns(DeprecationWarning, match="use batch_size")) + await fs._rm(["s3://b1/a", "s3://b2/b"], **{argument: 1}) + assert state["peak"] == 1 + assert sorted(r["Bucket"] for r in requests) == ["b1", "b2"] + assert all(set(r) == {"Bucket", "Delete"} for r in requests) + + @pytest.mark.asyncio + async def test_cp_file_fsspec_arguments(self): + # GH-1337: the recursive, maxdepth, and on_error arguments of + # _cp_file(), unused as with S3FileSystem.cp_file(), are not sent to + # CopyObject, which rejected them. + fs = AioS3FileSystem( + key="dummy", secret="dummy", region_name="us-east-1", skip_instance_cache=True + ) + _, requests = self._serve_keys(fs, {"key"}) + + await fs._cp_file( + "s3://bucket/key", "s3://bucket/dst", recursive=True, maxdepth=1, on_error="ignore" + ) + assert [r for name, r in requests if name == "copy_object"] == [ + {"CopySource": {"Bucket": "bucket", "Key": "key"}, "Bucket": "bucket", "Key": "dst"} + ] + @pytest.mark.parametrize(("mode", "open_mode"), [("overwrite", "wb"), ("create", "xb")]) def test_put_file_in_transaction_open_parameters(self, tmp_path, mode, open_mode): # GH-969: in a transaction, the open() parameters of put_file() go to - # open(), and the other parameters, also in s3_additional_kwargs, to + # the file, and the other parameters, also in s3_additional_kwargs, to # S3. GH-972: fsspec's mode argument selects the mode of the file. fs = AioS3FileSystem(connection=mock_connection(), skip_instance_cache=True) - fs.open = mock.MagicMock() + fs._open = mock.MagicMock() lpath = tmp_path / "data.csv" lpath.write_bytes(b"a") @@ -1802,17 +2360,19 @@ def test_put_file_in_transaction_open_parameters(self, tmp_path, mode, open_mode # Without them, the file gets the defaults of the filesystem. fs.put_file(str(lpath), "s3://bucket/key", mode=mode) - assert fs.open.call_args_list == [ + assert fs._open.call_args_list == [ mock.call( - "s3://bucket/key", + "bucket/key", open_mode, + autocommit=False, block_size=S3Core.MULTIPART_UPLOAD_MIN_PART_SIZE, max_workers=2, s3_additional_kwargs={"StorageClass": "STANDARD_IA", "ContentType": "text/csv"}, ), mock.call( - "s3://bucket/key", + "bucket/key", open_mode, + autocommit=False, block_size=fs._sync_fs.default_block_size, max_workers=fs._sync_fs.max_workers, s3_additional_kwargs={"ContentType": "text/csv"}, @@ -1852,8 +2412,9 @@ async def test_copy_version(self): @pytest.mark.asyncio async def test_expand_path_glob_lists_stem_prefix(self): - # Unversioned globs keep fsspec's async expansion, which lists only - # the keys that start with the stem before the first wildcard. + # A glob lists only the keys that start with the stem before the + # first wildcard, as fsspec's async expansion, which _expand_path() + # used before GH-1337, does. fs = AioS3FileSystem(connection=mock_connection(), skip_instance_cache=True) sync_fs = fs._sync_fs sync_fs._call = sync_fs._core.call = mock.MagicMock( diff --git a/tests/pyathena/filesystem/test_s3_executor.py b/tests/pyathena/filesystem/test_s3_executor.py index 6925f3767..558d84961 100644 --- a/tests/pyathena/filesystem/test_s3_executor.py +++ b/tests/pyathena/filesystem/test_s3_executor.py @@ -40,6 +40,7 @@ overlapping_calls, poll_until, retried, + run_bounded, run_in_forked_child, shutdown_bounded, sigint_when, @@ -155,10 +156,13 @@ async def main(): task = asyncio.ensure_future(run(work)) await asyncio.to_thread(operation.started.wait, EVENT_TIMEOUT) task.cancel() - with pytest.raises(asyncio.CancelledError): - await task + # Lets the cancellation reach the task, which waits for a running + # call unless it runs shielded. + await asyncio.sleep(0.1) operation.release.set() started = time.monotonic() + with pytest.raises(asyncio.CancelledError): + await task assert await asyncio.to_thread(finished.wait, EVENT_TIMEOUT) return time.monotonic() - started @@ -170,6 +174,42 @@ async def main(): assert operation.calls == (3 if shielded else 1) + @pytest.mark.parametrize("cancellations", [1, 2]) + def test_run_cancelled_waits_for_call(self, cancellations): + # GH-1337: a cancelled task raises once its running call has ended, + # so that nothing that the call does, such as the commit of a write, + # follows the cancellation. The task used to raise at once. A second + # cancellation stops the wait. + release = threading.Event() + started = threading.Event() + ended = threading.Event() + pool = S3ThreadPool(1) + + def work(): + started.set() + release.wait(EVENT_TIMEOUT) + ended.set() + + async def main(): + task = asyncio.ensure_future(pool.run(work)) + assert await asyncio.to_thread(started.wait, EVENT_TIMEOUT) + for _ in range(cancellations): + task.cancel() + done, _ = await asyncio.wait([task], timeout=0.1) + # Done before the call has ended only after the second one. + assert bool(done) == (cancellations == 2) + release.set() + with pytest.raises(asyncio.CancelledError): + await await_bounded(task) + if cancellations == 1: + assert ended.is_set() + + try: + run_bounded(main) + finally: + release.set() + pool.shutdown() + def test_run_cancelled_after_return(self): # GH-1203: a cancellation that arrives after the call has returned does # not stop the work that the call left running, such as the parts of a