Skip to content

Commit 446a6c8

Browse files
laughingman7743laughingman7743
authored
Merge pull request #740 from Liam3851/fix/739-cache-schema-catalog
Fix/739 cache schema catalog
2 parents 36f1d4e + fe66469 commit 446a6c8

4 files changed

Lines changed: 166 additions & 3 deletions

File tree

‎pyathena/aio/common.py‎

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -218,7 +218,11 @@ async def _find_previous_query_id( # type: ignore[override]
218218
):
219219
next_token = None
220220
break
221-
if execution.query == query:
221+
if (
222+
execution.query == query
223+
and execution.database == self._schema_name
224+
and (execution.catalog or "").lower() == (self._catalog_name or "").lower()
225+
):
222226
query_id = execution.query_id
223227
break
224228
if query_id or next_token is None:

‎pyathena/common.py‎

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -638,7 +638,11 @@ def _find_previous_query_id(
638638
):
639639
next_token = None
640640
break
641-
if execution.query == query:
641+
if (
642+
execution.query == query
643+
and execution.database == self._schema_name
644+
and (execution.catalog or "").lower() == (self._catalog_name or "").lower()
645+
):
642646
query_id = execution.query_id
643647
break
644648
if query_id or next_token is None:

‎tests/pyathena/aio/test_cursor.py‎

Lines changed: 82 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
import re
2-
from datetime import datetime
2+
from datetime import datetime, timezone
33
from unittest.mock import AsyncMock, MagicMock, patch
44

55
import pytest
@@ -130,6 +130,87 @@ async def test_execute_internal_legacy_kwargs_passthrough(self):
130130
cache_expiration_time=100,
131131
)
132132

133+
async def test_cache_size_different_schema(self):
134+
"""A cached result is only reused when it ran against the same schema (#739).
135+
136+
Mirrors the synchronous cursor test: identical SQL can resolve to different
137+
tables depending on the database it runs against, so a prior execution from
138+
another schema must not be a cache hit.
139+
"""
140+
query = "SELECT * FROM one_row"
141+
142+
def execution(schema):
143+
return AthenaQueryExecution(
144+
{
145+
"QueryExecution": {
146+
"QueryExecutionId": f"query_id_{schema}",
147+
"Query": query,
148+
"StatementType": AthenaQueryExecution.STATEMENT_TYPE_DML,
149+
"QueryExecutionContext": {"Database": schema},
150+
"Status": {
151+
"State": AthenaQueryExecution.STATE_SUCCEEDED,
152+
"CompletionDateTime": datetime.now(timezone.utc),
153+
},
154+
}
155+
}
156+
)
157+
158+
cursor = AioCursor.__new__(AioCursor) # bypass __init__ to avoid AWS calls
159+
cursor._catalog_name = None
160+
161+
with patch.object(
162+
AioCursor,
163+
"_list_query_executions",
164+
new_callable=AsyncMock,
165+
return_value=(None, [execution("other_schema")]),
166+
):
167+
cursor._schema_name = "this_schema"
168+
assert await cursor._find_previous_query_id(query, None, cache_size=100) is None
169+
cursor._schema_name = "other_schema"
170+
assert (
171+
await cursor._find_previous_query_id(query, None, cache_size=100)
172+
== "query_id_other_schema"
173+
)
174+
175+
async def test_cache_size_different_catalog(self):
176+
query = "SELECT * FROM one_row"
177+
schema = "this_schema"
178+
179+
def execution(catalog):
180+
return AthenaQueryExecution(
181+
{
182+
"QueryExecution": {
183+
"QueryExecutionId": f"query_id_{catalog}",
184+
"Query": query,
185+
"StatementType": AthenaQueryExecution.STATEMENT_TYPE_DML,
186+
"QueryExecutionContext": {"Database": schema, "Catalog": catalog},
187+
"Status": {
188+
"State": AthenaQueryExecution.STATE_SUCCEEDED,
189+
"CompletionDateTime": datetime.now(timezone.utc),
190+
},
191+
}
192+
}
193+
)
194+
195+
cursor = AioCursor.__new__(AioCursor)
196+
cursor._schema_name = schema
197+
198+
with patch.object(
199+
AioCursor,
200+
"_list_query_executions",
201+
new_callable=AsyncMock,
202+
return_value=(None, [execution("awsdatacatalog")]),
203+
):
204+
# A different catalog must not be a cache hit.
205+
cursor._catalog_name = "other_catalog"
206+
assert await cursor._find_previous_query_id(query, None, cache_size=100) is None
207+
# The same catalog, differing only in case, must still be a cache hit.
208+
cursor._catalog_name = "AwsDataCatalog"
209+
assert (
210+
await cursor._find_previous_query_id(query, None, cache_size=100)
211+
== "query_id_awsdatacatalog"
212+
)
213+
133214
async def test_no_result_set_raises(self, aio_cursor):
134215
with pytest.raises(ProgrammingError):
135216
await aio_cursor.fetchone()

