Trace S3 GET requests back to Athena queries.

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`.

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

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(