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 diff --git a/pyathena/common.py b/pyathena/common.py index 64ce0b0e..7cd4a6b1 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,25 @@ def _execute( cache_expiration_time=cache_expiration_time if cache_expiration_time else 0, ) if query_id is None: + query_trace_id = 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 + self._connection.client._client_config.user_agent_extra += f' QueryTraceId={query_trace_id}' + + 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(