Merge pull request #708 from pyathena-dev/feature/add-newer-query-execution-fields

feat: support newer GetQueryExecution response fields

Author: laughingman7743

pyathena/model.py

@@ -111,9 +111,13 @@ def __init__(self, response: dict[str, Any]) -> None:
         self._query_planning_time_in_millis: int | None = statistics.get(
             "QueryPlanningTimeInMillis", None
         )
+        self._service_pre_processing_time_in_millis: int | None = statistics.get(
+            "ServicePreProcessingTimeInMillis", None
+        )
         self._service_processing_time_in_millis: int | None = statistics.get(
             "ServiceProcessingTimeInMillis", None
         )
+        self._dpu_count: float | None = statistics.get("DpuCount")
         self._data_manifest_location: str | None = statistics.get("DataManifestLocation")
         reuse_info = statistics.get("ResultReuseInformation", {})
         self._reused_previous_result: bool | None = reuse_info.get("ReusedPreviousResult")
@@ -127,6 +131,24 @@ def __init__(self, response: dict[str, Any]) -> None:
         acl_conf = result_conf.get("AclConfiguration", {})
         self._s3_acl_option: str | None = acl_conf.get("S3AclOption")
 
+        managed_results_conf = query_execution.get("ManagedQueryResultsConfiguration", {})
+        self._managed_query_results_enabled: bool | None = managed_results_conf.get("Enabled")
+        managed_results_encryption_conf = managed_results_conf.get("EncryptionConfiguration", {})
+        self._managed_query_results_kms_key: str | None = managed_results_encryption_conf.get(
+            "KmsKey"
+        )
+
+        s3_access_grants_conf = query_execution.get("QueryResultsS3AccessGrantsConfiguration", {})
+        self._enable_s3_access_grants: bool | None = s3_access_grants_conf.get(
+            "EnableS3AccessGrants"
+        )
+        self._create_user_level_prefix: bool | None = s3_access_grants_conf.get(
+            "CreateUserLevelPrefix"
+        )
+        self._s3_access_grants_authentication_type: str | None = s3_access_grants_conf.get(
+            "AuthenticationType"
+        )
+
         engine_version = query_execution.get("EngineVersion", {})
         self._selected_engine_version: str | None = engine_version.get(
             "SelectedEngineVersion", None
@@ -224,10 +246,18 @@ def total_execution_time_in_millis(self) -> int | None:
     def query_planning_time_in_millis(self) -> int | None:
         return self._query_planning_time_in_millis
 
+    @property
+    def service_pre_processing_time_in_millis(self) -> int | None:
+        return self._service_pre_processing_time_in_millis
+
     @property
     def service_processing_time_in_millis(self) -> int | None:
         return self._service_processing_time_in_millis
 
+    @property
+    def dpu_count(self) -> float | None:
+        return self._dpu_count
+
     @property
     def output_location(self) -> str | None:
         return self._output_location
@@ -272,6 +302,26 @@ def result_reuse_enabled(self) -> bool | None:
     def result_reuse_minutes(self) -> int | None:
         return self._result_reuse_minutes
 
+    @property
+    def managed_query_results_enabled(self) -> bool | None:
+        return self._managed_query_results_enabled
+
+    @property
+    def managed_query_results_kms_key(self) -> str | None:
+        return self._managed_query_results_kms_key
+
+    @property
+    def enable_s3_access_grants(self) -> bool | None:
+        return self._enable_s3_access_grants
+
+    @property
+    def create_user_level_prefix(self) -> bool | None:
+        return self._create_user_level_prefix
+
+    @property
+    def s3_access_grants_authentication_type(self) -> str | None:
+        return self._s3_access_grants_authentication_type
+
 
 class AthenaCalculationExecutionStatus:
     """Status information for an Athena calculation execution.

pyproject.toml

@@ -5,8 +5,8 @@ authors = [
     {name = "laughingman7743", email = "laughingman7743@gmail.com"},
 ]
 dependencies = [
-    "boto3>=1.26.4",
-    "botocore>=1.29.4",
+    "boto3>=1.38.2",
+    "botocore>=1.41.2",
     "tenacity>=4.1.0",
     "fsspec",
     "python-dateutil",

tests/pyathena/test_model.py

@@ -34,6 +34,17 @@ def test_init(self):
                     "ResultReuseConfiguration": {
                         "ResultReuseByAgeConfiguration": {"Enabled": True, "MaxAgeInMinutes": 5}
                     },
+                    "ManagedQueryResultsConfiguration": {
+                        "Enabled": True,
+                        "EncryptionConfiguration": {
+                            "KmsKey": "test_managed_kms_key",
+                        },
+                    },
+                    "QueryResultsS3AccessGrantsConfiguration": {
+                        "EnableS3AccessGrants": True,
+                        "CreateUserLevelPrefix": True,
+                        "AuthenticationType": "DIRECTORY_IDENTITY",
+                    },
                     "QueryExecutionContext": {
                         "Database": "test_database",
                         "Catalog": "test_catalog",
@@ -57,7 +68,9 @@ def test_init(self):
                         "TotalExecutionTimeInMillis": 4567890,
                         "QueryQueueTimeInMillis": 34567890,
                         "QueryPlanningTimeInMillis": 567890,
+                        "ServicePreProcessingTimeInMillis": 12345,
                         "ServiceProcessingTimeInMillis": 67890,
+                        "DpuCount": 4.0,
                         "ResultReuseInformation": {"ReusedPreviousResult": True},
                     },
                     "WorkGroup": "test_work_group",
@@ -92,7 +105,9 @@ def test_init(self):
         assert actual.query_queue_time_in_millis == 34567890
         assert actual.total_execution_time_in_millis == 4567890
         assert actual.query_planning_time_in_millis == 567890
+        assert actual.service_pre_processing_time_in_millis == 12345
         assert actual.service_processing_time_in_millis == 67890
+        assert actual.dpu_count == 4.0
         assert actual.output_location == "s3://bucket/path/to/output/"
         assert actual.data_manifest_location == "s3://bucket/path/to/data_manifest/"
         assert actual.reused_previous_result
@@ -104,6 +119,29 @@ def test_init(self):
         assert actual.effective_engine_version == "Athena engine version 2"
         assert actual.result_reuse_enabled
         assert actual.result_reuse_minutes == 5
+        assert actual.managed_query_results_enabled
+        assert actual.managed_query_results_kms_key == "test_managed_kms_key"
+        assert actual.enable_s3_access_grants
+        assert actual.create_user_level_prefix
+        assert actual.s3_access_grants_authentication_type == "DIRECTORY_IDENTITY"
+
+    def test_init_without_optional_fields(self):
+        actual = AthenaQueryExecution(
+            {
+                "QueryExecution": {
+                    "QueryExecutionId": "12345678-90ab-cdef-1234-567890abcdef",
+                    "Query": "SELECT 1",
+                    "Status": {"State": "SUCCEEDED"},
+                }
+            }
+        )
+        assert actual.service_pre_processing_time_in_millis is None
+        assert actual.dpu_count is None
+        assert actual.managed_query_results_enabled is None
+        assert actual.managed_query_results_kms_key is None
+        assert actual.enable_s3_access_grants is None
+        assert actual.create_user_level_prefix is None
+        assert actual.s3_access_grants_authentication_type is None
 
 
 class TestAthenaCalculationExecution: