diff --git a/docs/usage.md b/docs/usage.md index 643c950e5..7cddaff3f 100644 --- a/docs/usage.md +++ b/docs/usage.md @@ -286,6 +286,7 @@ cursor.execute("SELECT * FROM one_row", cache_size=100, cache_expiration_time=36 Results will only be re-used from a succeeded DML query (the assumption being that you always want to re-run queries like `CREATE TABLE` and `DROP TABLE`) whose query string (with `pyformat` parameters substituted) matches *exactly*, and that ran with the same schema and catalog as the cursor. +The cache is not used for a `qmark` query with parameters. With `unload=True` on the pandas, Arrow, and Polars cursors, a query that is wrapped in `UNLOAD` is written to a new location each time, so the cache never matches it. The S3 staging directory is not checked, so it's possible that the location of the results is not in your provided `s3_staging_dir`. diff --git a/pyathena/aio/common.py b/pyathena/aio/common.py index 5b04cdec0..cec7e799a 100644 --- a/pyathena/aio/common.py +++ b/pyathena/aio/common.py @@ -48,7 +48,8 @@ async def _execute( # type: ignore[override] """Start a query execution, or find a previous one to reuse. The individual keyword arguments override the ``options`` field of the - same name unless None. + same name unless None. A query with execution parameters (``qmark``) + always starts a new execution. Args: operation: SQL query string. @@ -85,12 +86,16 @@ async def _execute( # type: ignore[override] paramstyle=paramstyle, ) query, request = self._build_execute_request(operation, parameters, options) - query_id = await self._find_previous_query_id( - query, - options.work_group, - cache_size=options.cache_size, - cache_expiration_time=options.cache_expiration_time, - ) + query_id = None + # Athena does not return the ExecutionParameters of earlier executions, + # so the cache cannot tell which parameters an execution ran with (#941). + if not request.get("ExecutionParameters"): + query_id = await self._find_previous_query_id( + query, + options.work_group, + cache_size=options.cache_size, + cache_expiration_time=options.cache_expiration_time, + ) if query_id is None: query_id = await self._start_execution(lambda: self._start_query_execution(request)) return query_id diff --git a/pyathena/common.py b/pyathena/common.py index 31c445db0..514eadab2 100644 --- a/pyathena/common.py +++ b/pyathena/common.py @@ -1280,7 +1280,8 @@ def _execute( """Start a query execution, or find a previous one to reuse. The individual keyword arguments override the ``options`` field of the - same name unless None. + same name unless None. A query with execution parameters (``qmark``) + always starts a new execution. Args: operation: SQL query string. @@ -1317,12 +1318,16 @@ def _execute( paramstyle=paramstyle, ) query, request = self._build_execute_request(operation, parameters, options) - query_id = self._find_previous_query_id( - query, - options.work_group, - cache_size=options.cache_size, - cache_expiration_time=options.cache_expiration_time, - ) + query_id = None + # Athena does not return the ExecutionParameters of earlier executions, + # so the cache cannot tell which parameters an execution ran with (#941). + if not request.get("ExecutionParameters"): + query_id = self._find_previous_query_id( + query, + options.work_group, + cache_size=options.cache_size, + cache_expiration_time=options.cache_expiration_time, + ) if query_id is None: query_id = self._start_execution(lambda: self._start_query_execution(request)) return query_id diff --git a/pyathena/options.py b/pyathena/options.py index 04a383521..58f4b3edb 100644 --- a/pyathena/options.py +++ b/pyathena/options.py @@ -48,6 +48,7 @@ class ExecuteOptions: caching. 0 (default) disables the cache lookup, unless ``cache_expiration_time`` is set to a positive value, in which case all queries within the expiration window are scanned. + A ``qmark`` query with parameters is never looked up. cache_expiration_time: Maximum age in seconds of a cached query result to consider for reuse. 0 (default) means no age limit. result_reuse_enable: Enable Athena server-side result reuse for this diff --git a/tests/pyathena/aio/test_cursor.py b/tests/pyathena/aio/test_cursor.py index 4561b4613..796a54add 100644 --- a/tests/pyathena/aio/test_cursor.py +++ b/tests/pyathena/aio/test_cursor.py @@ -242,6 +242,37 @@ async def test_execute_internal_legacy_kwargs_passthrough(self): cache_expiration_time=100, ) + 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" + } + cursor._retry_config = RetryConfig() + cursor._kill_on_interrupt = True + + 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, + ): + query_id = await cursor._execute( + "SELECT ?", ["'1'"], paramstyle="qmark", cache_size=10, cache_expiration_time=100 + ) + + assert query_id == "test_query_id" + assert request_mock.call_args.kwargs["execution_parameters"] == ["'1'"] + cache_mock.assert_not_awaited() + @pytest.mark.parametrize( "final_state", [AthenaQueryExecution.STATE_CANCELLED, AthenaQueryExecution.STATE_SUCCEEDED], diff --git a/tests/pyathena/test_cursor.py b/tests/pyathena/test_cursor.py index 22afd5781..f4a6b5528 100644 --- a/tests/pyathena/test_cursor.py +++ b/tests/pyathena/test_cursor.py @@ -199,6 +199,24 @@ def test_cache_size_with_work_group(self, cursor): assert first_query_id != second_query_id assert third_query_id in [first_query_id, second_query_id] + @pytest.mark.parametrize("cursor", [{"work_group": ENV.work_group}], indirect=["cursor"]) + def test_cache_size_with_qmark_parameters(self, cursor): + query = f"SELECT ? AS v -- {datetime.now(UTC)!s}" + + cursor.execute(query, ["'1'"], paramstyle="qmark") + first_query_id = cursor.query_id + + # Different parameters must not reuse the earlier execution (#941). + cursor.execute(query, ["'2'"], paramstyle="qmark", cache_size=100) + assert cursor.query_id != first_query_id + assert cursor.fetchall() == [("2",)] + + # Athena does not return the parameters of earlier executions, + # so even the same parameters run again. + cursor.execute(query, ["'1'"], paramstyle="qmark", cache_size=100) + assert cursor.query_id != first_query_id + assert cursor.fetchall() == [("1",)] + def test_cache_expiration_time(self, cursor): query = f"SELECT * FROM one_row -- {datetime.now(UTC)!s}" @@ -1303,6 +1321,32 @@ def test_execute_internal_legacy_kwargs_passthrough(self): cache_expiration_time=100, ) + def test_execute_qmark_parameters_skip_cache(self): + """A qmark query with parameters never searches the cache (no AWS, #941).""" + cursor = Cursor.__new__(Cursor) # 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( + Cursor, + "_build_start_query_execution_request", + return_value={"ExecutionParameters": ["'1'"]}, + ) as request_mock, + patch.object(Cursor, "_find_previous_query_id", return_value="cached") as cache_mock, + ): + query_id = cursor._execute( + "SELECT ?", ["'1'"], paramstyle="qmark", cache_size=10, cache_expiration_time=100 + ) + + assert query_id == "test_query_id" + assert request_mock.call_args.kwargs["execution_parameters"] == ["'1'"] + cache_mock.assert_not_called() + def test_connection_level_callback(self): """Test connection-level default callback.""" callback_results = []