Merge pull request #665 from pyathena-dev/fix/managed-query-result-storage

Support Athena managed query result storage

Author: laughingman7743

.github/workflows/test.yaml

@@ -18,6 +18,7 @@ jobs:
       AWS_ATHENA_S3_STAGING_DIR: s3://laughingman7743-pyathena/github/
       AWS_ATHENA_WORKGROUP: pyathena
       AWS_ATHENA_SPARK_WORKGROUP: pyathena-spark
+      AWS_ATHENA_MANAGED_WORKGROUP: pyathena-managed
 
     strategy:
       fail-fast: false

docs/testing.md

@@ -20,6 +20,15 @@ If primary is not available as the default workgroup, specify an alternative wor
 $ export AWS_ATHENA_DEFAULT_WORKGROUP=DEFAULT_WORKGROUP
 ```
 
+### Managed query result storage (optional)
+
+To test the managed query result storage feature, create a workgroup with managed storage enabled and set the `AWS_ATHENA_MANAGED_WORKGROUP` environment variable.
+If not set, managed storage tests will be skipped.
+
+```bash
+$ export AWS_ATHENA_MANAGED_WORKGROUP=pyathena-managed
+```
+
 ## Run test
 
 ```bash

docs/usage.md

@@ -14,6 +14,36 @@ print(cursor.description)
 print(cursor.fetchall())
 ```
 
+## Managed query result storage
+
+When using a workgroup with [managed query result storage](https://docs.aws.amazon.com/athena/latest/ug/managed-results.html) enabled,
+you don't need to specify an S3 staging directory.
+
+```python
+from pyathena import connect
+
+cursor = connect(work_group="YOUR_MANAGED_WORK_GROUP",
+                 region_name="us-west-2").cursor()
+cursor.execute("SELECT * FROM one_row")
+print(cursor.fetchall())
+```
+
+If the ``AWS_ATHENA_S3_STAGING_DIR`` environment variable is set, pass ``s3_staging_dir=""``
+to explicitly disable the fallback. Otherwise the API will reject the request because
+``ResultConfiguration`` and ``ManagedQueryResultsConfiguration`` cannot be set together.
+
+```python
+cursor = connect(work_group="YOUR_MANAGED_WORK_GROUP",
+                 s3_staging_dir="",
+                 region_name="us-west-2").cursor()
+```
+
+```{note}
+With managed query result storage, query results are retrieved via the `GetQueryResults` API
+(1000 rows per request) instead of reading S3 files directly. This may be slower for large
+result sets. For large datasets, consider using customer-managed storage or the `UNLOAD` statement.
+```
+
 ## Cursor iteration
 
 ```python
@@ -366,7 +396,7 @@ Support [Boto3 environment variables](https://boto3.amazonaws.com/v1/documentati
 ### Additional environment variables
 
 AWS_ATHENA_S3_STAGING_DIR
-: The S3 location where Athena automatically stores the query results and metadata information. Required if you have not set up workgroups. Not required if a workgroup has been set up.
+: The S3 location where Athena automatically stores the query results and metadata information. Required if you have not set up workgroups. Not required if a workgroup has been set up. When connecting to a workgroup with [managed query result storage](https://docs.aws.amazon.com/athena/latest/ug/managed-results.html), pass ``s3_staging_dir=""`` to explicitly disable this environment variable fallback (see [Managed query result storage](#managed-query-result-storage)).
 
 AWS_ATHENA_WORK_GROUP
 : The setting of the workgroup to execute the query.

pyathena/__init__.py

@@ -85,6 +85,9 @@ def connect(*args, **kwargs) -> "Connection[Any]":
     Args:
         s3_staging_dir: S3 location to store query results. Required if not
             using workgroups or if the workgroup doesn't have a result location.
+            Pass an empty string to explicitly disable S3 staging and skip
+            the ``AWS_ATHENA_S3_STAGING_DIR`` environment variable fallback
+            (required for workgroups with managed query result storage).
         region_name: AWS region name. If not specified, uses the default region
             from your AWS configuration.
         schema_name: Athena database/schema name. Defaults to "default".
@@ -109,7 +112,7 @@ def connect(*args, **kwargs) -> "Connection[Any]":
         A Connection object that can be used to create cursors and execute queries.
 
     Raises:
-        AssertionError: If neither s3_staging_dir nor work_group is provided.
+        ProgrammingError: If neither s3_staging_dir nor work_group is provided.
 
     Example:
         >>> import pyathena

pyathena/arrow/async_cursor.py

@@ -184,7 +184,8 @@ def execute(
     ) -> Tuple[str, "Future[Union[AthenaArrowResultSet, Any]]"]:
         if self._unload:
             s3_staging_dir = s3_staging_dir if s3_staging_dir else self._s3_staging_dir
-            assert s3_staging_dir, "If the unload option is used, s3_staging_dir is required."
+            if not s3_staging_dir:
+                raise ProgrammingError("If the unload option is used, s3_staging_dir is required.")
             operation, unload_location = self._formatter.wrap_unload(
                 operation,
                 s3_staging_dir=s3_staging_dir,

pyathena/arrow/cursor.py

@@ -209,7 +209,8 @@ def execute(
         self._reset_state()
         if self._unload:
             s3_staging_dir = s3_staging_dir if s3_staging_dir else self._s3_staging_dir
-            assert s3_staging_dir, "If the unload option is used, s3_staging_dir is required."
+            if not s3_staging_dir:
+                raise ProgrammingError("If the unload option is used, s3_staging_dir is required.")
             operation, unload_location = self._formatter.wrap_unload(
                 operation,
                 s3_staging_dir=s3_staging_dir,

pyathena/arrow/result_set.py

@@ -117,6 +117,8 @@ def __init__(
         self._fs = self.__s3_file_system()
         if self.state == AthenaQueryExecution.STATE_SUCCEEDED and self.output_location:
             self._table = self._as_arrow()
+        elif self.state == AthenaQueryExecution.STATE_SUCCEEDED:
+            self._table = self._as_arrow_from_api()
         else:
             import pyarrow as pa
 
@@ -346,6 +348,25 @@ def _as_arrow(self) -> "Table":
             table = self._read_csv()
         return table
 
+    def _as_arrow_from_api(self, converter: Optional[Converter] = None) -> "Table":
+        """Build an Arrow Table from GetQueryResults API.
+
+        Used as a fallback when ``output_location`` is not available
+        (e.g. managed query result storage).
+
+        Args:
+            converter: Type converter for result values. Defaults to
+                ``DefaultTypeConverter`` if not specified.
+        """
+        import pyarrow as pa
+
+        rows = self._fetch_all_rows(converter)
+        if not rows:
+            return pa.Table.from_pydict({})
+        description = self.description if self.description else []
+        columns = [d[0] for d in description]
+        return pa.table(self._rows_to_columnar(rows, columns))
+
     def as_arrow(self) -> "Table":
         return self._table
 

pyathena/common.py

@@ -210,15 +210,15 @@ def _build_start_query_execution_request(
         request: Dict[str, Any] = {
             "QueryString": query,
             "QueryExecutionContext": {},
-            "ResultConfiguration": {},
         }
         if self._schema_name:
             request["QueryExecutionContext"].update({"Database": self._schema_name})
         if self._catalog_name:
             request["QueryExecutionContext"].update({"Catalog": self._catalog_name})
+        result_configuration: Dict[str, Any] = {}
         if self._s3_staging_dir or s3_staging_dir:
-            request["ResultConfiguration"].update(
-                {"OutputLocation": s3_staging_dir if s3_staging_dir else self._s3_staging_dir}
+            result_configuration["OutputLocation"] = (
+                s3_staging_dir if s3_staging_dir else self._s3_staging_dir
             )
         if self._work_group or work_group:
             request.update({"WorkGroup": work_group if work_group else self._work_group})
@@ -228,7 +228,9 @@ def _build_start_query_execution_request(
             }
             if self._kms_key:
                 enc_conf.update({"KmsKey": self._kms_key})
-            request["ResultConfiguration"].update({"EncryptionConfiguration": enc_conf})
+            result_configuration["EncryptionConfiguration"] = enc_conf
+        if result_configuration:
+            request["ResultConfiguration"] = result_configuration
         if self._result_reuse_enable or result_reuse_enable:
             reuse_conf = {
                 "Enabled": result_reuse_enable

pyathena/connection.py

@@ -26,7 +26,7 @@
 from pyathena.common import BaseCursor, CursorIterator
 from pyathena.converter import Converter
 from pyathena.cursor import Cursor
-from pyathena.error import NotSupportedError
+from pyathena.error import NotSupportedError, ProgrammingError
 from pyathena.formatter import DefaultParameterFormatter, Formatter
 from pyathena.util import RetryConfig
 
@@ -77,7 +77,9 @@ class Connection(Generic[ConnectionCursor]):
     Note:
         Either s3_staging_dir or work_group must be specified. If using a
         workgroup, it must have a result location configured unless
-        s3_staging_dir is also provided.
+        s3_staging_dir is also provided. For workgroups with managed query
+        result storage, pass ``s3_staging_dir=""`` to skip the environment
+        variable fallback.
     """
 
     _ENV_S3_STAGING_DIR: str = "AWS_ATHENA_S3_STAGING_DIR"
@@ -198,6 +200,10 @@ def __init__(
         Args:
             s3_staging_dir: S3 location to store query results. Required if not
                 using workgroups or if workgroup doesn't have result location.
+                Pass an empty string to explicitly disable S3 staging and skip
+                the ``AWS_ATHENA_S3_STAGING_DIR`` environment variable fallback.
+                This is required when connecting to a workgroup with managed
+                query result storage enabled.
             region_name: AWS region name. Uses default region if not specified.
             schema_name: Default database/schema name. Defaults to "default".
             catalog_name: Data catalog name. Defaults to "awsdatacatalog".
@@ -226,12 +232,17 @@ def __init__(
             **kwargs: Additional arguments passed to boto3 Session and client.
 
         Raises:
-            AssertionError: If neither s3_staging_dir nor work_group is provided.
+            ProgrammingError: If neither s3_staging_dir nor work_group is provided.
 
         Note:
             Either s3_staging_dir or work_group must be specified. Environment
             variables AWS_ATHENA_S3_STAGING_DIR and AWS_ATHENA_WORK_GROUP are
             checked if parameters are not provided.
+
+            When using a workgroup with managed query result storage, pass
+            ``s3_staging_dir=""`` to prevent the environment variable fallback
+            from sending a ``ResultConfiguration`` that conflicts with
+            ``ManagedQueryResultsConfiguration``.
         """
         self._kwargs = {
             **kwargs,
@@ -241,8 +252,8 @@ def __init__(
             "serial_number": serial_number,
             "duration_seconds": duration_seconds,
         }
-        if s3_staging_dir:
-            self.s3_staging_dir: Optional[str] = s3_staging_dir
+        if s3_staging_dir is not None:
+            self.s3_staging_dir: Optional[str] = s3_staging_dir or None
         else:
             self.s3_staging_dir = os.getenv(self._ENV_S3_STAGING_DIR)
         self.region_name = region_name
@@ -258,9 +269,8 @@ def __init__(
         self.profile_name = profile_name
         self.config: Optional[Config] = config if config else Config()
 
-        assert self.s3_staging_dir or self.work_group, (
-            "Required argument `s3_staging_dir` or `work_group` not found."
-        )
+        if not self.s3_staging_dir and not self.work_group:
+            raise ProgrammingError("Required argument `s3_staging_dir` or `work_group` not found.")
 
         if self.s3_staging_dir and not self.s3_staging_dir.endswith("/"):
             self.s3_staging_dir = f"{self.s3_staging_dir}/"

pyathena/pandas/async_cursor.py

@@ -161,7 +161,8 @@ def execute(
     ) -> Tuple[str, "Future[Union[AthenaPandasResultSet, Any]]"]:
         if self._unload:
             s3_staging_dir = s3_staging_dir if s3_staging_dir else self._s3_staging_dir
-            assert s3_staging_dir, "If the unload option is used, s3_staging_dir is required."
+            if not s3_staging_dir:
+                raise ProgrammingError("If the unload option is used, s3_staging_dir is required.")
             operation, unload_location = self._formatter.wrap_unload(
                 operation,
                 s3_staging_dir=s3_staging_dir,