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.
Files changed (90) hide show
  1. {pyathena-3.34.0 → pyathena-3.35.1}/PKG-INFO +1 -1
  2. {pyathena-3.34.0 → pyathena-3.35.1}/pyathena/__init__.py +1 -0
  3. {pyathena-3.34.0 → pyathena-3.35.1}/pyathena/_version.py +2 -2
  4. {pyathena-3.34.0 → pyathena-3.35.1}/pyathena/aio/arrow/cursor.py +29 -6
  5. {pyathena-3.34.0 → pyathena-3.35.1}/pyathena/aio/common.py +24 -9
  6. {pyathena-3.34.0 → pyathena-3.35.1}/pyathena/aio/cursor.py +24 -6
  7. {pyathena-3.34.0 → pyathena-3.35.1}/pyathena/aio/pandas/cursor.py +29 -7
  8. {pyathena-3.34.0 → pyathena-3.35.1}/pyathena/aio/polars/cursor.py +29 -6
  9. {pyathena-3.34.0 → pyathena-3.35.1}/pyathena/aio/s3fs/cursor.py +28 -5
  10. {pyathena-3.34.0 → pyathena-3.35.1}/pyathena/arrow/async_cursor.py +38 -7
  11. {pyathena-3.34.0 → pyathena-3.35.1}/pyathena/arrow/cursor.py +20 -12
  12. {pyathena-3.34.0 → pyathena-3.35.1}/pyathena/async_cursor.py +17 -6
  13. {pyathena-3.34.0 → pyathena-3.35.1}/pyathena/common.py +39 -12
  14. {pyathena-3.34.0 → pyathena-3.35.1}/pyathena/cursor.py +19 -11
  15. pyathena-3.35.1/pyathena/options.py +102 -0
  16. {pyathena-3.34.0 → pyathena-3.35.1}/pyathena/pandas/async_cursor.py +41 -7
  17. {pyathena-3.34.0 → pyathena-3.35.1}/pyathena/pandas/cursor.py +20 -12
  18. {pyathena-3.34.0 → pyathena-3.35.1}/pyathena/polars/async_cursor.py +18 -7
  19. {pyathena-3.34.0 → pyathena-3.35.1}/pyathena/polars/cursor.py +20 -12
  20. {pyathena-3.34.0 → pyathena-3.35.1}/pyathena/s3fs/async_cursor.py +17 -6
  21. {pyathena-3.34.0 → pyathena-3.35.1}/pyathena/s3fs/cursor.py +19 -10
  22. {pyathena-3.34.0 → pyathena-3.35.1}/.gitignore +0 -0
  23. {pyathena-3.34.0 → pyathena-3.35.1}/LICENSE +0 -0
  24. {pyathena-3.34.0 → pyathena-3.35.1}/README.md +0 -0
  25. {pyathena-3.34.0 → pyathena-3.35.1}/pyathena/aio/__init__.py +0 -0
  26. {pyathena-3.34.0 → pyathena-3.35.1}/pyathena/aio/arrow/__init__.py +0 -0
  27. {pyathena-3.34.0 → pyathena-3.35.1}/pyathena/aio/connection.py +0 -0
  28. {pyathena-3.34.0 → pyathena-3.35.1}/pyathena/aio/pandas/__init__.py +0 -0
  29. {pyathena-3.34.0 → pyathena-3.35.1}/pyathena/aio/polars/__init__.py +0 -0
  30. {pyathena-3.34.0 → pyathena-3.35.1}/pyathena/aio/result_set.py +0 -0
  31. {pyathena-3.34.0 → pyathena-3.35.1}/pyathena/aio/s3fs/__init__.py +0 -0
  32. {pyathena-3.34.0 → pyathena-3.35.1}/pyathena/aio/spark/__init__.py +0 -0
  33. {pyathena-3.34.0 → pyathena-3.35.1}/pyathena/aio/spark/cursor.py +0 -0
  34. {pyathena-3.34.0 → pyathena-3.35.1}/pyathena/aio/sqlalchemy/__init__.py +0 -0
  35. {pyathena-3.34.0 → pyathena-3.35.1}/pyathena/aio/sqlalchemy/arrow.py +0 -0
  36. {pyathena-3.34.0 → pyathena-3.35.1}/pyathena/aio/sqlalchemy/base.py +0 -0
  37. {pyathena-3.34.0 → pyathena-3.35.1}/pyathena/aio/sqlalchemy/pandas.py +0 -0
  38. {pyathena-3.34.0 → pyathena-3.35.1}/pyathena/aio/sqlalchemy/polars.py +0 -0
  39. {pyathena-3.34.0 → pyathena-3.35.1}/pyathena/aio/sqlalchemy/rest.py +0 -0
  40. {pyathena-3.34.0 → pyathena-3.35.1}/pyathena/aio/sqlalchemy/s3fs.py +0 -0
  41. {pyathena-3.34.0 → pyathena-3.35.1}/pyathena/aio/util.py +0 -0
  42. {pyathena-3.34.0 → pyathena-3.35.1}/pyathena/arrow/__init__.py +0 -0
  43. {pyathena-3.34.0 → pyathena-3.35.1}/pyathena/arrow/converter.py +0 -0
  44. {pyathena-3.34.0 → pyathena-3.35.1}/pyathena/arrow/result_set.py +0 -0
  45. {pyathena-3.34.0 → pyathena-3.35.1}/pyathena/arrow/util.py +0 -0
  46. {pyathena-3.34.0 → pyathena-3.35.1}/pyathena/connection.py +0 -0
  47. {pyathena-3.34.0 → pyathena-3.35.1}/pyathena/converter.py +0 -0
  48. {pyathena-3.34.0 → pyathena-3.35.1}/pyathena/error.py +0 -0
  49. {pyathena-3.34.0 → pyathena-3.35.1}/pyathena/filesystem/__init__.py +0 -0
  50. {pyathena-3.34.0 → pyathena-3.35.1}/pyathena/filesystem/s3.py +0 -0
  51. {pyathena-3.34.0 → pyathena-3.35.1}/pyathena/filesystem/s3_async.py +0 -0
  52. {pyathena-3.34.0 → pyathena-3.35.1}/pyathena/filesystem/s3_errors.py +0 -0
  53. {pyathena-3.34.0 → pyathena-3.35.1}/pyathena/filesystem/s3_executor.py +0 -0
  54. {pyathena-3.34.0 → pyathena-3.35.1}/pyathena/filesystem/s3_object.py +0 -0
  55. {pyathena-3.34.0 → pyathena-3.35.1}/pyathena/formatter.py +0 -0
  56. {pyathena-3.34.0 → pyathena-3.35.1}/pyathena/model.py +0 -0
  57. {pyathena-3.34.0 → pyathena-3.35.1}/pyathena/pandas/__init__.py +0 -0
  58. {pyathena-3.34.0 → pyathena-3.35.1}/pyathena/pandas/converter.py +0 -0
  59. {pyathena-3.34.0 → pyathena-3.35.1}/pyathena/pandas/result_set.py +0 -0
  60. {pyathena-3.34.0 → pyathena-3.35.1}/pyathena/pandas/util.py +0 -0
  61. {pyathena-3.34.0 → pyathena-3.35.1}/pyathena/parser.py +0 -0
  62. {pyathena-3.34.0 → pyathena-3.35.1}/pyathena/polars/__init__.py +0 -0
  63. {pyathena-3.34.0 → pyathena-3.35.1}/pyathena/polars/converter.py +0 -0
  64. {pyathena-3.34.0 → pyathena-3.35.1}/pyathena/polars/result_set.py +0 -0
  65. {pyathena-3.34.0 → pyathena-3.35.1}/pyathena/polars/util.py +0 -0
  66. {pyathena-3.34.0 → pyathena-3.35.1}/pyathena/py.typed +0 -0
  67. {pyathena-3.34.0 → pyathena-3.35.1}/pyathena/result_set.py +0 -0
  68. {pyathena-3.34.0 → pyathena-3.35.1}/pyathena/s3fs/__init__.py +0 -0
  69. {pyathena-3.34.0 → pyathena-3.35.1}/pyathena/s3fs/converter.py +0 -0
  70. {pyathena-3.34.0 → pyathena-3.35.1}/pyathena/s3fs/reader.py +0 -0
  71. {pyathena-3.34.0 → pyathena-3.35.1}/pyathena/s3fs/result_set.py +0 -0
  72. {pyathena-3.34.0 → pyathena-3.35.1}/pyathena/spark/__init__.py +0 -0
  73. {pyathena-3.34.0 → pyathena-3.35.1}/pyathena/spark/async_cursor.py +0 -0
  74. {pyathena-3.34.0 → pyathena-3.35.1}/pyathena/spark/common.py +0 -0
  75. {pyathena-3.34.0 → pyathena-3.35.1}/pyathena/spark/cursor.py +0 -0
  76. {pyathena-3.34.0 → pyathena-3.35.1}/pyathena/sqlalchemy/__init__.py +0 -0
  77. {pyathena-3.34.0 → pyathena-3.35.1}/pyathena/sqlalchemy/arrow.py +0 -0
  78. {pyathena-3.34.0 → pyathena-3.35.1}/pyathena/sqlalchemy/base.py +0 -0
  79. {pyathena-3.34.0 → pyathena-3.35.1}/pyathena/sqlalchemy/compiler.py +0 -0
  80. {pyathena-3.34.0 → pyathena-3.35.1}/pyathena/sqlalchemy/constants.py +0 -0
  81. {pyathena-3.34.0 → pyathena-3.35.1}/pyathena/sqlalchemy/pandas.py +0 -0
  82. {pyathena-3.34.0 → pyathena-3.35.1}/pyathena/sqlalchemy/polars.py +0 -0
  83. {pyathena-3.34.0 → pyathena-3.35.1}/pyathena/sqlalchemy/preparer.py +0 -0
  84. {pyathena-3.34.0 → pyathena-3.35.1}/pyathena/sqlalchemy/requirements.py +0 -0
  85. {pyathena-3.34.0 → pyathena-3.35.1}/pyathena/sqlalchemy/rest.py +0 -0
  86. {pyathena-3.34.0 → pyathena-3.35.1}/pyathena/sqlalchemy/s3fs.py +0 -0
  87. {pyathena-3.34.0 → pyathena-3.35.1}/pyathena/sqlalchemy/types.py +0 -0
  88. {pyathena-3.34.0 → pyathena-3.35.1}/pyathena/sqlalchemy/util.py +0 -0
  89. {pyathena-3.34.0 → pyathena-3.35.1}/pyathena/util.py +0 -0
  90. {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.34.0
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/
@@ -4,6 +4,7 @@ import datetime
4
4
  from typing import TYPE_CHECKING, Any, overload
5
5
 
6
6
  from pyathena.error import * # noqa: F403
7
+ from pyathena.options import ExecuteOptions as ExecuteOptions
7
8
 
8
9
  if TYPE_CHECKING:
9
10
  from pyathena.aio.connection import AioConnection
@@ -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.34.0'
22
- __version_tuple__ = version_tuple = (3, 34, 0)
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 = 0,
87
- cache_expiration_time: int | None = 0,
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
- operation, unload_location = self._prepare_unload(operation, s3_staging_dir)
112
- self.query_id = await self._execute(
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 = 0,
33
- cache_expiration_time: int | None = 0,
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
- query, execution_parameters = self._prepare_query(operation, parameters, paramstyle)
39
-
40
- request = self._build_start_query_execution_request(
41
- query=query,
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 if cache_size else 0,
52
- cache_expiration_time=cache_expiration_time if cache_expiration_time else 0,
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 = 0,
78
- cache_expiration_time: int = 0,
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
- self.query_id = await self._execute(
107
- operation,
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 = 0,
101
- cache_expiration_time: int | None = 0,
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
- operation, unload_location = self._prepare_unload(operation, s3_staging_dir)
132
- self.query_id = await self._execute(
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 = 0,
93
- cache_expiration_time: int | None = 0,
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
- operation, unload_location = self._prepare_unload(operation, s3_staging_dir)
118
- self.query_id = await self._execute(
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 = 0,
85
- cache_expiration_time: int | None = 0,
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
- self.query_id = await self._execute(
110
- operation,
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 = 0,
180
- cache_expiration_time: int | None = 0,
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
- operation, unload_location = self._prepare_unload(operation, s3_staging_dir)
188
- query_id = self._execute(
189
- operation,
190
- parameters=parameters,
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
  ),