From 21c24ebd6b5557646ce64dd2a7605d07796004d1 Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Sat, 17 Oct 2026 22:00:42 +0900 Subject: [PATCH 1/4] Remove and rewrite unnecessary aio and Spark cursor tests Apply the rows of #1245 to the aio connection, cursor, result set, utility, and DataFrame cursor tests and to the Spark cursor tests. The aio keyword, callback, and arraysize tests become parameters of the synchronous connection tests, the aio cursors' fetch-without-result and cancel tests become offline tests over every aio cursor, and the copies that run through shared code are removed. The rewritten tests answer the Athena requests offline: the request parameters of execute() options, the cache search through QueryExecutionHistory, the Glue fallback thread, and the Spark calculations. A cancelled start request is checked after its executor thread ends instead of after a sleep, and the Spark session tests drive the cursor constructor instead of its private helpers. Co-Authored-By: Claude Opus 5.5 --- tests/pyathena/aio/arrow/test_cursor.py | 110 +--- tests/pyathena/aio/polars/test_cursor.py | 96 +--- tests/pyathena/aio/s3fs/test_cursor.py | 81 +-- tests/pyathena/aio/spark/test_cursor.py | 167 +----- tests/pyathena/aio/test_common.py | 22 +- tests/pyathena/aio/test_connection.py | 178 +------ tests/pyathena/aio/test_cursor.py | 598 +++++----------------- tests/pyathena/aio/test_result_set.py | 6 - tests/pyathena/aio/test_util.py | 44 +- tests/pyathena/spark/test_async_cursor.py | 65 +-- tests/pyathena/spark/test_common.py | 218 +++----- tests/pyathena/spark/test_spark_cursor.py | 28 +- tests/pyathena/test_connection.py | 96 +++- tests/pyathena/test_cursor.py | 23 +- tests/pyathena/util.py | 68 ++- 15 files changed, 418 insertions(+), 1382 deletions(-) diff --git a/tests/pyathena/aio/arrow/test_cursor.py b/tests/pyathena/aio/arrow/test_cursor.py index fa4e14e11..7ed96e6c0 100644 --- a/tests/pyathena/aio/arrow/test_cursor.py +++ b/tests/pyathena/aio/arrow/test_cursor.py @@ -10,125 +10,28 @@ import pytest from pyathena.aio.arrow.cursor import AioArrowCursor -from pyathena.arrow.result_set import AthenaArrowResultSet -from pyathena.error import ProgrammingError from pyathena.model import AthenaQueryExecution from pyathena.util import RetryConfig -from tests import ENV -from tests.pyathena.aio.conftest import _aio_connect class TestAioArrowCursor: - async def test_binary_null_vs_empty(self, aio_arrow_cursor): - query = """SELECT * FROM (VALUES - (1, CAST(NULL AS VARBINARY), 'null', CAST(NULL AS VARCHAR)), - (2, X'', 'empty', ''), - (3, X'00ff275c25', 'comma, quote" and' || chr(10) || 'newline', 'NULL') - ) AS t(id, value, label, text_value) ORDER BY id""" - await aio_arrow_cursor.execute(query) - result = aio_arrow_cursor - rows = await result.fetchall() - assert [row[:3] for row in rows] == [ - (1, None, "null"), - (2, b"", "empty"), - (3, b"\x00\xff'\\%", 'comma, quote" and\nnewline'), - ] - assert [row[3] for row in rows] == ["", "", "NULL"] - async def test_fetchone(self, aio_arrow_cursor): - await aio_arrow_cursor.execute("SELECT * FROM one_row") + assert await aio_arrow_cursor.execute("SELECT * FROM one_row") is aio_arrow_cursor assert aio_arrow_cursor.rownumber == 0 assert await aio_arrow_cursor.fetchone() == (1,) assert aio_arrow_cursor.rownumber == 1 assert await aio_arrow_cursor.fetchone() is None - async def test_fetchmany(self, aio_arrow_cursor): - await aio_arrow_cursor.execute("SELECT * FROM many_rows LIMIT 15") - assert len(await aio_arrow_cursor.fetchmany(10)) == 10 - assert len(await aio_arrow_cursor.fetchmany(10)) == 5 - - async def test_fetchall(self, aio_arrow_cursor): - await aio_arrow_cursor.execute("SELECT * FROM one_row") - assert await aio_arrow_cursor.fetchall() == [(1,)] - await aio_arrow_cursor.execute("SELECT a FROM many_rows ORDER BY a") - assert await aio_arrow_cursor.fetchall() == [(i,) for i in range(10000)] - async def test_as_arrow(self, aio_arrow_cursor): await aio_arrow_cursor.execute("SELECT * FROM one_row") table = aio_arrow_cursor.as_arrow() assert table.num_rows == 1 assert table.num_columns == 1 assert table.column_names == ["number_of_rows"] - - async def test_as_polars(self, aio_arrow_cursor): - await aio_arrow_cursor.execute("SELECT * FROM one_row") df = aio_arrow_cursor.as_polars() assert df.height == 1 assert df.width == 1 - async def test_execute_returns_self(self, aio_arrow_cursor): - result = await aio_arrow_cursor.execute("SELECT * FROM one_row") - assert result is aio_arrow_cursor - - async def test_no_result_set_raises(self, aio_arrow_cursor): - with pytest.raises(ProgrammingError): - await aio_arrow_cursor.fetchone() - with pytest.raises(ProgrammingError): - await aio_arrow_cursor.fetchmany() - with pytest.raises(ProgrammingError): - await aio_arrow_cursor.fetchall() - with pytest.raises(ProgrammingError): - aio_arrow_cursor.as_arrow() - with pytest.raises(ProgrammingError): - aio_arrow_cursor.as_polars() - - async def test_context_manager(self): - from pyathena.aio.arrow.cursor import AioArrowCursor - - conn = await _aio_connect(schema_name=ENV.schema, cursor_class=AioArrowCursor) - try: - async with conn.cursor() as cursor: - await cursor.execute("SELECT * FROM one_row") - assert await cursor.fetchone() == (1,) - finally: - conn.close() - - async def test_arraysize_default(self, aio_arrow_cursor): - assert aio_arrow_cursor.arraysize == AthenaArrowResultSet.DEFAULT_FETCH_SIZE - - async def test_invalid_arraysize(self, aio_arrow_cursor): - aio_arrow_cursor.arraysize = 10000 - assert aio_arrow_cursor.arraysize == 10000 - with pytest.raises(ProgrammingError): - aio_arrow_cursor.arraysize = -1 - - async def test_description(self, aio_arrow_cursor): - await aio_arrow_cursor.execute("SELECT CAST(1 AS INT) AS foobar FROM one_row") - assert await aio_arrow_cursor.fetchall() == [(1,)] - assert aio_arrow_cursor.description == [("foobar", "integer", None, None, 10, 0, "UNKNOWN")] - - async def test_description_initial(self, aio_arrow_cursor): - assert aio_arrow_cursor.description is None - - async def test_cancel_initial(self, aio_arrow_cursor): - with pytest.raises(ProgrammingError): - await aio_arrow_cursor.cancel() - - async def test_executemany_fetch(self, aio_arrow_cursor): - await aio_arrow_cursor.executemany( - "SELECT %(x)d FROM one_row", [{"x": i} for i in range(1, 2)] - ) - with pytest.raises(ProgrammingError): - await aio_arrow_cursor.fetchall() - with pytest.raises(ProgrammingError): - await aio_arrow_cursor.fetchmany() - with pytest.raises(ProgrammingError): - await aio_arrow_cursor.fetchone() - with pytest.raises(ProgrammingError): - aio_arrow_cursor.as_arrow() - with pytest.raises(ProgrammingError): - aio_arrow_cursor.as_polars() - @pytest.mark.parametrize( "aio_arrow_cursor", [{"cursor_kwargs": {"unload": True}}], @@ -136,19 +39,10 @@ async def test_executemany_fetch(self, aio_arrow_cursor): ) async def test_fetchone_unload(self, aio_arrow_cursor): await aio_arrow_cursor.execute("SELECT * FROM one_row") + assert aio_arrow_cursor.as_arrow().num_rows == 1 assert await aio_arrow_cursor.fetchone() == (1,) assert await aio_arrow_cursor.fetchone() is None - @pytest.mark.parametrize( - "aio_arrow_cursor", - [{"cursor_kwargs": {"unload": True}}], - indirect=["aio_arrow_cursor"], - ) - async def test_as_arrow_unload(self, aio_arrow_cursor): - await aio_arrow_cursor.execute("SELECT * FROM one_row") - table = aio_arrow_cursor.as_arrow() - assert table.num_rows == 1 - @pytest.mark.parametrize( "execute_kwargs", [{}, {"connect_timeout": 3.0, "request_timeout": 4.0}] ) diff --git a/tests/pyathena/aio/polars/test_cursor.py b/tests/pyathena/aio/polars/test_cursor.py index 9c9003c74..e45a9508b 100644 --- a/tests/pyathena/aio/polars/test_cursor.py +++ b/tests/pyathena/aio/polars/test_cursor.py @@ -10,110 +10,27 @@ import pytest from pyathena.aio.polars.cursor import AioPolarsCursor -from pyathena.error import ProgrammingError from pyathena.model import AthenaQueryExecution -from pyathena.polars.result_set import AthenaPolarsResultSet from pyathena.util import RetryConfig -from tests import ENV -from tests.pyathena.aio.conftest import _aio_connect class TestAioPolarsCursor: async def test_fetchone(self, aio_polars_cursor): - await aio_polars_cursor.execute("SELECT * FROM one_row") + assert await aio_polars_cursor.execute("SELECT * FROM one_row") is aio_polars_cursor assert aio_polars_cursor.rownumber == 0 assert await aio_polars_cursor.fetchone() == (1,) assert aio_polars_cursor.rownumber == 1 assert await aio_polars_cursor.fetchone() is None - async def test_fetchmany(self, aio_polars_cursor): - await aio_polars_cursor.execute("SELECT * FROM many_rows LIMIT 15") - assert len(await aio_polars_cursor.fetchmany(10)) == 10 - assert len(await aio_polars_cursor.fetchmany(10)) == 5 - - async def test_fetchall(self, aio_polars_cursor): - await aio_polars_cursor.execute("SELECT * FROM one_row") - assert await aio_polars_cursor.fetchall() == [(1,)] - await aio_polars_cursor.execute("SELECT a FROM many_rows ORDER BY a") - assert await aio_polars_cursor.fetchall() == [(i,) for i in range(10000)] - async def test_as_polars(self, aio_polars_cursor): await aio_polars_cursor.execute("SELECT * FROM one_row") df = aio_polars_cursor.as_polars() assert df.height == 1 assert df.width == 1 - - async def test_as_arrow(self, aio_polars_cursor): - await aio_polars_cursor.execute("SELECT * FROM one_row") table = aio_polars_cursor.as_arrow() assert table.num_rows == 1 assert table.num_columns == 1 - async def test_execute_returns_self(self, aio_polars_cursor): - result = await aio_polars_cursor.execute("SELECT * FROM one_row") - assert result is aio_polars_cursor - - async def test_no_result_set_raises(self, aio_polars_cursor): - with pytest.raises(ProgrammingError): - await aio_polars_cursor.fetchone() - with pytest.raises(ProgrammingError): - await aio_polars_cursor.fetchmany() - with pytest.raises(ProgrammingError): - await aio_polars_cursor.fetchall() - with pytest.raises(ProgrammingError): - aio_polars_cursor.as_polars() - with pytest.raises(ProgrammingError): - aio_polars_cursor.as_arrow() - - async def test_context_manager(self): - from pyathena.aio.polars.cursor import AioPolarsCursor - - conn = await _aio_connect(schema_name=ENV.schema, cursor_class=AioPolarsCursor) - try: - async with conn.cursor() as cursor: - await cursor.execute("SELECT * FROM one_row") - assert await cursor.fetchone() == (1,) - finally: - conn.close() - - async def test_arraysize_default(self, aio_polars_cursor): - assert aio_polars_cursor.arraysize == AthenaPolarsResultSet.DEFAULT_FETCH_SIZE - - async def test_invalid_arraysize(self, aio_polars_cursor): - aio_polars_cursor.arraysize = 10000 - assert aio_polars_cursor.arraysize == 10000 - with pytest.raises(ProgrammingError): - aio_polars_cursor.arraysize = -1 - - async def test_description(self, aio_polars_cursor): - await aio_polars_cursor.execute("SELECT CAST(1 AS INT) AS foobar FROM one_row") - assert await aio_polars_cursor.fetchall() == [(1,)] - assert aio_polars_cursor.description == [ - ("foobar", "integer", None, None, 10, 0, "UNKNOWN") - ] - - async def test_description_initial(self, aio_polars_cursor): - assert aio_polars_cursor.description is None - - async def test_cancel_initial(self, aio_polars_cursor): - with pytest.raises(ProgrammingError): - await aio_polars_cursor.cancel() - - async def test_executemany_fetch(self, aio_polars_cursor): - await aio_polars_cursor.executemany( - "SELECT %(x)d FROM one_row", [{"x": i} for i in range(1, 2)] - ) - with pytest.raises(ProgrammingError): - await aio_polars_cursor.fetchall() - with pytest.raises(ProgrammingError): - await aio_polars_cursor.fetchmany() - with pytest.raises(ProgrammingError): - await aio_polars_cursor.fetchone() - with pytest.raises(ProgrammingError): - aio_polars_cursor.as_polars() - with pytest.raises(ProgrammingError): - aio_polars_cursor.as_arrow() - @pytest.mark.parametrize( "aio_polars_cursor", [{"cursor_kwargs": {"unload": True}}], @@ -121,19 +38,10 @@ async def test_executemany_fetch(self, aio_polars_cursor): ) async def test_fetchone_unload(self, aio_polars_cursor): await aio_polars_cursor.execute("SELECT * FROM one_row") + assert aio_polars_cursor.as_polars().height == 1 assert await aio_polars_cursor.fetchone() == (1,) assert await aio_polars_cursor.fetchone() is None - @pytest.mark.parametrize( - "aio_polars_cursor", - [{"cursor_kwargs": {"unload": True}}], - indirect=["aio_polars_cursor"], - ) - async def test_as_polars_unload(self, aio_polars_cursor): - await aio_polars_cursor.execute("SELECT * FROM one_row") - df = aio_polars_cursor.as_polars() - assert df.height == 1 - @pytest.mark.parametrize( "execute_kwargs", [{}, {"block_size": 2048, "cache_type": "none", "s3_max_workers": 3, "chunksize": 20}], diff --git a/tests/pyathena/aio/s3fs/test_cursor.py b/tests/pyathena/aio/s3fs/test_cursor.py index 625f09628..eb8721af1 100644 --- a/tests/pyathena/aio/s3fs/test_cursor.py +++ b/tests/pyathena/aio/s3fs/test_cursor.py @@ -10,98 +10,19 @@ import pytest from pyathena.aio.s3fs.cursor import AioS3FSCursor -from pyathena.error import ProgrammingError from pyathena.model import AthenaQueryExecution from pyathena.s3fs.reader import AthenaCSVReader, EmptyStringAsNullCSVReader -from pyathena.s3fs.result_set import AthenaS3FSResultSet from pyathena.util import RetryConfig -from tests import ENV -from tests.pyathena.aio.conftest import _aio_connect class TestAioS3FSCursor: - @pytest.mark.parametrize( - ("aio_s3fs_cursor", "expected_empty"), - [ - ({}, ""), - ({"cursor_kwargs": {"csv_reader": EmptyStringAsNullCSVReader}}, None), - ], - indirect=["aio_s3fs_cursor"], - ) - async def test_null_vs_empty_string(self, aio_s3fs_cursor, expected_empty): - await aio_s3fs_cursor.execute( - "SELECT NULL AS null_col, '' AS empty_col, 'text' AS value_col" - ) - assert await aio_s3fs_cursor.fetchall() == [(None, expected_empty, "text")] - async def test_fetchone(self, aio_s3fs_cursor): - await aio_s3fs_cursor.execute("SELECT * FROM one_row") + assert await aio_s3fs_cursor.execute("SELECT * FROM one_row") is aio_s3fs_cursor assert aio_s3fs_cursor.rownumber == 0 assert await aio_s3fs_cursor.fetchone() == (1,) assert aio_s3fs_cursor.rownumber == 1 assert await aio_s3fs_cursor.fetchone() is None - async def test_fetchmany(self, aio_s3fs_cursor): - await aio_s3fs_cursor.execute("SELECT * FROM many_rows LIMIT 15") - assert len(await aio_s3fs_cursor.fetchmany(10)) == 10 - assert len(await aio_s3fs_cursor.fetchmany(10)) == 5 - - async def test_fetchall(self, aio_s3fs_cursor): - await aio_s3fs_cursor.execute("SELECT * FROM one_row") - assert await aio_s3fs_cursor.fetchall() == [(1,)] - await aio_s3fs_cursor.execute("SELECT a FROM many_rows ORDER BY a") - assert await aio_s3fs_cursor.fetchall() == [(i,) for i in range(10000)] - - async def test_description(self, aio_s3fs_cursor): - await aio_s3fs_cursor.execute("SELECT 1 AS foobar FROM one_row") - assert await aio_s3fs_cursor.fetchall() == [(1,)] - assert aio_s3fs_cursor.description == [("foobar", "integer", None, None, 10, 0, "UNKNOWN")] - - async def test_description_initial(self, aio_s3fs_cursor): - assert aio_s3fs_cursor.description is None - - async def test_cancel_initial(self, aio_s3fs_cursor): - with pytest.raises(ProgrammingError): - await aio_s3fs_cursor.cancel() - - async def test_executemany_fetch(self, aio_s3fs_cursor): - await aio_s3fs_cursor.executemany( - "SELECT %(x)d FROM one_row", [{"x": i} for i in range(1, 2)] - ) - with pytest.raises(ProgrammingError): - await aio_s3fs_cursor.fetchall() - with pytest.raises(ProgrammingError): - await aio_s3fs_cursor.fetchmany() - with pytest.raises(ProgrammingError): - await aio_s3fs_cursor.fetchone() - - async def test_arraysize_default(self, aio_s3fs_cursor): - assert aio_s3fs_cursor.arraysize == AthenaS3FSResultSet.DEFAULT_FETCH_SIZE - - async def test_invalid_arraysize(self, aio_s3fs_cursor): - aio_s3fs_cursor.arraysize = 10000 - assert aio_s3fs_cursor.arraysize == 10000 - 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 = [row async for row in aio_s3fs_cursor] - assert rows == [(1,)] - - async def test_context_manager(self): - conn = await _aio_connect(schema_name=ENV.schema) - try: - async with conn.cursor(AioS3FSCursor) as cursor: - await cursor.execute("SELECT * FROM one_row") - assert await cursor.fetchone() == (1,) - finally: - conn.close() - - async def test_execute_returns_self(self, aio_s3fs_cursor): - result = await aio_s3fs_cursor.execute("SELECT * FROM one_row") - assert result is aio_s3fs_cursor - @pytest.mark.parametrize("execute_kwargs", [{}, {"csv_reader": AthenaCSVReader}]) async def test_read_options(self, execute_kwargs): """The cursor's read options reach the result set, and execute() overrides them. diff --git a/tests/pyathena/aio/spark/test_cursor.py b/tests/pyathena/aio/spark/test_cursor.py index 8be05d274..8c376fd6b 100644 --- a/tests/pyathena/aio/spark/test_cursor.py +++ b/tests/pyathena/aio/spark/test_cursor.py @@ -16,14 +16,14 @@ from pyathena.aio.spark.cursor import AioSparkCursor from pyathena.error import DatabaseError, NotSupportedError, OperationalError -from pyathena.model import AthenaCalculationExecutionStatus, AthenaSessionStatus +from pyathena.model import AthenaCalculationExecutionStatus from pyathena.util import RetryConfig -from tests import ENV -from tests.pyathena.aio.conftest import _aio_connect from tests.pyathena.util import ( CANCELABLE_SPARK_JOB, + COMPLETED_CALCULATION_RESPONSES, + answered_requests, + offline_spark_cursor, wait_for_spark_job, - wait_for_spark_session_state, ) @@ -106,17 +106,7 @@ def start_calculation_execution(**kwargs): class TestAioSparkCursor: async def test_spark_dataframe(self, aio_spark_cursor): - await aio_spark_cursor.execute( - textwrap.dedent( - f""" - df = spark.read.format("csv") \\ - .option("header", "true") \\ - .option("inferSchema", "true") \\ - .load("{ENV.s3_staging_dir}{ENV.schema}/spark_group_by/spark_group_by.csv") - """ - ), - description="test description", - ) + await aio_spark_cursor.execute("print('hello')", description="test description") assert aio_spark_cursor.calculation_execution assert aio_spark_cursor.session_id assert aio_spark_cursor.calculation_id @@ -132,70 +122,7 @@ async def test_spark_dataframe(self, aio_spark_cursor): assert aio_spark_cursor.std_error_s3_uri assert aio_spark_cursor.result_s3_uri assert aio_spark_cursor.result_type - - await aio_spark_cursor.execute( - textwrap.dedent( - """ - from pyspark.sql.functions import sum - df_count = df.groupBy("name").agg(sum("count").alias("sum")) - df_count.show() - """ - ) - ) - assert ( - await aio_spark_cursor.get_std_out() - == textwrap.dedent( - """ - +----+---+ - |name|sum| - +----+---+ - | bar| 5| - | foo| 5| - +----+---+ - """ - ).strip() - ) - - await aio_spark_cursor.execute( - textwrap.dedent( - f""" - df_count.write.mode('overwrite') \\ - .format("parquet") \\ - .option("path", "{ENV.s3_staging_dir}{ENV.schema}/spark/group_by") \\ - .saveAsTable("{ENV.schema}.spark_group_by") - """ - ) - ) - - @pytest.mark.dependency(depends="test_spark_dataframe") - async def test_spark_sql(self, aio_spark_cursor): - await aio_spark_cursor.execute( - textwrap.dedent( - f""" - spark.sql("SELECT * FROM {ENV.schema}.one_row").show() - """ - ) - ) - assert ( - await aio_spark_cursor.get_std_out() - == textwrap.dedent( - """ - +--------------+ - |number_of_rows| - +--------------+ - | 1| - +--------------+ - """ - ).strip() - ) - - await aio_spark_cursor.execute( - textwrap.dedent( - f""" - spark.sql("DROP TABLE IF EXISTS {ENV.schema}.spark_group_by") - """ - ) - ) + assert await aio_spark_cursor.get_std_out() == "hello" async def test_failed(self, aio_spark_cursor): with pytest.raises(OperationalError): @@ -316,29 +243,6 @@ async def test_execute_cancellation_without_kill_on_interrupt(self): cancel.assert_not_awaited() assert cursor.calculation_execution is None - async def test_session_ownership(self, aio_spark_cursor): - client = aio_spark_cursor.connection.client - session_id = aio_spark_cursor.session_id - borrower = await asyncio.to_thread( - aio_spark_cursor.connection.cursor, AioSparkCursor, session_id=session_id - ) - async with borrower: - await borrower.execute("print(1)") - assert await borrower.get_std_out() == "1" - - # Closing a cursor that was given the session leaves the session running. - await aio_spark_cursor.execute("print(2)") - assert await aio_spark_cursor.get_std_out() == "2" - - # Closing the cursor that started the session terminates it. - await aio_spark_cursor.close() - await asyncio.to_thread( - wait_for_spark_session_state, - client, - session_id, - AthenaSessionStatus.STATE_TERMINATED, - ) - @pytest.mark.parametrize("kill_on_interrupt", [True, False]) async def test_calculate_reuses_generated_token_on_retry(self, kill_on_interrupt): cursor, _, _, release = _starting_cursor(kill_on_interrupt=kill_on_interrupt) @@ -412,24 +316,6 @@ async def test_execute_cancelled_before_request_is_sent(self): cancel.assert_not_awaited() assert cursor.calculation_id is None - async def test_execute_timeout_while_starting(self): - cursor, cancel, started, release = _starting_cursor() - task = asyncio.create_task(asyncio.wait_for(cursor.execute("code"), timeout=0.05)) - try: - assert await asyncio.to_thread(started.wait, _TIMEOUT) - # The timeout is due before this sleep ends, so the event loop handles it - # while the start request is still blocked. - await asyncio.sleep(0.1) - assert not task.done() - finally: - release.set() - with pytest.raises(asyncio.TimeoutError): - await task - - assert started.is_set() - cancel.assert_awaited_once_with("calculation_id") - assert cursor.calculation_id == "calculation_id" - @pytest.mark.parametrize("failing", ["start", "cancel", "wait"]) async def test_execute_cancelled_while_starting_failure(self, failing): error = OperationalError("failed") @@ -534,33 +420,18 @@ async def test_execute_rejects_work_group(self): ) async def test_execute_accepts_none_parameters(self, args, kwargs): """Omitted or None parameters run the calculation (no AWS).""" + cursor = offline_spark_cursor(AioSparkCursor) + with answered_requests( + cursor.connection.client, COMPLETED_CALCULATION_RESPONSES + ) as requests: + await cursor.execute("code", *args, **kwargs) + + start = dict(requests)["StartCalculationExecution"] + assert (start["SessionId"], start["CodeBlock"]) == ("session", "code") + assert cursor.calculation_id == "calculation-id" + assert cursor.state == AthenaCalculationExecutionStatus.STATE_COMPLETED + + async def test_executemany(self): cursor = AioSparkCursor.__new__(AioSparkCursor) # bypass __init__ to avoid AWS calls - cursor._session_id = "session_id" - cursor._calculate = AsyncMock(return_value="calculation_id") - cursor._poll = AsyncMock( - return_value=MagicMock(state=AthenaCalculationExecutionStatus.STATE_COMPLETED) - ) - - await cursor.execute("code", *args, **kwargs) - - cursor._calculate.assert_awaited_once_with( - session_id="session_id", - code_block="code", - description=None, - client_request_token=None, - ) - - 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): - 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" + await cursor.executemany("SELECT 1", []) diff --git a/tests/pyathena/aio/test_common.py b/tests/pyathena/aio/test_common.py index da9816739..2e31905a9 100644 --- a/tests/pyathena/aio/test_common.py +++ b/tests/pyathena/aio/test_common.py @@ -94,10 +94,7 @@ def record(*args): assert len(threads) == 3 assert threading.get_ident() not in threads - @pytest.mark.parametrize( - "cursor_class", - [AioArrowCursor, AioPandasCursor, AioPolarsCursor, AioS3FSCursor], - ) + @pytest.mark.parametrize("cursor_class", _ALL_CURSOR_CLASSES) async def test_fetch_without_result_set(self, cursor_class): cursor = cursor_class( connection=MagicMock(), converter=None, formatter=None, retry_config=None @@ -105,6 +102,19 @@ async def test_fetch_without_result_set(self, cursor_class): for fetch in (cursor.fetchone, cursor.fetchmany, cursor.fetchall): with pytest.raises(ProgrammingError, match=r"No result set\."): await fetch() + # The Arrow and Polars cursors also convert the whole result. + for name in ("as_arrow", "as_polars"): + if hasattr(cursor, name): + with pytest.raises(ProgrammingError): + getattr(cursor, name)() + + @pytest.mark.parametrize("cursor_class", _ALL_CURSOR_CLASSES) + async def test_cancel_without_query(self, cursor_class): + cursor = cursor_class( + connection=MagicMock(), converter=None, formatter=None, retry_config=None + ) + with pytest.raises(ProgrammingError): + await cursor.cancel() @pytest.mark.parametrize( ("method", "args"), [("fetchone", ()), ("fetchmany", (2,)), ("fetchall", ())] @@ -278,6 +288,7 @@ async def test_executemany_failure_closes_result_set_off_event_loop(self, cursor result_set.rowcount = 1 async def execute(operation, parameters, **kwargs): + cursor.query_id = "query_id" cursor.result_set = result_set def parameters(): @@ -290,3 +301,6 @@ def parameters(): assert closes == [True] assert cursor.result_set is None + assert cursor.rowcount == -1 + # The query ID of the last execution is kept. + assert cursor.query_id == "query_id" diff --git a/tests/pyathena/aio/test_connection.py b/tests/pyathena/aio/test_connection.py index 624cecaf1..99f0bb2b9 100644 --- a/tests/pyathena/aio/test_connection.py +++ b/tests/pyathena/aio/test_connection.py @@ -5,13 +5,11 @@ # # SPDX-License-Identifier: MIT -from inspect import Parameter, signature from typing import Any from unittest.mock import MagicMock, patch import pytest -import pyathena from pyathena.aio.arrow.cursor import AioArrowCursor from pyathena.aio.connection import AioConnection from pyathena.aio.cursor import AioCursor, AioDictCursor @@ -19,94 +17,21 @@ from pyathena.aio.polars.cursor import AioPolarsCursor from pyathena.aio.s3fs.cursor import AioS3FSCursor from pyathena.aio.spark.cursor import AioSparkCursor -from pyathena.converter import DefaultTypeConverter -from pyathena.error import ProgrammingError from pyathena.spark.common import SparkBaseCursor -# Cursors whose arraysize is the GetQueryResults page size, capped at 1000. -PAGED_CURSORS = [AioCursor, AioDictCursor] -# Cursors whose arraysize only sets the fetchmany() batch, with no upper limit. -UNCAPPED_CURSORS = [AioArrowCursor, AioPandasCursor, AioPolarsCursor, AioS3FSCursor] -CURSOR_CLASSES = PAGED_CURSORS + UNCAPPED_CURSORS + [AioSparkCursor] +CURSOR_CLASSES = [ + AioCursor, + AioDictCursor, + AioArrowCursor, + AioPandasCursor, + AioPolarsCursor, + AioS3FSCursor, + AioSparkCursor, +] # Cursors that pass execute() keyword arguments to a DataFrame reader. READER_CURSOR_CLASSES = [AioPandasCursor, AioPolarsCursor] -@pytest.mark.parametrize("factory", [pyathena.aio_connect, AioConnection.create]) -async def test_empty_work_group_disables_environment_fallback(factory, monkeypatch): - monkeypatch.setenv("AWS_ATHENA_WORK_GROUP", "environment-workgroup") - conn = await factory( - region_name="us-east-1", - s3_staging_dir="s3://bucket/path/", - work_group="", - aws_access_key_id="access_key", - aws_secret_access_key="secret_key", - ) - try: - assert conn.work_group is None - finally: - conn.close() - - -@pytest.mark.parametrize("factory", [pyathena.aio_connect, AioConnection.create]) -async def test_connect_keyword_only(factory): - assert all( - parameter.kind in (Parameter.KEYWORD_ONLY, Parameter.VAR_KEYWORD) - for parameter in signature(factory).parameters.values() - ) - with patch("pyathena.connection.Session") as session: - with pytest.raises(TypeError, match="positional"): - await factory("s3://bucket/path/", "us-east-1") - session.assert_not_called() - - -def test_connection_constructor_keyword_only(): - with patch("pyathena.connection.Session") as session: - with pytest.raises(TypeError, match="positional"): - AioConnection("s3://bucket/path/") - session.assert_not_called() - - -@pytest.mark.parametrize("cursor_class", CURSOR_CLASSES) -def test_cursor_constructor_keyword_only(cursor_class): - constructor = cursor_class.__init__ - assert [ - parameter.name - for parameter in signature(constructor).parameters.values() - if parameter.kind in (Parameter.POSITIONAL_ONLY, Parameter.POSITIONAL_OR_KEYWORD) - ] == ["self"] - with pytest.raises(TypeError, match="positional"): - constructor(object(), "positional-setting") - - -@pytest.mark.parametrize("cursor_class", CURSOR_CLASSES) -def test_cursor_execute_keyword_only(cursor_class): - execute_signature = signature(cursor_class.execute) - assert [ - parameter.name - for parameter in execute_signature.parameters.values() - if parameter.kind in (Parameter.POSITIONAL_ONLY, Parameter.POSITIONAL_OR_KEYWORD) - ] == ["self", "operation", "parameters"] - parameters = {"value": 1} - cursor = object.__new__(cursor_class) - bound = execute_signature.bind(cursor, "SELECT %(value)s", parameters) - assert bound.arguments["parameters"] is parameters - if cursor_class is AioSparkCursor: - assert "work_group" not in execute_signature.parameters - else: - bound = execute_signature.bind(cursor, "SELECT %(value)s", parameters, work_group="group") - assert bound.arguments["work_group"] == "group" - execute_signature.bind(cursor, operation="SELECT 1") - with pytest.raises(TypeError, match="positional"): - cursor.execute("SELECT %(value)s", parameters, "group") - - -class RecordingAioCursor(AioCursor): - def __init__(self, **kwargs: Any) -> None: - self.kwargs = kwargs - super().__init__(**kwargs) - - async def _connection(**kwargs: Any) -> AioConnection: return await AioConnection.create( region_name="us-east-1", @@ -142,26 +67,6 @@ def _open_cursor(conn: AioConnection, cursor_class: type[Any], **kwargs: Any) -> return conn.cursor(cursor_class, **kwargs) -@pytest.mark.parametrize("cursor_class", CURSOR_CLASSES) -async def test_cursor_unknown_constructor_keyword(offline_connection, cursor_class): - with pytest.raises(TypeError, match="unexpected keyword argument 'work_gruop'"): - offline_connection.cursor(cursor_class, work_gruop="typo") - offline_connection.cursor_kwargs = {"work_gruop": "typo"} - with pytest.raises(TypeError, match="unexpected keyword argument 'work_gruop'"): - offline_connection.cursor(cursor_class) - - -@pytest.mark.parametrize("cursor_class", CURSOR_CLASSES) -async def test_cursor_accepts_connection_callbacks(offline_connection, cursor_class): - on_start_query_execution, on_poll = MagicMock(), MagicMock() - offline_connection.on_start_query_execution = on_start_query_execution - offline_connection.on_poll = on_poll - - async with _open_cursor(offline_connection, cursor_class) as cursor: - assert cursor._on_start_query_execution is on_start_query_execution - assert cursor._on_poll is on_poll - - @pytest.mark.parametrize( "cursor_class", [c for c in CURSOR_CLASSES if c not in READER_CURSOR_CLASSES] ) @@ -202,68 +107,3 @@ async def test_cursor_execute_supported_keywords(offline_connection, cursor_clas pytest.raises(RuntimeError, match="query started"), ): await cursor.execute("SELECT 1", **kwargs) - - -class TestAioConnection: - async def test_cursor_arguments_override_cursor_kwargs(self): - conn = await _connection( - cursor_class=RecordingAioCursor, - cursor_kwargs={"kill_on_interrupt": False, "schema_name": "configured"}, - ) - - cursor = conn.cursor(kill_on_interrupt=True) - - assert cursor.kwargs["kill_on_interrupt"] is True - assert cursor.kwargs["schema_name"] == "configured" - assert conn.cursor_kwargs == {"kill_on_interrupt": False, "schema_name": "configured"} - - @pytest.mark.parametrize("cursor_class", PAGED_CURSORS + UNCAPPED_CURSORS) - async def test_cursor_arraysize(self, cursor_class): - conn = await _connection(cursor_kwargs={"arraysize": 25}) - - assert conn.cursor(cursor_class).arraysize == 25 - assert conn.cursor(cursor_class, arraysize=50).arraysize == 50 - - @pytest.mark.parametrize("cursor_class", PAGED_CURSORS) - async def test_cursor_arraysize_over_page_size(self, cursor_class): - conn = await _connection() - - with pytest.raises(ProgrammingError): - conn.cursor(cursor_class, arraysize=1001) - - @pytest.mark.parametrize("cursor_class", UNCAPPED_CURSORS) - async def test_cursor_arraysize_uncapped(self, cursor_class): - conn = await _connection() - - assert conn.cursor(cursor_class, arraysize=1001).arraysize == 1001 - - @pytest.mark.parametrize("cursor_class", PAGED_CURSORS + UNCAPPED_CURSORS) - async def test_cursor_arraysize_not_positive(self, cursor_class): - conn = await _connection() - - with pytest.raises(ProgrammingError): - conn.cursor(cursor_class, arraysize=0) - - async def test_internal_cursor_leaves_out_default_cursor_options(self): - cursor_kwargs = { - "unload": True, - "chunksize": 10, - "arraysize": 5000, - "schema_name": "configured", - } - conn = await _connection( - cursor_class=AioPandasCursor, - cursor_kwargs=cursor_kwargs, - converter=AioPandasCursor.get_default_converter(True), - ) - - cursor = conn._internal_cursor(AioCursor) - - assert type(cursor) is AioCursor - assert cursor._schema_name == "configured" - assert cursor.arraysize == AioCursor.DEFAULT_FETCH_SIZE - assert type(cursor._converter) is DefaultTypeConverter - assert conn.cursor_kwargs == cursor_kwargs - # Creating the cursor directly still passes the options it does not accept. - with pytest.raises(TypeError, match="unexpected keyword argument 'unload'"): - conn.cursor(AioCursor) diff --git a/tests/pyathena/aio/test_cursor.py b/tests/pyathena/aio/test_cursor.py index 4bb694e89..bf8cc21a5 100644 --- a/tests/pyathena/aio/test_cursor.py +++ b/tests/pyathena/aio/test_cursor.py @@ -1,32 +1,47 @@ import asyncio -import re import threading from collections import OrderedDict +from concurrent.futures import ThreadPoolExecutor from datetime import UTC, datetime, timedelta -from unittest.mock import AsyncMock, MagicMock, call, patch +from unittest.mock import AsyncMock, MagicMock import pytest from botocore.exceptions import ClientError -from pyathena import BINARY, Binary, ExecuteOptions +from pyathena import ExecuteOptions +from pyathena.aio.connection import AioConnection from pyathena.aio.cursor import AioCursor, AioDictCursor from pyathena.aio.result_set import AthenaAioDictResultSet -from pyathena.error import DatabaseError, OperationalError, ProgrammingError -from pyathena.glue import GlueMetadataClient +from pyathena.error import DatabaseError, OperationalError from pyathena.model import AthenaQueryExecution -from pyathena.result_set import AthenaResultSet from pyathena.util import RetryConfig from tests import ENV -from tests.pyathena.aio.conftest import _aio_connect from tests.pyathena.util import ( EVENT_TIMEOUT, + QueryExecutionHistory, ThrottledOperation, - succeeded_query_execution, + answered_requests, + cache_search_cursor, + succeeded_query_execution_data, throttle_metadata_api, wait_until, ) +def _offline_connection(): + """Build a connection whose Athena requests the test answers. + + Returns: + The connection, with ``AioCursor`` as its cursor class. + """ + return AioConnection( + region_name="us-east-1", + s3_staging_dir="s3://bucket/path/", + aws_access_key_id="access_key", + aws_secret_access_key="secret_key", + ) + + def _offline_cursor(kill_on_interrupt, final_state): """An AioCursor whose first status request blocks until the task is cancelled. @@ -110,26 +125,8 @@ def start_query_execution(**kwargs): class TestAioCursor: - @pytest.mark.parametrize( - ("value", "expected"), - [ - (b"", b""), - (b"\x00\xff'\\%", b"\x00\xff'\\%"), - (bytes(range(256)), bytes(range(256))), - (bytearray(b"\x00\xff"), b"\x00\xff"), - (memoryview(b"\x00\xff"), b"\x00\xff"), - (Binary(bytearray(b"abc")), b"abc"), - (None, None), - ], - ids=["empty", "special", "all_bytes", "bytearray", "memoryview", "dbapi_binary", "null"], - ) - async def test_binary_parameter(self, aio_cursor, value, expected): - await aio_cursor.execute("SELECT CAST(%(value)s AS VARBINARY)", {"value": value}) - assert await aio_cursor.fetchone() == (expected,) - assert aio_cursor.description[0][1] == BINARY - async def test_fetchone(self, aio_cursor): - await aio_cursor.execute("SELECT * FROM one_row") + assert await aio_cursor.execute("SELECT * FROM one_row") is aio_cursor assert aio_cursor.rowcount == -1 assert aio_cursor.rownumber == 0 assert await aio_cursor.fetchone() == (1,) @@ -173,15 +170,6 @@ async def test_fetchall(self, aio_cursor): await aio_cursor.execute("SELECT a FROM many_rows ORDER BY a") assert await aio_cursor.fetchall() == [(i,) for i in range(10000)] - async def test_async_iterator(self, aio_cursor): - await aio_cursor.execute("SELECT * FROM one_row") - rows = [row async for row in aio_cursor] - assert rows == [(1,)] - - async def test_execute_returns_self(self, aio_cursor): - result = await aio_cursor.execute("SELECT * FROM one_row") - assert result is aio_cursor - async def test_execute_with_callback(self, aio_cursor): callback_results = [] await aio_cursor.execute("SELECT 1", on_start_query_execution=callback_results.append) @@ -196,87 +184,49 @@ async def test_execute_with_options(self, aio_cursor): assert await aio_cursor.fetchone() == (1,) async def test_execute_internal_options_passthrough(self): - """The private _execute() forwards its options to the request and the cache (no AWS).""" - cursor = AioCursor.__new__(AioCursor) # bypass __init__ to avoid AWS calls - cursor._connection = MagicMock() - cursor._connection.client.start_query_execution.return_value = { - "QueryExecutionId": "test_query_id" - } - cursor._retry_config = RetryConfig() - cursor._kill_on_interrupt = True - - with ( - patch.object( - AioCursor, "_build_start_query_execution_request", return_value={} - ) as request_mock, - patch.object( - AioCursor, "_find_previous_query_id", new_callable=AsyncMock, return_value=None - ) as cache_mock, - ): - query_id = await cursor._execute( - "SELECT 1", - parameters=None, - options=ExecuteOptions( - work_group="test_work_group", - s3_staging_dir="s3://test-bucket/path/", - cache_size=10, - cache_expiration_time=100, - result_reuse_enable=True, - result_reuse_minutes=5, - paramstyle="qmark", - ), - ) - - assert query_id == "test_query_id" - request_mock.assert_called_once_with( - query="SELECT 1", + """The options of execute() reach its Athena requests (no AWS).""" + cursor = _offline_connection().cursor() + options = ExecuteOptions( work_group="test_work_group", s3_staging_dir="s3://test-bucket/path/", - result_reuse_enable=True, - result_reuse_minutes=5, - execution_parameters=None, - ) - cache_mock.assert_awaited_once_with( - "SELECT 1", - "test_work_group", cache_size=10, cache_expiration_time=100, + result_reuse_enable=True, + result_reuse_minutes=5, + paramstyle="qmark", ) + with ( + answered_requests(cursor.connection.client) as requests, + pytest.raises(OperationalError), + ): + await cursor.execute("SELECT 1", options=options) - async def test_execute_qmark_parameters_skip_cache(self): - """A qmark query with parameters never searches the cache (no AWS, #941). - - Mirrors the synchronous cursor test. - """ - cursor = AioCursor.__new__(AioCursor) # bypass __init__ to avoid AWS calls - cursor._connection = MagicMock() - cursor._connection.client.start_query_execution.return_value = { - "QueryExecutionId": "test_query_id" + parameters = dict(requests) + assert parameters["ListQueryExecutions"] == { + "WorkGroup": "test_work_group", + "MaxResults": 10, + } + start = parameters["StartQueryExecution"] + assert start["WorkGroup"] == "test_work_group" + assert start["ResultConfiguration"] == {"OutputLocation": "s3://test-bucket/path/"} + assert start["ResultReuseConfiguration"] == { + "ResultReuseByAgeConfiguration": {"Enabled": True, "MaxAgeInMinutes": 5} } - cursor._retry_config = RetryConfig() - cursor._kill_on_interrupt = True + async def test_execute_qmark_parameters_skip_cache(self): + """A qmark query with parameters never searches the cache (no AWS, #941).""" + cursor = _offline_connection().cursor() with ( - patch.object( - AioCursor, - "_build_start_query_execution_request", - return_value={"ExecutionParameters": ["'1'"]}, - ) as request_mock, - patch.object( - AioCursor, "_find_previous_query_id", new_callable=AsyncMock, return_value="cached" - ) as cache_mock, + answered_requests(cursor.connection.client) as requests, + pytest.raises(OperationalError), ): - query_id = await cursor._execute( - "SELECT ?", - ["'1'"], - options=ExecuteOptions( - paramstyle="qmark", cache_size=10, cache_expiration_time=100 - ), - ) + await cursor.execute("SELECT ?", ["'1'"], paramstyle="qmark", cache_size=10) - assert query_id == "test_query_id" - assert request_mock.call_args.kwargs["execution_parameters"] == ["'1'"] - cache_mock.assert_not_awaited() + assert [operation for operation, _ in requests] == [ + "StartQueryExecution", + "GetQueryExecution", + ] + assert requests[0][1]["ExecutionParameters"] == ["'1'"] @pytest.mark.parametrize( "final_state", @@ -303,16 +253,6 @@ async def test_execute_kill_on_interrupt(self, final_state): assert cursor.query_id == "query_id" assert cursor.result_set is None - async def test_execute_kill_on_interrupt_timeout(self): - """A timeout cancels the query and raises TimeoutError (no AWS).""" - cursor, cancel, _ = _offline_cursor( - kill_on_interrupt=True, final_state=AthenaQueryExecution.STATE_CANCELLED - ) - with pytest.raises(asyncio.TimeoutError): - await asyncio.wait_for(cursor.execute("SELECT 1"), timeout=0.01) - - cancel.assert_awaited_once_with("query_id") - @pytest.mark.parametrize("failing", ["cancel", "wait"]) async def test_execute_kill_on_interrupt_failure(self, failing): """A failure to cancel or wait becomes the cause of the cancellation (no AWS).""" @@ -404,6 +344,8 @@ def start_query_execution(**kwargs): async def test_execute_cancelled_again_while_start_is_retried(self): """GH-1202: a second cancellation stops the retries of the start request (no AWS).""" + executor = ThreadPoolExecutor() + asyncio.get_running_loop().set_default_executor(executor) cursor, cancel, _, _ = _starting_cursor() operation = ThrottledOperation(block=True, failures=1, operation_name="StartQueryExecution") @@ -424,8 +366,9 @@ def start_query_execution(**kwargs): await task finally: operation.release.set() - # A retry without the stop would follow the first attempt at once. - await asyncio.sleep(0.2) + # The thread of the start request ends once its retries stop; a retry + # would follow the first attempt at once. + executor.shutdown(wait=True) assert operation.calls == 1 cancel.assert_not_awaited() @@ -465,25 +408,6 @@ async def test_execute_cancelled_before_request_is_sent(self): cancel.assert_not_awaited() assert cursor.query_id is None - async def test_execute_timeout_while_starting(self): - """A timeout during the start request stops the query and raises TimeoutError (no AWS).""" - cursor, cancel, started, release = _starting_cursor() - task = asyncio.create_task(asyncio.wait_for(cursor.execute("SELECT 1"), timeout=0.05)) - try: - assert await asyncio.to_thread(started.wait, EVENT_TIMEOUT) - # The timeout is due before this sleep ends, so the event loop handles it - # while the start request is still blocked. - await asyncio.sleep(0.1) - assert not task.done() - finally: - release.set() - with pytest.raises(asyncio.TimeoutError): - await task - - assert started.is_set() - cancel.assert_awaited_once_with("query_id") - assert cursor.query_id == "query_id" - @pytest.mark.parametrize("failing", ["start", "cancel"]) async def test_execute_cancelled_while_starting_failure(self, failing): """A failure to start or cancel becomes the cause of the cancellation (no AWS).""" @@ -540,239 +464,71 @@ async def test_execute_cancelled_while_starting_without_kill_on_interrupt(self): cancel.assert_not_awaited() assert cursor.query_id is None - async def test_cache_size_different_schema(self): - """A cached result is only reused when it ran against the same schema (#739). - - Mirrors the synchronous cursor test: identical SQL can resolve to different - tables depending on the database it runs against, so a prior execution from - another schema must not be a cache hit. - """ - query = "SELECT * FROM one_row" - - def execution(schema): - return AthenaQueryExecution( - { - "QueryExecution": { - "QueryExecutionId": f"query_id_{schema}", - "Query": query, - "StatementType": AthenaQueryExecution.STATEMENT_TYPE_DML, - "QueryExecutionContext": {"Database": schema}, - "Status": { - "State": AthenaQueryExecution.STATE_SUCCEEDED, - "CompletionDateTime": datetime.now(UTC), - }, - } - } - ) - - cursor = AioCursor.__new__(AioCursor) # bypass __init__ to avoid AWS calls - cursor._catalog_name = None - - with patch.object( - AioCursor, - "_list_query_executions", - new_callable=AsyncMock, - return_value=(None, [execution("other_schema")]), - ): - cursor._schema_name = "this_schema" - assert await cursor._find_previous_query_id(query, None, cache_size=100) is None - cursor._schema_name = "other_schema" - assert ( - await cursor._find_previous_query_id(query, None, cache_size=100) - == "query_id_other_schema" - ) - - async def test_cache_size_different_catalog(self): - query = "SELECT * FROM one_row" - schema = "this_schema" - - def execution(catalog): - return AthenaQueryExecution( - { - "QueryExecution": { - "QueryExecutionId": f"query_id_{catalog}", - "Query": query, - "StatementType": AthenaQueryExecution.STATEMENT_TYPE_DML, - "QueryExecutionContext": {"Database": schema, "Catalog": catalog}, - "Status": { - "State": AthenaQueryExecution.STATE_SUCCEEDED, - "CompletionDateTime": datetime.now(UTC), - }, - } - } - ) - - cursor = AioCursor.__new__(AioCursor) - cursor._schema_name = schema - - with patch.object( - AioCursor, - "_list_query_executions", - new_callable=AsyncMock, - return_value=(None, [execution("awsdatacatalog")]), - ): - # A different catalog must not be a cache hit. - cursor._catalog_name = "other_catalog" - assert await cursor._find_previous_query_id(query, None, cache_size=100) is None - # The same catalog, differing only in case, must still be a cache hit. - cursor._catalog_name = "AwsDataCatalog" - assert ( - await cursor._find_previous_query_id(query, None, cache_size=100) - == "query_id_awsdatacatalog" - ) - - async def test_cache_search_stops_at_expired_execution(self, caplog): + async def test_cache_search_stops_at_expired_execution(self): query = "SELECT * FROM one_row" now = datetime.now(UTC) - cursor = AioCursor.__new__(AioCursor) - cursor._schema_name = "this_schema" - cursor._catalog_name = None - page = [ - succeeded_query_execution("expired", query, now - timedelta(hours=2)), - succeeded_query_execution("other", "SELECT 1", now), - ] - - # A page read after the expired execution would return this match. - next_page = (None, [succeeded_query_execution("next_page", query, now)]) - - with patch.object( - AioCursor, - "_list_query_executions", - new_callable=AsyncMock, - side_effect=[("next_token", page), next_page], - ) as list_mock: - # Without cache_size, cache_expiration_time alone bounds the search. - assert ( - await cursor._find_previous_query_id(query, None, cache_expiration_time=3600) - is None - ) - list_mock.assert_awaited_once_with(None, next_token=None, max_results=50) - # A failed search also returns None; the expiry must stop it without an error. - assert "Failed to check the cache" not in caplog.text - - async def test_cache_search_reads_pages_up_to_cache_size(self): - query = "SELECT * FROM one_row" - now = datetime.now(UTC) - cursor = AioCursor.__new__(AioCursor) - cursor._schema_name = "this_schema" - cursor._catalog_name = None - pages = [ - ("next_token", [succeeded_query_execution("other", "SELECT 1", now)]), - ("last_token", [succeeded_query_execution("match", query, now)]), - ] - - with patch.object( - AioCursor, "_list_query_executions", new_callable=AsyncMock, side_effect=pages - ) as list_mock: - assert await cursor._find_previous_query_id(query, "wg", cache_size=60) == "match" - assert list_mock.await_args_list == [ - call("wg", next_token=None, max_results=50), - call("wg", next_token="next_token", max_results=10), + page_size = AioCursor.LIST_QUERY_EXECUTIONS_MAX_RESULTS + executions = [ + succeeded_query_execution_data("other", "SELECT 1", now), + succeeded_query_execution_data("expired", query, now - timedelta(hours=2)), + *( + succeeded_query_execution_data(f"older_{i}", "SELECT 1", now - timedelta(hours=3)) + for i in range(page_size - 2) + ), + # A search that went on to the second page would return this match. + succeeded_query_execution_data("next_page", query, now), ] + history = QueryExecutionHistory(executions) + cursor = cache_search_cursor(AioCursor, history) + + # Without cache_size, cache_expiration_time alone bounds the search. + assert await cursor._find_previous_query_id(query, None, cache_expiration_time=3600) is None + assert len(history.list_requests) == 1 + # A failed search would also return None, but the same search reuses the + # execution once it is within cache_expiration_time. + assert ( + await cursor._find_previous_query_id(query, None, cache_expiration_time=3 * 3600) + == "expired" + ) - async def test_cache_search_prefers_latest_execution(self): + @pytest.mark.parametrize( + ("position", "expected"), + [(59, "match"), (60, None)], + ids=["last_searched", "after_cache_size"], + ) + async def test_cache_search_reads_pages_up_to_cache_size(self, position, expected): + """With cache_size=60, the search reads the 60 latest executions over two pages.""" query = "SELECT * FROM one_row" now = datetime.now(UTC) - cursor = AioCursor.__new__(AioCursor) - cursor._schema_name = "this_schema" - cursor._catalog_name = None - page = [ - succeeded_query_execution("older", query, now - timedelta(minutes=1)), - succeeded_query_execution("latest", query, now), + executions = [ + succeeded_query_execution_data(f"other_{i}", "SELECT 1", now) for i in range(70) ] + executions[position] = succeeded_query_execution_data("match", query, now) + history = QueryExecutionHistory(executions) + cursor = cache_search_cursor(AioCursor, history) - with patch.object( - AioCursor, - "_list_query_executions", - new_callable=AsyncMock, - return_value=(None, page), - ): - assert await cursor._find_previous_query_id(query, None, cache_size=10) == "latest" - - async def test_no_result_set_raises(self, aio_cursor): - with pytest.raises(ProgrammingError): - await aio_cursor.fetchone() - with pytest.raises(ProgrammingError): - await aio_cursor.fetchmany() - with pytest.raises(ProgrammingError): - await aio_cursor.fetchall() + assert await cursor._find_previous_query_id(query, "wg", cache_size=60) == expected + assert [request["WorkGroup"] for request in history.list_requests] == ["wg", "wg"] async def test_description(self, aio_cursor): await aio_cursor.execute("SELECT 1 AS foobar FROM one_row") assert await aio_cursor.fetchall() == [(1,)] assert aio_cursor.description == [("foobar", "integer", None, None, 10, 0, "UNKNOWN")] - async def test_description_initial(self, aio_cursor): - assert aio_cursor.description is None - async def test_bad_query(self, aio_cursor): with pytest.raises(DatabaseError): await aio_cursor.execute("SELECT does_not_exist FROM this_really_does_not_exist") - async def test_query_id(self, aio_cursor): - assert aio_cursor.query_id is None - await aio_cursor.execute("SELECT * FROM one_row") - expected_pattern = ( - r"^[0-9a-f]{8}-[0-9a-f]{4}-[4][0-9a-f]{3}-[89ab][0-9a-f]{3}-[0-9a-f]{12}$" - ) - assert re.match(expected_pattern, aio_cursor.query_id) - - async def test_query_execution_initial(self, aio_cursor): - assert not aio_cursor.has_result_set - assert aio_cursor.rownumber is None - assert aio_cursor.rowcount == -1 - assert aio_cursor.database is None - assert aio_cursor.catalog is None - assert aio_cursor.query_id is None - assert aio_cursor.query is None - assert aio_cursor.statement_type is None - assert aio_cursor.work_group is None - assert aio_cursor.state is None - assert aio_cursor.state_change_reason is None - assert aio_cursor.submission_date_time is None - assert aio_cursor.completion_date_time is None - assert aio_cursor.data_scanned_in_bytes is None - assert aio_cursor.engine_execution_time_in_millis is None - assert aio_cursor.query_queue_time_in_millis is None - assert aio_cursor.total_execution_time_in_millis is None - assert aio_cursor.query_planning_time_in_millis is None - assert aio_cursor.service_processing_time_in_millis is None - assert aio_cursor.output_location is None - assert aio_cursor.data_manifest_location is None - assert aio_cursor.encryption_option is None - assert aio_cursor.kms_key is None - assert aio_cursor.selected_engine_version is None - assert aio_cursor.effective_engine_version is None - - async def test_cancel_initial(self, aio_cursor): - with pytest.raises(ProgrammingError): - await aio_cursor.cancel() - - async def test_executemany(self, aio_cursor, empty_table): - rows = [(1, "foo"), (2, "bar"), (3, "jim o'rourke")] - await aio_cursor.executemany( - f"INSERT INTO {empty_table} (a, b) VALUES (%(a)d, %(b)s)", - [{"a": a, "b": b} for a, b in rows], - ) - assert aio_cursor.rowcount == len(rows) - await aio_cursor.execute(f"SELECT * FROM {empty_table}") - assert sorted(await aio_cursor.fetchall()) == list(rows) - + # The sum of the update counts does not depend on the statement; the synchronous + # test_iceberg_table checks the update count of INSERT, UPDATE, and DELETE. @pytest.mark.parametrize( ("operation", "expected"), [ - pytest.param( - "INSERT INTO {table} SELECT id+10, group_id+10, value " - "FROM {table} WHERE group_id=%(group_id)d", - [(1, 10), (2, 20), (3, 30), (11, 10), (12, 20), (13, 30)], - id="insert", - ), pytest.param( "UPDATE {table} SET value=value+1 WHERE group_id=%(group_id)d", [(1, 11), (2, 21), (3, 31)], id="update", ), - pytest.param("DELETE FROM {table} WHERE group_id=%(group_id)d", [], id="delete"), ], ) async def test_executemany_rowcount(self, aio_cursor, executemany_table, operation, expected): @@ -802,7 +558,9 @@ async def test_executemany_rowcount(self, aio_cursor, executemany_table, operati assert aio_cursor.rowcount == -1 assert await aio_cursor.fetchall() == expected - @pytest.mark.parametrize("failure_index", [0, 1]) + # A failure of the first statement reaches the same code; this case also shows + # that the statements before the failure are not rolled back. + @pytest.mark.parametrize("failure_index", [1]) async def test_executemany_failure(self, aio_cursor, executemany_table, failure_index): operation = ( f"UPDATE {executemany_table} SET value=value+1 " @@ -828,31 +586,6 @@ async def test_executemany_failure(self, aio_cursor, executemany_table, failure_ await aio_cursor.execute(operation, {"group_id": "2"}) assert aio_cursor.rowcount == 1 - async def test_executemany_parameter_iteration_failure(self, aio_cursor, executemany_table): - previous = None - query_id = None - - def parameters(): - nonlocal previous, query_id - yield {"group_id": 1} - previous = aio_cursor.result_set - query_id = aio_cursor.query_id - raise ValueError("invalid parameters") - - with pytest.raises(ValueError, match="invalid parameters"): - await aio_cursor.executemany( - f"UPDATE {executemany_table} SET value=value+1 WHERE group_id=%(group_id)d", - parameters(), - ) - assert aio_cursor.rowcount == -1 - assert aio_cursor.result_set is None - assert query_id is not None - assert aio_cursor.query_id == query_id - assert previous is not None - assert previous.is_closed - await aio_cursor.execute(f"SELECT id, value FROM {executemany_table} ORDER BY id") - assert await aio_cursor.fetchall() == [(1, 11), (2, 21), (3, 30)] - @pytest.mark.parametrize("aio_cursor", [{"kill_on_interrupt": False}], indirect=True) async def test_executemany_cancellation(self, aio_cursor, executemany_table): query_ids = [] @@ -883,31 +616,6 @@ def on_start(query_id): await aio_cursor._cancel(query_ids[-1]) await aio_cursor._poll(query_ids[-1]) - async def test_executemany_fetch(self, aio_cursor): - await aio_cursor.executemany("SELECT %(x)d FROM one_row", []) - assert aio_cursor.rowcount == 0 - await aio_cursor.executemany("SELECT %(x)d FROM one_row", [{"x": i} for i in range(1, 2)]) - assert aio_cursor.rowcount == -1 - with pytest.raises(ProgrammingError): - await aio_cursor.fetchall() - with pytest.raises(ProgrammingError): - await aio_cursor.fetchmany() - with pytest.raises(ProgrammingError): - await aio_cursor.fetchone() - - async def test_context_manager(self): - conn = await _aio_connect(schema_name=ENV.schema) - try: - async with conn.cursor() as cursor: - await cursor.execute("SELECT * FROM one_row") - assert await cursor.fetchone() == (1,) - finally: - conn.close() - - async def test_open_close(self): - conn = await _aio_connect() - conn.close() - async def test_aio_connect(self): from pyathena import aio_connect @@ -923,55 +631,23 @@ async def test_arraysize(self, aio_cursor): actual = await aio_cursor.fetchmany() assert len(actual) == 5 - async def test_arraysize_default(self, aio_cursor): - assert aio_cursor.arraysize == AthenaResultSet.DEFAULT_FETCH_SIZE + async def test_glue_request_runs_off_the_event_loop(self, monkeypatch): + """The Glue request of a throttled metadata lookup does not block the loop (no AWS).""" + conn = _offline_connection() + throttle_metadata_api(conn.client, monkeypatch) + threads = [] - async def test_invalid_arraysize(self, aio_cursor): - with pytest.raises(ProgrammingError): - aio_cursor.arraysize = 10000 - with pytest.raises(ProgrammingError): - aio_cursor.arraysize = -1 + def get_table(**kwargs): + threads.append(threading.get_ident()) + return {"Table": {"Name": kwargs["Name"], "StorageDescriptor": {"Columns": []}}} - async def test_glue_request_runs_off_the_event_loop(self, aio_cursor, monkeypatch): - throttle_metadata_api(aio_cursor.connection.client, monkeypatch) - loop_thread = threading.get_ident() - threads = [] - # Record where the real client is built and the real request is sent. - client = GlueMetadataClient.client - get_table = GlueMetadataClient.get_table - - def recorded_client(self): - threads.append(("client", threading.get_ident())) - return client.fget(self) - - def recorded_get_table(self, *args, **kwargs): - threads.append(("request", threading.get_ident())) - return get_table(self, *args, **kwargs) - - monkeypatch.setattr(GlueMetadataClient, "client", property(recorded_client)) - monkeypatch.setattr(GlueMetadataClient, "get_table", recorded_get_table) - - assert (await aio_cursor.get_table_metadata("one_row")).name == "one_row" - # Building the client and sending the request would block the loop. - assert [step for step, _ in threads] == ["request", "client"] - assert all(thread != loop_thread for _, thread in threads) - - async def test_list_databases(self, aio_cursor): - databases = await aio_cursor.list_databases(catalog_name="AwsDataCatalog") - assert len(databases) > 0 - database_names = [db.name for db in databases] - assert "default" in database_names - - async def test_get_table_metadata(self, aio_cursor): - metadata = await aio_cursor.get_table_metadata(table_name="one_row") - assert metadata.name == "one_row" - assert metadata.table_type - - async def test_list_table_metadata(self, aio_cursor): - metadata_list = await aio_cursor.list_table_metadata() - assert len(metadata_list) > 0 - table_names = [m.name for m in metadata_list] - assert "one_row" in table_names + glue = MagicMock() + glue.get_table.side_effect = get_table + monkeypatch.setattr(conn.session, "client", lambda *args, **kwargs: glue) + + assert (await conn.cursor().get_table_metadata("one_row")).name == "one_row" + assert len(threads) == 1 + assert threading.get_ident() not in threads async def test_throttled_metadata_reads_glue(self, aio_cursor, monkeypatch): def view(metadata): @@ -985,6 +661,10 @@ async def read(): ) expected = await read() + table, listing, has_schema = expected + assert table[0] == "one_row" + assert "one_row" in [name for name, _, _ in listing] + assert has_schema calls = throttle_metadata_api(aio_cursor.connection.client, monkeypatch) # The Glue request runs in a worker thread, as the Athena calls do. @@ -997,30 +677,7 @@ async def test_fetchone(self, aio_dict_cursor): await aio_dict_cursor.execute("SELECT * FROM one_row") assert await aio_dict_cursor.fetchone() == {"number_of_rows": 1} - async def test_fetchmany(self, aio_dict_cursor): - await aio_dict_cursor.execute("SELECT * FROM many_rows LIMIT 15") - actual1 = await aio_dict_cursor.fetchmany(10) - assert len(actual1) == 10 - assert actual1 == [{"a": i} for i in range(10)] - actual2 = await aio_dict_cursor.fetchmany(10) - assert len(actual2) == 5 - assert actual2 == [{"a": i} for i in range(10, 15)] - - async def test_fetchall(self, aio_dict_cursor): - await aio_dict_cursor.execute("SELECT * FROM one_row") - assert await aio_dict_cursor.fetchall() == [{"number_of_rows": 1}] - await aio_dict_cursor.execute("SELECT a FROM many_rows ORDER BY a") - assert await aio_dict_cursor.fetchall() == [{"a": i} for i in range(10000)] - async def test_dict_type(self, aio_dict_cursor): - async with aio_dict_cursor.connection.cursor(dict_type=OrderedDict) as ordered_cursor: - await ordered_cursor.execute("SELECT * FROM one_row") - assert type(await ordered_cursor.fetchone()) is OrderedDict - # dict_type of another cursor does not change the row type of this one. - await aio_dict_cursor.execute("SELECT * FROM one_row") - assert type(await aio_dict_cursor.fetchone()) is dict - - async def test_dict_type_custom_result_set(self, aio_dict_cursor): class CustomResultSet(AthenaAioDictResultSet): pass @@ -1031,7 +688,10 @@ def __init__(self, **kwargs): async with aio_dict_cursor.connection.cursor( CustomDictCursor, dict_type=OrderedDict - ) as cursor: - await cursor.execute("SELECT * FROM one_row") - assert isinstance(cursor.result_set, CustomResultSet) - assert type(await cursor.fetchone()) is OrderedDict + ) as ordered_cursor: + await ordered_cursor.execute("SELECT * FROM one_row") + assert isinstance(ordered_cursor.result_set, CustomResultSet) + assert type(await ordered_cursor.fetchone()) is OrderedDict + # dict_type of another cursor does not change the row type of this one. + await aio_dict_cursor.execute("SELECT * FROM one_row") + assert type(await aio_dict_cursor.fetchone()) is dict diff --git a/tests/pyathena/aio/test_result_set.py b/tests/pyathena/aio/test_result_set.py index 793c55c20..ee2fe890d 100644 --- a/tests/pyathena/aio/test_result_set.py +++ b/tests/pyathena/aio/test_result_set.py @@ -157,9 +157,3 @@ async def test_fetch_after_close(self): assert await result_set.fetchmany() == [] assert await result_set.fetchall() == [] assert result_set.rownumber is None - - async def test_get_query_results_after_close(self): - result_set = await _create_result_set(AthenaAioResultSet, [("1", "2"), ("3",)]) - result_set.close() - with pytest.raises(ProgrammingError, match=r"AthenaAioResultSet is closed\."): - await result_set._async_get_query_results(2, "token1") diff --git a/tests/pyathena/aio/test_util.py b/tests/pyathena/aio/test_util.py index 66a8957a0..6c8215255 100644 --- a/tests/pyathena/aio/test_util.py +++ b/tests/pyathena/aio/test_util.py @@ -25,20 +25,15 @@ _VALUE: contextvars.ContextVar[str] = contextvars.ContextVar("_VALUE", default="unset") -async def _cancel_once_started(task, operation, stop_in): - """Cancel a task once its operation has made its first attempt. +async def _cancel_once_started(task, operation): + """Cancel a task during the first attempt of its operation. Args: task: The task that runs the operation in a thread. operation: The blocking ``ThrottledOperation`` of the task. - stop_in: "request" to cancel during the first attempt, "backoff" to - cancel during the backoff wait after it. """ # Polled, since the default executor may have no free thread. await wait_until(operation.started.is_set) - if stop_in == "backoff": - # Lets the first attempt fail and the backoff wait begin. - await asyncio.sleep(0.1) task.cancel() with pytest.raises(asyncio.CancelledError): await task @@ -315,8 +310,8 @@ async def test_run_in_thread(): assert value == "caller" -@pytest.mark.parametrize("stop_in", ["request", "backoff"]) -async def test_async_retry_api_call_cancelled(stop_in): +# The synchronous test_retry_api_call_stopped_by_event checks the stop during the backoff wait. +async def test_async_retry_api_call_cancelled(): """GH-1202: cancelling the caller makes no further attempt and ends the backoff wait. Without the stop, the thread would wait at least EVENT_TIMEOUT seconds before @@ -325,12 +320,12 @@ async def test_async_retry_api_call_cancelled(stop_in): loop = asyncio.get_running_loop() executor = ThreadPoolExecutor(max_workers=1) loop.set_default_executor(executor) - operation = ThrottledOperation(block=stop_in == "request") + operation = ThrottledOperation(block=True) try: task = asyncio.create_task( async_retry_api_call(operation, RetryConfig(attempt=2, multiplier=EVENT_TIMEOUT)) ) - await _cancel_once_started(task, operation, stop_in) + await _cancel_once_started(task, operation) started = time.monotonic() # Runs once the thread of the cancelled call is free. await loop.run_in_executor(None, lambda: None) @@ -341,19 +336,6 @@ async def test_async_retry_api_call_cancelled(stop_in): executor.shutdown(wait=True) -async def test_async_retry_api_call_not_stoppable(): - """GH-1202: a cleanup request that is not stoppable keeps its retries after a cancellation.""" - operation = ThrottledOperation(block=True, failures=2) - task = asyncio.create_task( - async_retry_api_call(operation, RetryConfig(attempt=3, multiplier=0), stoppable=False) - ) - try: - await _cancel_once_started(task, operation, "request") - await wait_until(lambda: operation.calls == 3) - finally: - operation.release.set() - - async def test_async_retry_api_call_not_stoppable_in_stopped_scope(): """A request that is not stoppable ignores the stop scopes of its caller.""" stop = threading.Event() @@ -367,17 +349,9 @@ async def test_async_retry_api_call_not_stoppable_in_stopped_scope(): assert operation.calls == 3 -async def test_async_retry_api_call_retries(): - """Without a cancellation, the retries continue.""" - operation = ThrottledOperation(failures=2) - assert await async_retry_api_call(operation, RetryConfig(attempt=3, multiplier=0)) == "success" - assert operation.calls == 3 - - -@pytest.mark.parametrize("stop_in", ["request", "backoff"]) -async def test_create_in_thread_cancelled_stops_retries(stop_in): +async def test_create_in_thread_cancelled_stops_retries(): """GH-1202: cancelling the caller stops the retries of a running construction.""" - operation = ThrottledOperation(block=stop_in == "request") + operation = ThrottledOperation(block=True) finished = threading.Event() def factory(): @@ -388,7 +362,7 @@ def factory(): task = asyncio.create_task(create_in_thread(factory)) try: - await _cancel_once_started(task, operation, stop_in) + await _cancel_once_started(task, operation) started = time.monotonic() assert await asyncio.to_thread(finished.wait, EVENT_TIMEOUT) assert time.monotonic() - started < EVENT_TIMEOUT diff --git a/tests/pyathena/spark/test_async_cursor.py b/tests/pyathena/spark/test_async_cursor.py index bcc332ba1..6ad3c6d7a 100644 --- a/tests/pyathena/spark/test_async_cursor.py +++ b/tests/pyathena/spark/test_async_cursor.py @@ -13,13 +13,15 @@ import pytest from pyathena import NotSupportedError, OperationalError -from pyathena.model import AthenaCalculationExecutionStatus, AthenaSessionStatus +from pyathena.model import AthenaCalculationExecutionStatus from pyathena.spark.async_cursor import AsyncSparkCursor from tests import ENV from tests.pyathena.util import ( CANCELABLE_SPARK_JOB, + COMPLETED_CALCULATION_RESPONSES, + answered_requests, + offline_spark_cursor, wait_for_spark_job, - wait_for_spark_session_state, ) # Bounds how long the executor test task blocks when nothing releases it. @@ -154,25 +156,6 @@ def test_cancel(self, async_spark_cursor): == AthenaCalculationExecutionStatus.STATE_COMPLETED ) - def test_session_ownership(self, async_spark_cursor): - client = async_spark_cursor.connection.client - session_id = async_spark_cursor.session_id - with async_spark_cursor.connection.cursor( - AsyncSparkCursor, session_id=session_id - ) as borrower: - _, future = borrower.execute("print(1)") - calculation_execution = future.result() - assert borrower.get_std_out(calculation_execution).result() == "1" - - # Closing a cursor that was given the session leaves the session running. - _, future = async_spark_cursor.execute("print(2)") - calculation_execution = future.result() - assert async_spark_cursor.get_std_out(calculation_execution).result() == "2" - - # Closing the cursor that started the session terminates it. - async_spark_cursor.close() - wait_for_spark_session_state(client, session_id, AthenaSessionStatus.STATE_TERMINATED) - @staticmethod def _cursor_with_submitted_work(): """Build a cursor whose executor holds one running and one queued future. @@ -244,19 +227,6 @@ def test_close_no_wait_shuts_down_executor(self, fails): assert running.result(_TIMEOUT) == "running" assert queued.result(_TIMEOUT) == "queued" - def test_close_retries_termination_after_failure(self): - cursor = _owned_session_cursor() - cursor._terminate_session = MagicMock( - side_effect=[OperationalError("termination failed"), None] - ) - cursor._executor = ThreadPoolExecutor(max_workers=1) - - with pytest.raises(OperationalError, match="termination failed"): - cursor.close() - cursor.close() - - assert cursor._terminate_session.call_count == 2 - @pytest.mark.parametrize("parameters", [{}, [], {"value": 1}]) def test_execute_rejects_parameters(self, parameters): """Any parameters other than None raise before a calculation starts (no AWS).""" @@ -295,20 +265,15 @@ def test_execute_rejects_work_group(self): ) def test_execute_accepts_none_parameters(self, args, kwargs): """Omitted or None parameters start the calculation and its polling (no AWS).""" - cursor = AsyncSparkCursor.__new__(AsyncSparkCursor) # bypass __init__ to avoid AWS calls - cursor._session_id = "session_id" - cursor._executor = MagicMock() - - with patch.object( - AsyncSparkCursor, "_calculate", return_value="calculation_id" - ) as calculate: - query_id, future = cursor.execute("code", *args, **kwargs) + cursor = offline_spark_cursor(AsyncSparkCursor) + with answered_requests( + cursor.connection.client, COMPLETED_CALCULATION_RESPONSES + ) as requests: + calculation_id, future = cursor.execute("code", *args, **kwargs) + calculation_execution = future.result(_TIMEOUT) + cursor.close() - calculate.assert_called_once_with( - session_id="session_id", - code_block="code", - description=None, - client_request_token=None, - ) - assert query_id == "calculation_id" - assert future is cursor._executor.submit.return_value + start = dict(requests)["StartCalculationExecution"] + assert (start["SessionId"], start["CodeBlock"]) == ("session", "code") + assert calculation_id == "calculation-id" + assert calculation_execution.state == AthenaCalculationExecutionStatus.STATE_COMPLETED diff --git a/tests/pyathena/spark/test_common.py b/tests/pyathena/spark/test_common.py index ad83598b3..50923a4d3 100644 --- a/tests/pyathena/spark/test_common.py +++ b/tests/pyathena/spark/test_common.py @@ -6,7 +6,6 @@ # SPDX-License-Identifier: MIT import asyncio -import logging import threading import uuid from unittest.mock import MagicMock, PropertyMock, patch @@ -28,20 +27,6 @@ _TIMEOUT = 10 -def _session_status(state: str, reason: str | None = None) -> AthenaSessionStatus: - return AthenaSessionStatus( - {"SessionId": "session_id", "Status": {"State": state, "StateChangeReason": reason}} - ) - - -def _cursor() -> SparkCursor: - cursor = SparkCursor.__new__(SparkCursor) # bypass __init__ to avoid AWS calls - cursor._connection = MagicMock() - cursor._retry_config = RetryConfig() - cursor._poll_interval = 0 - return cursor - - def _connection(): connection = MagicMock() connection.client.start_session.return_value = {"SessionId": "new-session"} @@ -143,75 +128,40 @@ def _close(cursor) -> None: class TestSparkBaseCursor: - @pytest.mark.parametrize( - "state", - [ - AthenaSessionStatus.STATE_TERMINATED, - AthenaSessionStatus.STATE_DEGRADED, - AthenaSessionStatus.STATE_FAILED, - ], - ) - def test_wait_for_idle_session_raises_on_failure_state(self, state): - cursor = _cursor() - with ( - patch.object( - SparkCursor, - "_get_session_status", - return_value=_session_status(state, "session failure reason"), - ), - patch("pyathena.spark.common.time.sleep", side_effect=AssertionError("slept")), - pytest.raises( - OperationalError, - match=rf"^Session: session_id is {state}\. session failure reason$", - ), - ): - cursor._wait_for_idle_session("session_id") - - def test_wait_for_idle_session_raises_without_reason(self): - cursor = _cursor() - with ( - patch.object( - SparkCursor, - "_get_session_status", - return_value=_session_status(AthenaSessionStatus.STATE_TERMINATED), - ), - patch("pyathena.spark.common.time.sleep", side_effect=AssertionError("slept")), - pytest.raises(OperationalError, match=r"^Session: session_id is TERMINATED\.$"), - ): - cursor._wait_for_idle_session("session_id") - def test_wait_for_idle_session_waits_until_idle(self): - cursor = _cursor() - statuses = [ - _session_status(AthenaSessionStatus.STATE_CREATING), - _session_status(AthenaSessionStatus.STATE_BUSY), - _session_status(AthenaSessionStatus.STATE_IDLE), + connection = _connection() + connection.client.get_session_status.side_effect = [ + {"Status": {"State": state}} + for state in ( + AthenaSessionStatus.STATE_CREATING, + AthenaSessionStatus.STATE_BUSY, + AthenaSessionStatus.STATE_IDLE, + ) ] - with ( - patch.object(SparkCursor, "_get_session_status", side_effect=statuses) as get_status, - patch("pyathena.spark.common.time.sleep") as sleep, - ): - cursor._wait_for_idle_session("session_id") - assert get_status.call_count == 3 - assert sleep.call_count == 2 + + cursor = _init_cursor(SparkCursor, connection) + + assert cursor.session_id == "new-session" + # The status is requested until the session is idle. + assert connection.client.get_session_status.call_count == 3 + connection.client.terminate_session.assert_not_called() def test_exists_session_raises_on_failure_state(self): - cursor = _cursor() - with ( - patch.object( - SparkCursor, - "_get_session_status", - return_value=_session_status( - AthenaSessionStatus.STATE_TERMINATED, "session failure reason" - ), - ), - patch("pyathena.spark.common.time.sleep", side_effect=AssertionError("slept")), - pytest.raises(OperationalError, match="session failure reason"), - ): - cursor._exists_session("session_id") - cursor._connection.client.get_session.assert_called_once_with(SessionId="session_id") + connection = _connection() + connection.client.get_session_status.return_value = { + "Status": { + "State": AthenaSessionStatus.STATE_TERMINATED, + "StateChangeReason": "session failure reason", + } + } - @pytest.mark.parametrize("cursor_class", SYNC_SPARK_CURSOR_CLASSES) + with pytest.raises(OperationalError): + _init_cursor(SparkCursor, connection, session_id="supplied-session") + + connection.client.start_session.assert_not_called() + connection.client.terminate_session.assert_not_called() + + @pytest.mark.parametrize("cursor_class", [SparkCursor]) @pytest.mark.parametrize("kill_on_interrupt", [True, False]) def test_calculate_reuses_generated_token_on_retry(self, cursor_class, kill_on_interrupt): cursor = _calculation_cursor(cursor_class, kill_on_interrupt=kill_on_interrupt) @@ -238,7 +188,7 @@ def test_calculate_reuses_generated_token_on_retry(self, cursor_class, kill_on_i thread.join(_TIMEOUT) assert not thread.is_alive() - @pytest.mark.parametrize("cursor_class", SYNC_SPARK_CURSOR_CLASSES) + @pytest.mark.parametrize("cursor_class", [SparkCursor]) def test_calculate_generates_token_per_call(self, cursor_class): cursor = _calculation_cursor(cursor_class) client = cursor._connection.client @@ -252,7 +202,7 @@ def test_calculate_generates_token_per_call(self, cursor_class): ] assert tokens[0] != tokens[1] - @pytest.mark.parametrize("cursor_class", SYNC_SPARK_CURSOR_CLASSES) + @pytest.mark.parametrize("cursor_class", [SparkCursor]) def test_calculate_keeps_caller_token(self, cursor_class): cursor = _calculation_cursor(cursor_class) @@ -262,32 +212,7 @@ def test_calculate_keeps_caller_token(self, cursor_class): SessionId="session_id", CodeBlock="code", ClientRequestToken="token" ) - @pytest.mark.parametrize("cursor_class", SYNC_SPARK_CURSOR_CLASSES) - @pytest.mark.parametrize( - "final_state", - [ - AthenaCalculationExecutionStatus.STATE_CANCELED, - AthenaCalculationExecutionStatus.STATE_COMPLETED, - ], - ) - def test_calculate_interrupted_while_starting(self, cursor_class, final_state): - cursor = _calculation_cursor(cursor_class) - final_execution = MagicMock(state=final_state) - cursor._get_calculation_execution.return_value = final_execution - started, release = _block_start(cursor) - waits, raised = interrupt_start_waits(started, release) - - with waits, pytest.raises(KeyboardInterrupt) as exc_info: - cursor._calculate(session_id="session_id", code_block="code") - - assert exc_info.value is raised[0] - assert exc_info.value.__cause__ is None - cursor._connection.client.start_calculation_execution.assert_called_once() - cursor._cancel.assert_called_once_with("calculation_id") - assert cursor.calculation_id == "calculation_id" - assert cursor._calculation_execution is final_execution - - @pytest.mark.parametrize("cursor_class", SYNC_SPARK_CURSOR_CLASSES) + @pytest.mark.parametrize("cursor_class", [SparkCursor]) @pytest.mark.parametrize("failing", ["start", "cancel", "wait"]) def test_calculate_interrupted_while_starting_failure(self, cursor_class, failing): cursor = _calculation_cursor(cursor_class) @@ -321,7 +246,7 @@ def test_calculate_interrupted_while_starting_failure(self, cursor_class, failin cursor._cancel.assert_called_once_with("calculation_id") assert cursor.calculation_id == "calculation_id" - @pytest.mark.parametrize("cursor_class", SYNC_SPARK_CURSOR_CLASSES) + @pytest.mark.parametrize("cursor_class", [SparkCursor]) def test_calculate_second_interrupt_while_starting(self, cursor_class): cursor = _calculation_cursor(cursor_class) started, release = _block_start(cursor) @@ -337,7 +262,7 @@ def test_calculate_second_interrupt_while_starting(self, cursor_class): assert exc_info.value.__context__ is raised[0] cursor._cancel.assert_not_called() - @pytest.mark.parametrize("cursor_class", SYNC_SPARK_CURSOR_CLASSES) + @pytest.mark.parametrize("cursor_class", [SparkCursor]) @pytest.mark.parametrize("helper_starts", [False, True]) def test_calculate_interrupted_before_request_is_sent(self, cursor_class, helper_starts): cursor = _calculation_cursor(cursor_class) @@ -366,42 +291,32 @@ def start(self): cursor._cancel.assert_not_called() assert cursor.calculation_id is None - @pytest.mark.parametrize("cursor_class", SYNC_SPARK_CURSOR_CLASSES) - def test_calculate_interrupt_without_kill_on_interrupt(self, cursor_class): - cursor = _calculation_cursor(cursor_class, kill_on_interrupt=False) - cursor._connection.client.start_calculation_execution.side_effect = KeyboardInterrupt() - - with ( - patch("pyathena.common.threading.Thread") as thread, - pytest.raises(KeyboardInterrupt), - ): - cursor._calculate(session_id="session_id", code_block="code") - - thread.assert_not_called() - cursor._connection.client.start_calculation_execution.assert_called_once() - cursor._cancel.assert_not_called() - - def test_execute_interrupted_while_starting(self): + @pytest.mark.parametrize( + "final_state", + [ + AthenaCalculationExecutionStatus.STATE_CANCELED, + AthenaCalculationExecutionStatus.STATE_COMPLETED, + ], + ) + def test_execute_interrupted_while_starting(self, final_state): cursor = _calculation_cursor(SparkCursor) + final_execution = MagicMock(state=final_state) + cursor._get_calculation_execution.return_value = final_execution started, release = _block_start(cursor) - waits, _ = interrupt_start_waits(started, release) + waits, raised = interrupt_start_waits(started, release) - with waits, pytest.raises(KeyboardInterrupt): + with waits, pytest.raises(KeyboardInterrupt) as exc_info: cursor.execute("code") + assert exc_info.value is raised[0] + assert exc_info.value.__cause__ is None + cursor._connection.client.start_calculation_execution.assert_called_once() + cursor._cancel.assert_called_once_with("calculation_id") assert cursor.calculation_id == "calculation_id" - assert cursor.state == AthenaCalculationExecutionStatus.STATE_CANCELED + assert cursor.calculation_execution is final_execution + assert cursor.state == final_state - @pytest.mark.parametrize("cursor_class", SPARK_CURSOR_CLASSES) - def test_init_starts_session(self, cursor_class): - connection = _connection() - with patch.object(SparkBaseCursor, "_wait_for_idle_session"): - cursor = _init_cursor(cursor_class, connection) - - assert cursor.session_id == "new-session" - connection.client.terminate_session.assert_not_called() - - @pytest.mark.parametrize("cursor_class", SPARK_CURSOR_CLASSES) + @pytest.mark.parametrize("cursor_class", SYNC_SPARK_CURSOR_CLASSES) @pytest.mark.parametrize( "error", [OperationalError("Session did not become idle."), KeyboardInterrupt()], @@ -417,7 +332,7 @@ def test_init_terminates_new_session_that_does_not_become_idle(self, cursor_clas assert exc_info.value is error connection.client.terminate_session.assert_called_once_with(SessionId="new-session") - @pytest.mark.parametrize("cursor_class", SPARK_CURSOR_CLASSES) + @pytest.mark.parametrize("cursor_class", SYNC_SPARK_CURSOR_CLASSES) @pytest.mark.parametrize( "state", [ @@ -434,18 +349,14 @@ def test_init_terminates_new_session_in_failure_state(self, cursor_class, state) } with ( patch("pyathena.spark.common.time.sleep", side_effect=AssertionError("slept")), - pytest.raises( - OperationalError, - match=rf"^Session: new-session is {state}\. session failure reason$", - ), + pytest.raises(OperationalError, match="session failure reason"), ): _init_cursor(cursor_class, connection) - connection.client.get_session_status.assert_called_once_with(SessionId="new-session") connection.client.terminate_session.assert_called_once_with(SessionId="new-session") - @pytest.mark.parametrize("cursor_class", SPARK_CURSOR_CLASSES) - def test_init_keeps_original_error_when_cleanup_fails(self, cursor_class, caplog): + @pytest.mark.parametrize("cursor_class", SYNC_SPARK_CURSOR_CLASSES) + def test_init_keeps_original_error_when_cleanup_fails(self, cursor_class): connection = _connection() connection.client.terminate_session.side_effect = ClientError( {"Error": {"Code": "InternalServerException", "Message": "Cleanup failed."}}, @@ -453,7 +364,6 @@ def test_init_keeps_original_error_when_cleanup_fails(self, cursor_class, caplog ) error = OperationalError("Session did not become idle.") with ( - caplog.at_level(logging.ERROR, logger="pyathena.spark.common"), patch.object(SparkBaseCursor, "_wait_for_idle_session", side_effect=error), pytest.raises(OperationalError) as exc_info, ): @@ -461,9 +371,8 @@ def test_init_keeps_original_error_when_cleanup_fails(self, cursor_class, caplog assert exc_info.value is error connection.client.terminate_session.assert_called_once_with(SessionId="new-session") - assert "Failed to terminate session: new-session." in caplog.text - @pytest.mark.parametrize("cursor_class", SPARK_CURSOR_CLASSES) + @pytest.mark.parametrize("cursor_class", SYNC_SPARK_CURSOR_CLASSES) def test_init_does_not_terminate_supplied_session(self, cursor_class): connection = _connection() error = OperationalError("Session did not become idle.") @@ -477,7 +386,7 @@ def test_init_does_not_terminate_supplied_session(self, cursor_class): connection.client.start_session.assert_not_called() connection.client.terminate_session.assert_not_called() - @pytest.mark.parametrize("cursor_class", SPARK_CURSOR_CLASSES) + @pytest.mark.parametrize("cursor_class", SYNC_SPARK_CURSOR_CLASSES) def test_init_does_not_start_session_when_s3_client_fails(self, cursor_class): connection = _connection() type(connection).s3_client = PropertyMock( @@ -546,12 +455,3 @@ def test_close_retries_failed_termination(self, cursor_class): _close(cursor) assert connection.client.terminate_session.call_count == 2 - - def test_async_close_shuts_down_executor_without_terminating_session(self): - connection = _connection() - cursor = _init_cursor(AsyncSparkCursor, connection, session_id="supplied-session") - cursor.close() - - connection.client.terminate_session.assert_not_called() - with pytest.raises(RuntimeError): - cursor.calculation_execution("calculation_id") diff --git a/tests/pyathena/spark/test_spark_cursor.py b/tests/pyathena/spark/test_spark_cursor.py index cc66789fa..9da8f3255 100644 --- a/tests/pyathena/spark/test_spark_cursor.py +++ b/tests/pyathena/spark/test_spark_cursor.py @@ -20,7 +20,10 @@ from tests import ENV from tests.pyathena.util import ( CANCELABLE_SPARK_JOB, + COMPLETED_CALCULATION_RESPONSES, ThrottledOperation, + answered_requests, + offline_spark_cursor, wait_for_spark_job, wait_for_spark_session_state, ) @@ -325,25 +328,16 @@ def test_execute_rejects_work_group(self): ) def test_execute_accepts_none_parameters(self, args, kwargs): """Omitted or None parameters run the calculation (no AWS).""" - cursor = SparkCursor.__new__(SparkCursor) # bypass __init__ to avoid AWS calls - cursor._session_id = "session_id" - - with ( - patch.object(SparkCursor, "_calculate", return_value="calculation_id") as calculate, - patch.object( - SparkCursor, - "_poll", - return_value=MagicMock(state=AthenaCalculationExecutionStatus.STATE_COMPLETED), - ), - ): + cursor = offline_spark_cursor(SparkCursor) + with answered_requests( + cursor.connection.client, COMPLETED_CALCULATION_RESPONSES + ) as requests: cursor.execute("code", *args, **kwargs) - calculate.assert_called_once_with( - session_id="session_id", - code_block="code", - description=None, - client_request_token=None, - ) + start = dict(requests)["StartCalculationExecution"] + assert (start["SessionId"], start["CodeBlock"]) == ("session", "code") + assert cursor.calculation_id == "calculation-id" + assert cursor.state == AthenaCalculationExecutionStatus.STATE_COMPLETED def test_spark_on_poll_invoked_each_iteration(): diff --git a/tests/pyathena/test_connection.py b/tests/pyathena/test_connection.py index 8d25856ef..f54dece72 100644 --- a/tests/pyathena/test_connection.py +++ b/tests/pyathena/test_connection.py @@ -5,6 +5,8 @@ # # SPDX-License-Identifier: MIT +import contextlib +import inspect import os import threading from concurrent.futures import ThreadPoolExecutor @@ -16,6 +18,13 @@ from botocore.config import Config import pyathena +from pyathena.aio.arrow.cursor import AioArrowCursor +from pyathena.aio.connection import AioConnection +from pyathena.aio.cursor import AioCursor, AioDictCursor +from pyathena.aio.pandas.cursor import AioPandasCursor +from pyathena.aio.polars.cursor import AioPolarsCursor +from pyathena.aio.s3fs.cursor import AioS3FSCursor +from pyathena.aio.spark.cursor import AioSparkCursor from pyathena.arrow.async_cursor import AsyncArrowCursor from pyathena.arrow.cursor import ArrowCursor from pyathena.async_cursor import AsyncCursor, AsyncDictCursor @@ -38,7 +47,7 @@ from tests.pyathena.util import EVENT_TIMEOUT, answered_requests # Cursors whose arraysize is the GetQueryResults page size, capped at 1000. -PAGED_CURSORS = [Cursor, DictCursor, AsyncCursor, AsyncDictCursor] +PAGED_CURSORS = [Cursor, DictCursor, AsyncCursor, AsyncDictCursor, AioCursor, AioDictCursor] # Cursors whose arraysize only sets the fetchmany() batch, with no upper limit. UNCAPPED_CURSORS = [ ArrowCursor, @@ -49,10 +58,29 @@ AsyncPandasCursor, AsyncPolarsCursor, AsyncS3FSCursor, + AioArrowCursor, + AioPandasCursor, + AioPolarsCursor, + AioS3FSCursor, ] -CURSOR_CLASSES = PAGED_CURSORS + UNCAPPED_CURSORS + [SparkCursor, AsyncSparkCursor] +CURSOR_CLASSES = PAGED_CURSORS + UNCAPPED_CURSORS + [SparkCursor, AsyncSparkCursor, AioSparkCursor] # Cursors that run on_start_query_execution. -START_CALLBACK_CURSORS = [Cursor, DictCursor, ArrowCursor, PandasCursor, PolarsCursor, S3FSCursor] +START_CALLBACK_CURSORS = [ + Cursor, + DictCursor, + ArrowCursor, + PandasCursor, + PolarsCursor, + S3FSCursor, + AioCursor, + AioDictCursor, + AioArrowCursor, + AioPandasCursor, + AioPolarsCursor, + AioS3FSCursor, +] +# The native asyncio cursors, whose execute() is a coroutine function. +AIO_CURSOR_CLASSES = [c for c in CURSOR_CLASSES if c.__module__.startswith("pyathena.aio.")] # Cursors that pass execute() keyword arguments to a DataFrame reader. READER_CURSOR_CLASSES = [PandasCursor, PolarsCursor, AsyncPandasCursor, AsyncPolarsCursor] @@ -92,8 +120,12 @@ def test_empty_settings_disable_both_environment_fallbacks(monkeypatch): Connection(region_name="us-east-1", s3_staging_dir="", work_group="") -@pytest.mark.parametrize("factory", [pyathena.connect, Connection]) +@pytest.mark.parametrize( + "factory", + [pyathena.connect, Connection, pyathena.aio_connect, AioConnection.create, AioConnection], +) def test_connect_keyword_only(factory): + # aio_connect() and AioConnection.create() raise when they are called. with pytest.raises(TypeError): factory("s3://bucket/path/") @@ -154,6 +186,26 @@ def _open_cursor(conn: Connection[Any], cursor_class: type[Any], **kwargs: Any) return conn.cursor(cursor_class, **kwargs) +@contextlib.asynccontextmanager +async def _opened_cursor(conn: Connection[Any], cursor_class: type[Any]) -> Any: + """Open a cursor and close it, awaiting the close of an aio cursor. + + Args: + conn: The connection. + cursor_class: The cursor class. + + Yields: + The cursor. + """ + cursor = _open_cursor(conn, cursor_class) + try: + yield cursor + finally: + closed = cursor.close() + if inspect.isawaitable(closed): + await closed + + @pytest.mark.parametrize("cursor_class", CURSOR_CLASSES) def test_cursor_constructor_keyword_only(cursor_class): # Binds the arguments as a call does before the constructor runs. A call @@ -164,9 +216,11 @@ def test_cursor_constructor_keyword_only(cursor_class): @pytest.mark.parametrize("cursor_class", CURSOR_CLASSES) -def test_cursor_execute_keyword_only(offline_connection, cursor_class): - with _open_cursor(offline_connection, cursor_class) as cursor, pytest.raises(TypeError): - cursor.execute("SELECT %(value)s", {"value": 1}, "group") +async def test_cursor_execute_keyword_only(offline_connection, cursor_class): + async with _opened_cursor(offline_connection, cursor_class) as cursor: + # An aio execute() raises when it is called, before it is awaited. + with pytest.raises(TypeError): + cursor.execute("SELECT %(value)s", {"value": 1}, "group") @pytest.mark.parametrize("cursor_class", CURSOR_CLASSES) @@ -179,30 +233,32 @@ def test_cursor_unknown_constructor_keyword(offline_connection, cursor_class): @pytest.mark.parametrize("cursor_class", CURSOR_CLASSES) -def test_cursor_accepts_connection_callbacks(offline_connection, cursor_class): +async def test_cursor_accepts_connection_callbacks(offline_connection, cursor_class): started, polled = [], [] offline_connection.on_start_query_execution = started.append offline_connection.on_poll = polled.append - with ( - answered_requests(offline_connection.client), - _open_cursor(offline_connection, cursor_class) as cursor, - ): - try: - result = cursor.execute("SELECT 1") - except OperationalError: - pass - else: - # The Async cursors run the execution in a future. - result[1].result(timeout=EVENT_TIMEOUT) + with answered_requests(offline_connection.client): + async with _opened_cursor(offline_connection, cursor_class) as cursor: + try: + result = cursor.execute("SELECT 1") + if inspect.isawaitable(result): + await result + else: + # The Async cursors run the execution in a future. + result[1].result(timeout=EVENT_TIMEOUT) + except OperationalError: + pass assert [execution.state for execution in polled] == ["FAILED"] # The Async and Spark cursors do not run on_start_query_execution. assert started == (["query-id"] if cursor_class in START_CALLBACK_CURSORS else []) +# aio/test_connection.py tests the native asyncio cursors. @pytest.mark.parametrize( - "cursor_class", [c for c in CURSOR_CLASSES if c not in READER_CURSOR_CLASSES] + "cursor_class", + [c for c in CURSOR_CLASSES if c not in READER_CURSOR_CLASSES + AIO_CURSOR_CLASSES], ) def test_cursor_execute_unknown_keyword(offline_connection, cursor_class): with _open_cursor(offline_connection, cursor_class) as cursor: diff --git a/tests/pyathena/test_cursor.py b/tests/pyathena/test_cursor.py index 6401c07f6..a98ae904b 100644 --- a/tests/pyathena/test_cursor.py +++ b/tests/pyathena/test_cursor.py @@ -51,6 +51,7 @@ EVENT_TIMEOUT, QueryExecutionHistory, ThrottledOperation, + cache_search_cursor, interrupt_start_waits, succeeded_query_execution, succeeded_query_execution_data, @@ -91,24 +92,6 @@ def _offline_cursor(kill_on_interrupt, cursor_class=Cursor): return cursor, cancel -def _cache_search_cursor(client): - """A cursor whose result cache search sends its requests to a client stand-in. - - Args: - client: The Athena client stand-in, such as ``QueryExecutionHistory``. - - Returns: - The cursor, for the schema ``this_schema``. - """ - cursor = Cursor.__new__(Cursor) # bypass __init__ to avoid AWS calls - cursor._connection = MagicMock(_client=client) - cursor._retry_config = RetryConfig(attempt=1) - cursor._work_group = None - cursor._schema_name = "this_schema" - cursor._catalog_name = None - return cursor - - def _types(value): """The type of a value and of each value it contains, for comparing them. @@ -440,7 +423,7 @@ def test_cache_search_stops_at_expired_execution(self): succeeded_query_execution_data("next_page", query, now), ] history = QueryExecutionHistory(executions) - cursor = _cache_search_cursor(history) + cursor = cache_search_cursor(Cursor, history) # Without cache_size, cache_expiration_time alone bounds the search. assert cursor._find_previous_query_id(query, None, cache_expiration_time=3600) is None @@ -465,7 +448,7 @@ def test_cache_search_reads_pages_up_to_cache_size(self, position, expected): ] executions[position] = succeeded_query_execution_data("match", query, now) history = QueryExecutionHistory(executions) - cursor = _cache_search_cursor(history) + cursor = cache_search_cursor(Cursor, history) assert cursor._find_previous_query_id(query, "wg", cache_size=60) == expected assert [request["WorkGroup"] for request in history.list_requests] == ["wg", "wg"] diff --git a/tests/pyathena/util.py b/tests/pyathena/util.py index eed770d00..422c4da72 100644 --- a/tests/pyathena/util.py +++ b/tests/pyathena/util.py @@ -14,7 +14,7 @@ from datetime import UTC, datetime, timedelta, timezone from pathlib import Path from types import SimpleNamespace -from unittest.mock import patch +from unittest.mock import MagicMock, patch from zoneinfo import ZoneInfo import pytest @@ -26,6 +26,7 @@ from jinja2 import Environment, FileSystemLoader from sqlalchemy import cast, create_engine, engine_from_config, literal, types +from pyathena.connection import Connection from pyathena.error import ProgrammingError from pyathena.filesystem.s3 import S3FileSystem from pyathena.filesystem.s3_async import AioS3FileSystem @@ -33,7 +34,8 @@ from pyathena.formatter import DefaultParameterFormatter from pyathena.glue import GlueMetadataClient from pyathena.model import AthenaCalculationExecutionStatus, AthenaQueryExecution -from pyathena.util import override, retry_api_call +from pyathena.spark.common import SparkBaseCursor +from pyathena.util import RetryConfig, override, retry_api_call _queries = Environment( loader=FileSystemLoader(Path(__file__).parents[1].resolve() / "resources" / "queries") @@ -316,8 +318,29 @@ def succeeded_query_execution(query_id, query, completion_date_time, schema="thi ) -# Athena's answers to a query or calculation that fails at its first status request. +def cache_search_cursor(cursor_class, client): + """Build a cursor whose result cache search sends its requests to a client stand-in. + + Args: + cursor_class: The cursor class, such as ``Cursor`` or ``AioCursor``. + client: The Athena client stand-in, such as ``QueryExecutionHistory``. + + Returns: + The cursor, for the schema ``this_schema``. + """ + cursor = cursor_class.__new__(cursor_class) # bypass __init__ to avoid AWS calls + cursor._connection = MagicMock(_client=client) + cursor._retry_config = RetryConfig(attempt=1) + cursor._work_group = None + cursor._schema_name = "this_schema" + cursor._catalog_name = None + return cursor + + +# Athena's answers to a query or calculation that fails at its first status request, +# after a search that finds no earlier executions. FAILED_EXECUTION_RESPONSES = { + "ListQueryExecutions": {"QueryExecutionIds": []}, "StartQueryExecution": {"QueryExecutionId": "query-id"}, "GetQueryExecution": { "QueryExecution": { @@ -342,6 +365,45 @@ def succeeded_query_execution(query_id, query, completion_date_time, schema="thi } +# Athena's answers to a calculation that completes at its first status request. +COMPLETED_CALCULATION_RESPONSES = { + **FAILED_EXECUTION_RESPONSES, + "GetCalculationExecutionStatus": { + "Status": {"State": "COMPLETED"}, + "Statistics": {"DpuExecutionInMillis": 0}, + }, + "GetCalculationExecution": { + "CalculationExecutionId": "calculation-id", + "SessionId": "session", + "Status": {"State": "COMPLETED"}, + "Statistics": {"DpuExecutionInMillis": 0}, + "Result": {}, + }, +} + + +def offline_spark_cursor(cursor_class): + """Build a Spark cursor on a connection that sends no request to AWS on its own. + + The cursor reports a started session, ``session``, which closing it leaves + alone. + + Args: + cursor_class: The Spark cursor class. + + Returns: + The cursor. + """ + conn = Connection( + region_name="us-east-1", + s3_staging_dir="s3://bucket/path/", + aws_access_key_id="access_key", + aws_secret_access_key="secret_key", + ) + with patch.object(SparkBaseCursor, "_start_session", return_value="session"): + return conn.cursor(cursor_class, terminate_session_on_close=False) + + @contextmanager def answered_requests(client, responses=FAILED_EXECUTION_RESPONSES): """Answer each request of a boto3 client once, without AWS, and record them. From 18d208e4f807ab483fd9ebbe275b10cccc1384d5 Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Sat, 17 Oct 2026 22:33:44 +0900 Subject: [PATCH 2/4] Check the options, the Glue thread, and the session that the rewrites missed The aio options test lists an hour-old execution, so it fails if cache_expiration_time does not reach the cache search. The Glue test also records the thread that builds the client. The supplied-session test checks the GetSession request and the reason from Athena. An Async cursor's future is waited for outside the suppressed OperationalError. The placeholder connection and the aio cursor list are defined once, and the Arrow and Polars whole-result methods are named per class. Co-Authored-By: Claude Opus 5.5 --- tests/pyathena/aio/test_common.py | 12 ++++--- tests/pyathena/aio/test_cursor.py | 54 ++++++++++++++++------------- tests/pyathena/spark/test_common.py | 18 ++++++---- tests/pyathena/test_connection.py | 21 +++++++---- tests/pyathena/util.py | 26 ++++++++++---- 5 files changed, 83 insertions(+), 48 deletions(-) diff --git a/tests/pyathena/aio/test_common.py b/tests/pyathena/aio/test_common.py index 2e31905a9..a5508d745 100644 --- a/tests/pyathena/aio/test_common.py +++ b/tests/pyathena/aio/test_common.py @@ -29,6 +29,11 @@ AioS3FSCursor, ] +_WHOLE_RESULT_METHODS = { + AioArrowCursor: ("as_arrow", "as_polars"), + AioPolarsCursor: ("as_polars", "as_arrow"), +} + def _result_set_closing_after_loop_responds() -> tuple[MagicMock, list[bool]]: """Return a result set whose ``close()`` waits for a callback on the running loop. @@ -103,10 +108,9 @@ async def test_fetch_without_result_set(self, cursor_class): with pytest.raises(ProgrammingError, match=r"No result set\."): await fetch() # The Arrow and Polars cursors also convert the whole result. - for name in ("as_arrow", "as_polars"): - if hasattr(cursor, name): - with pytest.raises(ProgrammingError): - getattr(cursor, name)() + for name in _WHOLE_RESULT_METHODS.get(cursor_class, ()): + with pytest.raises(ProgrammingError, match=r"No result set\."): + getattr(cursor, name)() @pytest.mark.parametrize("cursor_class", _ALL_CURSOR_CLASSES) async def test_cancel_without_query(self, cursor_class): diff --git a/tests/pyathena/aio/test_cursor.py b/tests/pyathena/aio/test_cursor.py index bf8cc21a5..71f0951ea 100644 --- a/tests/pyathena/aio/test_cursor.py +++ b/tests/pyathena/aio/test_cursor.py @@ -18,30 +18,18 @@ from tests import ENV from tests.pyathena.util import ( EVENT_TIMEOUT, + FAILED_EXECUTION_RESPONSES, QueryExecutionHistory, ThrottledOperation, answered_requests, cache_search_cursor, + placeholder_connection, succeeded_query_execution_data, throttle_metadata_api, wait_until, ) -def _offline_connection(): - """Build a connection whose Athena requests the test answers. - - Returns: - The connection, with ``AioCursor`` as its cursor class. - """ - return AioConnection( - region_name="us-east-1", - s3_staging_dir="s3://bucket/path/", - aws_access_key_id="access_key", - aws_secret_access_key="secret_key", - ) - - def _offline_cursor(kill_on_interrupt, final_state): """An AioCursor whose first status request blocks until the task is cancelled. @@ -185,7 +173,7 @@ async def test_execute_with_options(self, aio_cursor): async def test_execute_internal_options_passthrough(self): """The options of execute() reach its Athena requests (no AWS).""" - cursor = _offline_connection().cursor() + cursor = placeholder_connection(AioConnection).cursor() options = ExecuteOptions( work_group="test_work_group", s3_staging_dir="s3://test-bucket/path/", @@ -195,8 +183,20 @@ async def test_execute_internal_options_passthrough(self): result_reuse_minutes=5, paramstyle="qmark", ) + # The same query an hour ago, too old for cache_expiration_time to reuse. + previous = { + **succeeded_query_execution_data( + "previous", "SELECT 1", datetime.now(UTC) - timedelta(hours=1), schema="default" + ), + "QueryExecutionContext": {"Database": "default", "Catalog": "awsdatacatalog"}, + } + responses = { + **FAILED_EXECUTION_RESPONSES, + "ListQueryExecutions": {"QueryExecutionIds": ["previous"]}, + "BatchGetQueryExecution": {"QueryExecutions": [previous]}, + } with ( - answered_requests(cursor.connection.client) as requests, + answered_requests(cursor.connection.client, responses) as requests, pytest.raises(OperationalError), ): await cursor.execute("SELECT 1", options=options) @@ -215,7 +215,7 @@ async def test_execute_internal_options_passthrough(self): async def test_execute_qmark_parameters_skip_cache(self): """A qmark query with parameters never searches the cache (no AWS, #941).""" - cursor = _offline_connection().cursor() + cursor = placeholder_connection(AioConnection).cursor() with ( answered_requests(cursor.connection.client) as requests, pytest.raises(OperationalError), @@ -632,22 +632,28 @@ async def test_arraysize(self, aio_cursor): assert len(actual) == 5 async def test_glue_request_runs_off_the_event_loop(self, monkeypatch): - """The Glue request of a throttled metadata lookup does not block the loop (no AWS).""" - conn = _offline_connection() + """The Glue client and request of a throttled lookup do not block the loop (no AWS).""" + conn = placeholder_connection(AioConnection) throttle_metadata_api(conn.client, monkeypatch) threads = [] def get_table(**kwargs): - threads.append(threading.get_ident()) + threads.append(("request", threading.get_ident())) return {"Table": {"Name": kwargs["Name"], "StorageDescriptor": {"Columns": []}}} glue = MagicMock() glue.get_table.side_effect = get_table - monkeypatch.setattr(conn.session, "client", lambda *args, **kwargs: glue) + + def build_client(service_name, **kwargs): + assert service_name == "glue" + threads.append(("client", threading.get_ident())) + return glue + + monkeypatch.setattr(conn.session, "client", build_client) assert (await conn.cursor().get_table_metadata("one_row")).name == "one_row" - assert len(threads) == 1 - assert threading.get_ident() not in threads + assert sorted(step for step, _ in threads) == ["client", "request"] + assert threading.get_ident() not in [thread for _, thread in threads] async def test_throttled_metadata_reads_glue(self, aio_cursor, monkeypatch): def view(metadata): @@ -667,7 +673,7 @@ async def read(): assert has_schema calls = throttle_metadata_api(aio_cursor.connection.client, monkeypatch) - # The Glue request runs in a worker thread, as the Athena calls do. + # Glue reports each table as Athena does, and nothing is retried first. assert await read() == expected assert sorted(calls) == ["get_table_metadata", "list_databases", "list_table_metadata"] diff --git a/tests/pyathena/spark/test_common.py b/tests/pyathena/spark/test_common.py index 50923a4d3..ca69c4c65 100644 --- a/tests/pyathena/spark/test_common.py +++ b/tests/pyathena/spark/test_common.py @@ -24,6 +24,9 @@ SPARK_CURSOR_CLASSES = [SparkCursor, AsyncSparkCursor, AioSparkCursor] SYNC_SPARK_CURSOR_CLASSES = [SparkCursor, AsyncSparkCursor] +# AsyncSparkCursor inherits _calculate and _start_execution from SparkCursor's bases; +# the parameter keeps the test IDs. +CALCULATE_CURSOR_CLASSES = [SparkCursor] _TIMEOUT = 10 @@ -155,13 +158,14 @@ def test_exists_session_raises_on_failure_state(self): } } - with pytest.raises(OperationalError): + with pytest.raises(OperationalError, match="session failure reason"): _init_cursor(SparkCursor, connection, session_id="supplied-session") + connection.client.get_session.assert_called_once_with(SessionId="supplied-session") connection.client.start_session.assert_not_called() connection.client.terminate_session.assert_not_called() - @pytest.mark.parametrize("cursor_class", [SparkCursor]) + @pytest.mark.parametrize("cursor_class", CALCULATE_CURSOR_CLASSES) @pytest.mark.parametrize("kill_on_interrupt", [True, False]) def test_calculate_reuses_generated_token_on_retry(self, cursor_class, kill_on_interrupt): cursor = _calculation_cursor(cursor_class, kill_on_interrupt=kill_on_interrupt) @@ -188,7 +192,7 @@ def test_calculate_reuses_generated_token_on_retry(self, cursor_class, kill_on_i thread.join(_TIMEOUT) assert not thread.is_alive() - @pytest.mark.parametrize("cursor_class", [SparkCursor]) + @pytest.mark.parametrize("cursor_class", CALCULATE_CURSOR_CLASSES) def test_calculate_generates_token_per_call(self, cursor_class): cursor = _calculation_cursor(cursor_class) client = cursor._connection.client @@ -202,7 +206,7 @@ def test_calculate_generates_token_per_call(self, cursor_class): ] assert tokens[0] != tokens[1] - @pytest.mark.parametrize("cursor_class", [SparkCursor]) + @pytest.mark.parametrize("cursor_class", CALCULATE_CURSOR_CLASSES) def test_calculate_keeps_caller_token(self, cursor_class): cursor = _calculation_cursor(cursor_class) @@ -212,7 +216,7 @@ def test_calculate_keeps_caller_token(self, cursor_class): SessionId="session_id", CodeBlock="code", ClientRequestToken="token" ) - @pytest.mark.parametrize("cursor_class", [SparkCursor]) + @pytest.mark.parametrize("cursor_class", CALCULATE_CURSOR_CLASSES) @pytest.mark.parametrize("failing", ["start", "cancel", "wait"]) def test_calculate_interrupted_while_starting_failure(self, cursor_class, failing): cursor = _calculation_cursor(cursor_class) @@ -246,7 +250,7 @@ def test_calculate_interrupted_while_starting_failure(self, cursor_class, failin cursor._cancel.assert_called_once_with("calculation_id") assert cursor.calculation_id == "calculation_id" - @pytest.mark.parametrize("cursor_class", [SparkCursor]) + @pytest.mark.parametrize("cursor_class", CALCULATE_CURSOR_CLASSES) def test_calculate_second_interrupt_while_starting(self, cursor_class): cursor = _calculation_cursor(cursor_class) started, release = _block_start(cursor) @@ -262,7 +266,7 @@ def test_calculate_second_interrupt_while_starting(self, cursor_class): assert exc_info.value.__context__ is raised[0] cursor._cancel.assert_not_called() - @pytest.mark.parametrize("cursor_class", [SparkCursor]) + @pytest.mark.parametrize("cursor_class", CALCULATE_CURSOR_CLASSES) @pytest.mark.parametrize("helper_starts", [False, True]) def test_calculate_interrupted_before_request_is_sent(self, cursor_class, helper_starts): cursor = _calculation_cursor(cursor_class) diff --git a/tests/pyathena/test_connection.py b/tests/pyathena/test_connection.py index f54dece72..94526da0f 100644 --- a/tests/pyathena/test_connection.py +++ b/tests/pyathena/test_connection.py @@ -80,7 +80,15 @@ AioS3FSCursor, ] # The native asyncio cursors, whose execute() is a coroutine function. -AIO_CURSOR_CLASSES = [c for c in CURSOR_CLASSES if c.__module__.startswith("pyathena.aio.")] +AIO_CURSOR_CLASSES = [ + AioCursor, + AioDictCursor, + AioArrowCursor, + AioPandasCursor, + AioPolarsCursor, + AioS3FSCursor, + AioSparkCursor, +] # Cursors that pass execute() keyword arguments to a DataFrame reader. READER_CURSOR_CLASSES = [PandasCursor, PolarsCursor, AsyncPandasCursor, AsyncPolarsCursor] @@ -240,15 +248,14 @@ async def test_cursor_accepts_connection_callbacks(offline_connection, cursor_cl with answered_requests(offline_connection.client): async with _opened_cursor(offline_connection, cursor_class) as cursor: - try: + result = None + with contextlib.suppress(OperationalError): result = cursor.execute("SELECT 1") if inspect.isawaitable(result): await result - else: - # The Async cursors run the execution in a future. - result[1].result(timeout=EVENT_TIMEOUT) - except OperationalError: - pass + if isinstance(result, tuple): + # The Async cursors run the execution in a future. + result[1].result(timeout=EVENT_TIMEOUT) assert [execution.state for execution in polled] == ["FAILED"] # The Async and Spark cursors do not run on_start_query_execution. diff --git a/tests/pyathena/util.py b/tests/pyathena/util.py index 422c4da72..794b0e3e3 100644 --- a/tests/pyathena/util.py +++ b/tests/pyathena/util.py @@ -382,6 +382,25 @@ def cache_search_cursor(cursor_class, client): } +def placeholder_connection(connection_class=Connection): + """Build a connection with placeholder settings and credentials. + + Building it sends no request to AWS. + + Args: + connection_class: The connection class, such as ``AioConnection``. + + Returns: + The connection. + """ + return connection_class( + region_name="us-east-1", + s3_staging_dir="s3://bucket/path/", + aws_access_key_id="access_key", + aws_secret_access_key="secret_key", + ) + + def offline_spark_cursor(cursor_class): """Build a Spark cursor on a connection that sends no request to AWS on its own. @@ -394,12 +413,7 @@ def offline_spark_cursor(cursor_class): Returns: The cursor. """ - conn = Connection( - region_name="us-east-1", - s3_staging_dir="s3://bucket/path/", - aws_access_key_id="access_key", - aws_secret_access_key="secret_key", - ) + conn = placeholder_connection() with patch.object(SparkBaseCursor, "_start_session", return_value="session"): return conn.cursor(cursor_class, terminate_session_on_close=False) From 078c13407a719ed6c939e6582b83fc5d02bbe078 Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Sat, 17 Oct 2026 22:45:12 +0900 Subject: [PATCH 3/4] Keep the checks that the independent review found unprotected Check offline that an executemany() with an unknown update count reports -1, keep the test that closing an AsyncSparkCursor on a supplied session shuts down its executor, and run a Spark job in the aio calculation test, so that its DPU time and progress are reported. Leave out the message of the whole-result errors, which is not documented. Co-Authored-By: Claude Opus 5.5 --- tests/pyathena/aio/spark/test_cursor.py | 7 +++++-- tests/pyathena/aio/test_common.py | 18 +++++++++++++++++- tests/pyathena/spark/test_common.py | 9 +++++++++ 3 files changed, 31 insertions(+), 3 deletions(-) diff --git a/tests/pyathena/aio/spark/test_cursor.py b/tests/pyathena/aio/spark/test_cursor.py index 8c376fd6b..ac90d4759 100644 --- a/tests/pyathena/aio/spark/test_cursor.py +++ b/tests/pyathena/aio/spark/test_cursor.py @@ -106,7 +106,10 @@ def start_calculation_execution(**kwargs): class TestAioSparkCursor: async def test_spark_dataframe(self, aio_spark_cursor): - await aio_spark_cursor.execute("print('hello')", description="test description") + # A Spark job, so that the calculation reports its DPU time and progress. + await aio_spark_cursor.execute( + "print(spark.range(3).count())", description="test description" + ) assert aio_spark_cursor.calculation_execution assert aio_spark_cursor.session_id assert aio_spark_cursor.calculation_id @@ -122,7 +125,7 @@ async def test_spark_dataframe(self, aio_spark_cursor): assert aio_spark_cursor.std_error_s3_uri assert aio_spark_cursor.result_s3_uri assert aio_spark_cursor.result_type - assert await aio_spark_cursor.get_std_out() == "hello" + assert await aio_spark_cursor.get_std_out() == "3" async def test_failed(self, aio_spark_cursor): with pytest.raises(OperationalError): diff --git a/tests/pyathena/aio/test_common.py b/tests/pyathena/aio/test_common.py index a5508d745..fd34279ea 100644 --- a/tests/pyathena/aio/test_common.py +++ b/tests/pyathena/aio/test_common.py @@ -109,7 +109,7 @@ async def test_fetch_without_result_set(self, cursor_class): await fetch() # The Arrow and Polars cursors also convert the whole result. for name in _WHOLE_RESULT_METHODS.get(cursor_class, ()): - with pytest.raises(ProgrammingError, match=r"No result set\."): + with pytest.raises(ProgrammingError): getattr(cursor, name)() @pytest.mark.parametrize("cursor_class", _ALL_CURSOR_CLASSES) @@ -283,6 +283,22 @@ async def test_executemany_closes_previous_result_set_off_event_loop(self, curso assert cursor.result_set is None assert cursor.rowcount == 0 + @pytest.mark.parametrize("cursor_class", _ALL_CURSOR_CLASSES) + async def test_executemany_unknown_rowcount(self, cursor_class): + """A statement without an update count, such as a SELECT, makes the sum unknown.""" + cursor = cursor_class( + connection=MagicMock(), converter=None, formatter=None, retry_config=None + ) + rowcounts = iter([1, -1, 1]) + + async def execute(operation, parameters, **kwargs): + cursor.result_set = MagicMock(is_closed=True, rowcount=next(rowcounts)) + + cursor.execute = execute + await cursor.executemany("SELECT 1", [None, None, None]) + + assert cursor.rowcount == -1 + @pytest.mark.parametrize("cursor_class", _ALL_CURSOR_CLASSES) async def test_executemany_failure_closes_result_set_off_event_loop(self, cursor_class): cursor = cursor_class( diff --git a/tests/pyathena/spark/test_common.py b/tests/pyathena/spark/test_common.py index ca69c4c65..1cd8bbbae 100644 --- a/tests/pyathena/spark/test_common.py +++ b/tests/pyathena/spark/test_common.py @@ -459,3 +459,12 @@ def test_close_retries_failed_termination(self, cursor_class): _close(cursor) assert connection.client.terminate_session.call_count == 2 + + def test_async_close_shuts_down_executor_without_terminating_session(self): + connection = _connection() + cursor = _init_cursor(AsyncSparkCursor, connection, session_id="supplied-session") + cursor.close() + + connection.client.terminate_session.assert_not_called() + with pytest.raises(RuntimeError): + cursor.calculation_execution("calculation_id") From a5b5b349b2086f5514670f6931258bce4592e320 Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Sun, 18 Oct 2026 09:18:40 +0900 Subject: [PATCH 4/4] Check the end of async with for AioSparkCursor, and the result set after executemany AioSparkCursor has its own __aexit__, which no test checked: an owned session is now terminated at the end of async with. The aio executemany() test also checks that no result set is left to fetch, which the removed test_executemany_fetch tests checked. Co-Authored-By: Claude Opus 5.5 --- tests/pyathena/aio/test_common.py | 6 +++++- tests/pyathena/spark/test_common.py | 8 ++++++++ 2 files changed, 13 insertions(+), 1 deletion(-) diff --git a/tests/pyathena/aio/test_common.py b/tests/pyathena/aio/test_common.py index fd34279ea..cf5d2e0b6 100644 --- a/tests/pyathena/aio/test_common.py +++ b/tests/pyathena/aio/test_common.py @@ -285,7 +285,10 @@ async def test_executemany_closes_previous_result_set_off_event_loop(self, curso @pytest.mark.parametrize("cursor_class", _ALL_CURSOR_CLASSES) async def test_executemany_unknown_rowcount(self, cursor_class): - """A statement without an update count, such as a SELECT, makes the sum unknown.""" + """A statement without an update count, such as a SELECT, makes the sum unknown. + + No result set is left to fetch. + """ cursor = cursor_class( connection=MagicMock(), converter=None, formatter=None, retry_config=None ) @@ -298,6 +301,7 @@ async def execute(operation, parameters, **kwargs): await cursor.executemany("SELECT 1", [None, None, None]) assert cursor.rowcount == -1 + assert cursor.result_set is None @pytest.mark.parametrize("cursor_class", _ALL_CURSOR_CLASSES) async def test_executemany_failure_closes_result_set_off_event_loop(self, cursor_class): diff --git a/tests/pyathena/spark/test_common.py b/tests/pyathena/spark/test_common.py index 1cd8bbbae..a5da2958c 100644 --- a/tests/pyathena/spark/test_common.py +++ b/tests/pyathena/spark/test_common.py @@ -433,6 +433,14 @@ def test_close_terminates_session_by_ownership( else: connection.client.terminate_session.assert_not_called() + async def test_aio_context_exit_terminates_owned_session(self): + """AioSparkCursor closes itself, and so its own session, at the end of async with.""" + connection = _connection() + async with _init_cursor(AioSparkCursor, connection): + connection.client.terminate_session.assert_not_called() + + connection.client.terminate_session.assert_called_once_with(SessionId="new-session") + @pytest.mark.parametrize("cursor_class", SPARK_CURSOR_CLASSES) def test_close_does_not_terminate_session_twice(self, cursor_class): connection = _connection()