Skip to content

Commit 064a0b3

Browse files
feat: support newer GetQueryExecution response fields
Add the following fields to AthenaQueryExecution, all introduced to the GetQueryExecution response over the past two years and not previously exposed by PyAthena: Statistics: - service_pre_processing_time_in_millis - dpu_count Top-level: - managed_query_results_enabled / managed_query_results_kms_key - enable_s3_access_grants / create_user_level_prefix / s3_access_grants_authentication_type Bump boto3 to >=1.38.2 and botocore to >=1.41.2 (Nov 2025) so that all new fields are actually populated by the SDK. Earlier versions would silently return None for the newest field (DpuCount), which would be misleading. Refs #707
1 parent 97ab886 commit 064a0b3

4 files changed

Lines changed: 101 additions & 13 deletions

File tree

pyathena/model.py

Lines changed: 50 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -111,9 +111,13 @@ def __init__(self, response: dict[str, Any]) -> None:
111111
self._query_planning_time_in_millis: int | None = statistics.get(
112112
"QueryPlanningTimeInMillis", None
113113
)
114+
self._service_pre_processing_time_in_millis: int | None = statistics.get(
115+
"ServicePreProcessingTimeInMillis", None
116+
)
114117
self._service_processing_time_in_millis: int | None = statistics.get(
115118
"ServiceProcessingTimeInMillis", None
116119
)
120+
self._dpu_count: float | None = statistics.get("DpuCount")
117121
self._data_manifest_location: str | None = statistics.get("DataManifestLocation")
118122
reuse_info = statistics.get("ResultReuseInformation", {})
119123
self._reused_previous_result: bool | None = reuse_info.get("ReusedPreviousResult")
@@ -127,6 +131,24 @@ def __init__(self, response: dict[str, Any]) -> None:
127131
acl_conf = result_conf.get("AclConfiguration", {})
128132
self._s3_acl_option: str | None = acl_conf.get("S3AclOption")
129133

