PyAthena 3.27.1__tar.gz → 3.28.0__tar.gz

{pyathena-3.27.1 → pyathena-3.28.0}/PKG-INFO

@@ -1,9 +1,11 @@
 Metadata-Version: 2.4
 Name: PyAthena
-Version: 3.27.1
+Version: 3.28.0
 Summary: Python DB API 2.0 (PEP 249) client for Amazon Athena
-Project-URL: homepage, https://github.com/laughingman7743/PyAthena/
-Project-URL: repository, https://github.com/laughingman7743/PyAthena/
+Project-URL: homepage, https://github.com/pyathena-dev/PyAthena/
+Project-URL: repository, https://github.com/pyathena-dev/PyAthena/
+Project-URL: documentation, https://pyathena.dev/
+Project-URL: issues, https://github.com/pyathena-dev/PyAthena/issues
 Author-email: laughingman7743 <laughingman7743@gmail.com>
 License: Copyright 2017 laughingman7743
         
@@ -57,14 +59,14 @@ Description-Content-Type: text/markdown
 
 <div align="center">
 
-<img src="https://raw.githubusercontent.com/laughingman7743/PyAthena/master/docs/_static/icon.png" alt="PyAthena logo" width="250">
+<img src="https://raw.githubusercontent.com/pyathena-dev/PyAthena/master/docs/_static/icon.png" alt="PyAthena logo" width="250">
 
 [![PyPI - Version](https://badge.fury.io/py/pyathena.svg)](https://badge.fury.io/py/pyathena)
 [![PyPI - Python Version](https://img.shields.io/pypi/pyversions/PyAthena.svg)](https://pypi.org/project/PyAthena/)
 [![PyPI - Downloads](https://static.pepy.tech/badge/pyathena/month)](https://pepy.tech/project/pyathena)
-[![CI - Test](https://github.com/laughingman7743/PyAthena/actions/workflows/test.yaml/badge.svg)](https://github.com/laughingman7743/PyAthena/actions/workflows/test.yaml)
-[![CD - Docs](https://github.com/laughingman7743/PyAthena/actions/workflows/docs.yaml/badge.svg)](https://github.com/laughingman7743/PyAthena/actions/workflows/docs.yaml)
-[![License - MIT](https://img.shields.io/pypi/l/PyAthena.svg)](https://github.com/laughingman7743/PyAthena/blob/master/LICENSE)
+[![CI - Test](https://github.com/pyathena-dev/PyAthena/actions/workflows/test.yaml/badge.svg)](https://github.com/pyathena-dev/PyAthena/actions/workflows/test.yaml)
+[![CD - Docs](https://github.com/pyathena-dev/PyAthena/actions/workflows/docs.yaml/badge.svg)](https://github.com/pyathena-dev/PyAthena/actions/workflows/docs.yaml)
+[![License - MIT](https://img.shields.io/pypi/l/PyAthena.svg)](https://github.com/pyathena-dev/PyAthena/blob/master/LICENSE)
 [![linting - Ruff](https://img.shields.io/endpoint?url=https://raw.githubusercontent.com/astral-sh/ruff/main/assets/badge/v2.json)](https://github.com/astral-sh/ruff)
 [![types - Mypy](https://www.mypy-lang.org/static/mypy_badge.svg)](https://mypy-lang.org/)
 
@@ -115,10 +117,10 @@ Many of the implementations in this library are based on [PyHive](https://github
 
 ## Links
 
-- Documentation: https://laughingman7743.github.io/PyAthena/
+- Documentation: https://pyathena.dev/
 - PyPI Releases: https://pypi.org/project/PyAthena/
-- Source Code: https://github.com/laughingman7743/PyAthena/
-- Issue Tracker: https://github.com/laughingman7743/PyAthena/issues
+- Source Code: https://github.com/pyathena-dev/PyAthena/
+- Issue Tracker: https://github.com/pyathena-dev/PyAthena/issues
 
 ## Logo
 

{pyathena-3.27.1 → pyathena-3.28.0}/README.md

@@ -2,14 +2,14 @@
 
 <div align="center">
 
-<img src="https://raw.githubusercontent.com/laughingman7743/PyAthena/master/docs/_static/icon.png" alt="PyAthena logo" width="250">
+<img src="https://raw.githubusercontent.com/pyathena-dev/PyAthena/master/docs/_static/icon.png" alt="PyAthena logo" width="250">
 
 [![PyPI - Version](https://badge.fury.io/py/pyathena.svg)](https://badge.fury.io/py/pyathena)
 [![PyPI - Python Version](https://img.shields.io/pypi/pyversions/PyAthena.svg)](https://pypi.org/project/PyAthena/)
 [![PyPI - Downloads](https://static.pepy.tech/badge/pyathena/month)](https://pepy.tech/project/pyathena)
-[![CI - Test](https://github.com/laughingman7743/PyAthena/actions/workflows/test.yaml/badge.svg)](https://github.com/laughingman7743/PyAthena/actions/workflows/test.yaml)
-[![CD - Docs](https://github.com/laughingman7743/PyAthena/actions/workflows/docs.yaml/badge.svg)](https://github.com/laughingman7743/PyAthena/actions/workflows/docs.yaml)
-[![License - MIT](https://img.shields.io/pypi/l/PyAthena.svg)](https://github.com/laughingman7743/PyAthena/blob/master/LICENSE)
+[![CI - Test](https://github.com/pyathena-dev/PyAthena/actions/workflows/test.yaml/badge.svg)](https://github.com/pyathena-dev/PyAthena/actions/workflows/test.yaml)
+[![CD - Docs](https://github.com/pyathena-dev/PyAthena/actions/workflows/docs.yaml/badge.svg)](https://github.com/pyathena-dev/PyAthena/actions/workflows/docs.yaml)
+[![License - MIT](https://img.shields.io/pypi/l/PyAthena.svg)](https://github.com/pyathena-dev/PyAthena/blob/master/LICENSE)
 [![linting - Ruff](https://img.shields.io/endpoint?url=https://raw.githubusercontent.com/astral-sh/ruff/main/assets/badge/v2.json)](https://github.com/astral-sh/ruff)
 [![types - Mypy](https://www.mypy-lang.org/static/mypy_badge.svg)](https://mypy-lang.org/)
 
@@ -60,10 +60,10 @@ Many of the implementations in this library are based on [PyHive](https://github
 
 ## Links
 
-- Documentation: https://laughingman7743.github.io/PyAthena/
+- Documentation: https://pyathena.dev/
 - PyPI Releases: https://pypi.org/project/PyAthena/
-- Source Code: https://github.com/laughingman7743/PyAthena/
-- Issue Tracker: https://github.com/laughingman7743/PyAthena/issues
+- Source Code: https://github.com/pyathena-dev/PyAthena/
+- Issue Tracker: https://github.com/pyathena-dev/PyAthena/issues
 
 ## Logo
 

{pyathena-3.27.1 → pyathena-3.28.0}/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-3.27.1 → pyathena-3.28.0}/pyathena/_version.py

@@ -28,7 +28,7 @@ version_tuple: VERSION_TUPLE
 commit_id: COMMIT_ID
 __commit_id__: COMMIT_ID
 
-__version__ = version = '3.27.1'
-__version_tuple__ = version_tuple = (3, 27, 1)
+__version__ = version = '3.28.0'
+__version_tuple__ = version_tuple = (3, 28, 0)
 
 __commit_id__ = commit_id = None

{pyathena-3.27.1 → pyathena-3.28.0}/pyathena/arrow/async_cursor.py

@@ -184,7 +184,8 @@ class AsyncArrowCursor(AsyncCursor):
     ) -> 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-3.27.1 → pyathena-3.28.0}/pyathena/arrow/cursor.py

@@ -209,7 +209,8 @@ class ArrowCursor(BaseCursor, CursorIterator, WithResultSet):
         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-3.27.1 → pyathena-3.28.0}/pyathena/arrow/result_set.py

@@ -117,6 +117,8 @@ class AthenaArrowResultSet(AthenaResultSet):
         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 @@ class AthenaArrowResultSet(AthenaResultSet):
             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-3.27.1 → pyathena-3.28.0}/pyathena/common.py

@@ -210,15 +210,15 @@ class BaseCursor(metaclass=ABCMeta):
         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 @@ class BaseCursor(metaclass=ABCMeta):
             }
             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-3.27.1 → pyathena-3.28.0}/pyathena/connection.py

@@ -26,7 +26,7 @@ import pyathena
 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 @@ class Connection(Generic[ConnectionCursor]):
         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 @@ class Connection(Generic[ConnectionCursor]):
             **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 @@ class Connection(Generic[ConnectionCursor]):
             "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 @@ class Connection(Generic[ConnectionCursor]):
         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-3.27.1 → pyathena-3.28.0}/pyathena/pandas/async_cursor.py

@@ -161,7 +161,8 @@ class AsyncPandasCursor(AsyncCursor):
     ) -> 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-3.27.1 → pyathena-3.28.0}/pyathena/pandas/cursor.py

@@ -238,7 +238,8 @@ class PandasCursor(BaseCursor, CursorIterator, WithResultSet):
         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-3.27.1 → pyathena-3.28.0}/pyathena/pandas/result_set.py

@@ -296,6 +296,9 @@ class AthenaPandasResultSet(AthenaResultSet):
             df = self._as_pandas()
             trunc_date = _no_trunc_date if self.is_unload else self._trunc_date
             self._df_iter = PandasDataFrameIterator(df, trunc_date)
+        elif self.state == AthenaQueryExecution.STATE_SUCCEEDED:
+            df = self._as_pandas_from_api()
+            self._df_iter = PandasDataFrameIterator(df, self._trunc_date)
         else:
             import pandas as pd
 
@@ -678,6 +681,25 @@ class AthenaPandasResultSet(AthenaResultSet):
             df = self._read_csv()
         return df
 
+    def _as_pandas_from_api(self, converter: Optional[Converter] = None) -> "DataFrame":
+        """Build a DataFrame 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 pandas as pd
+
+        rows = self._fetch_all_rows(converter)
+        if not rows:
+            return pd.DataFrame()
+        description = self.description if self.description else []
+        columns = [d[0] for d in description]
+        return pd.DataFrame(self._rows_to_columnar(rows, columns))
+
     def as_pandas(self) -> Union[PandasDataFrameIterator, "DataFrame"]:
         if self._chunksize is None:
             return next(self._df_iter)

{pyathena-3.27.1 → pyathena-3.28.0}/pyathena/polars/async_cursor.py

@@ -223,7 +223,8 @@ class AsyncPolarsCursor(AsyncCursor):
         """
         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-3.27.1 → pyathena-3.28.0}/pyathena/polars/cursor.py

@@ -250,7 +250,8 @@ class PolarsCursor(BaseCursor, CursorIterator, WithResultSet):
         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-3.27.1 → pyathena-3.28.0}/pyathena/polars/result_set.py

@@ -251,6 +251,9 @@ class AthenaPolarsResultSet(AthenaResultSet):
         # _metadata for unload queries, so we must cache column names AFTER this.
         if self.state == AthenaQueryExecution.STATE_SUCCEEDED and self.output_location:
             self._df_iter = self._create_dataframe_iterator()
+        elif self.state == AthenaQueryExecution.STATE_SUCCEEDED:
+            df = self._as_polars_from_api()
+            self._df_iter = PolarsDataFrameIterator(df, self.converters, self._get_column_names())
         else:
             import polars as pl
 
@@ -454,8 +457,8 @@ class AthenaPolarsResultSet(AthenaResultSet):
         if not self._is_csv_readable():
             return pl.DataFrame()
 
-        # After validation, output_location is guaranteed to be set
-        assert self.output_location is not None
+        if self.output_location is None:
+            raise ProgrammingError("output_location is not available.")
 
         separator, has_header, new_columns = self._get_csv_params()
 
@@ -489,8 +492,8 @@ class AthenaPolarsResultSet(AthenaResultSet):
         if not self._prepare_parquet_location():
             return pl.DataFrame()
 
-        # After preparation, unload_location is guaranteed to be set
-        assert self._unload_location is not None
+        if self._unload_location is None:
+            raise ProgrammingError("unload_location is not available.")
 
         try:
             return pl.read_parquet(
@@ -539,6 +542,25 @@ class AthenaPolarsResultSet(AthenaResultSet):
             df = self._read_csv()
         return df
 
+    def _as_polars_from_api(self, converter: Optional[Converter] = None) -> "pl.DataFrame":
+        """Build a Polars DataFrame 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 polars as pl
+
+        rows = self._fetch_all_rows(converter)
+        if not rows:
+            return pl.DataFrame()
+        description = self.description if self.description else []
+        columns = [d[0] for d in description]
+        return pl.DataFrame(self._rows_to_columnar(rows, columns))
+
     def as_polars(self) -> "pl.DataFrame":
         """Return query results as a Polars DataFrame.
 
@@ -618,8 +640,8 @@ class AthenaPolarsResultSet(AthenaResultSet):
         if not self._is_csv_readable():
             return
 
-        # After validation, output_location is guaranteed to be set
-        assert self.output_location is not None
+        if self.output_location is None:
+            raise ProgrammingError("output_location is not available.")
 
         separator, has_header, new_columns = self._get_csv_params()
 
@@ -656,8 +678,8 @@ class AthenaPolarsResultSet(AthenaResultSet):
         if not self._prepare_parquet_location():
             return
 
-        # After preparation, unload_location is guaranteed to be set
-        assert self._unload_location is not None
+        if self._unload_location is None:
+            raise ProgrammingError("unload_location is not available.")
 
         try:
             lazy_df = pl.scan_parquet(

{pyathena-3.27.1 → pyathena-3.28.0}/pyathena/result_set.py

@@ -19,7 +19,7 @@ from typing import (
 )
 
 from pyathena.common import CursorIterator
-from pyathena.converter import Converter
+from pyathena.converter import Converter, DefaultTypeConverter
 from pyathena.error import DataError, OperationalError, ProgrammingError
 from pyathena.model import AthenaQueryExecution
 from pyathena.util import RetryConfig, parse_output_location, retry_api_call
@@ -73,7 +73,8 @@ class AthenaResultSet(CursorIterator):
         self._connection: Optional["Connection[Any]"] = connection
         self._converter = converter
         self._query_execution: Optional[AthenaQueryExecution] = query_execution
-        assert self._query_execution, "Required argument `query_execution` not found."
+        if not self._query_execution:
+            raise ProgrammingError("Required argument `query_execution` not found.")
         self._retry_config = retry_config
         self._client = connection.session.client(
             "s3",
@@ -328,19 +329,21 @@ class AthenaResultSet(CursorIterator):
             raise ProgrammingError("AthenaResultSet is closed.")
         return cast("Connection[Any]", self._connection)
 
-    def __fetch(self, next_token: Optional[str] = None) -> Dict[str, Any]:
+    def __get_query_results(
+        self, max_results: int, next_token: Optional[str] = None
+    ) -> Dict[str, Any]:
         if not self.query_id:
             raise ProgrammingError("QueryExecutionId is none or empty.")
         if self.state != AthenaQueryExecution.STATE_SUCCEEDED:
             raise ProgrammingError("QueryExecutionState is not SUCCEEDED.")
         if self.is_closed:
             raise ProgrammingError("AthenaResultSet is closed.")
-        request = {
+        request: Dict[str, Any] = {
             "QueryExecutionId": self.query_id,
-            "MaxResults": self._arraysize,
+            "MaxResults": max_results,
         }
         if next_token:
-            request.update({"NextToken": next_token})
+            request["NextToken"] = next_token
         try:
             response = retry_api_call(
                 self.connection.client.get_query_results,
@@ -354,17 +357,23 @@ class AthenaResultSet(CursorIterator):
         else:
             return cast(Dict[str, Any], response)
 
+    def __fetch(self, next_token: Optional[str] = None) -> Dict[str, Any]:
+        return self.__get_query_results(self._arraysize, next_token)
+
     def _fetch(self) -> None:
         if not self._next_token:
             raise ProgrammingError("NextToken is none or empty.")
         response = self.__fetch(self._next_token)
-        self._process_rows(response)
+        rows, self._next_token = self._parse_result_rows(response)
+        self._process_rows(rows)
 
     def _pre_fetch(self) -> None:
         response = self.__fetch()
         self._process_metadata(response)
         self._process_update_count(response)
-        self._process_rows(response)
+        rows, self._next_token = self._parse_result_rows(response)
+        offset = 1 if rows and self._is_first_row_column_labels(rows) else 0
+        self._process_rows(rows, offset)
 
     def fetchone(
         self,
@@ -433,41 +442,131 @@ class AthenaResultSet(CursorIterator):
             self._rowcount = update_count
 
     def _get_rows(
-        self, offset: int, metadata: Tuple[Any, ...], rows: List[Dict[str, Any]]
+        self,
+        offset: int,
+        metadata: Tuple[Any, ...],
+        rows: List[Dict[str, Any]],
+        converter: Optional[Converter] = None,
     ) -> List[Union[Tuple[Optional[Any], ...], Dict[Any, Optional[Any]]]]:
+        conv = converter or self._converter
         return [
             tuple(
                 [
-                    self._converter.convert(meta.get("Type"), row.get("VarCharValue"))
+                    conv.convert(meta.get("Type"), row.get("VarCharValue"))
                     for meta, row in zip(metadata, rows[i].get("Data", []), strict=False)
                 ]
             )
             for i in range(offset, len(rows))
         ]
 
-    def _process_rows(self, response: Dict[str, Any]) -> None:
+    def _parse_result_rows(
+        self, response: Dict[str, Any]
+    ) -> Tuple[List[Dict[str, Any]], Optional[str]]:
+        """Parse a GetQueryResults response into raw rows and next token.
+
+        Handles response validation and pagination token extraction.
+        This is the shared parsing logic used by both ``_pre_fetch``
+        (normal path) and ``_fetch_all_rows`` (API fallback).
+
+        Args:
+            response: Raw response dict from ``GetQueryResults`` API.
+
+        Returns:
+            Tuple of (rows, next_token).
+        """
         result_set = response.get("ResultSet")
         if not result_set:
             raise DataError("KeyError `ResultSet`")
         rows = result_set.get("Rows")
         if rows is None:
             raise DataError("KeyError `Rows`")
-        processed_rows = []
-        if len(rows) > 0:
-            offset = 1 if not self._next_token and self._is_first_row_column_labels(rows) else 0
-            metadata = cast(Tuple[Any, ...], self._metadata)
-            processed_rows = self._get_rows(offset, metadata, rows)
-        self._rows.extend(processed_rows)
-        self._next_token = response.get("NextToken")
+        next_token = response.get("NextToken")
+        return rows, next_token
+
+    def _process_rows(self, rows: List[Dict[str, Any]], offset: int = 0) -> None:
+        if rows and self._metadata:
+            processed_rows = self._get_rows(offset, self._metadata, rows)
+            self._rows.extend(processed_rows)
 
     def _is_first_row_column_labels(self, rows: List[Dict[str, Any]]) -> bool:
         first_row_data = rows[0].get("Data", [])
-        metadata = cast(Tuple[Any, Any], self._metadata)
-        for meta, data in zip(metadata, first_row_data, strict=False):
+        for meta, data in zip(self._metadata or (), first_row_data, strict=False):
             if meta.get("Name") != data.get("VarCharValue"):
                 return False
         return True
 
+    def _fetch_all_rows(
+        self,
+        converter: Optional[Converter] = None,
+    ) -> List[Tuple[Optional[Any], ...]]:
+        """Fetch all rows via GetQueryResults API with type conversion.
+
+        Paginates through all results from the beginning using MaxResults=1000.
+        Defaults to ``DefaultTypeConverter`` for string-to-Python type conversion,
+        because subclass converters (e.g. Pandas/Arrow) are designed for S3 file
+        reading and may not handle API result strings.
+
+        This method is intended for use by subclass result sets that need to
+        fall back to the API when S3 output is not available (e.g., managed
+        query result storage).
+
+        Args:
+            converter: Type converter for result values. Defaults to
+                ``DefaultTypeConverter`` if not specified.
+
+        Returns:
+            List of converted row tuples.
+        """
+        if self._metadata is None:
+            raise ProgrammingError("Metadata is not available.")
+
+        _logger.warning(
+            "output_location is not available (e.g. managed query result storage). "
+            "Falling back to GetQueryResults API. "
+            "This may be slow for large result sets."
+        )
+
+        converter = converter or DefaultTypeConverter()
+        all_rows: List[Tuple[Optional[Any], ...]] = []
+        next_token: Optional[str] = None
+
+        while True:
+            response = self.__get_query_results(self.DEFAULT_FETCH_SIZE, next_token)
+            rows, next_token = self._parse_result_rows(response)
+
+            offset = 1 if rows and self._is_first_row_column_labels(rows) else 0
+            all_rows.extend(
+                cast(
+                    List[Tuple[Optional[Any], ...]],
+                    self._get_rows(offset, self._metadata, rows, converter),
+                )
+            )
+
+            if not next_token:
+                break
+
+        return all_rows
+
+    @staticmethod
+    def _rows_to_columnar(
+        rows: List[Tuple[Optional[Any], ...]],
+        columns: List[str],
+    ) -> Dict[str, List[Any]]:
+        """Convert row-oriented data to columnar format.
+
+        Args:
+            rows: List of row tuples from ``_fetch_all_rows()``.
+            columns: Column names in order.
+
+        Returns:
+            Dictionary mapping column names to lists of values.
+        """
+        columnar: Dict[str, List[Any]] = {col: [] for col in columns}
+        for row in rows:
+            for col, val in zip(columns, row, strict=False):
+                columnar[col].append(val)
+        return columnar
+
     def _get_content_length(self) -> int:
         if not self.output_location:
             raise ProgrammingError("OutputLocation is none or empty.")
@@ -530,14 +629,19 @@ class AthenaDictResultSet(AthenaResultSet):
     dict_type: Type[Any] = dict
 
     def _get_rows(
-        self, offset: int, metadata: Tuple[Any, ...], rows: List[Dict[str, Any]]
+        self,
+        offset: int,
+        metadata: Tuple[Any, ...],
+        rows: List[Dict[str, Any]],
+        converter: Optional[Converter] = None,
     ) -> List[Union[Tuple[Optional[Any], ...], Dict[Any, Optional[Any]]]]:
+        conv = converter or self._converter
         return [
             self.dict_type(
                 [
                     (
                         meta.get("Name"),
-                        self._converter.convert(meta.get("Type"), row.get("VarCharValue")),
+                        conv.convert(meta.get("Type"), row.get("VarCharValue")),
                     )
                     for meta, row in zip(metadata, rows[i].get("Data", []), strict=False)
                 ]

{pyathena-3.27.1 → pyathena-3.28.0}/pyathena/s3fs/result_set.py

@@ -82,10 +82,14 @@ class AthenaS3FSResultSet(AthenaResultSet):
 
         if self.state == AthenaQueryExecution.STATE_SUCCEEDED and self.output_location:
             self._init_csv_reader()
+        elif self.state == AthenaQueryExecution.STATE_SUCCEEDED:
+            # Managed query result storage: no output_location, use API
+            rows = self._fetch_all_rows()
+            self._rows.extend(rows)
 
         # If CSV reader was not initialized (e.g., CTAS, DDL),
         # fall back to pre-fetched data from Athena API
-        if not self._csv_reader and pre_fetched_rows:
+        if not self._csv_reader and not self._rows and pre_fetched_rows:
             self._rows.extend(pre_fetched_rows)
 
     def _create_s3_file_system(self) -> S3FileSystem:

{pyathena-3.27.1 → pyathena-3.28.0}/pyproject.toml

@@ -30,8 +30,10 @@ classifiers = [
 dynamic = ["version"]
 
 [project.urls]
-homepage = "https://github.com/laughingman7743/PyAthena/"
-repository = "https://github.com/laughingman7743/PyAthena/"
+homepage = "https://github.com/pyathena-dev/PyAthena/"
+repository = "https://github.com/pyathena-dev/PyAthena/"
+documentation = "https://pyathena.dev/"
+issues = "https://github.com/pyathena-dev/PyAthena/issues"
 
 [project.entry-points."sqlalchemy.dialects"]
 awsathena = "pyathena.sqlalchemy.base:AthenaDialect"