From 5bcf87d58907094952a1289ae9dfcf9a0fc039d7 Mon Sep 17 00:00:00 2001 From: David Krych Date: Fri, 10 Jul 2026 18:50:03 -0400 Subject: [PATCH 1/3] Check that the schema and catalog match when looking for cached results. --- pyathena/aio/common.py | 6 ++++- pyathena/common.py | 6 ++++- tests/pyathena/aio/test_cursor.py | 42 +++++++++++++++++++++++++++++++ tests/pyathena/test_cursor.py | 38 ++++++++++++++++++++++++++++ 4 files changed, 90 insertions(+), 2 deletions(-) diff --git a/pyathena/aio/common.py b/pyathena/aio/common.py index 19914164..6c9f8430 100644 --- a/pyathena/aio/common.py +++ b/pyathena/aio/common.py @@ -218,7 +218,11 @@ async def _find_previous_query_id( # type: ignore[override] ): next_token = None break - if execution.query == query: + if ( + execution.query == query + and (self._schema_name is None or execution.database == self._schema_name) + and (self._catalog_name is None or execution.catalog == self._catalog_name) + ): query_id = execution.query_id break if query_id or next_token is None: diff --git a/pyathena/common.py b/pyathena/common.py index 9364e394..b07efc59 100644 --- a/pyathena/common.py +++ b/pyathena/common.py @@ -638,7 +638,11 @@ def _find_previous_query_id( ): next_token = None break - if execution.query == query: + if ( + execution.query == query + and (self._schema_name is None or execution.database == self._schema_name) + and (self._catalog_name is None or execution.catalog == self._catalog_name) + ): query_id = execution.query_id break if query_id or next_token is None: diff --git a/tests/pyathena/aio/test_cursor.py b/tests/pyathena/aio/test_cursor.py index 520b8795..350001d5 100644 --- a/tests/pyathena/aio/test_cursor.py +++ b/tests/pyathena/aio/test_cursor.py @@ -130,6 +130,48 @@ async def test_execute_internal_legacy_kwargs_passthrough(self): cache_expiration_time=100, ) + 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(), + }, + } + } + ) + + 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_no_result_set_raises(self, aio_cursor): with pytest.raises(ProgrammingError): await aio_cursor.fetchone() diff --git a/tests/pyathena/test_cursor.py b/tests/pyathena/test_cursor.py index bfe9e82a..fd3c7b53 100644 --- a/tests/pyathena/test_cursor.py +++ b/tests/pyathena/test_cursor.py @@ -185,6 +185,44 @@ def test_cache_expiration_time_with_cache_size(self, cursor): assert query_id_7 != query_id_8 assert query_id_9 in [query_id_7, query_id_8] + def test_cache_size_different_schema(self): + """A cached result is only reused when it ran against the same schema (#739). + + 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(timezone.utc), + }, + } + } + ) + + cursor = Cursor.__new__(Cursor) # bypass __init__ to avoid AWS calls + cursor._catalog_name = None + + with patch.object( + Cursor, "_list_query_executions", return_value=(None, [execution("other_schema")]) + ): + cursor._schema_name = "this_schema" + assert cursor._find_previous_query_id(query, None, cache_size=100) is None + cursor._schema_name = "other_schema" + assert ( + cursor._find_previous_query_id(query, None, cache_size=100) + == "query_id_other_schema" + ) + @pytest.mark.parametrize( "cursor", [{"work_group": ENV.work_group, "result_reuse_enable": True, "result_reuse_minutes": 5}], From b59f1dca246cfbcf8d18dfc55d1afb8645217374 Mon Sep 17 00:00:00 2001 From: David Krych Date: Fri, 10 Jul 2026 18:59:07 -0400 Subject: [PATCH 2/3] Remove bad/unneeded none guard. --- pyathena/aio/common.py | 4 ++-- pyathena/common.py | 4 ++-- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/pyathena/aio/common.py b/pyathena/aio/common.py index 6c9f8430..15b4ebf1 100644 --- a/pyathena/aio/common.py +++ b/pyathena/aio/common.py @@ -220,8 +220,8 @@ async def _find_previous_query_id( # type: ignore[override] break if ( execution.query == query - and (self._schema_name is None or execution.database == self._schema_name) - and (self._catalog_name is None or execution.catalog == self._catalog_name) + and execution.database == self._schema_name + and execution.catalog == self._catalog_name ): query_id = execution.query_id break diff --git a/pyathena/common.py b/pyathena/common.py index b07efc59..6c9cbf9a 100644 --- a/pyathena/common.py +++ b/pyathena/common.py @@ -640,8 +640,8 @@ def _find_previous_query_id( break if ( execution.query == query - and (self._schema_name is None or execution.database == self._schema_name) - and (self._catalog_name is None or execution.catalog == self._catalog_name) + and execution.database == self._schema_name + and execution.catalog == self._catalog_name ): query_id = execution.query_id break From 61b8ba8940e4bde0a97f9c9296efe22585a56693 Mon Sep 17 00:00:00 2001 From: David Krych Date: Mon, 13 Jul 2026 11:31:04 -0400 Subject: [PATCH 3/3] Fix case sensitivity and improve tests. --- pyathena/aio/common.py | 2 +- pyathena/common.py | 2 +- tests/pyathena/aio/test_cursor.py | 43 +++++++++++++++++++++++++++++-- tests/pyathena/test_cursor.py | 36 ++++++++++++++++++++++++++ 4 files changed, 79 insertions(+), 4 deletions(-) diff --git a/pyathena/aio/common.py b/pyathena/aio/common.py index 15b4ebf1..e490a544 100644 --- a/pyathena/aio/common.py +++ b/pyathena/aio/common.py @@ -221,7 +221,7 @@ async def _find_previous_query_id( # type: ignore[override] if ( execution.query == query and execution.database == self._schema_name - and execution.catalog == self._catalog_name + and (execution.catalog or "").lower() == (self._catalog_name or "").lower() ): query_id = execution.query_id break diff --git a/pyathena/common.py b/pyathena/common.py index 6c9cbf9a..3a517694 100644 --- a/pyathena/common.py +++ b/pyathena/common.py @@ -641,7 +641,7 @@ def _find_previous_query_id( if ( execution.query == query and execution.database == self._schema_name - and execution.catalog == self._catalog_name + and (execution.catalog or "").lower() == (self._catalog_name or "").lower() ): query_id = execution.query_id break diff --git a/tests/pyathena/aio/test_cursor.py b/tests/pyathena/aio/test_cursor.py index 350001d5..85648a3f 100644 --- a/tests/pyathena/aio/test_cursor.py +++ b/tests/pyathena/aio/test_cursor.py @@ -1,5 +1,5 @@ import re -from datetime import datetime +from datetime import datetime, timezone from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -149,7 +149,7 @@ def execution(schema): "QueryExecutionContext": {"Database": schema}, "Status": { "State": AthenaQueryExecution.STATE_SUCCEEDED, - "CompletionDateTime": datetime.now(), + "CompletionDateTime": datetime.now(timezone.utc), }, } } @@ -172,6 +172,45 @@ def execution(schema): == "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(timezone.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_no_result_set_raises(self, aio_cursor): with pytest.raises(ProgrammingError): await aio_cursor.fetchone() diff --git a/tests/pyathena/test_cursor.py b/tests/pyathena/test_cursor.py index fd3c7b53..f7d84638 100644 --- a/tests/pyathena/test_cursor.py +++ b/tests/pyathena/test_cursor.py @@ -223,6 +223,42 @@ def execution(schema): == "query_id_other_schema" ) + 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(timezone.utc), + }, + } + } + ) + + cursor = Cursor.__new__(Cursor) + cursor._schema_name = schema + + with patch.object( + Cursor, "_list_query_executions", return_value=(None, [execution("awsdatacatalog")]) + ): + # A different catalog must not be a cache hit. + cursor._catalog_name = "other_catalog" + assert 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 ( + cursor._find_previous_query_id(query, None, cache_size=100) + == "query_id_awsdatacatalog" + ) + @pytest.mark.parametrize( "cursor", [{"work_group": ENV.work_group, "result_reuse_enable": True, "result_reuse_minutes": 5}],