[native] Pass extraCredentials to connector session properties

prestodb · Feb 17, 2025 · 83e9709 · 83e9709
1 parent 700fdc7
commit 83e9709
Show file tree

Hide file tree

Showing 3 changed files with 16 additions and 11 deletions.
diff --git a/presto-native-execution/presto_cpp/main/QueryContextManager.cpp b/presto-native-execution/presto_cpp/main/QueryContextManager.cpp
@@ -60,14 +60,17 @@ void updateFromSystemConfigs(
 }
 
 std::unordered_map<std::string, std::unordered_map<std::string, std::string>>
-toConnectorConfigs(const protocol::SessionRepresentation& session) {
+toConnectorConfigs(const protocol::TaskUpdateRequest& taskUpdateRequest) {
   std::unordered_map<std::string, std::unordered_map<std::string, std::string>>
       connectorConfigs;
-  for (const auto& entry : session.catalogProperties) {
-    connectorConfigs.insert(
-        {entry.first,
-         std::unordered_map<std::string, std::string>(
-             entry.second.begin(), entry.second.end())});
+  for (const auto& entry : taskUpdateRequest.session.catalogProperties) {
+    auto sessionProperties = std::unordered_map<std::string, std::string>(
+        entry.second.begin(), entry.second.end());
+    sessionProperties.insert(
+        taskUpdateRequest.extraCredentials.begin(),
+        taskUpdateRequest.extraCredentials.end());
+    sessionProperties.insert({"user", taskUpdateRequest.session.user});
+    connectorConfigs.insert({entry.first, sessionProperties});
   }
 
   return connectorConfigs;
@@ -106,9 +109,11 @@ QueryContextManager::QueryContextManager(
 std::shared_ptr<velox::core::QueryCtx>
 QueryContextManager::findOrCreateQueryCtx(
     const protocol::TaskId& taskId,
-    const protocol::SessionRepresentation& session) {
+    const protocol::TaskUpdateRequest& taskUpdateRequest) {
   return findOrCreateQueryCtx(
-      taskId, toVeloxConfigs(session), toConnectorConfigs(session));
+      taskId,
+      toVeloxConfigs(taskUpdateRequest.session),
+      toConnectorConfigs(taskUpdateRequest));
 }
 
 std::shared_ptr<core::QueryCtx> QueryContextManager::findOrCreateQueryCtx(

diff --git a/presto-native-execution/presto_cpp/main/QueryContextManager.h b/presto-native-execution/presto_cpp/main/QueryContextManager.h
@@ -107,7 +107,7 @@ class QueryContextManager {
 
   std::shared_ptr<velox::core::QueryCtx> findOrCreateQueryCtx(
       const protocol::TaskId& taskId,
-      const protocol::SessionRepresentation& session);
+      const protocol::TaskUpdateRequest& taskUpdateRequest);
 
   /// Calls the given functor for every present query context.
   void visitAllContexts(std::function<void(

diff --git a/presto-native-execution/presto_cpp/main/TaskResource.cpp b/presto-native-execution/presto_cpp/main/TaskResource.cpp
@@ -299,7 +299,7 @@ proxygen::RequestHandler* TaskResource::createOrUpdateBatchTask(
 
         auto queryCtx =
             taskManager_.getQueryContextManager()->findOrCreateQueryCtx(
-                taskId, updateRequest.session);
+                taskId, updateRequest);
 
         VeloxBatchQueryPlanConverter converter(
             shuffleName,
@@ -340,7 +340,7 @@ proxygen::RequestHandler* TaskResource::createOrUpdateTask(
 
           queryCtx =
               taskManager_.getQueryContextManager()->findOrCreateQueryCtx(
-                  taskId, updateRequest.session);
+                  taskId, updateRequest);
 
           VeloxInteractiveQueryPlanConverter converter(queryCtx.get(), pool_);
           planFragment = converter.toVeloxQueryPlan(