Fix AioSparkCursor bugs and add missing tests

- Change _calculate exception from OperationalError to DatabaseError
  to match sync SparkBaseCursor behavior
- Add retry logic to _read_s3_file_as_text via async_retry_api_call
  to match sync version's retry_api_call usage
- Remove incorrect # type: ignore[override] on name-mangled __poll
- Add test_async_iterator for AioS3FSCursor
- Add test_executemany and test_context_manager for AioSparkCursor
- Move runtime import to top-level in S3FS test file

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

Author: laughingman7743

pyathena/aio/spark/cursor.py

@@ -6,7 +6,7 @@
 from typing import Any, Dict, List, Optional, Union, cast
 
 from pyathena.aio.util import async_retry_api_call
-from pyathena.error import NotSupportedError, OperationalError, ProgrammingError
+from pyathena.error import DatabaseError, NotSupportedError, OperationalError, ProgrammingError
 from pyathena.model import (
     AthenaCalculationExecution,
     AthenaCalculationExecutionStatus,
@@ -123,10 +123,10 @@ async def _calculate(  # type: ignore[override]
             calculation_id = response.get("CalculationExecutionId")
         except Exception as e:
             _logger.exception("Failed to execute calculation.")
-            raise OperationalError(*e.args) from e
+            raise DatabaseError(*e.args) from e
         return cast(str, calculation_id)
 
-    async def __poll(  # type: ignore[override]
+    async def __poll(
         self, query_id: str
     ) -> Union[AthenaQueryExecution, AthenaCalculationExecution]:
         while True:
@@ -181,8 +181,10 @@ async def _terminate_session(self) -> None:  # type: ignore[override]
 
     async def _read_s3_file_as_text(self, uri) -> str:  # type: ignore[override]
         bucket, key = parse_output_location(uri)
-        response = await asyncio.to_thread(
+        response = await async_retry_api_call(
             self._client.get_object,
+            config=self._retry_config,
+            logger=_logger,
             Bucket=bucket,
             Key=key,
         )

tests/pyathena/aio/s3fs/test_cursor.py

@@ -1,6 +1,7 @@
 # -*- coding: utf-8 -*-
 import pytest
 
+from pyathena.aio.s3fs.cursor import AioS3FSCursor
 from pyathena.error import ProgrammingError
 from pyathena.s3fs.result_set import AthenaS3FSResultSet
 from tests import ENV
@@ -58,11 +59,16 @@ async def test_invalid_arraysize(self, aio_s3fs_cursor):
         with pytest.raises(ProgrammingError):
             aio_s3fs_cursor.arraysize = -1
 
+    async def test_async_iterator(self, aio_s3fs_cursor):
+        await aio_s3fs_cursor.execute("SELECT * FROM one_row")
+        rows = []
+        async for row in aio_s3fs_cursor:
+            rows.append(row)
+        assert rows == [(1,)]
+
     async def test_context_manager(self):
         conn = await _aio_connect(schema_name=ENV.schema)
         try:
-            from pyathena.aio.s3fs.cursor import AioS3FSCursor
-
             async with conn.cursor(AioS3FSCursor) as cursor:
                 await cursor.execute("SELECT * FROM one_row")
                 assert await cursor.fetchone() == (1,)

tests/pyathena/aio/spark/test_cursor.py

@@ -3,9 +3,10 @@
 
 import pytest
 
-from pyathena.error import DatabaseError, OperationalError
+from pyathena.error import DatabaseError, NotSupportedError, OperationalError
 from pyathena.model import AthenaCalculationExecutionStatus
 from tests import ENV
+from tests.pyathena.aio.conftest import _aio_connect
 
 
 class TestAioSparkCursor:
@@ -142,3 +143,22 @@ async def cancel_after_delay(c):
             )
 
         await task
+
+    async def test_executemany(self, aio_spark_cursor):
+        with pytest.raises(NotSupportedError):
+            await aio_spark_cursor.executemany("SELECT 1", [])
+
+    async def test_context_manager(self):
+        import asyncio
+
+        from pyathena.aio.spark.cursor import AioSparkCursor
+
+        conn = await _aio_connect(
+            schema_name=ENV.schema,
+            cursor_class=AioSparkCursor,
+            work_group=ENV.spark_work_group,
+        )
+        cursor = await asyncio.to_thread(conn.cursor)
+        async with cursor:
+            await cursor.execute("print('hello')")
+            assert await cursor.get_std_out() == "hello"