PyAthena 3.29.1__tar.gz → 3.29.3__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- {pyathena-3.29.1 → pyathena-3.29.3}/PKG-INFO +3 -3
- {pyathena-3.29.1 → pyathena-3.29.3}/README.md +2 -2
- {pyathena-3.29.1 → pyathena-3.29.3}/pyathena/__init__.py +2 -2
- {pyathena-3.29.1 → pyathena-3.29.3}/pyathena/_version.py +2 -2
- {pyathena-3.29.1 → pyathena-3.29.3}/pyathena/aio/arrow/cursor.py +1 -1
- {pyathena-3.29.1 → pyathena-3.29.3}/pyathena/aio/pandas/cursor.py +1 -1
- {pyathena-3.29.1 → pyathena-3.29.3}/pyathena/aio/polars/cursor.py +1 -1
- {pyathena-3.29.1 → pyathena-3.29.3}/pyathena/aio/s3fs/cursor.py +8 -5
- {pyathena-3.29.1 → pyathena-3.29.3}/pyathena/aio/spark/cursor.py +1 -1
- {pyathena-3.29.1 → pyathena-3.29.3}/pyathena/filesystem/s3.py +31 -22
- pyathena-3.29.3/pyathena/filesystem/s3_async.py +357 -0
- pyathena-3.29.3/pyathena/filesystem/s3_executor.py +89 -0
- {pyathena-3.29.1 → pyathena-3.29.3}/pyathena/s3fs/result_set.py +6 -2
- {pyathena-3.29.1 → pyathena-3.29.3}/pyproject.toml +10 -4
- {pyathena-3.29.1 → pyathena-3.29.3}/.gitignore +0 -0
- {pyathena-3.29.1 → pyathena-3.29.3}/LICENSE +0 -0
- {pyathena-3.29.1 → pyathena-3.29.3}/pyathena/aio/__init__.py +0 -0
- {pyathena-3.29.1 → pyathena-3.29.3}/pyathena/aio/arrow/__init__.py +0 -0
- {pyathena-3.29.1 → pyathena-3.29.3}/pyathena/aio/common.py +0 -0
- {pyathena-3.29.1 → pyathena-3.29.3}/pyathena/aio/connection.py +0 -0
- {pyathena-3.29.1 → pyathena-3.29.3}/pyathena/aio/cursor.py +0 -0
- {pyathena-3.29.1 → pyathena-3.29.3}/pyathena/aio/pandas/__init__.py +0 -0
- {pyathena-3.29.1 → pyathena-3.29.3}/pyathena/aio/polars/__init__.py +0 -0
- {pyathena-3.29.1 → pyathena-3.29.3}/pyathena/aio/result_set.py +0 -0
- {pyathena-3.29.1 → pyathena-3.29.3}/pyathena/aio/s3fs/__init__.py +0 -0
- {pyathena-3.29.1 → pyathena-3.29.3}/pyathena/aio/spark/__init__.py +0 -0
- {pyathena-3.29.1 → pyathena-3.29.3}/pyathena/aio/sqlalchemy/__init__.py +0 -0
- {pyathena-3.29.1 → pyathena-3.29.3}/pyathena/aio/sqlalchemy/arrow.py +0 -0
- {pyathena-3.29.1 → pyathena-3.29.3}/pyathena/aio/sqlalchemy/base.py +0 -0
- {pyathena-3.29.1 → pyathena-3.29.3}/pyathena/aio/sqlalchemy/pandas.py +0 -0
- {pyathena-3.29.1 → pyathena-3.29.3}/pyathena/aio/sqlalchemy/polars.py +0 -0
- {pyathena-3.29.1 → pyathena-3.29.3}/pyathena/aio/sqlalchemy/rest.py +0 -0
- {pyathena-3.29.1 → pyathena-3.29.3}/pyathena/aio/sqlalchemy/s3fs.py +0 -0
- {pyathena-3.29.1 → pyathena-3.29.3}/pyathena/aio/util.py +0 -0
- {pyathena-3.29.1 → pyathena-3.29.3}/pyathena/arrow/__init__.py +0 -0
- {pyathena-3.29.1 → pyathena-3.29.3}/pyathena/arrow/async_cursor.py +0 -0
- {pyathena-3.29.1 → pyathena-3.29.3}/pyathena/arrow/converter.py +0 -0
- {pyathena-3.29.1 → pyathena-3.29.3}/pyathena/arrow/cursor.py +0 -0
- {pyathena-3.29.1 → pyathena-3.29.3}/pyathena/arrow/result_set.py +0 -0
- {pyathena-3.29.1 → pyathena-3.29.3}/pyathena/arrow/util.py +0 -0
- {pyathena-3.29.1 → pyathena-3.29.3}/pyathena/async_cursor.py +0 -0
- {pyathena-3.29.1 → pyathena-3.29.3}/pyathena/common.py +0 -0
- {pyathena-3.29.1 → pyathena-3.29.3}/pyathena/connection.py +0 -0
- {pyathena-3.29.1 → pyathena-3.29.3}/pyathena/converter.py +0 -0
- {pyathena-3.29.1 → pyathena-3.29.3}/pyathena/cursor.py +0 -0
- {pyathena-3.29.1 → pyathena-3.29.3}/pyathena/error.py +0 -0
- {pyathena-3.29.1 → pyathena-3.29.3}/pyathena/filesystem/__init__.py +0 -0
- {pyathena-3.29.1 → pyathena-3.29.3}/pyathena/filesystem/s3_object.py +0 -0
- {pyathena-3.29.1 → pyathena-3.29.3}/pyathena/formatter.py +0 -0
- {pyathena-3.29.1 → pyathena-3.29.3}/pyathena/model.py +0 -0
- {pyathena-3.29.1 → pyathena-3.29.3}/pyathena/pandas/__init__.py +0 -0
- {pyathena-3.29.1 → pyathena-3.29.3}/pyathena/pandas/async_cursor.py +0 -0
- {pyathena-3.29.1 → pyathena-3.29.3}/pyathena/pandas/converter.py +0 -0
- {pyathena-3.29.1 → pyathena-3.29.3}/pyathena/pandas/cursor.py +0 -0
- {pyathena-3.29.1 → pyathena-3.29.3}/pyathena/pandas/result_set.py +0 -0
- {pyathena-3.29.1 → pyathena-3.29.3}/pyathena/pandas/util.py +0 -0
- {pyathena-3.29.1 → pyathena-3.29.3}/pyathena/polars/__init__.py +0 -0
- {pyathena-3.29.1 → pyathena-3.29.3}/pyathena/polars/async_cursor.py +0 -0
- {pyathena-3.29.1 → pyathena-3.29.3}/pyathena/polars/converter.py +0 -0
- {pyathena-3.29.1 → pyathena-3.29.3}/pyathena/polars/cursor.py +0 -0
- {pyathena-3.29.1 → pyathena-3.29.3}/pyathena/polars/result_set.py +0 -0
- {pyathena-3.29.1 → pyathena-3.29.3}/pyathena/polars/util.py +0 -0
- {pyathena-3.29.1 → pyathena-3.29.3}/pyathena/py.typed +0 -0
- {pyathena-3.29.1 → pyathena-3.29.3}/pyathena/result_set.py +0 -0
- {pyathena-3.29.1 → pyathena-3.29.3}/pyathena/s3fs/__init__.py +0 -0
- {pyathena-3.29.1 → pyathena-3.29.3}/pyathena/s3fs/async_cursor.py +0 -0
- {pyathena-3.29.1 → pyathena-3.29.3}/pyathena/s3fs/converter.py +0 -0
- {pyathena-3.29.1 → pyathena-3.29.3}/pyathena/s3fs/cursor.py +0 -0
- {pyathena-3.29.1 → pyathena-3.29.3}/pyathena/s3fs/reader.py +0 -0
- {pyathena-3.29.1 → pyathena-3.29.3}/pyathena/spark/__init__.py +0 -0
- {pyathena-3.29.1 → pyathena-3.29.3}/pyathena/spark/async_cursor.py +0 -0
- {pyathena-3.29.1 → pyathena-3.29.3}/pyathena/spark/common.py +0 -0
- {pyathena-3.29.1 → pyathena-3.29.3}/pyathena/spark/cursor.py +0 -0
- {pyathena-3.29.1 → pyathena-3.29.3}/pyathena/sqlalchemy/__init__.py +0 -0
- {pyathena-3.29.1 → pyathena-3.29.3}/pyathena/sqlalchemy/arrow.py +0 -0
- {pyathena-3.29.1 → pyathena-3.29.3}/pyathena/sqlalchemy/base.py +0 -0
- {pyathena-3.29.1 → pyathena-3.29.3}/pyathena/sqlalchemy/compiler.py +0 -0
- {pyathena-3.29.1 → pyathena-3.29.3}/pyathena/sqlalchemy/constants.py +0 -0
- {pyathena-3.29.1 → pyathena-3.29.3}/pyathena/sqlalchemy/pandas.py +0 -0
- {pyathena-3.29.1 → pyathena-3.29.3}/pyathena/sqlalchemy/polars.py +0 -0
- {pyathena-3.29.1 → pyathena-3.29.3}/pyathena/sqlalchemy/preparer.py +0 -0
- {pyathena-3.29.1 → pyathena-3.29.3}/pyathena/sqlalchemy/requirements.py +0 -0
- {pyathena-3.29.1 → pyathena-3.29.3}/pyathena/sqlalchemy/rest.py +0 -0
- {pyathena-3.29.1 → pyathena-3.29.3}/pyathena/sqlalchemy/s3fs.py +0 -0
- {pyathena-3.29.1 → pyathena-3.29.3}/pyathena/sqlalchemy/types.py +0 -0
- {pyathena-3.29.1 → pyathena-3.29.3}/pyathena/sqlalchemy/util.py +0 -0
- {pyathena-3.29.1 → pyathena-3.29.3}/pyathena/util.py +0 -0
|
@@ -1,6 +1,6 @@
|
|
|
1
1
|
Metadata-Version: 2.4
|
|
2
2
|
Name: PyAthena
|
|
3
|
-
Version: 3.29.
|
|
3
|
+
Version: 3.29.3
|
|
4
4
|
Summary: Python DB API 2.0 (PEP 249) client for Amazon Athena
|
|
5
5
|
Project-URL: homepage, https://github.com/pyathena-dev/PyAthena/
|
|
6
6
|
Project-URL: repository, https://github.com/pyathena-dev/PyAthena/
|
|
@@ -115,10 +115,10 @@ Native asyncio is also supported:
|
|
|
115
115
|
|
|
116
116
|
```python
|
|
117
117
|
import asyncio
|
|
118
|
-
from pyathena import
|
|
118
|
+
from pyathena import aio_connect
|
|
119
119
|
|
|
120
120
|
async def main():
|
|
121
|
-
async with await
|
|
121
|
+
async with await aio_connect(s3_staging_dir="s3://YOUR_S3_BUCKET/path/to/",
|
|
122
122
|
region_name="us-west-2") as conn:
|
|
123
123
|
cursor = conn.cursor()
|
|
124
124
|
await cursor.execute("SELECT 1")
|
|
@@ -56,10 +56,10 @@ Native asyncio is also supported:
|
|
|
56
56
|
|
|
57
57
|
```python
|
|
58
58
|
import asyncio
|
|
59
|
-
from pyathena import
|
|
59
|
+
from pyathena import aio_connect
|
|
60
60
|
|
|
61
61
|
async def main():
|
|
62
|
-
async with await
|
|
62
|
+
async with await aio_connect(s3_staging_dir="s3://YOUR_S3_BUCKET/path/to/",
|
|
63
63
|
region_name="us-west-2") as conn:
|
|
64
64
|
cursor = conn.cursor()
|
|
65
65
|
await cursor.execute("SELECT 1")
|
|
@@ -131,7 +131,7 @@ def connect(*args, **kwargs) -> "Connection[Any]":
|
|
|
131
131
|
return Connection(*args, **kwargs)
|
|
132
132
|
|
|
133
133
|
|
|
134
|
-
async def
|
|
134
|
+
async def aio_connect(*args, **kwargs) -> "AioConnection":
|
|
135
135
|
"""Create a new async database connection to Amazon Athena.
|
|
136
136
|
|
|
137
137
|
This is the async counterpart of :func:`connect`. It returns an
|
|
@@ -147,7 +147,7 @@ async def aconnect(*args, **kwargs) -> "AioConnection":
|
|
|
147
147
|
|
|
148
148
|
Example:
|
|
149
149
|
>>> import pyathena
|
|
150
|
-
>>> conn = await pyathena.
|
|
150
|
+
>>> conn = await pyathena.aio_connect(
|
|
151
151
|
... s3_staging_dir='s3://my-bucket/staging/',
|
|
152
152
|
... region_name='us-east-1',
|
|
153
153
|
... )
|
|
@@ -28,7 +28,7 @@ version_tuple: VERSION_TUPLE
|
|
|
28
28
|
commit_id: COMMIT_ID
|
|
29
29
|
__commit_id__: COMMIT_ID
|
|
30
30
|
|
|
31
|
-
__version__ = version = '3.29.
|
|
32
|
-
__version_tuple__ = version_tuple = (3, 29,
|
|
31
|
+
__version__ = version = '3.29.3'
|
|
32
|
+
__version_tuple__ = version_tuple = (3, 29, 3)
|
|
33
33
|
|
|
34
34
|
__commit_id__ = commit_id = None
|
|
@@ -29,7 +29,7 @@ class AioArrowCursor(WithAsyncFetch):
|
|
|
29
29
|
operations, keeping the event loop free.
|
|
30
30
|
|
|
31
31
|
Example:
|
|
32
|
-
>>> async with await pyathena.
|
|
32
|
+
>>> async with await pyathena.aio_connect(...) as conn:
|
|
33
33
|
... cursor = conn.cursor(AioArrowCursor)
|
|
34
34
|
... await cursor.execute("SELECT * FROM my_table")
|
|
35
35
|
... table = cursor.as_arrow()
|
|
@@ -40,7 +40,7 @@ class AioPandasCursor(WithAsyncFetch):
|
|
|
40
40
|
when ``chunksize`` is set, as fetch calls trigger lazy S3 reads.
|
|
41
41
|
|
|
42
42
|
Example:
|
|
43
|
-
>>> async with await pyathena.
|
|
43
|
+
>>> async with await pyathena.aio_connect(...) as conn:
|
|
44
44
|
... cursor = conn.cursor(AioPandasCursor)
|
|
45
45
|
... await cursor.execute("SELECT * FROM my_table")
|
|
46
46
|
... df = cursor.as_pandas()
|
|
@@ -31,7 +31,7 @@ class AioPolarsCursor(WithAsyncFetch):
|
|
|
31
31
|
when ``chunksize`` is set, as fetch calls trigger lazy S3 reads.
|
|
32
32
|
|
|
33
33
|
Example:
|
|
34
|
-
>>> async with await pyathena.
|
|
34
|
+
>>> async with await pyathena.aio_connect(...) as conn:
|
|
35
35
|
... cursor = conn.cursor(AioPolarsCursor)
|
|
36
36
|
... await cursor.execute("SELECT * FROM my_table")
|
|
37
37
|
... df = cursor.as_polars()
|
|
@@ -8,6 +8,7 @@ from typing import Any, Dict, List, Optional, Tuple, Union, cast
|
|
|
8
8
|
from pyathena.aio.common import WithAsyncFetch
|
|
9
9
|
from pyathena.common import CursorIterator
|
|
10
10
|
from pyathena.error import OperationalError, ProgrammingError
|
|
11
|
+
from pyathena.filesystem.s3_async import AioS3FileSystem
|
|
11
12
|
from pyathena.model import AthenaQueryExecution
|
|
12
13
|
from pyathena.s3fs.converter import DefaultS3FSTypeConverter
|
|
13
14
|
from pyathena.s3fs.result_set import AthenaS3FSResultSet, CSVReaderType
|
|
@@ -16,14 +17,15 @@ _logger = logging.getLogger(__name__) # type: ignore
|
|
|
16
17
|
|
|
17
18
|
|
|
18
19
|
class AioS3FSCursor(WithAsyncFetch):
|
|
19
|
-
"""Native asyncio cursor that reads CSV results via
|
|
20
|
+
"""Native asyncio cursor that reads CSV results via AioS3FileSystem.
|
|
20
21
|
|
|
21
|
-
Uses ``
|
|
22
|
-
|
|
23
|
-
|
|
22
|
+
Uses ``AioS3FileSystem`` for S3 operations, which replaces
|
|
23
|
+
``ThreadPoolExecutor`` parallelism with ``asyncio.gather`` +
|
|
24
|
+
``asyncio.to_thread``. Fetch operations are wrapped in
|
|
25
|
+
``asyncio.to_thread()`` because CSV reading is blocking I/O.
|
|
24
26
|
|
|
25
27
|
Example:
|
|
26
|
-
>>> async with await pyathena.
|
|
28
|
+
>>> async with await pyathena.aio_connect(...) as conn:
|
|
27
29
|
... cursor = conn.cursor(AioS3FSCursor)
|
|
28
30
|
... await cursor.execute("SELECT * FROM my_table")
|
|
29
31
|
... row = await cursor.fetchone()
|
|
@@ -127,6 +129,7 @@ class AioS3FSCursor(WithAsyncFetch):
|
|
|
127
129
|
arraysize=self.arraysize,
|
|
128
130
|
retry_config=self._retry_config,
|
|
129
131
|
csv_reader=self._csv_reader,
|
|
132
|
+
filesystem_class=AioS3FileSystem,
|
|
130
133
|
**kwargs,
|
|
131
134
|
)
|
|
132
135
|
else:
|
|
@@ -33,7 +33,7 @@ class AioSparkCursor(SparkBaseCursor, WithCalculationExecution):
|
|
|
33
33
|
|
|
34
34
|
Example:
|
|
35
35
|
>>> import asyncio
|
|
36
|
-
>>> async with await pyathena.
|
|
36
|
+
>>> async with await pyathena.aio_connect(
|
|
37
37
|
... work_group="spark-workgroup",
|
|
38
38
|
... cursor_class=AioSparkCursor,
|
|
39
39
|
... ) as conn:
|
|
@@ -1,13 +1,11 @@
|
|
|
1
1
|
# -*- coding: utf-8 -*-
|
|
2
2
|
from __future__ import annotations
|
|
3
3
|
|
|
4
|
-
import itertools
|
|
5
4
|
import logging
|
|
6
5
|
import mimetypes
|
|
7
6
|
import os.path
|
|
8
7
|
import re
|
|
9
8
|
from concurrent.futures import Future, as_completed
|
|
10
|
-
from concurrent.futures.thread import ThreadPoolExecutor
|
|
11
9
|
from copy import deepcopy
|
|
12
10
|
from datetime import datetime
|
|
13
11
|
from multiprocessing import cpu_count
|
|
@@ -23,6 +21,7 @@ from fsspec.spec import AbstractBufferedFile
|
|
|
23
21
|
from fsspec.utils import tokenize
|
|
24
22
|
|
|
25
23
|
import pyathena
|
|
24
|
+
from pyathena.filesystem.s3_executor import S3Executor, S3ThreadPoolExecutor
|
|
26
25
|
from pyathena.filesystem.s3_object import (
|
|
27
26
|
S3CompleteMultipartUpload,
|
|
28
27
|
S3MultipartUpload,
|
|
@@ -686,6 +685,20 @@ class S3FileSystem(AbstractFileSystem):
|
|
|
686
685
|
**request,
|
|
687
686
|
)
|
|
688
687
|
|
|
688
|
+
def _create_executor(self, max_workers: int) -> S3Executor:
|
|
689
|
+
"""Create an executor strategy for parallel operations.
|
|
690
|
+
|
|
691
|
+
Subclasses can override to provide alternative execution strategies
|
|
692
|
+
(e.g., asyncio-based execution).
|
|
693
|
+
|
|
694
|
+
Args:
|
|
695
|
+
max_workers: Maximum number of parallel workers.
|
|
696
|
+
|
|
697
|
+
Returns:
|
|
698
|
+
An S3Executor instance.
|
|
699
|
+
"""
|
|
700
|
+
return S3ThreadPoolExecutor(max_workers=max_workers)
|
|
701
|
+
|
|
689
702
|
def _delete_objects(
|
|
690
703
|
self, bucket: str, paths: List[str], max_workers: Optional[int] = None, **kwargs
|
|
691
704
|
) -> None:
|
|
@@ -703,7 +716,7 @@ class S3FileSystem(AbstractFileSystem):
|
|
|
703
716
|
object_.update({"VersionId": version_id})
|
|
704
717
|
delete_objects.append(object_)
|
|
705
718
|
|
|
706
|
-
with
|
|
719
|
+
with self._create_executor(max_workers=max_workers) as executor:
|
|
707
720
|
fs = []
|
|
708
721
|
for delete in [
|
|
709
722
|
delete_objects[i : i + self.DELETE_OBJECTS_MAX_KEYS]
|
|
@@ -861,7 +874,7 @@ class S3FileSystem(AbstractFileSystem):
|
|
|
861
874
|
**kwargs,
|
|
862
875
|
)
|
|
863
876
|
parts = []
|
|
864
|
-
with
|
|
877
|
+
with self._create_executor(max_workers=max_workers) as executor:
|
|
865
878
|
fs = [
|
|
866
879
|
executor.submit(
|
|
867
880
|
self._upload_part_copy,
|
|
@@ -1106,6 +1119,7 @@ class S3FileSystem(AbstractFileSystem):
|
|
|
1106
1119
|
mode,
|
|
1107
1120
|
version_id=None,
|
|
1108
1121
|
max_workers=max_workers,
|
|
1122
|
+
executor=self._create_executor(max_workers=max_workers),
|
|
1109
1123
|
block_size=block_size,
|
|
1110
1124
|
cache_type=cache_type,
|
|
1111
1125
|
autocommit=autocommit,
|
|
@@ -1256,6 +1270,7 @@ class S3File(AbstractBufferedFile):
|
|
|
1256
1270
|
mode: str = "rb",
|
|
1257
1271
|
version_id: Optional[str] = None,
|
|
1258
1272
|
max_workers: int = (cpu_count() or 1) * 5,
|
|
1273
|
+
executor: Optional[S3Executor] = None,
|
|
1259
1274
|
block_size: int = S3FileSystem.DEFAULT_BLOCK_SIZE,
|
|
1260
1275
|
cache_type: str = "bytes",
|
|
1261
1276
|
autocommit: bool = True,
|
|
@@ -1265,7 +1280,7 @@ class S3File(AbstractBufferedFile):
|
|
|
1265
1280
|
**kwargs,
|
|
1266
1281
|
) -> None:
|
|
1267
1282
|
self.max_workers = max_workers
|
|
1268
|
-
self._executor =
|
|
1283
|
+
self._executor: S3Executor = executor or S3ThreadPoolExecutor(max_workers=max_workers)
|
|
1269
1284
|
self.s3_additional_kwargs = s3_additional_kwargs if s3_additional_kwargs else {}
|
|
1270
1285
|
|
|
1271
1286
|
super().__init__(
|
|
@@ -1481,24 +1496,18 @@ class S3File(AbstractBufferedFile):
|
|
|
1481
1496
|
start, end, max_workers=self.max_workers, worker_block_size=self.blocksize
|
|
1482
1497
|
)
|
|
1483
1498
|
if len(ranges) > 1:
|
|
1484
|
-
|
|
1485
|
-
|
|
1486
|
-
self.
|
|
1487
|
-
|
|
1488
|
-
|
|
1489
|
-
|
|
1490
|
-
|
|
1491
|
-
|
|
1492
|
-
**kwargs,
|
|
1493
|
-
),
|
|
1494
|
-
itertools.repeat(self.bucket),
|
|
1495
|
-
itertools.repeat(self.key),
|
|
1496
|
-
ranges,
|
|
1497
|
-
itertools.repeat(self.version_id),
|
|
1498
|
-
itertools.repeat(self.s3_additional_kwargs),
|
|
1499
|
-
)
|
|
1499
|
+
futures = [
|
|
1500
|
+
self._executor.submit(
|
|
1501
|
+
self.fs._get_object,
|
|
1502
|
+
bucket=self.bucket,
|
|
1503
|
+
key=self.key,
|
|
1504
|
+
ranges=r,
|
|
1505
|
+
version_id=self.version_id,
|
|
1506
|
+
**self.s3_additional_kwargs,
|
|
1500
1507
|
)
|
|
1501
|
-
|
|
1508
|
+
for r in ranges
|
|
1509
|
+
]
|
|
1510
|
+
object_ = self._merge_objects([f.result() for f in as_completed(futures)])
|
|
1502
1511
|
else:
|
|
1503
1512
|
object_ = self.fs._get_object(
|
|
1504
1513
|
self.bucket,
|
|
@@ -0,0 +1,357 @@
|
|
|
1
|
+
# -*- coding: utf-8 -*-
|
|
2
|
+
from __future__ import annotations
|
|
3
|
+
|
|
4
|
+
import asyncio
|
|
5
|
+
import logging
|
|
6
|
+
from multiprocessing import cpu_count
|
|
7
|
+
from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union, cast
|
|
8
|
+
|
|
9
|
+
from fsspec.asyn import AsyncFileSystem
|
|
10
|
+
from fsspec.callbacks import _DEFAULT_CALLBACK
|
|
11
|
+
|
|
12
|
+
from pyathena.filesystem.s3 import S3File, S3FileSystem
|
|
13
|
+
from pyathena.filesystem.s3_executor import S3AioExecutor
|
|
14
|
+
from pyathena.filesystem.s3_object import S3Object
|
|
15
|
+
|
|
16
|
+
if TYPE_CHECKING:
|
|
17
|
+
from datetime import datetime
|
|
18
|
+
|
|
19
|
+
from pyathena.connection import Connection
|
|
20
|
+
|
|
21
|
+
_logger = logging.getLogger(__name__)
|
|
22
|
+
|
|
23
|
+
|
|
24
|
+
class AioS3FileSystem(AsyncFileSystem):
|
|
25
|
+
"""An async filesystem interface for Amazon S3 using fsspec's AsyncFileSystem.
|
|
26
|
+
|
|
27
|
+
This class wraps ``S3FileSystem`` to provide native asyncio support. Instead of
|
|
28
|
+
using ``ThreadPoolExecutor`` for parallel operations, it uses ``asyncio.gather``
|
|
29
|
+
with ``asyncio.to_thread`` for natural integration with the asyncio event loop.
|
|
30
|
+
|
|
31
|
+
The implementation uses composition: an internal ``S3FileSystem`` instance handles
|
|
32
|
+
all boto3 calls, while this class delegates to it via ``asyncio.to_thread()``.
|
|
33
|
+
This avoids diamond inheritance issues and keeps all boto3 logic in one place.
|
|
34
|
+
|
|
35
|
+
File handles created by ``_open`` use ``S3AioExecutor`` so that parallel
|
|
36
|
+
operations (range reads, multipart uploads) are dispatched via the event loop
|
|
37
|
+
instead of spawning additional threads.
|
|
38
|
+
|
|
39
|
+
Attributes:
|
|
40
|
+
_sync_fs: The internal synchronous S3FileSystem instance.
|
|
41
|
+
|
|
42
|
+
Example:
|
|
43
|
+
>>> from pyathena.filesystem.s3_async import AioS3FileSystem
|
|
44
|
+
>>> fs = AioS3FileSystem(asynchronous=True)
|
|
45
|
+
>>>
|
|
46
|
+
>>> # Use in async context
|
|
47
|
+
>>> files = await fs._ls('s3://my-bucket/data/')
|
|
48
|
+
>>>
|
|
49
|
+
>>> # Sync wrappers also available (auto-generated by fsspec)
|
|
50
|
+
>>> files = fs.ls('s3://my-bucket/data/')
|
|
51
|
+
"""
|
|
52
|
+
|
|
53
|
+
# https://docs.aws.amazon.com/AmazonS3/latest/API/API_DeleteObjects.html
|
|
54
|
+
DELETE_OBJECTS_MAX_KEYS: int = 1000
|
|
55
|
+
|
|
56
|
+
protocol = ("s3", "s3a")
|
|
57
|
+
mirror_sync_methods = True
|
|
58
|
+
async_impl = True
|
|
59
|
+
_extra_tokenize_attributes = ("default_block_size",)
|
|
60
|
+
|
|
61
|
+
def __init__(
|
|
62
|
+
self,
|
|
63
|
+
connection: Optional["Connection[Any]"] = None,
|
|
64
|
+
default_block_size: Optional[int] = None,
|
|
65
|
+
default_cache_type: Optional[str] = None,
|
|
66
|
+
max_workers: int = (cpu_count() or 1) * 5,
|
|
67
|
+
s3_additional_kwargs: Optional[Dict[str, Any]] = None,
|
|
68
|
+
asynchronous: bool = False,
|
|
69
|
+
loop: Optional[Any] = None,
|
|
70
|
+
batch_size: Optional[int] = None,
|
|
71
|
+
**kwargs,
|
|
72
|
+
) -> None:
|
|
73
|
+
super().__init__(
|
|
74
|
+
asynchronous=asynchronous,
|
|
75
|
+
loop=loop,
|
|
76
|
+
batch_size=batch_size,
|
|
77
|
+
**kwargs,
|
|
78
|
+
)
|
|
79
|
+
self._sync_fs = S3FileSystem(
|
|
80
|
+
connection=connection,
|
|
81
|
+
default_block_size=default_block_size,
|
|
82
|
+
default_cache_type=default_cache_type,
|
|
83
|
+
max_workers=max_workers,
|
|
84
|
+
s3_additional_kwargs=s3_additional_kwargs,
|
|
85
|
+
**kwargs,
|
|
86
|
+
)
|
|
87
|
+
# Share dircache for cache coherence between async and sync instances
|
|
88
|
+
self.dircache = self._sync_fs.dircache
|
|
89
|
+
|
|
90
|
+
@staticmethod
|
|
91
|
+
def parse_path(path: str) -> Tuple[str, Optional[str], Optional[str]]:
|
|
92
|
+
return S3FileSystem.parse_path(path)
|
|
93
|
+
|
|
94
|
+
async def _info(self, path: str, **kwargs) -> S3Object:
|
|
95
|
+
return await asyncio.to_thread(self._sync_fs.info, path, **kwargs)
|
|
96
|
+
|
|
97
|
+
async def _ls(
|
|
98
|
+
self, path: str, detail: bool = False, **kwargs
|
|
99
|
+
) -> Union[List[S3Object], List[str]]:
|
|
100
|
+
return await asyncio.to_thread(self._sync_fs.ls, path, detail=detail, **kwargs)
|
|
101
|
+
|
|
102
|
+
async def _cat_file(
|
|
103
|
+
self, path: str, start: Optional[int] = None, end: Optional[int] = None, **kwargs
|
|
104
|
+
) -> bytes:
|
|
105
|
+
return await asyncio.to_thread(self._sync_fs.cat_file, path, start=start, end=end, **kwargs)
|
|
106
|
+
|
|
107
|
+
async def _exists(self, path: str, **kwargs) -> bool:
|
|
108
|
+
return await asyncio.to_thread(self._sync_fs.exists, path, **kwargs)
|
|
109
|
+
|
|
110
|
+
async def _rm_file(self, path: str, **kwargs) -> None:
|
|
111
|
+
await asyncio.to_thread(self._sync_fs.rm_file, path, **kwargs)
|
|
112
|
+
|
|
113
|
+
async def _pipe_file(self, path: str, value: bytes, **kwargs) -> None:
|
|
114
|
+
await asyncio.to_thread(self._sync_fs.pipe_file, path, value, **kwargs)
|
|
115
|
+
|
|
116
|
+
async def _put_file(self, lpath: str, rpath: str, callback=_DEFAULT_CALLBACK, **kwargs) -> None:
|
|
117
|
+
await asyncio.to_thread(self._sync_fs.put_file, lpath, rpath, callback=callback, **kwargs)
|
|
118
|
+
|
|
119
|
+
async def _get_file(self, rpath: str, lpath: str, callback=_DEFAULT_CALLBACK, **kwargs) -> None:
|
|
120
|
+
await asyncio.to_thread(self._sync_fs.get_file, rpath, lpath, callback=callback, **kwargs)
|
|
121
|
+
|
|
122
|
+
async def _mkdir(self, path: str, create_parents: bool = True, **kwargs) -> None:
|
|
123
|
+
await asyncio.to_thread(self._sync_fs.mkdir, path, create_parents=create_parents, **kwargs)
|
|
124
|
+
|
|
125
|
+
async def _makedirs(self, path: str, exist_ok: bool = False) -> None:
|
|
126
|
+
await asyncio.to_thread(self._sync_fs.makedirs, path, exist_ok=exist_ok)
|
|
127
|
+
|
|
128
|
+
async def _rm(self, path: Union[str, List[str]], recursive: bool = False, **kwargs) -> None:
|
|
129
|
+
"""Remove files or directories using async parallel batch deletion.
|
|
130
|
+
|
|
131
|
+
For multiple paths, chunks into batches of 1000 (S3 API limit) and uses
|
|
132
|
+
``asyncio.gather`` with ``asyncio.to_thread`` instead of ThreadPoolExecutor.
|
|
133
|
+
"""
|
|
134
|
+
if isinstance(path, str):
|
|
135
|
+
path = [path]
|
|
136
|
+
|
|
137
|
+
bucket, _, _ = self.parse_path(path[0])
|
|
138
|
+
|
|
139
|
+
expand_paths: List[str] = []
|
|
140
|
+
for p in path:
|
|
141
|
+
expanded = await asyncio.to_thread(self._sync_fs.expand_path, p, recursive=recursive)
|
|
142
|
+
expand_paths.extend(expanded)
|
|
143
|
+
|
|
144
|
+
if not expand_paths:
|
|
145
|
+
return
|
|
146
|
+
|
|
147
|
+
quiet = kwargs.pop("Quiet", True)
|
|
148
|
+
delete_objects: List[Dict[str, Any]] = []
|
|
149
|
+
for p in expand_paths:
|
|
150
|
+
_, key, version_id = self.parse_path(p)
|
|
151
|
+
if key:
|
|
152
|
+
object_: Dict[str, Any] = {"Key": key}
|
|
153
|
+
if version_id:
|
|
154
|
+
object_["VersionId"] = version_id
|
|
155
|
+
delete_objects.append(object_)
|
|
156
|
+
|
|
157
|
+
if not delete_objects:
|
|
158
|
+
return
|
|
159
|
+
|
|
160
|
+
chunks = [
|
|
161
|
+
delete_objects[i : i + self.DELETE_OBJECTS_MAX_KEYS]
|
|
162
|
+
for i in range(0, len(delete_objects), self.DELETE_OBJECTS_MAX_KEYS)
|
|
163
|
+
]
|
|
164
|
+
|
|
165
|
+
async def _delete_chunk(chunk: List[Dict[str, Any]]) -> None:
|
|
166
|
+
request = {
|
|
167
|
+
"Bucket": bucket,
|
|
168
|
+
"Delete": {
|
|
169
|
+
"Objects": chunk,
|
|
170
|
+
"Quiet": quiet,
|
|
171
|
+
},
|
|
172
|
+
}
|
|
173
|
+
await asyncio.to_thread(
|
|
174
|
+
self._sync_fs._call, self._sync_fs._client.delete_objects, **request
|
|
175
|
+
)
|
|
176
|
+
|
|
177
|
+
await asyncio.gather(*[_delete_chunk(chunk) for chunk in chunks])
|
|
178
|
+
|
|
179
|
+
for p in expand_paths:
|
|
180
|
+
self._sync_fs.invalidate_cache(p)
|
|
181
|
+
|
|
182
|
+
async def _cp_file(self, path1: str, path2: str, **kwargs) -> None:
|
|
183
|
+
"""Copy an S3 object, using async parallel multipart upload for large files."""
|
|
184
|
+
kwargs.pop("onerror", None)
|
|
185
|
+
bucket1, key1, version_id1 = self.parse_path(path1)
|
|
186
|
+
bucket2, key2, version_id2 = self.parse_path(path2)
|
|
187
|
+
if version_id2:
|
|
188
|
+
raise ValueError("Cannot copy to a versioned file.")
|
|
189
|
+
if not key1 or not key2:
|
|
190
|
+
raise ValueError("Cannot copy buckets.")
|
|
191
|
+
|
|
192
|
+
info1 = await self._info(path1)
|
|
193
|
+
size1 = info1.get("size", 0)
|
|
194
|
+
if size1 <= S3FileSystem.MULTIPART_UPLOAD_MAX_PART_SIZE:
|
|
195
|
+
await asyncio.to_thread(
|
|
196
|
+
self._sync_fs._copy_object,
|
|
197
|
+
bucket1=bucket1,
|
|
198
|
+
key1=key1,
|
|
199
|
+
version_id1=version_id1,
|
|
200
|
+
bucket2=bucket2,
|
|
201
|
+
key2=key2,
|
|
202
|
+
**kwargs,
|
|
203
|
+
)
|
|
204
|
+
else:
|
|
205
|
+
await self._copy_object_with_multipart_upload(
|
|
206
|
+
bucket1=bucket1,
|
|
207
|
+
key1=key1,
|
|
208
|
+
version_id1=version_id1,
|
|
209
|
+
size1=size1,
|
|
210
|
+
bucket2=bucket2,
|
|
211
|
+
key2=key2,
|
|
212
|
+
**kwargs,
|
|
213
|
+
)
|
|
214
|
+
self._sync_fs.invalidate_cache(path2)
|
|
215
|
+
|
|
216
|
+
async def _copy_object_with_multipart_upload(
|
|
217
|
+
self,
|
|
218
|
+
bucket1: str,
|
|
219
|
+
key1: str,
|
|
220
|
+
size1: int,
|
|
221
|
+
bucket2: str,
|
|
222
|
+
key2: str,
|
|
223
|
+
block_size: Optional[int] = None,
|
|
224
|
+
version_id1: Optional[str] = None,
|
|
225
|
+
**kwargs,
|
|
226
|
+
) -> None:
|
|
227
|
+
block_size = block_size if block_size else S3FileSystem.MULTIPART_UPLOAD_MAX_PART_SIZE
|
|
228
|
+
if (
|
|
229
|
+
block_size < S3FileSystem.MULTIPART_UPLOAD_MIN_PART_SIZE
|
|
230
|
+
or block_size > S3FileSystem.MULTIPART_UPLOAD_MAX_PART_SIZE
|
|
231
|
+
):
|
|
232
|
+
raise ValueError("Block size must be greater than 5MiB and less than 5GiB.")
|
|
233
|
+
|
|
234
|
+
copy_source: Dict[str, Any] = {
|
|
235
|
+
"Bucket": bucket1,
|
|
236
|
+
"Key": key1,
|
|
237
|
+
}
|
|
238
|
+
if version_id1:
|
|
239
|
+
copy_source["VersionId"] = version_id1
|
|
240
|
+
|
|
241
|
+
ranges = S3File._get_ranges(
|
|
242
|
+
0,
|
|
243
|
+
size1,
|
|
244
|
+
self._sync_fs.max_workers,
|
|
245
|
+
block_size,
|
|
246
|
+
)
|
|
247
|
+
multipart_upload = await asyncio.to_thread(
|
|
248
|
+
self._sync_fs._create_multipart_upload,
|
|
249
|
+
bucket=bucket2,
|
|
250
|
+
key=key2,
|
|
251
|
+
**kwargs,
|
|
252
|
+
)
|
|
253
|
+
|
|
254
|
+
async def _upload_part(i: int, range_: Tuple[int, int]) -> Dict[str, Any]:
|
|
255
|
+
result = await asyncio.to_thread(
|
|
256
|
+
self._sync_fs._upload_part_copy,
|
|
257
|
+
bucket=bucket2,
|
|
258
|
+
key=key2,
|
|
259
|
+
copy_source=copy_source,
|
|
260
|
+
upload_id=cast(str, multipart_upload.upload_id),
|
|
261
|
+
part_number=i + 1,
|
|
262
|
+
copy_source_ranges=range_,
|
|
263
|
+
)
|
|
264
|
+
return {
|
|
265
|
+
"ETag": result.etag,
|
|
266
|
+
"PartNumber": result.part_number,
|
|
267
|
+
}
|
|
268
|
+
|
|
269
|
+
parts = await asyncio.gather(*[_upload_part(i, r) for i, r in enumerate(ranges)])
|
|
270
|
+
parts_list = sorted(parts, key=lambda x: x["PartNumber"])
|
|
271
|
+
|
|
272
|
+
await asyncio.to_thread(
|
|
273
|
+
self._sync_fs._complete_multipart_upload,
|
|
274
|
+
bucket=bucket2,
|
|
275
|
+
key=key2,
|
|
276
|
+
upload_id=cast(str, multipart_upload.upload_id),
|
|
277
|
+
parts=parts_list,
|
|
278
|
+
)
|
|
279
|
+
|
|
280
|
+
async def _find(
|
|
281
|
+
self,
|
|
282
|
+
path: str,
|
|
283
|
+
maxdepth: Optional[int] = None,
|
|
284
|
+
withdirs: bool = False,
|
|
285
|
+
**kwargs,
|
|
286
|
+
) -> Union[Dict[str, S3Object], List[str]]:
|
|
287
|
+
detail = kwargs.pop("detail", False)
|
|
288
|
+
files = await asyncio.to_thread(
|
|
289
|
+
self._sync_fs._find, path, maxdepth=maxdepth, withdirs=withdirs, **kwargs
|
|
290
|
+
)
|
|
291
|
+
if detail:
|
|
292
|
+
return {f.name: f for f in files}
|
|
293
|
+
return [f.name for f in files]
|
|
294
|
+
|
|
295
|
+
def _open(
|
|
296
|
+
self,
|
|
297
|
+
path: str,
|
|
298
|
+
mode: str = "rb",
|
|
299
|
+
block_size: Optional[int] = None,
|
|
300
|
+
cache_type: Optional[str] = None,
|
|
301
|
+
autocommit: bool = True,
|
|
302
|
+
cache_options: Optional[Dict[Any, Any]] = None,
|
|
303
|
+
**kwargs,
|
|
304
|
+
) -> "AioS3File":
|
|
305
|
+
if block_size is None:
|
|
306
|
+
block_size = self._sync_fs.default_block_size
|
|
307
|
+
if cache_type is None:
|
|
308
|
+
cache_type = self._sync_fs.default_cache_type
|
|
309
|
+
max_workers = kwargs.pop("max_worker", self._sync_fs.max_workers)
|
|
310
|
+
s3_additional_kwargs = kwargs.pop("s3_additional_kwargs", {})
|
|
311
|
+
s3_additional_kwargs.update(self._sync_fs.s3_additional_kwargs)
|
|
312
|
+
|
|
313
|
+
return AioS3File(
|
|
314
|
+
self._sync_fs,
|
|
315
|
+
path,
|
|
316
|
+
mode,
|
|
317
|
+
version_id=None,
|
|
318
|
+
max_workers=max_workers,
|
|
319
|
+
executor=S3AioExecutor(loop=self._loop),
|
|
320
|
+
block_size=block_size,
|
|
321
|
+
cache_type=cache_type,
|
|
322
|
+
autocommit=autocommit,
|
|
323
|
+
cache_options=cache_options,
|
|
324
|
+
s3_additional_kwargs=s3_additional_kwargs,
|
|
325
|
+
**kwargs,
|
|
326
|
+
)
|
|
327
|
+
|
|
328
|
+
def sign(self, path: str, expiration: int = 3600, **kwargs) -> str:
|
|
329
|
+
return cast(str, self._sync_fs.sign(path, expiration=expiration, **kwargs))
|
|
330
|
+
|
|
331
|
+
def checksum(self, path: str, **kwargs) -> int:
|
|
332
|
+
return cast(int, self._sync_fs.checksum(path, **kwargs))
|
|
333
|
+
|
|
334
|
+
def created(self, path: str) -> "datetime":
|
|
335
|
+
return self._sync_fs.created(path)
|
|
336
|
+
|
|
337
|
+
def modified(self, path: str) -> "datetime":
|
|
338
|
+
return self._sync_fs.modified(path)
|
|
339
|
+
|
|
340
|
+
def invalidate_cache(self, path: Optional[str] = None) -> None:
|
|
341
|
+
self._sync_fs.invalidate_cache(path)
|
|
342
|
+
|
|
343
|
+
async def _touch(self, path: str, truncate: bool = True, **kwargs) -> None:
|
|
344
|
+
await asyncio.to_thread(self._sync_fs.touch, path, truncate=truncate, **kwargs)
|
|
345
|
+
|
|
346
|
+
|
|
347
|
+
class AioS3File(S3File):
|
|
348
|
+
"""Async-aware S3 file handle using ``S3AioExecutor``.
|
|
349
|
+
|
|
350
|
+
Functionally identical to ``S3File``; exists as a distinct type for
|
|
351
|
+
``isinstance`` checks and to document the async execution model.
|
|
352
|
+
All parallel operations (range reads, multipart uploads) are dispatched
|
|
353
|
+
through the ``S3Executor`` interface — the ``S3AioExecutor``
|
|
354
|
+
provided by ``AioS3FileSystem`` uses the event loop instead of threads.
|
|
355
|
+
"""
|
|
356
|
+
|
|
357
|
+
pass
|
|
@@ -0,0 +1,89 @@
|
|
|
1
|
+
# -*- coding: utf-8 -*-
|
|
2
|
+
from __future__ import annotations
|
|
3
|
+
|
|
4
|
+
import asyncio
|
|
5
|
+
from abc import ABCMeta, abstractmethod
|
|
6
|
+
from concurrent.futures import Future
|
|
7
|
+
from concurrent.futures.thread import ThreadPoolExecutor
|
|
8
|
+
from typing import Any, Callable, Optional, TypeVar
|
|
9
|
+
|
|
10
|
+
T = TypeVar("T")
|
|
11
|
+
|
|
12
|
+
|
|
13
|
+
class S3Executor(metaclass=ABCMeta):
|
|
14
|
+
"""Abstract executor for parallel S3 operations.
|
|
15
|
+
|
|
16
|
+
Defines the interface used by ``S3File`` and ``S3FileSystem`` for submitting
|
|
17
|
+
work to run in parallel and for shutting down the executor when done.
|
|
18
|
+
Both ``submit`` and ``shutdown`` mirror the ``concurrent.futures.Executor``
|
|
19
|
+
interface so that ``as_completed()`` and ``Future.cancel()`` work unchanged.
|
|
20
|
+
"""
|
|
21
|
+
|
|
22
|
+
@abstractmethod
|
|
23
|
+
def submit(self, fn: Callable[..., T], *args: Any, **kwargs: Any) -> Future[T]:
|
|
24
|
+
"""Submit a callable for execution and return a Future."""
|
|
25
|
+
...
|
|
26
|
+
|
|
27
|
+
@abstractmethod
|
|
28
|
+
def shutdown(self, wait: bool = True) -> None:
|
|
29
|
+
"""Shut down the executor, freeing any resources."""
|
|
30
|
+
...
|
|
31
|
+
|
|
32
|
+
def __enter__(self) -> "S3Executor":
|
|
33
|
+
return self
|
|
34
|
+
|
|
35
|
+
def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
|
|
36
|
+
self.shutdown(wait=True)
|
|
37
|
+
|
|
38
|
+
|
|
39
|
+
class S3ThreadPoolExecutor(S3Executor):
|
|
40
|
+
"""Executor that delegates to a ``ThreadPoolExecutor``.
|
|
41
|
+
|
|
42
|
+
This is the default executor used by ``S3File`` and ``S3FileSystem``
|
|
43
|
+
for synchronous parallel operations.
|
|
44
|
+
"""
|
|
45
|
+
|
|
46
|
+
def __init__(self, max_workers: int) -> None:
|
|
47
|
+
self._executor = ThreadPoolExecutor(max_workers=max_workers)
|
|
48
|
+
|
|
49
|
+
def submit(self, fn: Callable[..., T], *args: Any, **kwargs: Any) -> Future[T]:
|
|
50
|
+
return self._executor.submit(fn, *args, **kwargs)
|
|
51
|
+
|
|
52
|
+
def shutdown(self, wait: bool = True) -> None:
|
|
53
|
+
self._executor.shutdown(wait=wait)
|
|
54
|
+
|
|
55
|
+
|
|
56
|
+
class S3AioExecutor(S3Executor):
|
|
57
|
+
"""Executor that schedules work on an asyncio event loop.
|
|
58
|
+
|
|
59
|
+
Uses ``asyncio.run_coroutine_threadsafe(asyncio.to_thread(fn), loop)`` to
|
|
60
|
+
dispatch blocking functions onto the event loop's thread pool, returning
|
|
61
|
+
``concurrent.futures.Future`` objects that are compatible with
|
|
62
|
+
``as_completed()`` and ``Future.cancel()``.
|
|
63
|
+
|
|
64
|
+
This avoids thread-in-thread nesting when ``S3File`` is used from within
|
|
65
|
+
``asyncio.to_thread()`` calls (the pattern used by ``AioS3FileSystem``).
|
|
66
|
+
|
|
67
|
+
Args:
|
|
68
|
+
loop: A running asyncio event loop.
|
|
69
|
+
|
|
70
|
+
Raises:
|
|
71
|
+
RuntimeError: If the event loop is not running when ``submit`` is called.
|
|
72
|
+
"""
|
|
73
|
+
|
|
74
|
+
def __init__(self, loop: Optional[asyncio.AbstractEventLoop] = None) -> None:
|
|
75
|
+
self._loop = loop
|
|
76
|
+
|
|
77
|
+
def submit(self, fn: Callable[..., T], *args: Any, **kwargs: Any) -> Future[T]:
|
|
78
|
+
if self._loop is not None and self._loop.is_running():
|
|
79
|
+
return asyncio.run_coroutine_threadsafe(
|
|
80
|
+
asyncio.to_thread(fn, *args, **kwargs), self._loop
|
|
81
|
+
)
|
|
82
|
+
raise RuntimeError(
|
|
83
|
+
"S3AioExecutor requires a running event loop. "
|
|
84
|
+
"Use S3ThreadPoolExecutor for synchronous usage."
|
|
85
|
+
)
|
|
86
|
+
|
|
87
|
+
def shutdown(self, wait: bool = True) -> None:
|
|
88
|
+
# No resources to release — work is dispatched to the event loop.
|
|
89
|
+
pass
|
|
@@ -5,6 +5,8 @@ import logging
|
|
|
5
5
|
from io import TextIOWrapper
|
|
6
6
|
from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Type, Union
|
|
7
7
|
|
|
8
|
+
from fsspec import AbstractFileSystem
|
|
9
|
+
|
|
8
10
|
from pyathena.converter import Converter
|
|
9
11
|
from pyathena.error import OperationalError, ProgrammingError
|
|
10
12
|
from pyathena.filesystem.s3 import S3FileSystem
|
|
@@ -62,6 +64,7 @@ class AthenaS3FSResultSet(AthenaResultSet):
|
|
|
62
64
|
retry_config: RetryConfig,
|
|
63
65
|
block_size: Optional[int] = None,
|
|
64
66
|
csv_reader: Optional[CSVReaderType] = None,
|
|
67
|
+
filesystem_class: Optional[Type[AbstractFileSystem]] = None,
|
|
65
68
|
**kwargs,
|
|
66
69
|
) -> None:
|
|
67
70
|
super().__init__(
|
|
@@ -77,6 +80,7 @@ class AthenaS3FSResultSet(AthenaResultSet):
|
|
|
77
80
|
self._arraysize = arraysize
|
|
78
81
|
self._block_size = block_size if block_size else self.DEFAULT_BLOCK_SIZE
|
|
79
82
|
self._csv_reader_class: CSVReaderType = csv_reader or AthenaCSVReader
|
|
83
|
+
self._filesystem_class: Type[AbstractFileSystem] = filesystem_class or S3FileSystem
|
|
80
84
|
self._fs = self._create_s3_file_system()
|
|
81
85
|
self._csv_reader: Optional[Any] = None
|
|
82
86
|
|
|
@@ -92,9 +96,9 @@ class AthenaS3FSResultSet(AthenaResultSet):
|
|
|
92
96
|
if not self._csv_reader and not self._rows and pre_fetched_rows:
|
|
93
97
|
self._rows.extend(pre_fetched_rows)
|
|
94
98
|
|
|
95
|
-
def _create_s3_file_system(self) ->
|
|
99
|
+
def _create_s3_file_system(self) -> AbstractFileSystem:
|
|
96
100
|
"""Create S3FileSystem using connection settings."""
|
|
97
|
-
return
|
|
101
|
+
return self._filesystem_class(
|
|
98
102
|
connection=self.connection,
|
|
99
103
|
default_block_size=self._block_size,
|
|
100
104
|
)
|
|
@@ -178,7 +178,7 @@ exclude = [
|
|
|
178
178
|
legacy_tox_ini = """
|
|
179
179
|
[tox]
|
|
180
180
|
isolated_build = true
|
|
181
|
-
envlist = py{310,311,312,313,314}
|
|
181
|
+
envlist = py{310,311,312,313,314}-{pyathena,sqla,sqla_async}
|
|
182
182
|
|
|
183
183
|
[gh-actions]
|
|
184
184
|
python =
|
|
@@ -188,6 +188,12 @@ python =
|
|
|
188
188
|
3.13: py313
|
|
189
189
|
3.14: py314
|
|
190
190
|
|
|
191
|
+
[gh-actions:env]
|
|
192
|
+
TEST_TYPE =
|
|
193
|
+
pyathena: pyathena
|
|
194
|
+
sqla: sqla
|
|
195
|
+
sqla_async: sqla_async
|
|
196
|
+
|
|
191
197
|
[testenv]
|
|
192
198
|
allowlist_externals =
|
|
193
199
|
uv
|
|
@@ -195,9 +201,9 @@ allowlist_externals =
|
|
|
195
201
|
make
|
|
196
202
|
commands =
|
|
197
203
|
uv sync --group dev
|
|
198
|
-
make test
|
|
199
|
-
make test-sqla
|
|
200
|
-
make test-sqla-async
|
|
204
|
+
pyathena: make test
|
|
205
|
+
sqla: make test-sqla
|
|
206
|
+
sqla_async: make test-sqla-async
|
|
201
207
|
passenv =
|
|
202
208
|
TOXENV
|
|
203
209
|
AWS_*
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|