Add tests, docs, CI for managed query result storage; replace assert with ProgrammingError

- Add test_fetch_all_rows to all cursor test files (Cursor, Pandas, Arrow, Polars, S3FS)
  parameterized by workgroup (default vs managed) using fixture indirect
- Support s3_staging_dir="" to explicitly disable env var fallback, required
  for managed workgroups where ResultConfiguration conflicts with
  ManagedQueryResultsConfiguration
- Replace all assert statements with ProgrammingError exceptions across
  connection.py, cursor files, result_set files (11 occurrences)
- Add AWS_ATHENA_MANAGED_WORKGROUP env var to tests/Env, GitHub Actions, docs
- Add ManagedWorkGroup resource to CloudFormation template

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

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/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/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,

pyathena/pandas/cursor.py

@@ -238,7 +238,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/polars/async_cursor.py

@@ -223,7 +223,8 @@ def execute(
         """
         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/polars/cursor.py

@@ -250,7 +250,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,