‎tests/pyathena/test_cursor.py‎

Lines changed: 74 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -185,6 +185,80 @@ def test_cache_expiration_time_with_cache_size(self, cursor):
185185
assert query_id_7 != query_id_8
186186
assert query_id_9 in [query_id_7, query_id_8]
187187

188+
def test_cache_size_different_schema(self):
189+
"""A cached result is only reused when it ran against the same schema (#739).
190+
191+
Identical SQL can resolve to different tables depending on the database it
192+
runs against, so a prior execution from another schema must not be a cache hit.
193+
"""
194+
query = "SELECT * FROM one_row"
195+
196+
def execution(schema):
197+
return AthenaQueryExecution(
198+
{
199+
"QueryExecution": {
200+
"QueryExecutionId": f"query_id_{schema}",
201+
"Query": query,
202+
"StatementType": AthenaQueryExecution.STATEMENT_TYPE_DML,
203+
"QueryExecutionContext": {"Database": schema},
204+
"Status": {
205+
"State": AthenaQueryExecution.STATE_SUCCEEDED,
206+
"CompletionDateTime": datetime.now(timezone.utc),
207+
},
208+
}
209+
}
210+
)
211+
212+
cursor = Cursor.__new__(Cursor) # bypass __init__ to avoid AWS calls
213+
cursor._catalog_name = None
214+
215+
with patch.object(
216+
Cursor, "_list_query_executions", return_value=(None, [execution("other_schema")])
217+
):
218+
cursor._schema_name = "this_schema"
219+
assert cursor._find_previous_query_id(query, None, cache_size=100) is None
220+
cursor._schema_name = "other_schema"
221+
assert (
222+
cursor._find_previous_query_id(query, None, cache_size=100)
223+
== "query_id_other_schema"
224+
)
225+
226+
def test_cache_size_different_catalog(self):
227+
query = "SELECT * FROM one_row"
228+
schema = "this_schema"
229+
230+
def execution(catalog):
231+
return AthenaQueryExecution(
232+
{
233+
"QueryExecution": {
234+
"QueryExecutionId": f"query_id_{catalog}",
235+
"Query": query,
236+
"StatementType": AthenaQueryExecution.STATEMENT_TYPE_DML,
237+
"QueryExecutionContext": {"Database": schema, "Catalog": catalog},
238+
"Status": {
239+
"State": AthenaQueryExecution.STATE_SUCCEEDED,
240+
"CompletionDateTime": datetime.now(timezone.utc),
241+
},
242+
}
243+
}
244+
)
245+
246+
cursor = Cursor.__new__(Cursor)
247+
cursor._schema_name = schema
248+
249+
with patch.object(
250+
Cursor, "_list_query_executions", return_value=(None, [execution("awsdatacatalog")])
251+
):
252+
# A different catalog must not be a cache hit.
253+
cursor._catalog_name = "other_catalog"
254+
assert cursor._find_previous_query_id(query, None, cache_size=100) is None
255+
# The same catalog, differing only in case, must still be a cache hit
256+
cursor._catalog_name = "AwsDataCatalog"
257+
assert (
258+
cursor._find_previous_query_id(query, None, cache_size=100)
259+
== "query_id_awsdatacatalog"
260+
)
261+
188262
@pytest.mark.parametrize(
189263
"cursor",
190264
[{"work_group": ENV.work_group, "result_reuse_enable": True, "result_reuse_minutes": 5}],

0 commit comments

Comments
 (0)