feat: Athena - add client_request_token (#2474)

* feat: Athena - add client_request_token

Signed-off-by: Anton Kukushkin <[email protected]>

* refactor: Remove cache defaults copy-paste

Signed-off-by: Anton Kukushkin <[email protected]>

* test: Add client_request_token test case

Signed-off-by: Anton Kukushkin <[email protected]>

* PR feedback

Signed-off-by: Anton Kukushkin <[email protected]>

* Fix local cache test case

Signed-off-by: Anton Kukushkin <[email protected]>

* Drop coverage limit slightly

Signed-off-by: Anton Kukushkin <[email protected]>

---------

Signed-off-by: Anton Kukushkin <[email protected]>

Author: kukushking

awswrangler/athena/_cache.py

@@ -8,7 +8,7 @@
 
 import boto3
 
-from awswrangler import _utils
+from awswrangler import _utils, typing
 
 if TYPE_CHECKING:
     from mypy_boto3_athena.type_defs import QueryExecutionTypeDef
@@ -170,15 +170,23 @@ def _check_for_cached_results(
     sql: str,
     boto3_session: Optional[boto3.Session],
     workgroup: Optional[str],
-    max_cache_seconds: int,
-    max_cache_query_inspections: int,
-    max_remote_cache_entries: int,
+    athena_cache_settings: Optional[typing.AthenaCacheSettings] = None,
 ) -> _CacheInfo:
     """
     Check whether `sql` has been run before, within the `max_cache_seconds` window, by the `workgroup`.
 
     If so, returns a dict with Athena's `query_execution_info` and the data format.
     """
+    athena_cache_settings = athena_cache_settings or {}
+
+    max_cache_seconds = athena_cache_settings.get("max_cache_seconds", 0)
+    max_cache_query_inspections = athena_cache_settings.get("max_cache_query_inspections", 50)
+    max_remote_cache_entries = athena_cache_settings.get("max_remote_cache_entries", 50)
+    max_local_cache_entries = athena_cache_settings.get("max_local_cache_entries", 100)
+    max_remote_cache_entries = min(max_remote_cache_entries, max_local_cache_entries)
+
+    _cache_manager.max_cache_size = max_local_cache_entries
+
     if max_cache_seconds <= 0:
         return _CacheInfo(has_valid_cache=False)
     num_executions_inspected: int = 0

awswrangler/athena/_executions.py

@@ -17,7 +17,7 @@
 from awswrangler import _utils, exceptions, typing
 from awswrangler._config import apply_configs
 
-from ._cache import _cache_manager, _CacheInfo, _check_for_cached_results
+from ._cache import _CacheInfo, _check_for_cached_results
 from ._utils import (
     _QUERY_FINAL_STATES,
     _QUERY_WAIT_POLLING_DELAY,
@@ -41,6 +41,7 @@ def start_query_execution(
     params: Union[Dict[str, Any], List[str], None] = None,
     paramstyle: Literal["qmark", "named"] = "named",
     boto3_session: Optional[boto3.Session] = None,
+    client_request_token: Optional[str] = None,
     athena_cache_settings: Optional[typing.AthenaCacheSettings] = None,
     athena_query_wait_polling_delay: float = _QUERY_WAIT_POLLING_DELAY,
     data_source: Optional[str] = None,
@@ -88,6 +89,13 @@ def start_query_execution(
         - ``qmark``
     boto3_session : boto3.Session(), optional
         Boto3 Session. The default boto3 session will be used if boto3_session receive None.
+    client_request_token : str, optional
+        A unique case-sensitive string used to ensure the request to create the query is idempotent (executes only once).
+        If another StartQueryExecution request is received, the same response is returned and another query is not created.
+        If a parameter has changed, for example, the QueryString , an error is returned.
+        If you pass the same client_request_token value with different parameters the query fails with error
+        message "Idempotent parameters do not match". Use this only with ctas_approach=False and unload_approach=False
+        and disabled cache.
     athena_cache_settings: typing.AthenaCacheSettings, optional
         Parameters of the Athena cache settings such as max_cache_seconds, max_cache_query_inspections,
         max_remote_cache_entries, and max_local_cache_entries.
@@ -125,26 +133,16 @@ def start_query_execution(
     sql, execution_params = _apply_formatter(sql, params, paramstyle)
     _logger.debug("Executing query:\n%s", sql)
 
-    athena_cache_settings = athena_cache_settings if athena_cache_settings else {}
-    max_cache_seconds = athena_cache_settings.get("max_cache_seconds", 0)
-    max_cache_query_inspections = athena_cache_settings.get("max_cache_query_inspections", 50)
-    max_remote_cache_entries = athena_cache_settings.get("max_remote_cache_entries", 50)
-    max_local_cache_entries = athena_cache_settings.get("max_local_cache_entries", 100)
-
-    max_remote_cache_entries = min(max_remote_cache_entries, max_local_cache_entries)
-
-    _cache_manager.max_cache_size = max_local_cache_entries
-    cache_info: _CacheInfo = _check_for_cached_results(
-        sql=sql,
-        boto3_session=boto3_session,
-        workgroup=workgroup,
-        max_cache_seconds=max_cache_seconds,
-        max_cache_query_inspections=max_cache_query_inspections,
-        max_remote_cache_entries=max_remote_cache_entries,
-    )
-    _logger.debug("Cache info:\n%s", cache_info)
+    if not client_request_token:
+        cache_info: _CacheInfo = _check_for_cached_results(
+            sql=sql,
+            boto3_session=boto3_session,
+            workgroup=workgroup,
+            athena_cache_settings=athena_cache_settings,
+        )
+        _logger.debug("Cache info:\n%s", cache_info)
 
-    if cache_info.has_valid_cache and cache_info.query_execution_id is not None:
+    if not client_request_token and cache_info.has_valid_cache and cache_info.query_execution_id is not None:
         query_execution_id = cache_info.query_execution_id
         _logger.debug("Valid cache found. Retrieving...")
     else:
@@ -159,6 +157,7 @@ def start_query_execution(
             encryption=encryption,
             kms_key=kms_key,
             execution_params=execution_params,
+            client_request_token=client_request_token,
             boto3_session=boto3_session,
         )
     if wait:

awswrangler/athena/_read.py

@@ -427,6 +427,7 @@ def _resolve_query_without_cache_regular(
     boto3_session: Optional[boto3.Session],
     execution_params: Optional[List[str]] = None,
     dtype_backend: Literal["numpy_nullable", "pyarrow"] = "numpy_nullable",
+    client_request_token: Optional[str] = None,
 ) -> Union[pd.DataFrame, Iterator[pd.DataFrame]]:
     wg_config: _WorkGroupConfig = _get_workgroup_config(session=boto3_session, workgroup=workgroup)
     s3_output = _get_s3_output(s3_output=s3_output, wg_config=wg_config, boto3_session=boto3_session)
@@ -442,6 +443,7 @@ def _resolve_query_without_cache_regular(
         encryption=encryption,
         kms_key=kms_key,
         execution_params=execution_params,
+        client_request_token=client_request_token,
         boto3_session=boto3_session,
     )
     _logger.debug("Query id: %s", query_id)
@@ -490,6 +492,7 @@ def _resolve_query_without_cache(
     pyarrow_additional_kwargs: Optional[Dict[str, Any]] = None,
     execution_params: Optional[List[str]] = None,
     dtype_backend: Literal["numpy_nullable", "pyarrow"] = "numpy_nullable",
+    client_request_token: Optional[str] = None,
 ) -> Union[pd.DataFrame, Iterator[pd.DataFrame]]:
     """
     Execute a query in Athena and returns results as DataFrame, back to `read_sql_query`.
@@ -570,6 +573,7 @@ def _resolve_query_without_cache(
         boto3_session=boto3_session,
         execution_params=execution_params,
         dtype_backend=dtype_backend,
+        client_request_token=client_request_token,
     )
 
 
@@ -776,6 +780,7 @@ def read_sql_query(  # pylint: disable=too-many-arguments,too-many-locals
     keep_files: bool = True,
     use_threads: Union[bool, int] = True,
     boto3_session: Optional[boto3.Session] = None,
+    client_request_token: Optional[str] = None,
     athena_cache_settings: Optional[typing.AthenaCacheSettings] = None,
     data_source: Optional[str] = None,
     athena_query_wait_polling_delay: float = _QUERY_WAIT_POLLING_DELAY,
@@ -938,6 +943,13 @@ def read_sql_query(  # pylint: disable=too-many-arguments,too-many-locals
         If integer is provided, specified number is used.
     boto3_session : boto3.Session(), optional
         Boto3 Session. The default boto3 session will be used if boto3_session receive None.
+    client_request_token : str, optional
+        A unique case-sensitive string used to ensure the request to create the query is idempotent (executes only once).
+        If another StartQueryExecution request is received, the same response is returned and another query is not created.
+        If a parameter has changed, for example, the QueryString , an error is returned.
+        If you pass the same client_request_token value with different parameters the query fails with error
+        message "Idempotent parameters do not match". Use this only with ctas_approach=False and unload_approach=False
+        and disabled cache.
     athena_cache_settings: typing.AthenaCacheSettings, optional
         Parameters of the Athena cache settings such as max_cache_seconds, max_cache_query_inspections,
         max_remote_cache_entries, and max_local_cache_entries.
@@ -1022,46 +1034,44 @@ def read_sql_query(  # pylint: disable=too-many-arguments,too-many-locals
         raise exceptions.InvalidArgumentCombination("Only one of ctas_approach=True or unload_approach=True is allowed")
     if unload_parameters and unload_parameters.get("file_format") not in (None, "PARQUET"):
         raise exceptions.InvalidArgumentCombination("Only PARQUET file format is supported if unload_approach=True")
+    if client_request_token and athena_cache_settings:
+        raise exceptions.InvalidArgumentCombination(
+            "Only one of `client_request_token` or `athena_cache_settings` is allowed."
+        )
+    if client_request_token and (ctas_approach or unload_approach):
+        raise exceptions.InvalidArgumentCombination(
+            "Using `client_request_token` is only allowed when `ctas_approach=False` and `unload_approach=False`."
+        )
     chunksize = sys.maxsize if ctas_approach is False and chunksize is True else chunksize
 
     # Substitute query parameters if applicable
     sql, execution_params = _apply_formatter(sql, params, paramstyle)
 
-    athena_cache_settings = athena_cache_settings if athena_cache_settings else {}
-    max_cache_seconds = athena_cache_settings.get("max_cache_seconds", 0)
-    max_cache_query_inspections = athena_cache_settings.get("max_cache_query_inspections", 50)
-    max_remote_cache_entries = athena_cache_settings.get("max_remote_cache_entries", 50)
-    max_local_cache_entries = athena_cache_settings.get("max_local_cache_entries", 100)
-
-    max_remote_cache_entries = min(max_remote_cache_entries, max_local_cache_entries)
-
-    _cache_manager.max_cache_size = max_local_cache_entries
-    cache_info: _CacheInfo = _check_for_cached_results(
-        sql=sql,
-        boto3_session=boto3_session,
-        workgroup=workgroup,
-        max_cache_seconds=max_cache_seconds,
-        max_cache_query_inspections=max_cache_query_inspections,
-        max_remote_cache_entries=max_remote_cache_entries,
-    )
-    _logger.debug("Cache info:\n%s", cache_info)
-    if cache_info.has_valid_cache is True:
-        _logger.debug("Valid cache found. Retrieving...")
-        try:
-            return _resolve_query_with_cache(
-                cache_info=cache_info,
-                categories=categories,
-                chunksize=chunksize,
-                use_threads=use_threads,
-                session=boto3_session,
-                athena_query_wait_polling_delay=athena_query_wait_polling_delay,
-                s3_additional_kwargs=s3_additional_kwargs,
-                pyarrow_additional_kwargs=pyarrow_additional_kwargs,
-                dtype_backend=dtype_backend,
-            )
-        except Exception as e:  # pylint: disable=broad-except
-            _logger.error(e)  # if there is anything wrong with the cache, just fallback to the usual path
-            _logger.debug("Corrupted cache. Continuing to execute query...")
+    if not client_request_token:
+        cache_info: _CacheInfo = _check_for_cached_results(
+            sql=sql,
+            boto3_session=boto3_session,
+            workgroup=workgroup,
+            athena_cache_settings=athena_cache_settings,
+        )
+        _logger.debug("Cache info:\n%s", cache_info)
+        if cache_info.has_valid_cache is True:
+            _logger.debug("Valid cache found. Retrieving...")
+            try:
+                return _resolve_query_with_cache(
+                    cache_info=cache_info,
+                    categories=categories,
+                    chunksize=chunksize,
+                    use_threads=use_threads,
+                    session=boto3_session,
+                    athena_query_wait_polling_delay=athena_query_wait_polling_delay,
+                    s3_additional_kwargs=s3_additional_kwargs,
+                    pyarrow_additional_kwargs=pyarrow_additional_kwargs,
+                    dtype_backend=dtype_backend,
+                )
+            except Exception as e:  # pylint: disable=broad-except
+                _logger.error(e)  # if there is anything wrong with the cache, just fallback to the usual path
+                _logger.debug("Corrupted cache. Continuing to execute query...")
 
     ctas_parameters = ctas_parameters if ctas_parameters else {}
     ctas_database = ctas_parameters.get("database")
@@ -1094,6 +1104,7 @@ def read_sql_query(  # pylint: disable=too-many-arguments,too-many-locals
         pyarrow_additional_kwargs=pyarrow_additional_kwargs,
         execution_params=execution_params,
         dtype_backend=dtype_backend,
+        client_request_token=client_request_token,
     )
 
 
@@ -1117,6 +1128,7 @@ def read_sql_table(
     keep_files: bool = True,
     use_threads: Union[bool, int] = True,
     boto3_session: Optional[boto3.Session] = None,
+    client_request_token: Optional[str] = None,
     athena_cache_settings: Optional[typing.AthenaCacheSettings] = None,
     data_source: Optional[str] = None,
     dtype_backend: Literal["numpy_nullable", "pyarrow"] = "numpy_nullable",
@@ -1274,6 +1286,13 @@ def read_sql_table(
         If integer is provided, specified number is used.
     boto3_session : boto3.Session(), optional
         Boto3 Session. The default boto3 session will be used if boto3_session receive None.
+    client_request_token : str, optional
+        A unique case-sensitive string used to ensure the request to create the query is idempotent (executes only once).
+        If another StartQueryExecution request is received, the same response is returned and another query is not created.
+        If a parameter has changed, for example, the QueryString , an error is returned.
+        If you pass the same client_request_token value with different parameters the query fails with error
+        message "Idempotent parameters do not match". Use this only with ctas_approach=False and unload_approach=False
+        and disabled cache.
     athena_cache_settings: typing.AthenaCacheSettings, optional
         Parameters of the Athena cache settings such as max_cache_seconds, max_cache_query_inspections,
         max_remote_cache_entries, and max_local_cache_entries.
@@ -1327,6 +1346,7 @@ def read_sql_table(
         keep_files=keep_files,
         use_threads=use_threads,
         boto3_session=boto3_session,
+        client_request_token=client_request_token,
         athena_cache_settings=athena_cache_settings,
         data_source=data_source,
         s3_additional_kwargs=s3_additional_kwargs,

awswrangler/athena/_utils.py

@@ -86,6 +86,7 @@ def _start_query_execution(
     encryption: Optional[str] = None,
     kms_key: Optional[str] = None,
     execution_params: Optional[List[str]] = None,
+    client_request_token: Optional[str] = None,
     boto3_session: Optional[boto3.Session] = None,
 ) -> str:
     args: Dict[str, Any] = {"QueryString": sql}
@@ -116,6 +117,9 @@ def _start_query_execution(
     if workgroup is not None:
         args["WorkGroup"] = workgroup
 
+    if client_request_token:
+        args["ClientRequestToken"] = client_request_token
+
     if execution_params:
         args["ExecutionParameters"] = execution_params
 

tests/unit/test_athena_cache.py

@@ -175,6 +175,9 @@ def test_local_cache(wr, path, glue_database, glue_table):
     df = pd.DataFrame({"c0": [0, None]}, dtype="Int64")
     wr.s3.to_parquet(df=df, path=path, dataset=True, mode="overwrite", database=glue_database, table=glue_table)
 
+    # Set max cache size because it is supposed to be set in the patched method below
+    wr.athena._cache._cache_manager.max_cache_size = 1
+
     with patch(
         "awswrangler.athena._read._check_for_cached_results",
         return_value=wr.athena._read._CacheInfo(has_valid_cache=False),
@@ -189,7 +192,7 @@ def test_local_cache(wr, path, glue_database, glue_table):
         assert df.shape == df2.shape
         assert df.c0.sum() == df2.c0.sum()
         first_query_id = df2.query_metadata["QueryExecutionId"]
-        assert first_query_id in wr.athena._read._cache_manager
+        assert first_query_id in wr.athena._cache._cache_manager
 
         df3 = wr.athena.read_sql_query(
             f"SELECT * FROM {glue_table}",
@@ -202,8 +205,8 @@ def test_local_cache(wr, path, glue_database, glue_table):
         assert df.c0.sum() == df3.c0.sum()
         second_query_id = df3.query_metadata["QueryExecutionId"]
 
-        assert first_query_id not in wr.athena._read._cache_manager
-        assert second_query_id in wr.athena._read._cache_manager
+        assert first_query_id not in wr.athena._cache._cache_manager
+        assert second_query_id in wr.athena._cache._cache_manager
 
 
 def test_paginated_remote_cache(wr, path, glue_database, glue_table, workgroup1):
@@ -260,3 +263,30 @@ def test_cache_start_query(wr, path, glue_database, glue_table, data_source):
         )
         internal_start_query.assert_not_called()
         assert query_id == query_id_2
+
+
+def test_start_query_client_request_token(wr, path, glue_database, glue_table):
+    df = pd.DataFrame({"c0": [0, None]}, dtype="Int64")
+
+    wr.s3.to_parquet(
+        df=df,
+        path=path,
+        dataset=True,
+        mode="overwrite",
+        database=glue_database,
+        table=glue_table,
+    )
+
+    client_request_token = f"token-{glue_database}-{glue_table}-1"
+    query_id_1 = wr.athena.start_query_execution(
+        sql=f"SELECT * FROM {glue_table}",
+        database=glue_database,
+        client_request_token=client_request_token,
+    )
+    query_id_2 = wr.athena.start_query_execution(
+        sql=f"SELECT * FROM {glue_table}",
+        database=glue_database,
+        client_request_token=client_request_token,
+    )
+
+    assert query_id_1 == query_id_2

tox.ini

@@ -10,7 +10,7 @@ passenv =
        AWS_SECRET_ACCESS_KEY
        AWS_SESSION_TOKEN
 setenv =
-       COV_FAIL_UNDER = 87.00
+       COV_FAIL_UNDER = 82.00
 allowlist_externals = poetry
 commands_pre =
        poetry install --no-root --sync --extras "deltalake gremlin mysql opencypher opensearch oracle postgres redshift sparql sqlserver geopandas"