134+
managed_results_conf = query_execution.get("ManagedQueryResultsConfiguration", {})
135+
self._managed_query_results_enabled: bool | None = managed_results_conf.get("Enabled")
136+
managed_results_encryption_conf = managed_results_conf.get("EncryptionConfiguration", {})
137+
self._managed_query_results_kms_key: str | None = managed_results_encryption_conf.get(
138+
"KmsKey"
139+
)
140+
141+
s3_access_grants_conf = query_execution.get("QueryResultsS3AccessGrantsConfiguration", {})
142+
self._enable_s3_access_grants: bool | None = s3_access_grants_conf.get(
143+
"EnableS3AccessGrants"
144+
)
145+
self._create_user_level_prefix: bool | None = s3_access_grants_conf.get(
146+
"CreateUserLevelPrefix"
147+
)
148+
self._s3_access_grants_authentication_type: str | None = s3_access_grants_conf.get(
149+
"AuthenticationType"
150+
)
151+
130152
engine_version = query_execution.get("EngineVersion", {})
131153
self._selected_engine_version: str | None = engine_version.get(
132154
"SelectedEngineVersion", None
@@ -224,10 +246,18 @@ def total_execution_time_in_millis(self) -> int | None:
224246
def query_planning_time_in_millis(self) -> int | None:
225247
return self._query_planning_time_in_millis
226248

249+
@property
250+
def service_pre_processing_time_in_millis(self) -> int | None:
251+
return self._service_pre_processing_time_in_millis
252+
227253
@property
228254
def service_processing_time_in_millis(self) -> int | None:
229255
return self._service_processing_time_in_millis
230256

257+
@property
258+
def dpu_count(self) -> float | None:
259+
return self._dpu_count
260+
231261
@property
232262
def output_location(self) -> str | None:
233263
return self._output_location
@@ -272,6 +302,26 @@ def result_reuse_enabled(self) -> bool | None:
272302
def result_reuse_minutes(self) -> int | None:
273303
return self._result_reuse_minutes
274304

305+
@property
306+
def managed_query_results_enabled(self) -> bool | None:
307+
return self._managed_query_results_enabled
308+
309+
@property
310+
def managed_query_results_kms_key(self) -> str | None:
311+
return self._managed_query_results_kms_key
312+
313+
@property
314+
def enable_s3_access_grants(self) -> bool | None:
315+
return self._enable_s3_access_grants
316+
317+
@property
318+
def create_user_level_prefix(self) -> bool | None:
319+
return self._create_user_level_prefix
320+
321+
@property
322+
def s3_access_grants_authentication_type(self) -> str | None:
323+
return self._s3_access_grants_authentication_type
324+
275325

276326
class AthenaCalculationExecutionStatus:
277327
"""Status information for an Athena calculation execution.

pyproject.toml

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -5,8 +5,8 @@ authors = [
55
{name = "laughingman7743", email = "laughingman7743@gmail.com"},
66
]
77
dependencies = [
8-
"boto3>=1.26.4",
9-
"botocore>=1.29.4",
8+
"boto3>=1.38.2",
9+
"botocore>=1.41.2",
1010
"tenacity>=4.1.0",
1111
"fsspec",
1212
"python-dateutil",

tests/pyathena/test_model.py

Lines changed: 38 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -34,6 +34,17 @@ def test_init(self):
3434
"ResultReuseConfiguration": {
3535
"ResultReuseByAgeConfiguration": {"Enabled": True, "MaxAgeInMinutes": 5}
3636
},
37+
"ManagedQueryResultsConfiguration": {
38+
"Enabled": True,
39+
"EncryptionConfiguration": {
40+
"KmsKey": "test_managed_kms_key",
41+
},
42+
},
43+
"QueryResultsS3AccessGrantsConfiguration": {
44+
"EnableS3AccessGrants": True,
45+
"CreateUserLevelPrefix": True,
46+
"AuthenticationType": "DIRECTORY_IDENTITY",
47+
},
3748
"QueryExecutionContext": {
3849
"Database": "test_database",
3950
"Catalog": "test_catalog",
@@ -57,7 +68,9 @@ def test_init(self):
5768
"TotalExecutionTimeInMillis": 4567890,
5869
"QueryQueueTimeInMillis": 34567890,
5970
"QueryPlanningTimeInMillis": 567890,
71+
"ServicePreProcessingTimeInMillis": 12345,
6072
"ServiceProcessingTimeInMillis": 67890,
73+
"DpuCount": 4.0,
6174
"ResultReuseInformation": {"ReusedPreviousResult": True},
6275
},
6376
"WorkGroup": "test_work_group",
@@ -92,7 +105,9 @@ def test_init(self):
92105
assert actual.query_queue_time_in_millis == 34567890
93106
assert actual.total_execution_time_in_millis == 4567890
94107
assert actual.query_planning_time_in_millis == 567890
108+
assert actual.service_pre_processing_time_in_millis == 12345
95109
assert actual.service_processing_time_in_millis == 67890
110+
assert actual.dpu_count == 4.0
96111
assert actual.output_location == "s3://bucket/path/to/output/"
97112
assert actual.data_manifest_location == "s3://bucket/path/to/data_manifest/"
98113
assert actual.reused_previous_result
@@ -104,6 +119,29 @@ def test_init(self):
104119
assert actual.effective_engine_version == "Athena engine version 2"
105120
assert actual.result_reuse_enabled
106121
assert actual.result_reuse_minutes == 5
122+
assert actual.managed_query_results_enabled
123+
assert actual.managed_query_results_kms_key == "test_managed_kms_key"
124+
assert actual.enable_s3_access_grants
125+
assert actual.create_user_level_prefix
126+
assert actual.s3_access_grants_authentication_type == "DIRECTORY_IDENTITY"
127+
128+
def test_init_without_optional_fields(self):
129+
actual = AthenaQueryExecution(
130+
{
131+
"QueryExecution": {
132+
"QueryExecutionId": "12345678-90ab-cdef-1234-567890abcdef",
133+
"Query": "SELECT 1",
134+
"Status": {"State": "SUCCEEDED"},
135+
}
136+
}
137+
)
138+
assert actual.service_pre_processing_time_in_millis is None
139+
assert actual.dpu_count is None
140+
assert actual.managed_query_results_enabled is None
141+
assert actual.managed_query_results_kms_key is None
142+
assert actual.enable_s3_access_grants is None
143+
assert actual.create_user_level_prefix is None
144+
assert actual.s3_access_grants_authentication_type is None
107145

108146

109147
class TestAthenaCalculationExecution:

uv.lock

Lines changed: 11 additions & 11 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

0 commit comments

Comments
 (0)