diff --git a/querybook/server/datasources/query_execution.py b/querybook/server/datasources/query_execution.py index 1c0c72876..97ab2cd8d 100644 --- a/querybook/server/datasources/query_execution.py +++ b/querybook/server/datasources/query_execution.py @@ -1,6 +1,6 @@ from datetime import datetime -from flask import abort, Response, redirect +from flask import abort, Response, redirect, request from flask_login import current_user from app.flask_app import socketio @@ -72,6 +72,10 @@ def create_query_execution( query=query, engine_id=engine_id, uid=uid, session=session ) + metadata = metadata or {} + 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 diff --git a/querybook/server/lib/query_executor/base_executor.py b/querybook/server/lib/query_executor/base_executor.py index 0bfb45ca4..6ef2f9e2e 100644 --- a/querybook/server/lib/query_executor/base_executor.py +++ b/querybook/server/lib/query_executor/base_executor.py @@ -528,6 +528,18 @@ def __init__( self._client = None self._cursor = None + query_execution_metadata = ( + qe_logic.get_query_execution_metadata_by_execution_id( + self._query_execution_id + ) + ) + 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 del self._cursor