Skip to content

Commit 83e9709

Browse files
majetideepakBryanCutler
authored andcommitted
[native] Pass extraCredentials to connector session properties
1 parent 700fdc7 commit 83e9709

File tree

3 files changed

+16
-11
lines changed

3 files changed

+16
-11
lines changed

presto-native-execution/presto_cpp/main/QueryContextManager.cpp

Lines changed: 13 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -60,14 +60,17 @@ void updateFromSystemConfigs(
6060
}
6161

6262
std::unordered_map<std::string, std::unordered_map<std::string, std::string>>
63-
toConnectorConfigs(const protocol::SessionRepresentation& session) {
63+
toConnectorConfigs(const protocol::TaskUpdateRequest& taskUpdateRequest) {
6464
std::unordered_map<std::string, std::unordered_map<std::string, std::string>>
6565
connectorConfigs;
66-
for (const auto& entry : session.catalogProperties) {
67-
connectorConfigs.insert(
68-
{entry.first,
69-
std::unordered_map<std::string, std::string>(
70-
entry.second.begin(), entry.second.end())});
66+
for (const auto& entry : taskUpdateRequest.session.catalogProperties) {
67+
auto sessionProperties = std::unordered_map<std::string, std::string>(
68+
entry.second.begin(), entry.second.end());
69+
sessionProperties.insert(
70+
taskUpdateRequest.extraCredentials.begin(),
71+
taskUpdateRequest.extraCredentials.end());
72+
sessionProperties.insert({"user", taskUpdateRequest.session.user});
73+
connectorConfigs.insert({entry.first, sessionProperties});
7174
}
7275

7376
return connectorConfigs;
@@ -106,9 +109,11 @@ QueryContextManager::QueryContextManager(
106109
std::shared_ptr<velox::core::QueryCtx>
107110
QueryContextManager::findOrCreateQueryCtx(
108111
const protocol::TaskId& taskId,
109-
const protocol::SessionRepresentation& session) {
112+
const protocol::TaskUpdateRequest& taskUpdateRequest) {
110113
return findOrCreateQueryCtx(
111-
taskId, toVeloxConfigs(session), toConnectorConfigs(session));
114+
taskId,
115+
toVeloxConfigs(taskUpdateRequest.session),
116+
toConnectorConfigs(taskUpdateRequest));
112117
}
113118

114119
std::shared_ptr<core::QueryCtx> QueryContextManager::findOrCreateQueryCtx(

presto-native-execution/presto_cpp/main/QueryContextManager.h

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -107,7 +107,7 @@ class QueryContextManager {
107107

108108
std::shared_ptr<velox::core::QueryCtx> findOrCreateQueryCtx(
109109
const protocol::TaskId& taskId,
110-
const protocol::SessionRepresentation& session);
110+
const protocol::TaskUpdateRequest& taskUpdateRequest);
111111

112112
/// Calls the given functor for every present query context.
113113
void visitAllContexts(std::function<void(

presto-native-execution/presto_cpp/main/TaskResource.cpp

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -299,7 +299,7 @@ proxygen::RequestHandler* TaskResource::createOrUpdateBatchTask(
299299

300300
auto queryCtx =
301301
taskManager_.getQueryContextManager()->findOrCreateQueryCtx(
302-
taskId, updateRequest.session);
302+
taskId, updateRequest);
303303

304304
VeloxBatchQueryPlanConverter converter(
305305
shuffleName,
@@ -340,7 +340,7 @@ proxygen::RequestHandler* TaskResource::createOrUpdateTask(
340340

341341
queryCtx =
342342
taskManager_.getQueryContextManager()->findOrCreateQueryCtx(
343-
taskId, updateRequest.session);
343+
taskId, updateRequest);
344344

345345
VeloxInteractiveQueryPlanConverter converter(queryCtx.get(), pool_);
346346
planFragment = converter.toVeloxQueryPlan(

0 commit comments

Comments
 (0)