diff --git a/pyathena/aio/common.py b/pyathena/aio/common.py index 19914164..e490a544 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 execution.database == self._schema_name + and (execution.catalog or "").lower() == (self._catalog_name or "").lower() + ): 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..3a517694 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 execution.database == self._schema_name + and (execution.catalog or "").lower() == (self._catalog_name or "").lower() + ): 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..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 @@ -130,6 +130,87 @@ 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(timezone.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(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 bfe9e82a..f7d84638 100644 --- a/tests/pyathena/test_cursor.py +++ b/tests/pyathena/test_cursor.py @@ -185,6 +185,80 @@ 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" + ) + + 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}],