diff --git a/benchmarks/pyathena_bench/filesystem.py b/benchmarks/pyathena_bench/filesystem.py index a88203af7..00eff7f35 100644 --- a/benchmarks/pyathena_bench/filesystem.py +++ b/benchmarks/pyathena_bench/filesystem.py @@ -19,6 +19,7 @@ from botocore.config import Config from pyathena import connect +from pyathena.filesystem import s3_filesystem from pyathena.filesystem.s3 import S3FileSystem from pyathena.filesystem.s3_async import AioS3FileSystem from pyathena.util import parse_output_location @@ -237,8 +238,8 @@ def filesystem(settings: Settings, case: Case, session_: Any, connection: Any) - fs.connect() return fs filesystem_class = AioS3FileSystem if case.api == "aio" else S3FileSystem - return filesystem_class( - connection=connection, max_workers=case.max_workers, skip_instance_cache=True + return s3_filesystem( + connection, filesystem_class, max_workers=case.max_workers, skip_instance_cache=True ) diff --git a/benchmarks/tests/test_filesystem.py b/benchmarks/tests/test_filesystem.py index 50b941302..7143108f9 100644 --- a/benchmarks/tests/test_filesystem.py +++ b/benchmarks/tests/test_filesystem.py @@ -296,7 +296,7 @@ class TestRunFilesystem: def s3_client(self, monkeypatch): """Replace the connection and filesystems; the fake client answers HeadObject.""" client = FakeS3Client(10) - connection = SimpleNamespace(s3_client=client, close=lambda: None) + connection = SimpleNamespace(s3_client=client, retry_config=None, close=lambda: None) monkeypatch.setattr("pyathena_bench.filesystem.connect", lambda **kwargs: connection) for name in ("S3FileSystem", "AioS3FileSystem"): monkeypatch.setattr( diff --git a/docs/api/filesystem.rst b/docs/api/filesystem.rst index 12852bf7a..3b7dc202d 100644 --- a/docs/api/filesystem.rst +++ b/docs/api/filesystem.rst @@ -16,6 +16,8 @@ This section covers S3 filesystem integration and object management. S3 FileSystem ------------- +.. autofunction:: pyathena.filesystem.s3_filesystem + .. autoclass:: pyathena.filesystem.s3.S3FileSystem :members: diff --git a/docs/filesystem.md b/docs/filesystem.md index a69364ad9..3c6f9cee9 100644 --- a/docs/filesystem.md +++ b/docs/filesystem.md @@ -48,17 +48,24 @@ fsspec.register_implementation("s3", s3fs.S3FileSystem, clobber=True) ## Basic usage -The filesystem can be constructed from a PyAthena connection, whose S3 client it -then uses (see "S3 client" in [Usage](usage.md)), or directly with -s3fs-compatible credential arguments: +The filesystem can be created from a PyAthena connection with `s3_filesystem()`, which +gives it the S3 client and the retry policy of the connection (see "S3 client" in +[Usage](usage.md)), with a boto3 S3 client, or with s3fs-compatible credential +arguments: ```python +import boto3 + from pyathena import connect +from pyathena.filesystem import s3_filesystem from pyathena.filesystem.s3 import S3FileSystem -fs = S3FileSystem(connect(s3_staging_dir="s3://YOUR_S3_BUCKET/path/to/", +fs = s3_filesystem(connect(s3_staging_dir="s3://YOUR_S3_BUCKET/path/to/", region_name="us-west-2")) +# Or with a boto3 S3 client, and optionally a pyathena.util.RetryConfig. +fs = S3FileSystem(client=boto3.client("s3", region_name="us-west-2")) + # Or with direct credentials (s3fs-compatible arguments). fs = S3FileSystem(key="YOUR_ACCESS_KEY", secret="YOUR_SECRET_KEY") @@ -69,10 +76,23 @@ fs = S3FileSystem(profile="YOUR_PROFILE") fs = S3FileSystem(anon=True) ``` +The filesystem uses a given client as it is and does not close it. Combining `client` +with an argument from which the filesystem builds a client, such as `key`, `secret`, +`token`, `profile`, `anon`, `endpoint_url`, `client_kwargs`, `config_kwargs`, or +`region_name`, raises `TypeError`; such settings from fsspec's configuration, such as +the `FSSPEC_S3_ENDPOINT_URL` environment variable, are ignored. `requester_pays`, `max_workers`, the cache options, +and the arguments of fsspec can be combined with it. `s3_filesystem()` takes the other +arguments of the filesystem, and `AioS3FileSystem` as its second argument. + +The `connection` argument of `S3FileSystem` and `AioS3FileSystem` is deprecated and +emits a `DeprecationWarning`; it will be removed in the next major version. It is +subject to the same restrictions as `client`. + The S3 client computes checksums of uploads and validates those of downloads by botocore's `request_checksum_calculation` and `response_checksum_validation` options, -set with `s3_config` of the connection or with `config_kwargs`; see "Checksums" under -"S3 client" in [Usage](usage.md) for their settings and their effect on data integrity. +set with `s3_config` of the connection, with `config_kwargs`, or in the configuration of +a given client; see "Checksums" under "S3 client" in [Usage](usage.md) for their +settings and their effect on data integrity. Standard fsspec operations work as expected: @@ -161,12 +181,24 @@ reads, parts, part copies, the ACLs of a recursive `chmod`, the DeleteObjects ba filesystem, with `max_workers` threads. Each file or operation runs at most its own `max_workers` requests at once in that pool. The `max_workers` of the filesystem defaults to the `max_pool_connections` of its S3 client: 10 unless configured, or that -of the client of the connection given with `connection`. A client that the filesystem -builds takes its `max_pool_connections` from `max_workers` unless `config_kwargs` sets -it. For more parallel requests, raise `max_workers`, and with `connection`, also the -`max_pool_connections` of its S3 client, such as with its `s3_config`. `close()`, or the end of a `with` block, releases the threads of the pool once they +of the given `client`, such as the client of the connection. A client that the +filesystem builds takes its `max_pool_connections` from `max_workers` unless +`config_kwargs` sets it; `max_workers` does not change the connection pool of a given +client. For more parallel requests, raise `max_workers`, and with a given client, also +its `max_pool_connections`, such as with the `s3_config` of the connection. `close()`, or the end of a `with` block, releases the threads of the pool once they are idle, and the filesystem stays usable. +(filesystem-pickling)= + +### Pickling and processes + +A filesystem created with s3fs-compatible arguments can be pickled, as for Dask, Ray, or +`multiprocessing` workers: the unpickled filesystem builds its own client from the same +arguments. A filesystem created with `client`, `connection`, or `s3_filesystem()`, and +the files that it opens for reading, raise `TypeError` when they are pickled, as a boto3 +client cannot be pickled; as in fsspec, a file opened for writing cannot be pickled. To use a filesystem in such workers, create it with s3fs-compatible +arguments, such as `profile` or `region_name`, or create a filesystem in each worker. + In a forked child process, create a new filesystem. A file that the parent opened raises `RuntimeError` there from its reads, writes, and flushes, `commit()`, and `discard()`, and closing it only marks it closed, without uploading or aborting anything. diff --git a/docs/s3fs.md b/docs/s3fs.md index d5091ed4a..02f2a060f 100644 --- a/docs/s3fs.md +++ b/docs/s3fs.md @@ -255,6 +255,8 @@ The reader is used only for result files; managed query results are read through S3FSCursor reads a CSV result file in 32 MiB blocks with the `"background"` {ref}`fsspec cache `, which reads the next block while rows are converted and keeps two blocks. With `block_size` given to `execute()`, the file is read in blocks of that size with the `"bytes"` cache. A `filesystem_class` other than PyAthena's `S3FileSystem` and `AioS3FileSystem`, including a subclass, reads the file with its own default cache type. +PyAthena's filesystems are created with the S3 client of the connection; another `filesystem_class` is created with the `connection`, `default_block_size`, and `skip_instance_cache=True` arguments. +In the next major version, it will be created with `client` and `retry_config` in place of `connection`. ### Limitations diff --git a/docs/usage.md b/docs/usage.md index 9519cdff8..27ebd7efb 100644 --- a/docs/usage.md +++ b/docs/usage.md @@ -761,7 +761,7 @@ The Glue request does not carry the connection's workgroup; turn the fallback of ## S3 client -A connection builds one S3 client on first use and shares it with the result sets of its cursors, the `S3FileSystem` instances created with `connection=`, and the Spark cursors. +A connection builds one S3 client on first use and shares it with the result sets of its cursors, the filesystems created with `pyathena.filesystem.s3_filesystem()`, and the Spark cursors. The result sets of `Cursor`, `DictCursor`, and their asynchronous versions do not use it, so these cursors build no S3 client. `ArrowCursor`, and `PolarsCursor` for Parquet (`unload=True`) and chunked results, read the query results through their libraries' own S3 clients and use the shared client for their other S3 requests. @@ -807,7 +807,7 @@ conn = connect( ) ``` -An `S3FileSystem` built without a connection takes them in `config_kwargs`: +An `S3FileSystem` built from s3fs-style arguments takes them in `config_kwargs`, and one given `client` uses the configuration of that client: ```python from pyathena.filesystem.s3 import S3FileSystem diff --git a/pyathena/filesystem/__init__.py b/pyathena/filesystem/__init__.py index bcdea01f1..9477feb9d 100644 --- a/pyathena/filesystem/__init__.py +++ b/pyathena/filesystem/__init__.py @@ -7,14 +7,65 @@ """fsspec filesystem implementations for Amazon S3.""" +from __future__ import annotations + import logging +from typing import TYPE_CHECKING, Any, TypeVar, overload import fsspec +from fsspec import AbstractFileSystem from pyathena.filesystem.s3 import S3FileSystem +if TYPE_CHECKING: + from pyathena.connection import Connection + _logger = logging.getLogger(__name__) +_FileSystemT = TypeVar("_FileSystemT", bound=AbstractFileSystem) + + +@overload +def s3_filesystem(connection: Connection[Any], **kwargs: Any) -> S3FileSystem: ... + + +@overload +def s3_filesystem( + connection: Connection[Any], filesystem_class: type[_FileSystemT], **kwargs: Any +) -> _FileSystemT: ... + + +def s3_filesystem( + connection: Connection[Any], + filesystem_class: type[AbstractFileSystem] = S3FileSystem, + **kwargs: Any, +) -> AbstractFileSystem: + """Create an S3 filesystem with the S3 client and retry policy of a connection. + + The filesystem borrows ``connection.s3_client``, which the connection + shares with its result sets, and uses ``connection.retry_config``. + Closing the filesystem does not close the client. A filesystem created + this way cannot be pickled; see :class:`~pyathena.filesystem.s3.S3FileSystem`. + + Example: + >>> from pyathena import connect + >>> from pyathena.filesystem import s3_filesystem + >>> fs = s3_filesystem(connect(s3_staging_dir="s3://my-bucket/path/")) + + Args: + connection: The PyAthena connection. + filesystem_class: ``S3FileSystem``, ``AioS3FileSystem``, or a subclass of + either. + **kwargs: The other arguments of the filesystem, except the arguments + from which it builds a client. + + Returns: + The filesystem. + """ + return filesystem_class( + client=connection.s3_client, retry_config=connection.retry_config, **kwargs + ) + def register_s3_filesystem() -> None: """Register PyAthena's S3 filesystem as fsspec's "s3" / "s3a" protocols. diff --git a/pyathena/filesystem/s3.py b/pyathena/filesystem/s3.py index c7e7c0966..644e40b58 100644 --- a/pyathena/filesystem/s3.py +++ b/pyathena/filesystem/s3.py @@ -8,6 +8,7 @@ import logging import mimetypes import os.path +import sys import threading import time import warnings @@ -17,22 +18,23 @@ from datetime import datetime from glob import has_magic from io import BytesIO -from typing import Any, BinaryIO, cast +from typing import TYPE_CHECKING, Any, BinaryIO, ClassVar, cast from urllib.parse import unquote_plus +import fsspec from boto3 import Session from botocore import UNSIGNED from botocore.client import BaseClient, Config from fsspec import AbstractFileSystem from fsspec.callbacks import _DEFAULT_CALLBACK, Callback from fsspec.compression import compr +from fsspec.config import apply_config 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, other_paths, tokenize import pyathena -from pyathena.connection import Connection from pyathena.filesystem.s3_core import S3Core, S3DeleteBatch, S3DeleteResult from pyathena.filesystem.s3_errors import S3ClientError from pyathena.filesystem.s3_executor import ( @@ -57,8 +59,26 @@ from pyathena.filesystem.s3_writer import _S3WriteSession from pyathena.util import RetryConfig, override, raise_if_stopped, unstoppable_retries +if TYPE_CHECKING: + from pyathena.connection import Connection + _logger = logging.getLogger(__name__) +# The message of the TypeError that a filesystem with an injected client, and +# the files that it opens for reading, raise when they are pickled. +_UNPICKLABLE_CLIENT_MESSAGE = ( + "{cls} built with client= or connection= cannot be pickled, as its boto3 client " + "cannot be. To pickle it, for example to send it to Dask, Ray, or multiprocessing " + "workers, build it from s3fs-style arguments such as profile=, region_name=, or " + "key=, secret=, and token=." +) +# The DeprecationWarning of the connection argument of the filesystems. +_CONNECTION_DEPRECATION_MESSAGE = ( + "The connection argument of {cls} is deprecated and will be removed in the next " + "major version. Use pyathena.filesystem.s3_filesystem(connection), or pass " + "client=connection.s3_client and retry_config=connection.retry_config." +) + # The request parameters on which the authorization of a lookup (HeadObject, # HeadBucket, or ListObjectsV2) depends. _LOOKUP_REQUEST_PARAMETERS = frozenset( @@ -273,6 +293,48 @@ class S3FileSystem(AbstractFileSystem): protocol = ("s3", "s3a") _extra_tokenize_attributes = ("default_block_size",) + # The boto3 session arguments of the client that the filesystem builds. + _SESSION_PASSING_ARGS: ClassVar[frozenset[str]] = frozenset( + { + "aws_access_key_id", + "aws_secret_access_key", + "aws_session_token", + "region_name", + "botocore_session", + "profile_name", + } + ) + # The arguments of ``Session.client()`` taken from ``client_kwargs``. + _CLIENT_PASSING_ARGS: ClassVar[frozenset[str]] = frozenset( + { + "aws_access_key_id", + "aws_secret_access_key", + "aws_session_token", + "api_version", + "use_ssl", + "verify", + "endpoint_url", + "region_name", + "config", + } + ) + # The arguments from which the filesystem builds its client, which an + # injected client (``client`` or ``connection``) does not take. + _CLIENT_BUILDING_ARGS: ClassVar[frozenset[str]] = _SESSION_PASSING_ARGS | { + "key", + "username", + "secret", + "password", + "token", + "profile", + "anon", + "use_ssl", + "endpoint_url", + "connect_timeout", + "read_timeout", + "client_kwargs", + "config_kwargs", + } def __init__( self, @@ -286,15 +348,23 @@ def __init__( version_aware: bool = False, *args, batch_size: int | None = None, + client: BaseClient | None = None, + retry_config: RetryConfig | None = None, **kwargs, ) -> None: """Create a filesystem for Amazon S3. + The filesystem sends its requests with the boto3 S3 client given as + ``client``, or with a client that it builds from s3fs-compatible + arguments in ``kwargs``. :func:`pyathena.filesystem.s3_filesystem` + creates one with the S3 client and retry policy of a PyAthena + connection. + Args: - connection: A PyAthena connection whose S3 client - (``Connection.s3_client``) and retry policy the filesystem uses. - Without one, the client is built from s3fs-compatible arguments - in ``kwargs``. + connection: Deprecated; use ``client`` and ``retry_config``, or + :func:`pyathena.filesystem.s3_filesystem`. A PyAthena connection + whose S3 client (``Connection.s3_client``) and retry policy the + filesystem uses, with a ``DeprecationWarning``. default_block_size: The block size for reads and writes; defaults to ``DEFAULT_BLOCK_SIZE``. default_cache_type: The fsspec cache type for reads, the name of a @@ -313,7 +383,8 @@ def __init__( None, the ``max_pool_connections`` of the S3 client (10 unless configured). A client that the filesystem builds gets ``max_pool_connections`` from ``max_workers`` unless - ``config_kwargs`` sets it. + ``config_kwargs`` sets it; the connection pool of a given + ``client`` keeps its size. s3_additional_kwargs: Extra arguments for the object requests of ``open()``, ``pipe_file()`` and ``put_file()``. Their authorization parameters (``RequestPayer``, @@ -333,22 +404,33 @@ def __init__( (``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``. + client: The boto3 S3 client that the filesystem sends its requests + with. The filesystem does not close it. A filesystem with a + given client cannot be pickled. + retry_config: The retry policy of the requests; defaults to + ``RetryConfig()``. + **kwargs: Passed to ``fsspec.AbstractFileSystem``; without + ``client`` or ``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. + TypeError: If ``connection`` is given with ``client`` or + ``retry_config``, or ``client`` or ``connection`` with an + argument from which the filesystem builds a client, such as + ``key``, ``profile``, ``anon``, ``endpoint_url``, + ``client_kwargs``, or ``region_name``, other than one that + fsspec's configuration supplies. """ super().__init__(*args, **kwargs) self.batch_size = self._check_batch_size(batch_size) - if connection: - client = connection.s3_client - retry_config = connection.retry_config - else: + self._client_injected = client is not None or connection is not None + client, retry_config = self._resolve_injected_client( + type(self), connection, client, retry_config, kwargs + ) + if client is None: client = self._get_client_compatible_with_s3fs(max_workers=max_workers, **kwargs) - retry_config = RetryConfig() self.default_block_size = ( default_block_size if default_block_size else self.DEFAULT_BLOCK_SIZE ) @@ -372,10 +454,94 @@ def __init__( request_kwargs.setdefault("RequestPayer", "requester") self._core = S3Core( client, - retry_config=retry_config, + retry_config=retry_config if retry_config is not None else RetryConfig(), request_kwargs={**request_kwargs, **S3Core.copy_from_source_params(request_kwargs)}, ) + @classmethod + def _resolve_injected_client( + cls, + filesystem_class: type[AbstractFileSystem], + connection: Connection[Any] | None, + client: BaseClient | None, + retry_config: RetryConfig | None, + kwargs: Mapping[str, Any], + ) -> tuple[BaseClient | None, RetryConfig | None]: + """Check the injected client of a filesystem and resolve it from a connection. + + Warns that ``connection`` is deprecated, and rejects the arguments that + conflict with an injected client before the client of the connection is + used. ``AioS3FileSystem`` resolves its arguments with it too. + + Args: + filesystem_class: The class of the filesystem being created, whose + name the warning shows and whose fsspec configuration applies. + connection: The deprecated ``connection`` argument. + client: The ``client`` argument. + retry_config: The ``retry_config`` argument. + kwargs: The other keyword arguments of the filesystem, merged with + fsspec's configuration for the protocol. + + Returns: + The injected client and retry policy, or the given ones if + ``connection`` is None. + + Raises: + TypeError: If the arguments conflict with an injected client. + """ + if connection is not None: + message = _CONNECTION_DEPRECATION_MESSAGE.format(cls=filesystem_class.__name__) + if sys.version_info >= (3, 12): + # Attributed to the first caller outside fsspec and the filesystems, + # such as through fsspec.filesystem() or a subclass. + warnings.warn( + message, + DeprecationWarning, + skip_file_prefixes=( + os.path.dirname(fsspec.__file__) + os.sep, + os.path.dirname(__file__) + os.sep, + ), + ) + else: + # Past this method, __init__, and fsspec's metaclass. + warnings.warn(message, DeprecationWarning, stacklevel=4) + if client is not None or retry_config is not None: + raise TypeError("connection cannot be combined with client or retry_config.") + if client is not None or connection is not None: + # The settings of fsspec's configuration for the protocol, such as + # the FSSPEC_S3_* environment variables, apply to every filesystem + # and are ignored instead. + configured = apply_config(filesystem_class, {}) + building_args = sorted( + k + for k in cls._CLIENT_BUILDING_ARGS + if kwargs.get(k) is not None and (k not in configured or kwargs[k] != configured[k]) + ) + if building_args: + raise TypeError( + f"{', '.join(building_args)} cannot be combined with client or " + "connection, whose client the filesystem uses as it is." + ) + if connection is not None: + return connection.s3_client, connection.retry_config + return client, retry_config + + @override + def __reduce__(self) -> tuple[Any, ...]: + """Reduce the filesystem for pickling, as fsspec does. + + Returns: + The reduction of ``fsspec.AbstractFileSystem``, which creates the + filesystem again from its arguments. + + Raises: + TypeError: If the filesystem was built with ``client`` or + ``connection``, whose boto3 client cannot be pickled. + """ + if self._client_injected: + raise TypeError(_UNPICKLABLE_CLIENT_MESSAGE.format(cls=type(self).__name__)) + return cast(tuple[Any, ...], super().__reduce__()) + @property def core(self) -> S3Core: """The typed S3 operations that the filesystem sends its requests with.""" @@ -476,13 +642,11 @@ def _get_client_compatible_with_s3fs( if profile := kwargs.pop("profile", None): kwargs.setdefault("profile_name", profile) - session = Session( - **{k: v for k, v in kwargs.items() if k in Connection._SESSION_PASSING_ARGS} - ) + session = Session(**{k: v for k, v in kwargs.items() if k in self._SESSION_PASSING_ARGS}) return session.client( "s3", config=Config(**config_kwargs), - **{k: v for k, v in client_kwargs.items() if k in Connection._CLIENT_PASSING_ARGS}, + **{k: v for k, v in client_kwargs.items() if k in self._CLIENT_PASSING_ARGS}, ) @staticmethod diff --git a/pyathena/filesystem/s3_async.py b/pyathena/filesystem/s3_async.py index 6c9880b1d..dc3a0e8db 100644 --- a/pyathena/filesystem/s3_async.py +++ b/pyathena/filesystem/s3_async.py @@ -30,7 +30,13 @@ from fsspec.utils import check_contained from pyathena.aio.util import create_in_thread -from pyathena.filesystem.s3 import CompressedBuffer, S3File, S3FileSystem, S3FileTarget +from pyathena.filesystem.s3 import ( + _UNPICKLABLE_CLIENT_MESSAGE, + CompressedBuffer, + S3File, + S3FileSystem, + S3FileTarget, +) from pyathena.filesystem.s3_core import S3Core from pyathena.filesystem.s3_executor import ItemT, S3Executor, S3MapResult, S3ThreadPool, T from pyathena.filesystem.s3_object import ( @@ -44,12 +50,15 @@ 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 +from pyathena.util import override, retries_stopped if TYPE_CHECKING: from datetime import datetime + from botocore.client import BaseClient + from pyathena.connection import Connection + from pyathena.util import RetryConfig _logger = logging.getLogger(__name__) @@ -300,12 +309,18 @@ def __init__( asynchronous: bool = False, loop: Any | None = None, batch_size: int | None = None, + *, + client: BaseClient | None = None, + retry_config: RetryConfig | None = None, **kwargs, ) -> None: """Initialize the filesystem and its internal ``S3FileSystem``. Args: - connection: Passed to the internal ``S3FileSystem``. + connection: Deprecated; use ``client`` and ``retry_config``, or + :func:`pyathena.filesystem.s3_filesystem` with + ``AioS3FileSystem``. A PyAthena connection whose S3 client and + retry policy the filesystem uses, with a ``DeprecationWarning``. default_block_size: Passed to the internal ``S3FileSystem``. default_cache_type: Passed to the internal ``S3FileSystem``. max_workers: The number of threads that run the blocking calls of @@ -321,8 +336,16 @@ def __init__( asynchronous: Passed to ``fsspec.asyn.AsyncFileSystem``. loop: Passed to ``fsspec.asyn.AsyncFileSystem``. batch_size: Passed to ``fsspec.asyn.AsyncFileSystem``. + client: Passed to the internal ``S3FileSystem``. A filesystem with + a given client cannot be pickled. + retry_config: Passed to the internal ``S3FileSystem``. **kwargs: Passed to both ``fsspec.asyn.AsyncFileSystem`` and the internal ``S3FileSystem``. + + Raises: + TypeError: If ``connection`` is given with ``client`` or + ``retry_config``, or the internal ``S3FileSystem`` rejects + the arguments. """ super().__init__( asynchronous=asynchronous, @@ -330,8 +353,13 @@ def __init__( batch_size=batch_size, **kwargs, ) + # Checked before the client of the connection is used. + client, retry_config = S3FileSystem._resolve_injected_client( + type(self), connection, client, retry_config, kwargs + ) self._sync_fs = S3FileSystem( - connection=connection, + client=client, + retry_config=retry_config, default_block_size=default_block_size, default_cache_type=default_cache_type, max_workers=max_workers, @@ -351,6 +379,22 @@ def __init__( # Share dircache for cache coherence between async and sync instances self.dircache = self._sync_fs.dircache + @override + def __reduce__(self) -> tuple[Any, ...]: + """Reduce the filesystem for pickling, as fsspec does. + + Returns: + The reduction of ``fsspec.AbstractFileSystem``, which creates the + filesystem again from its arguments. + + Raises: + TypeError: If the filesystem was built with ``client`` or + ``connection``, whose boto3 client cannot be pickled. + """ + if self._sync_fs._client_injected: + raise TypeError(_UNPICKLABLE_CLIENT_MESSAGE.format(cls=type(self).__name__)) + return cast(tuple[Any, ...], super().__reduce__()) + @property def core(self) -> S3Core: """The typed S3 operations of the wrapped ``S3FileSystem``.""" diff --git a/pyathena/pandas/result_set.py b/pyathena/pandas/result_set.py index 548cb7604..feb6e052f 100644 --- a/pyathena/pandas/result_set.py +++ b/pyathena/pandas/result_set.py @@ -31,6 +31,7 @@ from pyathena.arrow.util import to_column_info from pyathena.converter import Converter, _parse_time_zone, _to_datetime_with_tz, _to_time from pyathena.error import DataError, ProgrammingError +from pyathena.filesystem import s3_filesystem from pyathena.filesystem.s3 import S3FileSystem from pyathena.model import AthenaQueryExecution from pyathena.pandas.reader import _BINARY_NULL, BinaryCSVReader @@ -919,8 +920,8 @@ def filesystem(self) -> S3FileSystem: if self._filesystem is None: # Not cached by fsspec so that the connection and the dircache are released # with the result set. - self._filesystem = S3FileSystem( - connection=self._connection, + self._filesystem = s3_filesystem( + self._connection, default_block_size=self._block_size, default_cache_type=self._cache_type, max_workers=self._max_workers, diff --git a/pyathena/polars/result_set.py b/pyathena/polars/result_set.py index 108ddcb6d..c728ba5dd 100644 --- a/pyathena/polars/result_set.py +++ b/pyathena/polars/result_set.py @@ -34,6 +34,7 @@ from pyathena.arrow.reader import ParquetChunkReader from pyathena.converter import Converter, _to_default from pyathena.error import DataError, ProgrammingError +from pyathena.filesystem import s3_filesystem from pyathena.filesystem.s3 import S3FileSystem from pyathena.model import AthenaQueryExecution from pyathena.polars.converter import _to_datetimes @@ -609,15 +610,15 @@ def _read_kwargs( def _filesystem_kwargs(self) -> dict[str, Any]: """Get the arguments of the S3FileSystem that reads the result files. - ``_read_csv()`` opens a CSV result file with an ``S3FileSystem`` built from - them and passes the open file to Polars; the chunked read of UNLOAD results - with pyarrow reads its files with one. + ``_read_csv()`` opens a CSV result file with an ``S3FileSystem`` that + ``s3_filesystem()`` creates with them and the S3 client of the connection, + and passes the open file to Polars; the chunked read of UNLOAD results with + pyarrow reads its files with one. Returns: - The ``S3FileSystem`` arguments. + The ``S3FileSystem`` arguments other than the client of the connection. """ return { - "connection": self.connection, "default_block_size": self._block_size, "default_cache_type": self._cache_type, "max_workers": self._max_workers, @@ -884,7 +885,7 @@ def _read_csv(self) -> pl.DataFrame: fsspec.open, self.output_location, "rb", **(storage_options or {}) ) else: - filesystem = S3FileSystem(**storage_options) + filesystem = s3_filesystem(self.connection, S3FileSystem, **storage_options) open_csv = partial(filesystem.open, self.output_location, "rb") schema_overrides, timestamp_dtypes = self._get_csv_schema_overrides(has_header) strict = not kwargs.get("ignore_errors", False) @@ -1250,7 +1251,7 @@ def _iter_parquet_chunks(self, schema: pl.Schema) -> Iterator[pl.DataFrame]: if not self._kwargs and importlib.util.find_spec("pyarrow") is not None: return self._read_parquet_chunks( - S3FileSystem(**self._filesystem_kwargs), self._data_manifest + s3_filesystem(self.connection, **self._filesystem_kwargs), self._data_manifest ) return self._scan_parquet_chunks(self._unload_location, self.connection, schema) diff --git a/pyathena/s3fs/result_set.py b/pyathena/s3fs/result_set.py index b6518e268..4d2432fdd 100644 --- a/pyathena/s3fs/result_set.py +++ b/pyathena/s3fs/result_set.py @@ -18,6 +18,7 @@ from pyathena.converter import Converter from pyathena.error import OperationalError, ProgrammingError +from pyathena.filesystem import s3_filesystem from pyathena.filesystem.s3 import S3FileSystem from pyathena.filesystem.s3_async import AioS3FileSystem from pyathena.model import AthenaQueryExecution @@ -139,14 +140,24 @@ def __init__( self._rows.extend(pre_fetched_rows) def _create_s3_file_system(self) -> AbstractFileSystem: - """Create S3FileSystem using connection settings.""" + """Create the filesystem that reads the results with the connection's S3 client. + + PyAthena's ``S3FileSystem`` and ``AioS3FileSystem`` are created with + ``s3_filesystem()``. Another ``filesystem_class``, including a subclass, + still receives ``connection`` until the next major version. + + Returns: + The filesystem. + """ # Not cached by fsspec so that the connection and the dircache are # released with the result set. - return self._filesystem_class( - connection=self.connection, - default_block_size=self._block_size or self.DEFAULT_BLOCK_SIZE, - skip_instance_cache=True, - ) + kwargs: dict[str, Any] = { + "default_block_size": self._block_size or self.DEFAULT_BLOCK_SIZE, + "skip_instance_cache": True, + } + if self._filesystem_class in (S3FileSystem, AioS3FileSystem): + return s3_filesystem(self.connection, self._filesystem_class, **kwargs) + return self._filesystem_class(connection=self.connection, **kwargs) def _get_csv_open_options(self) -> dict[str, Any]: """Get the options to open the CSV result file with the filesystem. diff --git a/tests/pyathena/filesystem/test_fsspec_conformance.py b/tests/pyathena/filesystem/test_fsspec_conformance.py index 93106b04d..a76647ad6 100644 --- a/tests/pyathena/filesystem/test_fsspec_conformance.py +++ b/tests/pyathena/filesystem/test_fsspec_conformance.py @@ -14,7 +14,7 @@ import fsspec.tests.abstract as abstract import pytest -from pyathena.filesystem.s3 import S3FileSystem +from pyathena.filesystem import s3_filesystem from pyathena.filesystem.s3_async import AioS3FileSystem from tests import ENV from tests.pyathena.conftest import connect @@ -23,7 +23,7 @@ class S3Fixtures(abstract.AbstractFixtures): @pytest.fixture(scope="class") def fs(self): - return S3FileSystem(connect()) + return s3_filesystem(connect()) @pytest.fixture def fs_join(self): @@ -54,7 +54,7 @@ def supports_empty_directories(self): class AioS3Fixtures(S3Fixtures): @pytest.fixture(scope="class") def fs(self): - fs = AioS3FileSystem(connection=connect()) + fs = s3_filesystem(connect(), AioS3FileSystem) yield fs fs.close() diff --git a/tests/pyathena/filesystem/test_s3.py b/tests/pyathena/filesystem/test_s3.py index b48f339b3..33e629759 100644 --- a/tests/pyathena/filesystem/test_s3.py +++ b/tests/pyathena/filesystem/test_s3.py @@ -10,6 +10,7 @@ import itertools import lzma import os +import pickle import re import sys import tempfile @@ -42,7 +43,7 @@ from fsspec.implementations.memory import MemoryFileSystem import pyathena -from pyathena.filesystem import register_s3_filesystem +from pyathena.filesystem import register_s3_filesystem, s3_filesystem from pyathena.filesystem.s3 import CompressedBuffer, S3File, S3FileSystem from pyathena.filesystem.s3_core import S3Core, S3DeleteBatch, S3MultipartCopyPlan from pyathena.filesystem.s3_errors import S3ClientError @@ -62,7 +63,7 @@ ThrottledOperation, forking, interrupt_each_opcode, - mock_connection, + mock_client, overlapping_calls, poll_until, retried, @@ -3630,12 +3631,154 @@ def test_request_pool_size(self, max_workers, config_kwargs, pool_size, pool_con assert fs._request_pool.max_workers == pool_size assert fs._client.meta.config.max_pool_connections == pool_connections - def test_request_pool_size_of_connection(self): - # GH-1332: with a connection, the request pool takes the connection - # pool of its S3 client. - fs = S3FileSystem(connection=mock_connection(7), skip_instance_cache=True) + def test_request_pool_size_of_client(self): + # GH-1332: with a given client, the request pool takes the connection + # pool of the client. + fs = S3FileSystem(client=mock_client(7), skip_instance_cache=True) assert fs.max_workers == 7 + def test_client(self): + # GH-1363: the filesystem sends its requests with the given client and + # retry policy, which retries SlowDown, unlike the default one, and + # borrows the client: close() does not close it. + client = boto3.client( + "s3", region_name="us-east-1", aws_access_key_id="dummy", aws_secret_access_key="dummy" + ) + retry_config = RetryConfig(exceptions=["SlowDown"], attempt=2, multiplier=0) + fs = S3FileSystem(client=client, retry_config=retry_config, skip_instance_cache=True) + with Stubber(client) as stubber: + for _ in range(2): + stubber.add_client_error( + "head_object", service_error_code="SlowDown", http_status_code=503 + ) + with pytest.raises(OSError, match="Errno 16"): + fs.info("bucket/key") + stubber.assert_no_pending_responses() + with mock.patch.object(client, "close") as close: + fs.close() + close.assert_not_called() + + def test_s3_filesystem(self): + # GH-1363: s3_filesystem() creates a filesystem with the S3 client and + # retry policy of a connection, without the deprecated argument. + connection = mock.MagicMock() + connection.s3_client = mock_client(7) + with warnings.catch_warnings(): + warnings.simplefilter("error", DeprecationWarning) + fs = s3_filesystem(connection, skip_instance_cache=True) + assert type(fs) is S3FileSystem + assert fs._client is connection.s3_client + assert fs.core.retry_config is connection.retry_config + assert fs.max_workers == 7 + + def test_connection_deprecated(self): + # GH-1363: connection= still uses the S3 client and retry policy of the + # connection, with a DeprecationWarning that points at the caller. + connection = mock.MagicMock() + connection.s3_client = mock_client() + with pytest.warns(DeprecationWarning, match="s3_filesystem") as record: + fs = S3FileSystem(connection=connection, skip_instance_cache=True) + assert record[0].filename == __file__ + assert fs._client is connection.s3_client + assert fs.core.retry_config is connection.retry_config + + @pytest.mark.skipif(sys.version_info < (3, 12), reason="skip_file_prefixes of Python 3.12") + def test_connection_deprecated_through_fsspec(self): + # GH-1363: a filesystem created through fsspec.filesystem() attributes the + # warning to its caller, not to fsspec. + connection = mock.MagicMock() + connection.s3_client = mock_client() + with pytest.warns(DeprecationWarning, match="s3_filesystem") as record: + fsspec.filesystem("s3", connection=connection, skip_instance_cache=True) + assert record[0].filename == __file__ + + @pytest.mark.parametrize( + "kwargs", + [ + {"key": "dummy"}, + {"username": "dummy"}, + {"secret": "dummy"}, + {"password": "dummy"}, + {"token": "dummy"}, + {"profile": "dummy"}, + {"profile_name": "dummy"}, + {"anon": True}, + {"use_ssl": False}, + {"endpoint_url": "http://localhost:4566"}, + {"connect_timeout": 1}, + {"read_timeout": 1}, + {"client_kwargs": {}}, + {"config_kwargs": {}}, + {"region_name": "us-east-1"}, + {"aws_access_key_id": "dummy"}, + {"aws_secret_access_key": "dummy"}, + {"aws_session_token": "dummy"}, + {"botocore_session": mock.sentinel.botocore_session}, + ], + ids=lambda kwargs: next(iter(kwargs)), + ) + @pytest.mark.parametrize("injected", ["client", "connection"]) + def test_client_building_args_rejected(self, injected, kwargs): + # GH-1339: a client-building argument given with a client used to be + # ignored silently. + client = mock_client() + connection = mock.MagicMock() + s3_client = mock.PropertyMock(return_value=client) + type(connection).s3_client = s3_client + injected_kwargs = {injected: client if injected == "client" else connection} + with ( + warnings.catch_warnings(action="ignore", category=DeprecationWarning), + pytest.raises(TypeError, match=f"{next(iter(kwargs))} cannot be combined"), + ): + S3FileSystem(**injected_kwargs, **kwargs) + # Checked before the client of the connection is used. + s3_client.assert_not_called() + + @pytest.mark.parametrize("kwargs", [{"client": mock_client()}, {"retry_config": RetryConfig()}]) + def test_connection_with_client_rejected(self, kwargs): + with ( + pytest.warns(DeprecationWarning, match="connection argument"), + pytest.raises(TypeError, match="connection cannot be combined"), + ): + S3FileSystem(connection=mock.MagicMock(), **kwargs) + + def test_pickle_client_rejected(self): + # GH-1339: a boto3 client cannot be pickled; the filesystem names the + # arguments that can. + fs = S3FileSystem(client=S3_CLIENT, default_cache_type="none", skip_instance_cache=True) + with pytest.raises(TypeError, match=r"cannot be pickled.*profile="): + pickle.dumps(fs) + fs.info = mock.MagicMock(return_value=self._file_object("key")) + fs.info.return_value.size = 4 + with fs.open("s3://bucket/key", "rb") as f, pytest.raises(TypeError, match="profile="): + pickle.dumps(f) + + def test_pickle_s3fs_args(self): + fs = S3FileSystem( + key="dummy", + secret="dummy", + region_name="us-east-1", + default_block_size=2**20, + skip_instance_cache=True, + ) + unpickled = pickle.loads(pickle.dumps(fs)) + assert unpickled.default_block_size == 2**20 + assert unpickled._client.meta.region_name == "us-east-1" + + def test_client_ignores_fsspec_config(self, monkeypatch): + # GH-1363: settings of fsspec's configuration, such as the + # FSSPEC_S3_ENDPOINT_URL environment variable, apply to every filesystem, + # including the ones of the result sets, and do not conflict with a + # client. An explicit argument with another value still does. + monkeypatch.setitem(fsspec.config.conf, "s3", {"endpoint_url": "http://localhost:9000"}) + client = mock_client() + fs = S3FileSystem(client=client, skip_instance_cache=True) + assert fs._client is client + with pytest.raises(TypeError, match="endpoint_url cannot be combined"): + S3FileSystem( + client=client, endpoint_url="http://localhost:9001", skip_instance_cache=True + ) + def test_files_share_request_pool(self): # GH-1332: the reads of the files of a filesystem run in its request # pool, which bounds them together. Each file used to have a pool of @@ -5971,7 +6114,7 @@ def test_multipart_copy_uses_creation_algorithm(self, algorithm, checksum_type): def fs(self, request): if not hasattr(request, "param"): request.param = {} - return S3FileSystem(connect(), **request.param) + return s3_filesystem(connect(), **request.param) @pytest.mark.parametrize( ("algorithm", "checksum_type"), diff --git a/tests/pyathena/filesystem/test_s3_async.py b/tests/pyathena/filesystem/test_s3_async.py index 60bf10e1e..f6de38530 100644 --- a/tests/pyathena/filesystem/test_s3_async.py +++ b/tests/pyathena/filesystem/test_s3_async.py @@ -6,6 +6,7 @@ import gzip import io import os +import pickle import threading import time import traceback @@ -18,10 +19,12 @@ from types import SimpleNamespace from unittest import mock +import boto3 import pytest from botocore.stub import Stubber from fsspec.callbacks import Callback +from pyathena.filesystem import s3_filesystem from pyathena.filesystem.s3 import S3FileSystem from pyathena.filesystem.s3_async import ( AioS3File, @@ -54,7 +57,7 @@ cancel_when_started, forking, interrupt_each_opcode, - mock_connection, + mock_client, overlapping_calls, poll_until, retried, @@ -147,7 +150,7 @@ def test_read_in_forked_child(self): # parent opens reads its ranges in parallel. Its open() raised there # since GH-1336, and before that, its reads waited forever for the # event loop of the parent. - fs = AioS3FileSystem(connection=mock_connection(), skip_instance_cache=True) + fs = AioS3FileSystem(client=mock_client(), skip_instance_cache=True) block_size = 2**10 size = block_size * 4 fs.core.get_object = mock.MagicMock( @@ -260,7 +263,7 @@ def test_mv_source_lookup_and_delete_params(self): @pytest.mark.asyncio async def test_cp_file_failed_multipart_copy_invalidates_cache(self): # GH-973: see TestS3FileSystem.test_cp_file_failed_multipart_copy_invalidates_cache. - fs = AioS3FileSystem(connection=mock_connection(), skip_instance_cache=True) + fs = AioS3FileSystem(client=mock_client(), skip_instance_cache=True) fs._info = mock.AsyncMock( return_value=S3Object( init={"ContentLength": S3Core.MULTIPART_UPLOAD_MAX_PART_SIZE + 1}, @@ -282,7 +285,7 @@ async def test_copy_object_with_multipart_upload_max_workers(self): # whatever the size of the default executor of the event loop, in # which they used to run. asyncio.get_running_loop().set_default_executor(ThreadPoolExecutor(max_workers=1)) - fs = AioS3FileSystem(connection=mock_connection(), skip_instance_cache=True) + fs = AioS3FileSystem(client=mock_client(), skip_instance_cache=True) sync_fs = fs._sync_fs sync_fs.core.create_multipart_upload = mock.MagicMock( return_value=S3MultipartUpload( @@ -337,7 +340,7 @@ async def test_copy_object_with_multipart_upload_shares_threads(self): # GH-1159: concurrent copies share the threads of the filesystem, so # that a bulk copy of large objects does not start max_workers threads # per object. - fs = AioS3FileSystem(connection=mock_connection(), max_workers=2, skip_instance_cache=True) + fs = AioS3FileSystem(client=mock_client(), max_workers=2, skip_instance_cache=True) sync_fs = fs._sync_fs sync_fs.core.create_multipart_upload = mock.MagicMock( return_value=S3MultipartUpload( @@ -380,7 +383,7 @@ async def test_blocking_calls_bypass_default_executor(self): # does not hold them. They used to run in it. loop = asyncio.get_running_loop() loop.set_default_executor(ThreadPoolExecutor(max_workers=1)) - fs = AioS3FileSystem(connection=mock_connection(), skip_instance_cache=True) + fs = AioS3FileSystem(client=mock_client(), skip_instance_cache=True) fs._sync_fs.info = mock.MagicMock(return_value={"name": "bucket/key"}) release = threading.Event() # Held until it is released, after the lookup has had to finish. @@ -663,9 +666,7 @@ async def test_max_workers_none(self): # the filesystem, its request pool, and as the default max_workers of # its files and copies (GH-1332; it used to take the default of # ThreadPoolExecutor). The files and copies used to fail with it. - fs = AioS3FileSystem( - connection=mock_connection(7), max_workers=None, skip_instance_cache=True - ) + fs = AioS3FileSystem(client=mock_client(7), max_workers=None, skip_instance_cache=True) sync_fs = fs._sync_fs assert fs._thread_pool.max_workers == sync_fs._request_pool.max_workers == 7 sync_fs.info = mock.MagicMock(return_value={"name": "bucket/key"}) @@ -710,7 +711,7 @@ async def test_copy_object_with_multipart_upload_cancelled(self, cancellations): # aborts the upload, and is re-raised, as S3FileSystem does on an # interrupt. A repeated cancellation returns without stopping the # cleanup. - fs = AioS3FileSystem(connection=mock_connection(), max_workers=2, skip_instance_cache=True) + fs = AioS3FileSystem(client=mock_client(), max_workers=2, skip_instance_cache=True) sync_fs = fs._sync_fs sync_fs.core.create_multipart_upload = mock.MagicMock( return_value=S3MultipartUpload( @@ -788,7 +789,7 @@ def abort_multipart_upload(*args, **kwargs): async def test_close_releases_threads(self, exit_with): # GH-1209: close(), also at the end of a with block, ends the threads # of the filesystem, which stays usable with new ones. - fs = AioS3FileSystem(connection=mock_connection(), skip_instance_cache=True) + fs = AioS3FileSystem(client=mock_client(), skip_instance_cache=True) threads = [] def exists(path, **kwargs): @@ -815,7 +816,7 @@ def exists(path, **kwargs): def test_close_releases_request_threads(self): # GH-1332: close() also ends the threads of the request pool of the # internal S3FileSystem, which the files and parallel requests share. - fs = AioS3FileSystem(connection=mock_connection(), skip_instance_cache=True) + fs = AioS3FileSystem(client=mock_client(), skip_instance_cache=True) threads = [] fs._sync_fs._request_pool.submit(lambda: threads.append(threading.current_thread())).result( EVENT_TIMEOUT @@ -829,7 +830,7 @@ def test_close_releases_request_threads(self): async def test_close_during_multipart_copy(self, cancel): # GH-1209: a copy that runs at close() still schedules its remaining # requests and its completion or abort. - fs = AioS3FileSystem(connection=mock_connection(), max_workers=1, skip_instance_cache=True) + fs = AioS3FileSystem(client=mock_client(), max_workers=1, skip_instance_cache=True) sync_fs = fs._sync_fs sync_fs.core.create_multipart_upload = mock.MagicMock( return_value=S3MultipartUpload( @@ -892,7 +893,7 @@ def upload_part_copy(**kw): async def test_copy_object_with_multipart_upload_cancelled_completion(self, completion_fails): # GH-1046: a cancellation during CompleteMultipartUpload waits for it, # aborts the upload only if it failed, and is re-raised. - fs = AioS3FileSystem(connection=mock_connection(), skip_instance_cache=True) + fs = AioS3FileSystem(client=mock_client(), skip_instance_cache=True) sync_fs = fs._sync_fs sync_fs.core.create_multipart_upload = mock.MagicMock( return_value=S3MultipartUpload( @@ -971,7 +972,7 @@ async def test_copy_object_with_multipart_upload_cancelled_creation(self, creati # A cancellation during CreateMultipartUpload waits for it, aborts # the upload that it created before any part is copied, and is # re-raised. The created upload used to be left incomplete. - fs = AioS3FileSystem(connection=mock_connection(), skip_instance_cache=True) + fs = AioS3FileSystem(client=mock_client(), skip_instance_cache=True) sync_fs = fs._sync_fs events = [] started = threading.Event() @@ -1041,7 +1042,7 @@ async def test_copy_object_with_multipart_upload_cancelled_queued_completion(sel # GH-1186: a cancellation cancels a CompleteMultipartUpload request # that waits for a thread of the filesystem, and aborts the upload. The # request used to run after the cancellation and complete the copy. - fs = AioS3FileSystem(connection=mock_connection(), max_workers=1, skip_instance_cache=True) + fs = AioS3FileSystem(client=mock_client(), max_workers=1, skip_instance_cache=True) sync_fs = fs._sync_fs sync_fs.core.create_multipart_upload = mock.MagicMock( return_value=S3MultipartUpload( @@ -1125,7 +1126,7 @@ def test_copy_object_with_multipart_upload_interrupted_while_scheduling_completi # completion is being scheduled, before or after it is queued, keeps it # from running and aborts the upload, instead of leaving the upload # incomplete with a completion that nothing sends. - fs = AioS3FileSystem(connection=mock_connection(), max_workers=1, skip_instance_cache=True) + fs = AioS3FileSystem(client=mock_client(), max_workers=1, skip_instance_cache=True) sync_fs = fs._sync_fs sync_fs.core.create_multipart_upload = mock.MagicMock( return_value=S3MultipartUpload( @@ -1185,7 +1186,7 @@ async def test_copy_object_with_multipart_upload_failed_part_with_threads_taken( # thread of the request pool for retries; a completion queued behind it # would notice the failure only after its backoff wait of at least # EVENT_TIMEOUT seconds. - fs = AioS3FileSystem(connection=mock_connection(), max_workers=1, skip_instance_cache=True) + fs = AioS3FileSystem(client=mock_client(), max_workers=1, skip_instance_cache=True) sync_fs = fs._sync_fs plan = S3MultipartCopyPlan( source=S3Path("bucket", "src"), @@ -1241,7 +1242,7 @@ async def test_copy_object_with_multipart_upload_failed_annotation_stops_running # GH-1203: a failed annotation copy stops the retries of the running # ones, whose backoff wait would take at least EVENT_TIMEOUT seconds, # and its error is raised, whichever annotation comes first. - fs = AioS3FileSystem(connection=mock_connection(), skip_instance_cache=True) + fs = AioS3FileSystem(client=mock_client(), skip_instance_cache=True) sync_fs = fs._sync_fs plan = S3MultipartCopyPlan( source=S3Path("bucket", "src"), @@ -1281,7 +1282,7 @@ def copy_object_annotation(name, *args, **kwargs): async def test_copy_object_with_multipart_upload_cancelled_annotation_copy(self): # GH-1332 (B6): a cancelled copy raises once its running annotation # copy has ended, which used to run on after the cancellation. - fs = AioS3FileSystem(connection=mock_connection(), skip_instance_cache=True) + fs = AioS3FileSystem(client=mock_client(), skip_instance_cache=True) sync_fs = fs._sync_fs plan = S3MultipartCopyPlan( source=S3Path("bucket", "src"), @@ -1329,7 +1330,7 @@ def copy_object_annotation(name, *args, **kwargs): async def test_copy_object_with_multipart_upload_failed_annotation_skips_queued_copies(self): # GH-1203: the annotation copies that wait for a thread of the filesystem # when one fails are not sent. - fs = AioS3FileSystem(connection=mock_connection(), max_workers=1, skip_instance_cache=True) + fs = AioS3FileSystem(client=mock_client(), max_workers=1, skip_instance_cache=True) sync_fs = fs._sync_fs plan = S3MultipartCopyPlan( source=S3Path("bucket", "src"), @@ -1373,7 +1374,7 @@ async def test_copy_object_with_multipart_upload_cancelled_queued_parts(self): # files of the filesystem, before it waits for the running ones and # aborts the upload. The queued part copies used to run before the # abort. - fs = AioS3FileSystem(connection=mock_connection(), max_workers=1, skip_instance_cache=True) + fs = AioS3FileSystem(client=mock_client(), max_workers=1, skip_instance_cache=True) sync_fs = fs._sync_fs sync_fs.core.create_multipart_upload = mock.MagicMock( return_value=S3MultipartUpload( @@ -1456,7 +1457,7 @@ def test_copy_object_with_multipart_upload_loop_shutdown(self, stage, expected, # also cancelled the tasks of the requests and of the cleanup, so the # cleanup left an upload created during the cancellation, aborted # before the running parts ended, or aborted during the completion. - fs = AioS3FileSystem(connection=mock_connection(), skip_instance_cache=True) + fs = AioS3FileSystem(client=mock_client(), skip_instance_cache=True) sync_fs = fs._sync_fs events = [] lock = threading.Lock() @@ -1565,7 +1566,7 @@ def write(): def test_transaction_pipe_put_file_create_existing(self, tmp_path): # GH-972: in a transaction, put_file(mode="create") also raises, when # the file is opened, for an existing object. - fs = AioS3FileSystem(connection=mock_connection(), skip_instance_cache=True) + fs = AioS3FileSystem(client=mock_client(), skip_instance_cache=True) fs._sync_fs._object_exists = mock.MagicMock(return_value=True) fs._sync_fs._call = fs._sync_fs._core.call = mock.MagicMock() local = tmp_path / "local" @@ -1590,7 +1591,7 @@ def test_transaction_pipe_file_compression(self, path, compression): # GH-1037: the value is compressed before it is uploaded, also in a # transaction of this filesystem, whose pipe_file() is its own; outside # one, it is the synchronous filesystem's. - fs = AioS3FileSystem(connection=mock_connection(), skip_instance_cache=True) + fs = AioS3FileSystem(client=mock_client(), skip_instance_cache=True) put_object = fs._sync_fs.core.put_object = mock.MagicMock() with fs.transaction: @@ -1605,7 +1606,7 @@ def test_transaction_pipe_file_write(self): # and a failed write does not replace the object with an empty one # when the transaction commits. GH-1211: the commit of the failed # write raises, and fsspec discards the rest of the transaction. - fs = AioS3FileSystem(connection=mock_connection(), skip_instance_cache=True) + fs = AioS3FileSystem(client=mock_client(), skip_instance_cache=True) put_object = fs._sync_fs.core.put_object = mock.MagicMock() def write_in_transaction(): @@ -1634,7 +1635,7 @@ def test_transaction_put_file_failed_write(self, tmp_path, error): # GH-1211: the commit of the failed write raises, and fsspec discards # the rest of the transaction. A local file that cannot be read opens # no remote file. - fs = AioS3FileSystem(connection=mock_connection(), skip_instance_cache=True) + fs = AioS3FileSystem(client=mock_client(), skip_instance_cache=True) put_object = fs._sync_fs.core.put_object = mock.MagicMock() local = tmp_path / "local" local.write_bytes(b"a") @@ -1669,7 +1670,7 @@ def test_transaction_pipe_put_file_exceeding_max_parts(self, tmp_path, kwargs): # GH-953: in a transaction, as outside one, pipe_file() and put_file() # reject data that does not fit in the maximum number of parts before # opening the file. - fs = AioS3FileSystem(connection=mock_connection(), skip_instance_cache=True) + fs = AioS3FileSystem(client=mock_client(), skip_instance_cache=True) 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() @@ -1690,7 +1691,7 @@ async def test_rm_requests(self): # GH-971: _rm() sent the keys of every bucket to the bucket of the # first path and dropped the DeleteObjects parameters. The async # filesystem sends the request of each bucket in parallel. - fs = AioS3FileSystem(connection=mock_connection(), skip_instance_cache=True) + fs = AioS3FileSystem(client=mock_client(), skip_instance_cache=True) sync_fs = fs._sync_fs sync_fs._call = sync_fs._core.call = mock.MagicMock(return_value={}) @@ -1705,7 +1706,7 @@ async def test_rm_requests(self): @pytest.mark.asyncio async def test_rm_request_error_keeps_errors(self): - fs = AioS3FileSystem(connection=mock_connection(), skip_instance_cache=True) + fs = AioS3FileSystem(client=mock_client(), skip_instance_cache=True) def call(method, **request): if request["Bucket"] == "b1": @@ -1725,7 +1726,7 @@ async def test_rm_cancel_waits_for_running_requests(self): # have ended, and each request invalidates the cache of its objects # when it ends. The requests used to run on after the cancellation was # raised. - fs = AioS3FileSystem(connection=mock_connection(), skip_instance_cache=True) + fs = AioS3FileSystem(client=mock_client(), skip_instance_cache=True) started = {"b1": threading.Event(), "b2": threading.Event()} release = {"b1": threading.Event(), "b2": threading.Event()} @@ -1876,7 +1877,7 @@ def test_batch_cancelled_waits_for_calls(self, tmp_path, method, cancellation): # 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) + fs = AioS3FileSystem(client=mock_client(), skip_instance_cache=True) sync_fs = fs._sync_fs state = {"started": 0, "ended": 0, "stopped": 0} release = threading.Event() @@ -1950,7 +1951,7 @@ def test_batch_failure_waits_for_calls(self, tmp_path, method): # 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) + fs = AioS3FileSystem(client=mock_client(), skip_instance_cache=True) sync_fs = fs._sync_fs state = {"started": 0, "ended": 0, "stopped": 0} release = threading.Event() @@ -2303,7 +2304,7 @@ async def test_rm_batch_size(self, argument): # 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) + fs = AioS3FileSystem(client=mock_client(), skip_instance_cache=True) overlap, state = overlapping_calls(2, timeout=0.2) requests = [] @@ -2343,7 +2344,7 @@ 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 # 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 = AioS3FileSystem(client=mock_client(), skip_instance_cache=True) fs._open = mock.MagicMock() lpath = tmp_path / "data.csv" lpath.write_bytes(b"a") @@ -2383,7 +2384,7 @@ def test_put_file_in_transaction_open_parameters(self, tmp_path, mode, open_mode async def test_cp_file_directory(self): # GH-1008: recursive copy() passes the directories, which used to be # sent to CopyObject and fail with NoSuchKey. - fs = AioS3FileSystem(connection=mock_connection(), skip_instance_cache=True) + fs = AioS3FileSystem(client=mock_client(), skip_instance_cache=True) fs._info = mock.AsyncMock(return_value=S3FileSystem._directory_object("bucket", "src")) fs._sync_fs._call = fs._sync_fs._core.call = mock.MagicMock() @@ -2415,7 +2416,7 @@ async def test_expand_path_glob_lists_stem_prefix(self): # 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) + fs = AioS3FileSystem(client=mock_client(), skip_instance_cache=True) sync_fs = fs._sync_fs sync_fs._call = sync_fs._core.call = mock.MagicMock( side_effect=lambda method, **request: {"Contents": [], "KeyCount": 0} @@ -2434,7 +2435,7 @@ async def test_expand_path_glob_lists_stem_prefix(self): async def test_expand_path_version_maxdepth(self): # A maxdepth below 1 is rejected, also for a version, which fsspec's # expansion does not see. - fs = AioS3FileSystem(connection=mock_connection(), skip_instance_cache=True) + fs = AioS3FileSystem(client=mock_client(), skip_instance_cache=True) with pytest.raises(ValueError): # noqa: PT011 await fs._expand_path("s3://bucket/b?versionId=v1", maxdepth=0) @@ -2453,7 +2454,7 @@ async def test_expand_path_version_maxdepth(self): ) @pytest.mark.asyncio async def test_get_version(self, tmp_path, destination, expected): - fs = AioS3FileSystem(connection=mock_connection(), skip_instance_cache=True) + fs = AioS3FileSystem(client=mock_client(), skip_instance_cache=True) fs._sync_fs.isdir = mock.MagicMock(return_value=False) fs._get_file = mock.AsyncMock() @@ -2468,7 +2469,7 @@ async def test_get_version(self, tmp_path, destination, expected): async def test_mv(self): # GH-1008: the files are copied in parallel, the directories are # skipped, and only the copied sources are deleted. - fs = AioS3FileSystem(connection=mock_connection(), skip_instance_cache=True) + fs = AioS3FileSystem(client=mock_client(), skip_instance_cache=True) started = asyncio.Event() copies = [] @@ -2504,7 +2505,7 @@ async def copy_file(path1, path2, **kwargs): async def test_mv_copy_failure(self): # A failed copy is raised after the other copies have finished, and # nothing is deleted. - fs = AioS3FileSystem(connection=mock_connection(), skip_instance_cache=True) + fs = AioS3FileSystem(client=mock_client(), skip_instance_cache=True) finished = [] async def copy_file(path1, path2, **kwargs): @@ -2601,23 +2602,94 @@ def upload_part_copy(**kw): == "requester" ) + def test_client(self): + # GH-1363: the requests go through the given client with the given retry + # policy, which retries SlowDown, unlike the default one. + client = boto3.client( + "s3", region_name="us-east-1", aws_access_key_id="dummy", aws_secret_access_key="dummy" + ) + retry_config = RetryConfig(exceptions=["SlowDown"], attempt=2, multiplier=0) + fs = AioS3FileSystem(client=client, retry_config=retry_config, skip_instance_cache=True) + with Stubber(client) as stubber: + for _ in range(2): + stubber.add_client_error( + "head_object", service_error_code="SlowDown", http_status_code=503 + ) + with pytest.raises(OSError, match="Errno 16"): + fs.info("bucket/key") + stubber.assert_no_pending_responses() + + def test_s3_filesystem(self): + connection = mock.MagicMock() + connection.s3_client = mock_client() + with warnings.catch_warnings(): + warnings.simplefilter("error", DeprecationWarning) + fs = s3_filesystem(connection, AioS3FileSystem, skip_instance_cache=True) + assert type(fs) is AioS3FileSystem + assert fs._sync_fs._client is connection.s3_client + assert fs.core.retry_config is connection.retry_config + + def test_connection_deprecated(self): + # GH-1363: one DeprecationWarning, which points at the caller. + connection = mock.MagicMock() + connection.s3_client = mock_client() + with pytest.warns(DeprecationWarning, match="s3_filesystem") as record: + fs = AioS3FileSystem(connection=connection, skip_instance_cache=True) + assert len(record) == 1 + assert record[0].filename == __file__ + assert fs._sync_fs._client is connection.s3_client + assert fs.core.retry_config is connection.retry_config + + @pytest.mark.parametrize( + ("kwargs", "match"), + [ + ({"client": mock_client()}, "connection cannot be combined"), + ({"retry_config": RetryConfig()}, "connection cannot be combined"), + ({"profile": "dummy"}, "profile cannot be combined"), + ], + ids=["client", "retry_config", "profile"], + ) + def test_connection_conflicts_rejected(self, kwargs, match): + connection = mock.MagicMock() + s3_client = mock.PropertyMock(return_value=mock_client()) + type(connection).s3_client = s3_client + with ( + pytest.warns(DeprecationWarning, match="connection argument"), + pytest.raises(TypeError, match=match), + ): + AioS3FileSystem(connection=connection, skip_instance_cache=True, **kwargs) + # Checked before the client of the connection is used. + s3_client.assert_not_called() + + def test_pickle(self): + # GH-1339: a filesystem with a given client raises TypeError, which names + # the s3fs-style arguments; one built from them pickles. + fs = AioS3FileSystem(client=mock_client(), skip_instance_cache=True) + with pytest.raises(TypeError, match=r"AioS3FileSystem.*cannot be pickled.*profile="): + pickle.dumps(fs) + fs = AioS3FileSystem( + key="dummy", secret="dummy", region_name="us-east-1", skip_instance_cache=True + ) + unpickled = pickle.loads(pickle.dumps(fs)) + assert unpickled._sync_fs._client.meta.region_name == "us-east-1" + def test_internal_file_system_not_cached(self): # GH-978: the internal S3FileSystem was kept in the fsspec instance # cache, so skip_instance_cache=True instances shared it. - connection = mock_connection() - fs1 = AioS3FileSystem(connection=connection, skip_instance_cache=True) - fs2 = AioS3FileSystem(connection=connection, skip_instance_cache=True) + client = mock_client() + fs1 = AioS3FileSystem(client=client, skip_instance_cache=True) + fs2 = AioS3FileSystem(client=client, skip_instance_cache=True) assert fs1._sync_fs is not fs2._sync_fs assert fs1.dircache is not fs2.dircache - fs3 = AioS3FileSystem(connection=connection) - assert AioS3FileSystem(connection=connection) is fs3 + fs3 = AioS3FileSystem(client=client) + assert AioS3FileSystem(client=client) is fs3 for fs in (fs1, fs2, fs3): assert fs._sync_fs not in S3FileSystem._cache.values() @pytest.fixture(scope="class") def fs(self): - fs = AioS3FileSystem(connection=connect()) + fs = s3_filesystem(connect(), AioS3FileSystem) yield fs fs.close() @@ -2746,7 +2818,7 @@ class TestAioS3File: def test_open_max_workers(self, max_workers): # GH-1158: None takes the default of the filesystem. It used to raise # TypeError when the executor of the file was created. - fs = AioS3FileSystem(connection=mock_connection(), skip_instance_cache=True) + fs = AioS3FileSystem(client=mock_client(), skip_instance_cache=True) with fs.open("s3://bucket/key", "wb", max_workers=max_workers) as f: assert isinstance(f, AioS3File) assert f.max_workers == (max_workers or fs._sync_fs.max_workers) @@ -2758,7 +2830,7 @@ async def test_open_parallel_requests(self, asynchronous): # reads. GH-977: a filesystem created with asynchronous=True has no # event loop of its own and used to fail to run them. fs = AioS3FileSystem( - connection=mock_connection(), asynchronous=asynchronous, skip_instance_cache=True + client=mock_client(), asynchronous=asynchronous, skip_instance_cache=True ) block_size = S3Core.MULTIPART_UPLOAD_MIN_PART_SIZE size = block_size * 4 @@ -2823,14 +2895,14 @@ def read(): assert state["peak"] == 2 def test_open_invalid_max_workers(self): - fs = AioS3FileSystem(connection=mock_connection(), skip_instance_cache=True) + fs = AioS3FileSystem(client=mock_client(), skip_instance_cache=True) with pytest.raises(ValueError, match="max_workers must be greater than 0"): fs.open("s3://bucket/key", "wb", max_workers=0) def test_open_lookup_parameters(self): # GH-1004: the lookup of the file sends its lookup parameters, and # looks up the version that is opened. - fs = AioS3FileSystem(connection=mock_connection(), skip_instance_cache=True) + fs = AioS3FileSystem(client=mock_client(), skip_instance_cache=True) fs._sync_fs.info = mock.MagicMock( return_value=S3Object( init={"ContentLength": 4}, @@ -2855,7 +2927,7 @@ def test_open_lookup_parameters(self): class TestAioS3StreamedFile: @pytest.fixture def fs(self): - fs = AioS3FileSystem(connection=mock_connection(), skip_instance_cache=True) + fs = AioS3FileSystem(client=mock_client(), skip_instance_cache=True) core = fs._sync_fs.core core.create_multipart_upload = mock.MagicMock( return_value=S3MultipartUpload( @@ -3396,9 +3468,7 @@ def test_close_after_interrupted_cancel(self): # request that a close waits for, including one that skips the # callbacks of its future, does not keep the close waiting. def case(interrupt): - fs = AioS3FileSystem( - connection=mock_connection(), max_workers=1, skip_instance_cache=True - ) + fs = AioS3FileSystem(client=mock_client(), max_workers=1, skip_instance_cache=True) info = S3Object( init={"Key": "key", "ETag": '"etag"'}, type=S3ObjectType.S3_OBJECT_TYPE_FILE, diff --git a/tests/pyathena/pandas/test_cursor.py b/tests/pyathena/pandas/test_cursor.py index fd70756da..2f3c2cf92 100644 --- a/tests/pyathena/pandas/test_cursor.py +++ b/tests/pyathena/pandas/test_cursor.py @@ -89,7 +89,7 @@ def _offline_cursor(tmp_path, columns, rows, **cursor_kwargs): connection.client.get_query_results.return_value = query_results_response(columns) connection.s3_client.head_object.return_value = {"ContentLength": len(data)} with ( - patch("pyathena.pandas.result_set.S3FileSystem", lambda **_: fs), + patch("pyathena.pandas.result_set.s3_filesystem", lambda *_, **__: fs), PandasCursor( connection=connection, converter=DefaultPandasTypeConverter(), @@ -367,7 +367,11 @@ def test_csv_storage_options(self, pandas_cursor, query, expected): """ pandas_cursor.execute( query, - storage_options={"connection": pandas_cursor.connection, "skip_instance_cache": True}, + storage_options={ + "client": pandas_cursor.connection.s3_client, + "retry_config": pandas_cursor.connection.retry_config, + "skip_instance_cache": True, + }, ) assert pandas_cursor.fetchall() == expected with pytest.raises(OperationalError): @@ -1522,7 +1526,11 @@ def test_callback(query_id: str): @pytest.mark.parametrize("option", ["filesystem", "storage_options"]) def test_unload_with_filesystem_options(self, pandas_cursor, option): """filesystem or storage_options given to execute() read the UNLOAD result.""" - fs_kwargs = {"connection": pandas_cursor.connection, "skip_instance_cache": True} + fs_kwargs = { + "client": pandas_cursor.connection.s3_client, + "retry_config": pandas_cursor.connection.retry_config, + "skip_instance_cache": True, + } kwargs = ( {"filesystem": S3FileSystem(**fs_kwargs)} if option == "filesystem" diff --git a/tests/pyathena/pandas/test_result_set.py b/tests/pyathena/pandas/test_result_set.py index d5b73e9e2..4b2e27554 100644 --- a/tests/pyathena/pandas/test_result_set.py +++ b/tests/pyathena/pandas/test_result_set.py @@ -1662,7 +1662,7 @@ def _local_result_set( connection.s3_client.meta.config.max_pool_connections = 2 connection.s3_client.head_object.return_value = {"ContentLength": length} connection.client.get_query_results.return_value = query_results_response(types) - with patch("pyathena.pandas.result_set.S3FileSystem", return_value=fs): + with patch("pyathena.pandas.result_set.s3_filesystem", return_value=fs): if describe: return AthenaPandasResultSetMetadataReader( connection, query_execution, RetryConfig(), unload=unload, **kwargs @@ -2006,7 +2006,7 @@ def _fetching_result_set(result, arraysize): ) with ( patch.object(AthenaPandasResultSet, "_load_metadata"), - patch("pyathena.pandas.result_set.S3FileSystem"), + patch("pyathena.pandas.result_set.s3_filesystem"), patch.object(AthenaPandasResultSet, "_as_pandas", return_value=result), patch.object( AthenaPandasResultSet, "is_unload", new_callable=PropertyMock, return_value=True diff --git a/tests/pyathena/polars/test_cursor.py b/tests/pyathena/polars/test_cursor.py index defbb2d5e..4940ada97 100644 --- a/tests/pyathena/polars/test_cursor.py +++ b/tests/pyathena/polars/test_cursor.py @@ -95,7 +95,8 @@ def test_as_polars_with_read_kwargs(self, polars_cursor): "SELECT * FROM one_row", schema_overrides={"number_of_rows": pl.Utf8}, storage_options={ - "connection": polars_cursor.connection, + "client": polars_cursor.connection.s3_client, + "retry_config": polars_cursor.connection.retry_config, "skip_instance_cache": True, }, ).as_polars() diff --git a/tests/pyathena/polars/test_result_set.py b/tests/pyathena/polars/test_result_set.py index dca16725a..535eaf1e2 100644 --- a/tests/pyathena/polars/test_result_set.py +++ b/tests/pyathena/polars/test_result_set.py @@ -102,8 +102,8 @@ def _local_result_set( # The HEAD request parses the output location, which is a local path here. patch.object(AthenaPolarsResultSet, "_get_content_length", return_value=1), patch( - "pyathena.polars.result_set.S3FileSystem", - file_system or (lambda **_: _LocalS3FileSystem()), + "pyathena.polars.result_set.s3_filesystem", + file_system or (lambda *_, **__: _LocalS3FileSystem()), ), ): return AthenaPolarsResultSet( diff --git a/tests/pyathena/s3fs/test_result_set.py b/tests/pyathena/s3fs/test_result_set.py index 9ae446165..f2e81d3e1 100644 --- a/tests/pyathena/s3fs/test_result_set.py +++ b/tests/pyathena/s3fs/test_result_set.py @@ -5,6 +5,7 @@ # # SPDX-License-Identifier: MIT +import warnings from collections.abc import Iterator from io import BytesIO from typing import Any @@ -124,6 +125,8 @@ class TestAthenaS3FSResultSet: ], ids=["default", "aio", "block_size", "subclass"], ) + # The subclass receives the deprecated connection argument (see test_filesystem_client). + @pytest.mark.filterwarnings("ignore:The connection argument:DeprecationWarning") def test_csv_open_options(self, block_size, filesystem_class, prefetch): """PyAthena's filesystems read the CSV result in 32 MiB blocks with the background cache, which keeps two blocks, unless block_size is given. @@ -149,6 +152,31 @@ def test_csv_open_options(self, block_size, filesystem_class, prefetch): assert "cache_type" not in options assert "block_size" not in options + @pytest.mark.parametrize( + ("filesystem_class", "deprecated"), + [(S3FileSystem, False), (AioS3FileSystem, False), (_UserS3FileSystem, True)], + ids=["default", "aio", "subclass"], + ) + def test_filesystem_client(self, filesystem_class, deprecated): + """PyAthena's filesystems read the results with the connection's S3 client, + without the deprecated connection argument. + + Another filesystem_class still receives the connection until the next + major version. + """ + data = "null_col,empty_col,value_col\n,,a\n" + with ( + patch.object( + filesystem_class, "_open", autospec=True, return_value=BytesIO(data.encode()) + ), + warnings.catch_warnings(record=True) as record, + ): + warnings.simplefilter("always", DeprecationWarning) + result_set, _ = _result_set(data, None, filesystem_class=filesystem_class) + with result_set: + assert result_set._fs.core.client is result_set.connection.s3_client + assert any(issubclass(w.category, DeprecationWarning) for w in record) is deprecated + def test_reader_closes_at_end_of_results(self): """The reader is closed once every row has been read, before close(). diff --git a/tests/pyathena/test_connection.py b/tests/pyathena/test_connection.py index 30fd1cf24..e7ae455bb 100644 --- a/tests/pyathena/test_connection.py +++ b/tests/pyathena/test_connection.py @@ -38,7 +38,7 @@ from pyathena.converter import DefaultTypeConverter from pyathena.cursor import Cursor, DictCursor from pyathena.error import OperationalError, ProgrammingError -from pyathena.filesystem.s3 import S3FileSystem +from pyathena.filesystem import s3_filesystem from pyathena.pandas.async_cursor import AsyncPandasCursor from pyathena.pandas.cursor import PandasCursor from pyathena.polars.async_cursor import AsyncPolarsCursor @@ -530,9 +530,10 @@ def test_s3_client_uses_s3_endpoint_setting(self, isolated_aws_config, monkeypat def test_s3_filesystem_uses_connection_s3_client(self): conn = _connection() - fs = S3FileSystem(connection=conn, skip_instance_cache=True) + fs = s3_filesystem(conn, skip_instance_cache=True) assert fs._client is conn.s3_client + assert fs.core.retry_config is conn.retry_config def test_s3_config_defaults_to_config(self): conn = _connection(config=Config(max_pool_connections=20)) diff --git a/tests/pyathena/util.py b/tests/pyathena/util.py index 6de9141a9..5e8fbc506 100644 --- a/tests/pyathena/util.py +++ b/tests/pyathena/util.py @@ -500,12 +500,26 @@ def interrupt_each_opcode(target, case, limit=2000): return n - 1 -def mock_connection(max_pool_connections=10): - """Return a mock connection whose S3 client has a connection pool of the given size. +def mock_client(max_pool_connections=10): + """Return a mock S3 client with a connection pool of the given size. The filesystems size their request pool from the ``max_pool_connections`` of the S3 client when ``max_workers`` is not given. + Args: + max_pool_connections: The connection pool size of the S3 client. + + Returns: + The mock S3 client. + """ + client = MagicMock() + client.meta.config.max_pool_connections = max_pool_connections + return client + + +def mock_connection(max_pool_connections=10): + """Return a mock connection whose S3 client has a connection pool of the given size. + Args: max_pool_connections: The connection pool size of the S3 client. @@ -513,7 +527,7 @@ def mock_connection(max_pool_connections=10): The mock connection. """ connection = MagicMock() - connection.s3_client.meta.config.max_pool_connections = max_pool_connections + connection.s3_client = mock_client(max_pool_connections) return connection