Inject the filesystem's S3 client with client= and retry_config= (GH-1363)

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
     )
 
 

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(

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:
 

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.

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 <filesystem-read-caching>`, 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
 

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

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.