diff --git a/querybook/server/datasources/query_execution.py b/querybook/server/datasources/query_execution.py index 57571d60a..97ab2cd8d 100644 --- a/querybook/server/datasources/query_execution.py +++ b/querybook/server/datasources/query_execution.py @@ -72,14 +72,14 @@ def create_query_execution( query=query, engine_id=engine_id, uid=uid, session=session ) - api_access_token = ( - True if request.headers.get("api-access-token", None) else False - ) metadata = metadata or {} - metadata["api_access_token"] = api_access_token - logic.create_query_execution_metadata( - query_execution.id, metadata, session=session - ) + used_api_token = request.headers.get("api-access-token") is not None + if used_api_token: + metadata["used_api_token"] = used_api_token + if metadata: + logic.create_query_execution_metadata( + query_execution.id, metadata, session=session + ) data_doc = None if data_cell_id: diff --git a/querybook/server/lib/query_executor/base_executor.py b/querybook/server/lib/query_executor/base_executor.py index 805815e2c..6ef2f9e2e 100644 --- a/querybook/server/lib/query_executor/base_executor.py +++ b/querybook/server/lib/query_executor/base_executor.py @@ -528,15 +528,17 @@ def __init__( self._client = None self._cursor = None - with DBSession() as session: - query_execution_metadata = ( - qe_logic.get_query_execution_metadata_by_execution_id( - self._query_execution_id, session=session - ).execution_metadata + query_execution_metadata = ( + qe_logic.get_query_execution_metadata_by_execution_id( + self._query_execution_id ) - self._api_access_token = query_execution_metadata.get( - "api_access_token", False + ) + if query_execution_metadata: + self._used_api_token = query_execution_metadata.execution_metadata.get( + "used_api_token", False ) + else: + self._used_api_token = False def __del__(self): del self._logger