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..ac90d4759 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,16 +106,9 @@ def start_calculation_execution(**kwargs): class TestAioSparkCursor: async def test_spark_dataframe(self, aio_spark_cursor): + # A Spark job, so that the calculation reports its DPU time and progress. 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", + "print(spark.range(3).count())", description="test description" ) assert aio_spark_cursor.calculation_execution assert aio_spark_cursor.session_id @@ -132,70 +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 - - 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() == "3" async def test_failed(self, aio_spark_cursor): with pytest.raises(OperationalError): @@ -316,29 +246,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 +319,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 +423,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..cf5d2e0b6 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. @@ -94,10 +99,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 +107,18 @@ 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 _WHOLE_RESULT_METHODS.get(cursor_class, ()): + 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", ())] @@ -269,6 +283,26 @@ 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. + + No result set is left to fetch. + """ + 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 + 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): cursor = cursor_class( @@ -278,6 +312,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 +325,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..71f0951ea 100644 --- a/tests/pyathena/aio/test_cursor.py +++ b/tests/pyathena/aio/test_cursor.py @@ -1,27 +1,30 @@ 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, + FAILED_EXECUTION_RESPONSES, + QueryExecutionHistory, ThrottledOperation, - succeeded_query_execution, + answered_requests, + cache_search_cursor, + placeholder_connection, + succeeded_query_execution_data, throttle_metadata_api, wait_until, ) @@ -110,26 +113,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 +158,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 +172,61 @@ 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 = placeholder_connection(AioConnection).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", ) + # 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, responses) 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 = placeholder_connection(AioConnection).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): - 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): + 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 - 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,29 @@ 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_invalid_arraysize(self, aio_cursor): - with pytest.raises(ProgrammingError): - aio_cursor.arraysize = 10000 - with pytest.raises(ProgrammingError): - aio_cursor.arraysize = -1 - - 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() + async def test_glue_request_runs_off_the_event_loop(self, monkeypatch): + """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 = [] - # 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): + def get_table(**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) + return {"Table": {"Name": kwargs["Name"], "StorageDescriptor": {"Columns": []}}} - 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) + glue = MagicMock() + glue.get_table.side_effect = get_table - 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 + def build_client(service_name, **kwargs): + assert service_name == "glue" + threads.append(("client", threading.get_ident())) + return glue - 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 + monkeypatch.setattr(conn.session, "client", build_client) - 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 + assert (await conn.cursor().get_table_metadata("one_row")).name == "one_row" + 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): @@ -985,9 +667,13 @@ 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. + # 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"] @@ -997,30 +683,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 +694,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..a5da2958c 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 @@ -25,23 +24,12 @@ 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 -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 +131,41 @@ 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, 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", 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) @@ -238,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", SYNC_SPARK_CURSOR_CLASSES) + @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 @@ -252,7 +206,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", CALCULATE_CURSOR_CLASSES) def test_calculate_keeps_caller_token(self, cursor_class): cursor = _calculation_cursor(cursor_class) @@ -262,32 +216,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", 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) @@ -321,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", SYNC_SPARK_CURSOR_CLASSES) + @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) @@ -337,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", SYNC_SPARK_CURSOR_CLASSES) + @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) @@ -366,42 +295,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 +336,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 +353,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 +368,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 +375,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 +390,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( @@ -520,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() 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..94526da0f 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,37 @@ 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 = [ + AioCursor, + AioDictCursor, + AioArrowCursor, + AioPandasCursor, + AioPolarsCursor, + AioS3FSCursor, + AioSparkCursor, +] # Cursors that pass execute() keyword arguments to a DataFrame reader. READER_CURSOR_CLASSES = [PandasCursor, PolarsCursor, AsyncPandasCursor, AsyncPolarsCursor] @@ -92,8 +128,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 +194,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 +224,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 +241,31 @@ 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: + result = None + with contextlib.suppress(OperationalError): + result = cursor.execute("SELECT 1") + if inspect.isawaitable(result): + await result + 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. 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..794b0e3e3 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,59 @@ 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 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. + + The cursor reports a started session, ``session``, which closing it leaves + alone. + + Args: + cursor_class: The Spark cursor class. + + Returns: + The cursor. + """ + conn = placeholder_connection() + 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.