From ed379a52341cfe2911f68500c1f1131f91922675 Mon Sep 17 00:00:00 2001 From: Antony Southworth Date: Fri, 12 Jul 2024 09:58:32 +1200 Subject: [PATCH 1/3] Inject ClientRequestToken into User-Agent for StartQueryExecution --- pyathena/common.py | 29 ++++++++++++++++++++--------- 1 file changed, 20 insertions(+), 9 deletions(-) diff --git a/pyathena/common.py b/pyathena/common.py index 64ce0b0e..8b2d5104 100644 --- a/pyathena/common.py +++ b/pyathena/common.py @@ -7,6 +7,7 @@ from abc import ABCMeta, abstractmethod from datetime import datetime, timedelta, timezone from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union, cast +from uuid import uuid4 from pyathena.converter import Converter, DefaultTypeConverter from pyathena.error import DatabaseError, OperationalError, ProgrammingError @@ -571,16 +572,26 @@ def _execute( cache_expiration_time=cache_expiration_time if cache_expiration_time else 0, ) if query_id is None: + client_request_token = str(uuid4()) + old_user_agent_extra = self._connection.client._client_config.user_agent_extra try: - query_id = retry_api_call( - self._connection.client.start_query_execution, - config=self._retry_config, - logger=_logger, - **request, - ).get("QueryExecutionId") - except Exception as e: - _logger.exception("Failed to execute query.") - raise DatabaseError(*e.args) from e + request['ClientRequestToken'] = client_request_token + self._connection.client._client_config.user_agent_extra += f' ClientRequestToken={client_request_token}' + + try: + query_id = retry_api_call( + self._connection.client.start_query_execution, + config=self._retry_config, + logger=_logger, + **request, + ).get("QueryExecutionId") + except Exception as e: + _logger.exception("Failed to execute query.") + raise DatabaseError(*e.args) from e + + finally: + self._connection.client._client_config.user_agent_extra = old_user_agent_extra + return query_id def _calculate( From 95b46c5834e3d1e2c86418eb12ce10d9c8543c4b Mon Sep 17 00:00:00 2001 From: Antony Southworth Date: Fri, 12 Jul 2024 10:16:35 +1200 Subject: [PATCH 2/3] Actually, it doesn't need to be the ClientRequestToken --- pyathena/common.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/pyathena/common.py b/pyathena/common.py index 8b2d5104..7cd4a6b1 100644 --- a/pyathena/common.py +++ b/pyathena/common.py @@ -572,11 +572,10 @@ def _execute( cache_expiration_time=cache_expiration_time if cache_expiration_time else 0, ) if query_id is None: - client_request_token = str(uuid4()) + query_trace_id = str(uuid4()) old_user_agent_extra = self._connection.client._client_config.user_agent_extra try: - request['ClientRequestToken'] = client_request_token - self._connection.client._client_config.user_agent_extra += f' ClientRequestToken={client_request_token}' + self._connection.client._client_config.user_agent_extra += f' QueryTraceId={query_trace_id}' try: query_id = retry_api_call( From 9ddd20b209cbf2b61dc1dc9e1f716b179fa6e161 Mon Sep 17 00:00:00 2001 From: Antony Southworth Date: Fri, 12 Jul 2024 11:06:04 +1200 Subject: [PATCH 3/3] Add basic CONTRIBUTING. --- CONTRIBUTING.md | 7 +++++++ pyathena/__init__.py | 2 +- 2 files changed, 8 insertions(+), 1 deletion(-) create mode 100644 CONTRIBUTING.md diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md new file mode 100644 index 00000000..a6d7760c --- /dev/null +++ b/CONTRIBUTING.md @@ -0,0 +1,7 @@ +# Contributing + +This project uses [`hatch`](https://hatch.pypa.io/latest/) to manage building, running, and testing. + +Install `hatch` with `pip install hatch` then use the commands in the `Makefile`. + +To build the wheel locally, use `hatch build`. diff --git a/pyathena/__init__.py b/pyathena/__init__.py index 01f86976..dd6bfc14 100644 --- a/pyathena/__init__.py +++ b/pyathena/__init__.py @@ -10,7 +10,7 @@ from pyathena.connection import Connection, ConnectionCursor from pyathena.cursor import Cursor -__version__ = "3.8.3" +__version__ = "3.8.4+tonybeta" user_agent_extra: str = f"PyAthena/{__version__}" # Globals https://www.python.org/dev/peps/pep-0249/#globals