Skip to content

Commit

Permalink
Merge pull request #434 from laughingman7743/support_requester_payments
Browse files Browse the repository at this point in the history
Support for requester payments with s3fs
  • Loading branch information
laughingman7743 committed Apr 23, 2023
1 parent f714fde commit 0419f22
Showing 1 changed file with 7 additions and 0 deletions.
7 changes: 7 additions & 0 deletions pyathena/filesystem/s3.py
Original file line number Diff line number Diff line change
Expand Up @@ -62,6 +62,9 @@ def __init__(
self.default_cache_type = default_cache_type if default_cache_type else "bytes"
self.max_workers = max_workers

requester_pays = kwargs.pop("requester_pays", False)
self.request_kwargs = {"RequestPayer": "requester"} if requester_pays else {}

def _get_client_compatible_with_s3fs(self, **kwargs) -> BaseClient:
"""
https://github.com/fsspec/s3fs/blob/2023.4.0/s3fs/core.py#L457-L535
Expand Down Expand Up @@ -159,6 +162,7 @@ def _head_object(self, path: str, refresh: bool = False) -> Optional[Dict[Any, A
logger=_logger,
Bucket=bucket,
Key=key,
**self.request_kwargs,
)
except botocore.exceptions.ClientError as e:
if e.response["Error"]["Code"] in ["NoSuchKey", "NoSuchBucket", "404"]:
Expand Down Expand Up @@ -236,6 +240,7 @@ def _ls_dirs(
config=self._retry_config,
logger=_logger,
**request,
**self.request_kwargs,
)
files.extend(
S3Object(
Expand Down Expand Up @@ -327,6 +332,7 @@ def info(self, path, **kwargs):
Prefix=f"{key.rstrip('/')}/" if key else "",
Delimiter="/",
MaxKeys=1,
**self.request_kwargs,
)
if (
response.get("KeyCount", 0) > 0
Expand Down Expand Up @@ -457,6 +463,7 @@ def _get_object(
logger=_logger,
**request,
**kwargs,
**self.request_kwargs,
)
return ranges[0], cast(bytes, response["Body"].read())

Expand Down

0 comments on commit 0419f22

Please sign in to comment.