address comments

pinterest · Jul 8, 2024 · 58867a4 · 58867a4
1 parent 8697311
commit 58867a4
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 14 deletions.
diff --git 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
@@ -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