PyAthena 3.34.0__tar.gz → 3.35.1__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.34.0 → pyathena-3.35.1}/PKG-INFO +1 -1
- {pyathena-3.34.0 → pyathena-3.35.1}/pyathena/__init__.py +1 -0
- {pyathena-3.34.0 → pyathena-3.35.1}/pyathena/_version.py +2 -2
- {pyathena-3.34.0 → pyathena-3.35.1}/pyathena/aio/arrow/cursor.py +29 -6
- {pyathena-3.34.0 → pyathena-3.35.1}/pyathena/aio/common.py +24 -9
- {pyathena-3.34.0 → pyathena-3.35.1}/pyathena/aio/cursor.py +24 -6
- {pyathena-3.34.0 → pyathena-3.35.1}/pyathena/aio/pandas/cursor.py +29 -7
- {pyathena-3.34.0 → pyathena-3.35.1}/pyathena/aio/polars/cursor.py +29 -6
- {pyathena-3.34.0 → pyathena-3.35.1}/pyathena/aio/s3fs/cursor.py +28 -5
- {pyathena-3.34.0 → pyathena-3.35.1}/pyathena/arrow/async_cursor.py +38 -7
- {pyathena-3.34.0 → pyathena-3.35.1}/pyathena/arrow/cursor.py +20 -12
- {pyathena-3.34.0 → pyathena-3.35.1}/pyathena/async_cursor.py +17 -6
- {pyathena-3.34.0 → pyathena-3.35.1}/pyathena/common.py +39 -12
- {pyathena-3.34.0 → pyathena-3.35.1}/pyathena/cursor.py +19 -11
- pyathena-3.35.1/pyathena/options.py +102 -0
- {pyathena-3.34.0 → pyathena-3.35.1}/pyathena/pandas/async_cursor.py +41 -7
- {pyathena-3.34.0 → pyathena-3.35.1}/pyathena/pandas/cursor.py +20 -12
- {pyathena-3.34.0 → pyathena-3.35.1}/pyathena/polars/async_cursor.py +18 -7
- {pyathena-3.34.0 → pyathena-3.35.1}/pyathena/polars/cursor.py +20 -12
- {pyathena-3.34.0 → pyathena-3.35.1}/pyathena/s3fs/async_cursor.py +17 -6
- {pyathena-3.34.0 → pyathena-3.35.1}/pyathena/s3fs/cursor.py +19 -10
- {pyathena-3.34.0 → pyathena-3.35.1}/.gitignore +0 -0
- {pyathena-3.34.0 → pyathena-3.35.1}/LICENSE +0 -0
- {pyathena-3.34.0 → pyathena-3.35.1}/README.md +0 -0
- {pyathena-3.34.0 → pyathena-3.35.1}/pyathena/aio/__init__.py +0 -0
- {pyathena-3.34.0 → pyathena-3.35.1}/pyathena/aio/arrow/__init__.py +0 -0
- {pyathena-3.34.0 → pyathena-3.35.1}/pyathena/aio/connection.py +0 -0
- {pyathena-3.34.0 → pyathena-3.35.1}/pyathena/aio/pandas/__init__.py +0 -0
- {pyathena-3.34.0 → pyathena-3.35.1}/pyathena/aio/polars/__init__.py +0 -0
- {pyathena-3.34.0 → pyathena-3.35.1}/pyathena/aio/result_set.py +0 -0
- {pyathena-3.34.0 → pyathena-3.35.1}/pyathena/aio/s3fs/__init__.py +0 -0
- {pyathena-3.34.0 → pyathena-3.35.1}/pyathena/aio/spark/__init__.py +0 -0
- {pyathena-3.34.0 → pyathena-3.35.1}/pyathena/aio/spark/cursor.py +0 -0
- {pyathena-3.34.0 → pyathena-3.35.1}/pyathena/aio/sqlalchemy/__init__.py +0 -0
- {pyathena-3.34.0 → pyathena-3.35.1}/pyathena/aio/sqlalchemy/arrow.py +0 -0
- {pyathena-3.34.0 → pyathena-3.35.1}/pyathena/aio/sqlalchemy/base.py +0 -0
- {pyathena-3.34.0 → pyathena-3.35.1}/pyathena/aio/sqlalchemy/pandas.py +0 -0
- {pyathena-3.34.0 → pyathena-3.35.1}/pyathena/aio/sqlalchemy/polars.py +0 -0
- {pyathena-3.34.0 → pyathena-3.35.1}/pyathena/aio/sqlalchemy/rest.py +0 -0
- {pyathena-3.34.0 → pyathena-3.35.1}/pyathena/aio/sqlalchemy/s3fs.py +0 -0
- {pyathena-3.34.0 → pyathena-3.35.1}/pyathena/aio/util.py +0 -0
- {pyathena-3.34.0 → pyathena-3.35.1}/pyathena/arrow/__init__.py +0 -0
- {pyathena-3.34.0 → pyathena-3.35.1}/pyathena/arrow/converter.py +0 -0
- {pyathena-3.34.0 → pyathena-3.35.1}/pyathena/arrow/result_set.py +0 -0
- {pyathena-3.34.0 → pyathena-3.35.1}/pyathena/arrow/util.py +0 -0
- {pyathena-3.34.0 → pyathena-3.35.1}/pyathena/connection.py +0 -0
- {pyathena-3.34.0 → pyathena-3.35.1}/pyathena/converter.py +0 -0
- {pyathena-3.34.0 → pyathena-3.35.1}/pyathena/error.py +0 -0
- {pyathena-3.34.0 → pyathena-3.35.1}/pyathena/filesystem/__init__.py +0 -0
- {pyathena-3.34.0 → pyathena-3.35.1}/pyathena/filesystem/s3.py +0 -0
- {pyathena-3.34.0 → pyathena-3.35.1}/pyathena/filesystem/s3_async.py +0 -0
- {pyathena-3.34.0 → pyathena-3.35.1}/pyathena/filesystem/s3_errors.py +0 -0
- {pyathena-3.34.0 → pyathena-3.35.1}/pyathena/filesystem/s3_executor.py +0 -0
- {pyathena-3.34.0 → pyathena-3.35.1}/pyathena/filesystem/s3_object.py +0 -0
- {pyathena-3.34.0 → pyathena-3.35.1}/pyathena/formatter.py +0 -0
- {pyathena-3.34.0 → pyathena-3.35.1}/pyathena/model.py +0 -0
- {pyathena-3.34.0 → pyathena-3.35.1}/pyathena/pandas/__init__.py +0 -0
- {pyathena-3.34.0 → pyathena-3.35.1}/pyathena/pandas/converter.py +0 -0
- {pyathena-3.34.0 → pyathena-3.35.1}/pyathena/pandas/result_set.py +0 -0
- {pyathena-3.34.0 → pyathena-3.35.1}/pyathena/pandas/util.py +0 -0
- {pyathena-3.34.0 → pyathena-3.35.1}/pyathena/parser.py +0 -0
- {pyathena-3.34.0 → pyathena-3.35.1}/pyathena/polars/__init__.py +0 -0
- {pyathena-3.34.0 → pyathena-3.35.1}/pyathena/polars/converter.py +0 -0
- {pyathena-3.34.0 → pyathena-3.35.1}/pyathena/polars/result_set.py +0 -0
- {pyathena-3.34.0 → pyathena-3.35.1}/pyathena/polars/util.py +0 -0
- {pyathena-3.34.0 → pyathena-3.35.1}/pyathena/py.typed +0 -0
- {pyathena-3.34.0 → pyathena-3.35.1}/pyathena/result_set.py +0 -0
- {pyathena-3.34.0 → pyathena-3.35.1}/pyathena/s3fs/__init__.py +0 -0
- {pyathena-3.34.0 → pyathena-3.35.1}/pyathena/s3fs/converter.py +0 -0
- {pyathena-3.34.0 → pyathena-3.35.1}/pyathena/s3fs/reader.py +0 -0
- {pyathena-3.34.0 → pyathena-3.35.1}/pyathena/s3fs/result_set.py +0 -0
- {pyathena-3.34.0 → pyathena-3.35.1}/pyathena/spark/__init__.py +0 -0
- {pyathena-3.34.0 → pyathena-3.35.1}/pyathena/spark/async_cursor.py +0 -0
- {pyathena-3.34.0 → pyathena-3.35.1}/pyathena/spark/common.py +0 -0
- {pyathena-3.34.0 → pyathena-3.35.1}/pyathena/spark/cursor.py +0 -0
- {pyathena-3.34.0 → pyathena-3.35.1}/pyathena/sqlalchemy/__init__.py +0 -0
- {pyathena-3.34.0 → pyathena-3.35.1}/pyathena/sqlalchemy/arrow.py +0 -0
- {pyathena-3.34.0 → pyathena-3.35.1}/pyathena/sqlalchemy/base.py +0 -0
- {pyathena-3.34.0 → pyathena-3.35.1}/pyathena/sqlalchemy/compiler.py +0 -0
- {pyathena-3.34.0 → pyathena-3.35.1}/pyathena/sqlalchemy/constants.py +0 -0
- {pyathena-3.34.0 → pyathena-3.35.1}/pyathena/sqlalchemy/pandas.py +0 -0
- {pyathena-3.34.0 → pyathena-3.35.1}/pyathena/sqlalchemy/polars.py +0 -0
- {pyathena-3.34.0 → pyathena-3.35.1}/pyathena/sqlalchemy/preparer.py +0 -0
- {pyathena-3.34.0 → pyathena-3.35.1}/pyathena/sqlalchemy/requirements.py +0 -0
- {pyathena-3.34.0 → pyathena-3.35.1}/pyathena/sqlalchemy/rest.py +0 -0
- {pyathena-3.34.0 → pyathena-3.35.1}/pyathena/sqlalchemy/s3fs.py +0 -0
- {pyathena-3.34.0 → pyathena-3.35.1}/pyathena/sqlalchemy/types.py +0 -0
- {pyathena-3.34.0 → pyathena-3.35.1}/pyathena/sqlalchemy/util.py +0 -0
- {pyathena-3.34.0 → pyathena-3.35.1}/pyathena/util.py +0 -0
- {pyathena-3.34.0 → pyathena-3.35.1}/pyproject.toml +0 -0
|
@@ -1,6 +1,6 @@
|
|
|
1
1
|
Metadata-Version: 2.4
|
|
2
2
|
Name: PyAthena
|
|
3
|
-
Version: 3.
|
|
3
|
+
Version: 3.35.1
|
|
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/
|
|
@@ -18,7 +18,7 @@ version_tuple: tuple[int | str, ...]
|
|
|
18
18
|
commit_id: str | None
|
|
19
19
|
__commit_id__: str | None
|
|
20
20
|
|
|
21
|
-
__version__ = version = '3.
|
|
22
|
-
__version_tuple__ = version_tuple = (3,
|
|
21
|
+
__version__ = version = '3.35.1'
|
|
22
|
+
__version_tuple__ = version_tuple = (3, 35, 1)
|
|
23
23
|
|
|
24
24
|
__commit_id__ = commit_id = None
|
|
@@ -2,6 +2,7 @@ from __future__ import annotations
|
|
|
2
2
|
|
|
3
3
|
import asyncio
|
|
4
4
|
import logging
|
|
5
|
+
from collections.abc import Callable
|
|
5
6
|
from typing import TYPE_CHECKING, Any, cast
|
|
6
7
|
|
|
7
8
|
from pyathena.aio.common import WithAsyncFetch
|
|
@@ -13,6 +14,7 @@ from pyathena.arrow.result_set import AthenaArrowResultSet
|
|
|
13
14
|
from pyathena.common import CursorIterator
|
|
14
15
|
from pyathena.error import OperationalError, ProgrammingError
|
|
15
16
|
from pyathena.model import AthenaQueryExecution
|
|
17
|
+
from pyathena.options import ExecuteOptions
|
|
16
18
|
|
|
17
19
|
if TYPE_CHECKING:
|
|
18
20
|
import polars as pl
|
|
@@ -83,11 +85,15 @@ class AioArrowCursor(WithAsyncFetch):
|
|
|
83
85
|
parameters: dict[str, Any] | list[str] | None = None,
|
|
84
86
|
work_group: str | None = None,
|
|
85
87
|
s3_staging_dir: str | None = None,
|
|
86
|
-
cache_size: int | None =
|
|
87
|
-
cache_expiration_time: int | None =
|
|
88
|
+
cache_size: int | None = None,
|
|
89
|
+
cache_expiration_time: int | None = None,
|
|
88
90
|
result_reuse_enable: bool | None = None,
|
|
89
91
|
result_reuse_minutes: int | None = None,
|
|
90
92
|
paramstyle: str | None = None,
|
|
93
|
+
on_start_query_execution: Callable[[str], None] | None = None,
|
|
94
|
+
result_set_type_hints: dict[str | int, str] | None = None,
|
|
95
|
+
*,
|
|
96
|
+
options: ExecuteOptions | None = None,
|
|
91
97
|
**kwargs,
|
|
92
98
|
) -> AioArrowCursor:
|
|
93
99
|
"""Execute a SQL query asynchronously and return results as Arrow Tables.
|
|
@@ -102,16 +108,21 @@ class AioArrowCursor(WithAsyncFetch):
|
|
|
102
108
|
result_reuse_enable: Enable Athena result reuse for this query.
|
|
103
109
|
result_reuse_minutes: Minutes to reuse cached results.
|
|
104
110
|
paramstyle: Parameter style ('qmark' or 'pyformat').
|
|
111
|
+
on_start_query_execution: Callback called when query starts.
|
|
112
|
+
result_set_type_hints: Optional dictionary mapping column names to
|
|
113
|
+
Athena DDL type signatures for precise type conversion within
|
|
114
|
+
complex types.
|
|
115
|
+
options: Shared execution options as an
|
|
116
|
+
:class:`~pyathena.options.ExecuteOptions` instance. Individual
|
|
117
|
+
keyword arguments take precedence over ``options`` fields.
|
|
105
118
|
**kwargs: Additional execution parameters.
|
|
106
119
|
|
|
107
120
|
Returns:
|
|
108
121
|
Self reference for method chaining.
|
|
109
122
|
"""
|
|
110
123
|
self._reset_state()
|
|
111
|
-
|
|
112
|
-
|
|
113
|
-
operation,
|
|
114
|
-
parameters=parameters,
|
|
124
|
+
options = ExecuteOptions.resolve(
|
|
125
|
+
options,
|
|
115
126
|
work_group=work_group,
|
|
116
127
|
s3_staging_dir=s3_staging_dir,
|
|
117
128
|
cache_size=cache_size,
|
|
@@ -119,8 +130,19 @@ class AioArrowCursor(WithAsyncFetch):
|
|
|
119
130
|
result_reuse_enable=result_reuse_enable,
|
|
120
131
|
result_reuse_minutes=result_reuse_minutes,
|
|
121
132
|
paramstyle=paramstyle,
|
|
133
|
+
on_start_query_execution=on_start_query_execution,
|
|
134
|
+
result_set_type_hints=result_set_type_hints,
|
|
135
|
+
)
|
|
136
|
+
operation, unload_location = self._prepare_unload(operation, options.s3_staging_dir)
|
|
137
|
+
self.query_id = await self._execute(
|
|
138
|
+
operation,
|
|
139
|
+
parameters=parameters,
|
|
140
|
+
options=options,
|
|
122
141
|
)
|
|
123
142
|
|
|
143
|
+
# Call user callbacks immediately after start_query_execution
|
|
144
|
+
self._call_on_start_query_execution(self.query_id, options)
|
|
145
|
+
|
|
124
146
|
query_execution = await self._poll(self.query_id)
|
|
125
147
|
if query_execution.state == AthenaQueryExecution.STATE_SUCCEEDED:
|
|
126
148
|
self.result_set = await asyncio.to_thread(
|
|
@@ -134,6 +156,7 @@ class AioArrowCursor(WithAsyncFetch):
|
|
|
134
156
|
unload_location=unload_location,
|
|
135
157
|
connect_timeout=self._connect_timeout,
|
|
136
158
|
request_timeout=self._request_timeout,
|
|
159
|
+
result_set_type_hints=options.result_set_type_hints,
|
|
137
160
|
**kwargs,
|
|
138
161
|
)
|
|
139
162
|
else:
|
|
@@ -10,6 +10,7 @@ from pyathena.aio.util import async_retry_api_call
|
|
|
10
10
|
from pyathena.common import BaseCursor, CursorIterator
|
|
11
11
|
from pyathena.error import DatabaseError, OperationalError, ProgrammingError
|
|
12
12
|
from pyathena.model import AthenaDatabase, AthenaQueryExecution, AthenaTableMetadata
|
|
13
|
+
from pyathena.options import ExecuteOptions
|
|
13
14
|
from pyathena.result_set import AthenaResultSet, WithResultSet
|
|
14
15
|
|
|
15
16
|
_logger = logging.getLogger(__name__)
|
|
@@ -29,27 +30,41 @@ class AioBaseCursor(BaseCursor):
|
|
|
29
30
|
parameters: dict[str, Any] | list[str] | None = None,
|
|
30
31
|
work_group: str | None = None,
|
|
31
32
|
s3_staging_dir: str | None = None,
|
|
32
|
-
cache_size: int | None =
|
|
33
|
-
cache_expiration_time: int | None =
|
|
33
|
+
cache_size: int | None = None,
|
|
34
|
+
cache_expiration_time: int | None = None,
|
|
34
35
|
result_reuse_enable: bool | None = None,
|
|
35
36
|
result_reuse_minutes: int | None = None,
|
|
36
37
|
paramstyle: str | None = None,
|
|
38
|
+
options: ExecuteOptions | None = None,
|
|
37
39
|
) -> str:
|
|
38
|
-
|
|
39
|
-
|
|
40
|
-
|
|
41
|
-
|
|
40
|
+
# The individual keyword arguments are retained for backward compatibility
|
|
41
|
+
# with external callers that predate ExecuteOptions, mirroring
|
|
42
|
+
# BaseCursor._execute().
|
|
43
|
+
options = ExecuteOptions.resolve(
|
|
44
|
+
options,
|
|
42
45
|
work_group=work_group,
|
|
43
46
|
s3_staging_dir=s3_staging_dir,
|
|
47
|
+
cache_size=cache_size,
|
|
48
|
+
cache_expiration_time=cache_expiration_time,
|
|
44
49
|
result_reuse_enable=result_reuse_enable,
|
|
45
50
|
result_reuse_minutes=result_reuse_minutes,
|
|
51
|
+
paramstyle=paramstyle,
|
|
52
|
+
)
|
|
53
|
+
query, execution_parameters = self._prepare_query(operation, parameters, options.paramstyle)
|
|
54
|
+
|
|
55
|
+
request = self._build_start_query_execution_request(
|
|
56
|
+
query=query,
|
|
57
|
+
work_group=options.work_group,
|
|
58
|
+
s3_staging_dir=options.s3_staging_dir,
|
|
59
|
+
result_reuse_enable=options.result_reuse_enable,
|
|
60
|
+
result_reuse_minutes=options.result_reuse_minutes,
|
|
46
61
|
execution_parameters=execution_parameters,
|
|
47
62
|
)
|
|
48
63
|
query_id = await self._find_previous_query_id(
|
|
49
64
|
query,
|
|
50
|
-
work_group,
|
|
51
|
-
cache_size=cache_size
|
|
52
|
-
cache_expiration_time=cache_expiration_time
|
|
65
|
+
options.work_group,
|
|
66
|
+
cache_size=options.cache_size,
|
|
67
|
+
cache_expiration_time=options.cache_expiration_time,
|
|
53
68
|
)
|
|
54
69
|
if query_id is None:
|
|
55
70
|
try:
|
|
@@ -1,6 +1,7 @@
|
|
|
1
1
|
from __future__ import annotations
|
|
2
2
|
|
|
3
3
|
import logging
|
|
4
|
+
from collections.abc import Callable
|
|
4
5
|
from typing import Any, cast
|
|
5
6
|
|
|
6
7
|
from pyathena.aio.common import WithAsyncFetch
|
|
@@ -8,6 +9,7 @@ from pyathena.aio.result_set import AthenaAioDictResultSet, AthenaAioResultSet
|
|
|
8
9
|
from pyathena.common import CursorIterator
|
|
9
10
|
from pyathena.error import OperationalError, ProgrammingError
|
|
10
11
|
from pyathena.model import AthenaQueryExecution
|
|
12
|
+
from pyathena.options import ExecuteOptions
|
|
11
13
|
|
|
12
14
|
_logger = logging.getLogger(__name__)
|
|
13
15
|
|
|
@@ -74,12 +76,15 @@ class AioCursor(WithAsyncFetch):
|
|
|
74
76
|
parameters: dict[str, Any] | list[str] | None = None,
|
|
75
77
|
work_group: str | None = None,
|
|
76
78
|
s3_staging_dir: str | None = None,
|
|
77
|
-
cache_size: int =
|
|
78
|
-
cache_expiration_time: int =
|
|
79
|
+
cache_size: int | None = None,
|
|
80
|
+
cache_expiration_time: int | None = None,
|
|
79
81
|
result_reuse_enable: bool | None = None,
|
|
80
82
|
result_reuse_minutes: int | None = None,
|
|
81
83
|
paramstyle: str | None = None,
|
|
84
|
+
on_start_query_execution: Callable[[str], None] | None = None,
|
|
82
85
|
result_set_type_hints: dict[str | int, str] | None = None,
|
|
86
|
+
*,
|
|
87
|
+
options: ExecuteOptions | None = None,
|
|
83
88
|
**kwargs,
|
|
84
89
|
) -> AioCursor:
|
|
85
90
|
"""Execute a SQL query asynchronously.
|
|
@@ -94,18 +99,21 @@ class AioCursor(WithAsyncFetch):
|
|
|
94
99
|
result_reuse_enable: Enable result reuse (optional).
|
|
95
100
|
result_reuse_minutes: Result reuse duration in minutes (optional).
|
|
96
101
|
paramstyle: Parameter style to use (optional).
|
|
102
|
+
on_start_query_execution: Callback called when query starts.
|
|
97
103
|
result_set_type_hints: Optional dictionary mapping column names to
|
|
98
104
|
Athena DDL type signatures for precise type conversion within
|
|
99
105
|
complex types.
|
|
106
|
+
options: Shared execution options as an
|
|
107
|
+
:class:`~pyathena.options.ExecuteOptions` instance. Individual
|
|
108
|
+
keyword arguments take precedence over ``options`` fields.
|
|
100
109
|
**kwargs: Additional execution parameters.
|
|
101
110
|
|
|
102
111
|
Returns:
|
|
103
112
|
Self reference for method chaining.
|
|
104
113
|
"""
|
|
105
114
|
self._reset_state()
|
|
106
|
-
|
|
107
|
-
|
|
108
|
-
parameters=parameters,
|
|
115
|
+
options = ExecuteOptions.resolve(
|
|
116
|
+
options,
|
|
109
117
|
work_group=work_group,
|
|
110
118
|
s3_staging_dir=s3_staging_dir,
|
|
111
119
|
cache_size=cache_size,
|
|
@@ -113,8 +121,18 @@ class AioCursor(WithAsyncFetch):
|
|
|
113
121
|
result_reuse_enable=result_reuse_enable,
|
|
114
122
|
result_reuse_minutes=result_reuse_minutes,
|
|
115
123
|
paramstyle=paramstyle,
|
|
124
|
+
on_start_query_execution=on_start_query_execution,
|
|
125
|
+
result_set_type_hints=result_set_type_hints,
|
|
126
|
+
)
|
|
127
|
+
self.query_id = await self._execute(
|
|
128
|
+
operation,
|
|
129
|
+
parameters=parameters,
|
|
130
|
+
options=options,
|
|
116
131
|
)
|
|
117
132
|
|
|
133
|
+
# Call user callbacks immediately after start_query_execution
|
|
134
|
+
self._call_on_start_query_execution(self.query_id, options)
|
|
135
|
+
|
|
118
136
|
query_execution = await self._poll(self.query_id)
|
|
119
137
|
if query_execution.state == AthenaQueryExecution.STATE_SUCCEEDED:
|
|
120
138
|
self.result_set = await self._result_set_class.create(
|
|
@@ -123,7 +141,7 @@ class AioCursor(WithAsyncFetch):
|
|
|
123
141
|
query_execution,
|
|
124
142
|
self.arraysize,
|
|
125
143
|
self._retry_config,
|
|
126
|
-
result_set_type_hints=result_set_type_hints,
|
|
144
|
+
result_set_type_hints=options.result_set_type_hints,
|
|
127
145
|
)
|
|
128
146
|
else:
|
|
129
147
|
raise OperationalError(query_execution.state_change_reason)
|
|
@@ -2,7 +2,7 @@ from __future__ import annotations
|
|
|
2
2
|
|
|
3
3
|
import asyncio
|
|
4
4
|
import logging
|
|
5
|
-
from collections.abc import Iterable
|
|
5
|
+
from collections.abc import Callable, Iterable
|
|
6
6
|
from multiprocessing import cpu_count
|
|
7
7
|
from typing import (
|
|
8
8
|
TYPE_CHECKING,
|
|
@@ -14,6 +14,7 @@ from pyathena.aio.common import WithAsyncFetch
|
|
|
14
14
|
from pyathena.common import CursorIterator
|
|
15
15
|
from pyathena.error import OperationalError, ProgrammingError
|
|
16
16
|
from pyathena.model import AthenaQueryExecution
|
|
17
|
+
from pyathena.options import ExecuteOptions
|
|
17
18
|
from pyathena.pandas.converter import (
|
|
18
19
|
DefaultPandasTypeConverter,
|
|
19
20
|
DefaultPandasUnloadTypeConverter,
|
|
@@ -97,14 +98,18 @@ class AioPandasCursor(WithAsyncFetch):
|
|
|
97
98
|
parameters: dict[str, Any] | list[str] | None = None,
|
|
98
99
|
work_group: str | None = None,
|
|
99
100
|
s3_staging_dir: str | None = None,
|
|
100
|
-
cache_size: int | None =
|
|
101
|
-
cache_expiration_time: int | None =
|
|
101
|
+
cache_size: int | None = None,
|
|
102
|
+
cache_expiration_time: int | None = None,
|
|
102
103
|
result_reuse_enable: bool | None = None,
|
|
103
104
|
result_reuse_minutes: int | None = None,
|
|
104
105
|
paramstyle: str | None = None,
|
|
105
106
|
keep_default_na: bool = False,
|
|
106
107
|
na_values: Iterable[str] | None = ("",),
|
|
107
108
|
quoting: int = 1,
|
|
109
|
+
on_start_query_execution: Callable[[str], None] | None = None,
|
|
110
|
+
result_set_type_hints: dict[str | int, str] | None = None,
|
|
111
|
+
*,
|
|
112
|
+
options: ExecuteOptions | None = None,
|
|
108
113
|
**kwargs,
|
|
109
114
|
) -> AioPandasCursor:
|
|
110
115
|
"""Execute a SQL query asynchronously and return results as pandas DataFrames.
|
|
@@ -122,16 +127,21 @@ class AioPandasCursor(WithAsyncFetch):
|
|
|
122
127
|
keep_default_na: Whether to keep default pandas NA values.
|
|
123
128
|
na_values: Additional values to treat as NA.
|
|
124
129
|
quoting: CSV quoting behavior (pandas csv.QUOTE_* constants).
|
|
130
|
+
on_start_query_execution: Callback called when query starts.
|
|
131
|
+
result_set_type_hints: Optional dictionary mapping column names to
|
|
132
|
+
Athena DDL type signatures for precise type conversion within
|
|
133
|
+
complex types.
|
|
134
|
+
options: Shared execution options as an
|
|
135
|
+
:class:`~pyathena.options.ExecuteOptions` instance. Individual
|
|
136
|
+
keyword arguments take precedence over ``options`` fields.
|
|
125
137
|
**kwargs: Additional pandas read_csv/read_parquet parameters.
|
|
126
138
|
|
|
127
139
|
Returns:
|
|
128
140
|
Self reference for method chaining.
|
|
129
141
|
"""
|
|
130
142
|
self._reset_state()
|
|
131
|
-
|
|
132
|
-
|
|
133
|
-
operation,
|
|
134
|
-
parameters=parameters,
|
|
143
|
+
options = ExecuteOptions.resolve(
|
|
144
|
+
options,
|
|
135
145
|
work_group=work_group,
|
|
136
146
|
s3_staging_dir=s3_staging_dir,
|
|
137
147
|
cache_size=cache_size,
|
|
@@ -139,8 +149,19 @@ class AioPandasCursor(WithAsyncFetch):
|
|
|
139
149
|
result_reuse_enable=result_reuse_enable,
|
|
140
150
|
result_reuse_minutes=result_reuse_minutes,
|
|
141
151
|
paramstyle=paramstyle,
|
|
152
|
+
on_start_query_execution=on_start_query_execution,
|
|
153
|
+
result_set_type_hints=result_set_type_hints,
|
|
154
|
+
)
|
|
155
|
+
operation, unload_location = self._prepare_unload(operation, options.s3_staging_dir)
|
|
156
|
+
self.query_id = await self._execute(
|
|
157
|
+
operation,
|
|
158
|
+
parameters=parameters,
|
|
159
|
+
options=options,
|
|
142
160
|
)
|
|
143
161
|
|
|
162
|
+
# Call user callbacks immediately after start_query_execution
|
|
163
|
+
self._call_on_start_query_execution(self.query_id, options)
|
|
164
|
+
|
|
144
165
|
query_execution = await self._poll(self.query_id)
|
|
145
166
|
if query_execution.state == AthenaQueryExecution.STATE_SUCCEEDED:
|
|
146
167
|
self.result_set = await asyncio.to_thread(
|
|
@@ -161,6 +182,7 @@ class AioPandasCursor(WithAsyncFetch):
|
|
|
161
182
|
cache_type=kwargs.pop("cache_type", self._cache_type),
|
|
162
183
|
max_workers=kwargs.pop("max_workers", self._max_workers),
|
|
163
184
|
auto_optimize_chunksize=self._auto_optimize_chunksize,
|
|
185
|
+
result_set_type_hints=options.result_set_type_hints,
|
|
164
186
|
**kwargs,
|
|
165
187
|
)
|
|
166
188
|
else:
|
|
@@ -2,6 +2,7 @@ from __future__ import annotations
|
|
|
2
2
|
|
|
3
3
|
import asyncio
|
|
4
4
|
import logging
|
|
5
|
+
from collections.abc import Callable
|
|
5
6
|
from multiprocessing import cpu_count
|
|
6
7
|
from typing import TYPE_CHECKING, Any, cast
|
|
7
8
|
|
|
@@ -9,6 +10,7 @@ from pyathena.aio.common import WithAsyncFetch
|
|
|
9
10
|
from pyathena.common import CursorIterator
|
|
10
11
|
from pyathena.error import OperationalError, ProgrammingError
|
|
11
12
|
from pyathena.model import AthenaQueryExecution
|
|
13
|
+
from pyathena.options import ExecuteOptions
|
|
12
14
|
from pyathena.polars.converter import (
|
|
13
15
|
DefaultPolarsTypeConverter,
|
|
14
16
|
DefaultPolarsUnloadTypeConverter,
|
|
@@ -89,11 +91,15 @@ class AioPolarsCursor(WithAsyncFetch):
|
|
|
89
91
|
parameters: dict[str, Any] | list[str] | None = None,
|
|
90
92
|
work_group: str | None = None,
|
|
91
93
|
s3_staging_dir: str | None = None,
|
|
92
|
-
cache_size: int | None =
|
|
93
|
-
cache_expiration_time: int | None =
|
|
94
|
+
cache_size: int | None = None,
|
|
95
|
+
cache_expiration_time: int | None = None,
|
|
94
96
|
result_reuse_enable: bool | None = None,
|
|
95
97
|
result_reuse_minutes: int | None = None,
|
|
96
98
|
paramstyle: str | None = None,
|
|
99
|
+
on_start_query_execution: Callable[[str], None] | None = None,
|
|
100
|
+
result_set_type_hints: dict[str | int, str] | None = None,
|
|
101
|
+
*,
|
|
102
|
+
options: ExecuteOptions | None = None,
|
|
97
103
|
**kwargs,
|
|
98
104
|
) -> AioPolarsCursor:
|
|
99
105
|
"""Execute a SQL query asynchronously and return results as Polars DataFrames.
|
|
@@ -108,16 +114,21 @@ class AioPolarsCursor(WithAsyncFetch):
|
|
|
108
114
|
result_reuse_enable: Enable Athena result reuse for this query.
|
|
109
115
|
result_reuse_minutes: Minutes to reuse cached results.
|
|
110
116
|
paramstyle: Parameter style ('qmark' or 'pyformat').
|
|
117
|
+
on_start_query_execution: Callback called when query starts.
|
|
118
|
+
result_set_type_hints: Optional dictionary mapping column names to
|
|
119
|
+
Athena DDL type signatures for precise type conversion within
|
|
120
|
+
complex types.
|
|
121
|
+
options: Shared execution options as an
|
|
122
|
+
:class:`~pyathena.options.ExecuteOptions` instance. Individual
|
|
123
|
+
keyword arguments take precedence over ``options`` fields.
|
|
111
124
|
**kwargs: Additional execution parameters passed to Polars read functions.
|
|
112
125
|
|
|
113
126
|
Returns:
|
|
114
127
|
Self reference for method chaining.
|
|
115
128
|
"""
|
|
116
129
|
self._reset_state()
|
|
117
|
-
|
|
118
|
-
|
|
119
|
-
operation,
|
|
120
|
-
parameters=parameters,
|
|
130
|
+
options = ExecuteOptions.resolve(
|
|
131
|
+
options,
|
|
121
132
|
work_group=work_group,
|
|
122
133
|
s3_staging_dir=s3_staging_dir,
|
|
123
134
|
cache_size=cache_size,
|
|
@@ -125,8 +136,19 @@ class AioPolarsCursor(WithAsyncFetch):
|
|
|
125
136
|
result_reuse_enable=result_reuse_enable,
|
|
126
137
|
result_reuse_minutes=result_reuse_minutes,
|
|
127
138
|
paramstyle=paramstyle,
|
|
139
|
+
on_start_query_execution=on_start_query_execution,
|
|
140
|
+
result_set_type_hints=result_set_type_hints,
|
|
141
|
+
)
|
|
142
|
+
operation, unload_location = self._prepare_unload(operation, options.s3_staging_dir)
|
|
143
|
+
self.query_id = await self._execute(
|
|
144
|
+
operation,
|
|
145
|
+
parameters=parameters,
|
|
146
|
+
options=options,
|
|
128
147
|
)
|
|
129
148
|
|
|
149
|
+
# Call user callbacks immediately after start_query_execution
|
|
150
|
+
self._call_on_start_query_execution(self.query_id, options)
|
|
151
|
+
|
|
130
152
|
query_execution = await self._poll(self.query_id)
|
|
131
153
|
if query_execution.state == AthenaQueryExecution.STATE_SUCCEEDED:
|
|
132
154
|
self.result_set = await asyncio.to_thread(
|
|
@@ -142,6 +164,7 @@ class AioPolarsCursor(WithAsyncFetch):
|
|
|
142
164
|
cache_type=self._cache_type,
|
|
143
165
|
max_workers=self._max_workers,
|
|
144
166
|
chunksize=self._chunksize,
|
|
167
|
+
result_set_type_hints=options.result_set_type_hints,
|
|
145
168
|
**kwargs,
|
|
146
169
|
)
|
|
147
170
|
else:
|
|
@@ -2,6 +2,7 @@ from __future__ import annotations
|
|
|
2
2
|
|
|
3
3
|
import asyncio
|
|
4
4
|
import logging
|
|
5
|
+
from collections.abc import Callable
|
|
5
6
|
from typing import Any, cast
|
|
6
7
|
|
|
7
8
|
from pyathena.aio.common import WithAsyncFetch
|
|
@@ -9,6 +10,7 @@ from pyathena.common import CursorIterator
|
|
|
9
10
|
from pyathena.error import OperationalError, ProgrammingError
|
|
10
11
|
from pyathena.filesystem.s3_async import AioS3FileSystem
|
|
11
12
|
from pyathena.model import AthenaQueryExecution
|
|
13
|
+
from pyathena.options import ExecuteOptions
|
|
12
14
|
from pyathena.s3fs.converter import DefaultS3FSTypeConverter
|
|
13
15
|
from pyathena.s3fs.result_set import AthenaS3FSResultSet, CSVReaderType
|
|
14
16
|
|
|
@@ -81,11 +83,15 @@ class AioS3FSCursor(WithAsyncFetch):
|
|
|
81
83
|
parameters: dict[str, Any] | list[str] | None = None,
|
|
82
84
|
work_group: str | None = None,
|
|
83
85
|
s3_staging_dir: str | None = None,
|
|
84
|
-
cache_size: int | None =
|
|
85
|
-
cache_expiration_time: int | None =
|
|
86
|
+
cache_size: int | None = None,
|
|
87
|
+
cache_expiration_time: int | None = None,
|
|
86
88
|
result_reuse_enable: bool | None = None,
|
|
87
89
|
result_reuse_minutes: int | None = None,
|
|
88
90
|
paramstyle: str | None = None,
|
|
91
|
+
on_start_query_execution: Callable[[str], None] | None = None,
|
|
92
|
+
result_set_type_hints: dict[str | int, str] | None = None,
|
|
93
|
+
*,
|
|
94
|
+
options: ExecuteOptions | None = None,
|
|
89
95
|
**kwargs,
|
|
90
96
|
) -> AioS3FSCursor:
|
|
91
97
|
"""Execute a SQL query asynchronously via S3FileSystem CSV reader.
|
|
@@ -100,15 +106,21 @@ class AioS3FSCursor(WithAsyncFetch):
|
|
|
100
106
|
result_reuse_enable: Enable Athena result reuse for this query.
|
|
101
107
|
result_reuse_minutes: Minutes to reuse cached results.
|
|
102
108
|
paramstyle: Parameter style ('qmark' or 'pyformat').
|
|
109
|
+
on_start_query_execution: Callback called when query starts.
|
|
110
|
+
result_set_type_hints: Optional dictionary mapping column names to
|
|
111
|
+
Athena DDL type signatures for precise type conversion within
|
|
112
|
+
complex types.
|
|
113
|
+
options: Shared execution options as an
|
|
114
|
+
:class:`~pyathena.options.ExecuteOptions` instance. Individual
|
|
115
|
+
keyword arguments take precedence over ``options`` fields.
|
|
103
116
|
**kwargs: Additional execution parameters.
|
|
104
117
|
|
|
105
118
|
Returns:
|
|
106
119
|
Self reference for method chaining.
|
|
107
120
|
"""
|
|
108
121
|
self._reset_state()
|
|
109
|
-
|
|
110
|
-
|
|
111
|
-
parameters=parameters,
|
|
122
|
+
options = ExecuteOptions.resolve(
|
|
123
|
+
options,
|
|
112
124
|
work_group=work_group,
|
|
113
125
|
s3_staging_dir=s3_staging_dir,
|
|
114
126
|
cache_size=cache_size,
|
|
@@ -116,8 +128,18 @@ class AioS3FSCursor(WithAsyncFetch):
|
|
|
116
128
|
result_reuse_enable=result_reuse_enable,
|
|
117
129
|
result_reuse_minutes=result_reuse_minutes,
|
|
118
130
|
paramstyle=paramstyle,
|
|
131
|
+
on_start_query_execution=on_start_query_execution,
|
|
132
|
+
result_set_type_hints=result_set_type_hints,
|
|
133
|
+
)
|
|
134
|
+
self.query_id = await self._execute(
|
|
135
|
+
operation,
|
|
136
|
+
parameters=parameters,
|
|
137
|
+
options=options,
|
|
119
138
|
)
|
|
120
139
|
|
|
140
|
+
# Call user callbacks immediately after start_query_execution
|
|
141
|
+
self._call_on_start_query_execution(self.query_id, options)
|
|
142
|
+
|
|
121
143
|
query_execution = await self._poll(self.query_id)
|
|
122
144
|
if query_execution.state == AthenaQueryExecution.STATE_SUCCEEDED:
|
|
123
145
|
self.result_set = await asyncio.to_thread(
|
|
@@ -129,6 +151,7 @@ class AioS3FSCursor(WithAsyncFetch):
|
|
|
129
151
|
retry_config=self._retry_config,
|
|
130
152
|
csv_reader=self._csv_reader,
|
|
131
153
|
filesystem_class=AioS3FileSystem,
|
|
154
|
+
result_set_type_hints=options.result_set_type_hints,
|
|
132
155
|
**kwargs,
|
|
133
156
|
)
|
|
134
157
|
else:
|
|
@@ -14,6 +14,7 @@ from pyathena.arrow.result_set import AthenaArrowResultSet
|
|
|
14
14
|
from pyathena.async_cursor import AsyncCursor
|
|
15
15
|
from pyathena.common import CursorIterator
|
|
16
16
|
from pyathena.model import AthenaQueryExecution
|
|
17
|
+
from pyathena.options import ExecuteOptions
|
|
17
18
|
|
|
18
19
|
_logger = logging.getLogger(__name__)
|
|
19
20
|
|
|
@@ -176,18 +177,41 @@ class AsyncArrowCursor(AsyncCursor):
|
|
|
176
177
|
parameters: dict[str, Any] | list[str] | None = None,
|
|
177
178
|
work_group: str | None = None,
|
|
178
179
|
s3_staging_dir: str | None = None,
|
|
179
|
-
cache_size: int | None =
|
|
180
|
-
cache_expiration_time: int | None =
|
|
180
|
+
cache_size: int | None = None,
|
|
181
|
+
cache_expiration_time: int | None = None,
|
|
181
182
|
result_reuse_enable: bool | None = None,
|
|
182
183
|
result_reuse_minutes: int | None = None,
|
|
183
184
|
paramstyle: str | None = None,
|
|
184
185
|
result_set_type_hints: dict[str | int, str] | None = None,
|
|
186
|
+
*,
|
|
187
|
+
options: ExecuteOptions | None = None,
|
|
185
188
|
**kwargs,
|
|
186
189
|
) -> tuple[str, Future[AthenaArrowResultSet | Any]]:
|
|
187
|
-
|
|
188
|
-
|
|
189
|
-
|
|
190
|
-
|
|
190
|
+
"""Execute a SQL query asynchronously and return results as Arrow Tables.
|
|
191
|
+
|
|
192
|
+
Args:
|
|
193
|
+
operation: SQL query string to execute.
|
|
194
|
+
parameters: Query parameters for parameterized queries.
|
|
195
|
+
work_group: Athena workgroup to use for this query.
|
|
196
|
+
s3_staging_dir: S3 location for query results.
|
|
197
|
+
cache_size: Number of queries to check for result caching.
|
|
198
|
+
cache_expiration_time: Cache expiration time in seconds.
|
|
199
|
+
result_reuse_enable: Enable Athena result reuse for this query.
|
|
200
|
+
result_reuse_minutes: Minutes to reuse cached results.
|
|
201
|
+
paramstyle: Parameter style ('qmark' or 'pyformat').
|
|
202
|
+
result_set_type_hints: Optional dictionary mapping column names to
|
|
203
|
+
Athena DDL type signatures for precise type conversion within
|
|
204
|
+
complex types.
|
|
205
|
+
options: Shared execution options as an
|
|
206
|
+
:class:`~pyathena.options.ExecuteOptions` instance. Individual
|
|
207
|
+
keyword arguments take precedence over ``options`` fields.
|
|
208
|
+
**kwargs: Additional execution parameters.
|
|
209
|
+
|
|
210
|
+
Returns:
|
|
211
|
+
Tuple of (query_id, future) where future resolves to AthenaArrowResultSet.
|
|
212
|
+
"""
|
|
213
|
+
options = ExecuteOptions.resolve(
|
|
214
|
+
options,
|
|
191
215
|
work_group=work_group,
|
|
192
216
|
s3_staging_dir=s3_staging_dir,
|
|
193
217
|
cache_size=cache_size,
|
|
@@ -195,13 +219,20 @@ class AsyncArrowCursor(AsyncCursor):
|
|
|
195
219
|
result_reuse_enable=result_reuse_enable,
|
|
196
220
|
result_reuse_minutes=result_reuse_minutes,
|
|
197
221
|
paramstyle=paramstyle,
|
|
222
|
+
result_set_type_hints=result_set_type_hints,
|
|
223
|
+
)
|
|
224
|
+
operation, unload_location = self._prepare_unload(operation, options.s3_staging_dir)
|
|
225
|
+
query_id = self._execute(
|
|
226
|
+
operation,
|
|
227
|
+
parameters=parameters,
|
|
228
|
+
options=options,
|
|
198
229
|
)
|
|
199
230
|
return (
|
|
200
231
|
query_id,
|
|
201
232
|
self._executor.submit(
|
|
202
233
|
self._collect_result_set,
|
|
203
234
|
query_id,
|
|
204
|
-
result_set_type_hints,
|
|
235
|
+
options.result_set_type_hints,
|
|
205
236
|
unload_location,
|
|
206
237
|
kwargs,
|
|
207
238
|
),
|