[native] Support using custom storage access username based on source and client tags (prestodb#25827)

## Description
To support using custom storage access username based on source and
client tags
- Pass source and client tag from session to QueryConfig and QueryCtx
- Abstract functions in PrestoServer, TaskManager, and QueryTextManager

## Release Notes
Please follow [release notes
guidelines](https://github.com/prestodb/presto/wiki/Release-Notes-Guidelines)
and fill in the release notes below.

```
== NO RELEASE NOTE ==
```

Author: vhsu14

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

@@ -271,8 +271,8 @@ void PrestoServer::run() {
     exit(EXIT_FAILURE);
   }
 
-  registerFileSinks();
   registerFileSystems();
+  registerFileSinks();
   registerFileReadersAndWriters();
   registerMemoryArbitrators();
   registerShuffleInterfaceFactories();
@@ -434,8 +434,7 @@ void PrestoServer::run() {
   nativeWorkerPool_ = velox::memory::MemoryManager::getInstance()->addLeafPool(
       "PrestoNativeWorker");
 
-  taskManager_ = std::make_unique<TaskManager>(
-      driverExecutor_.get(), httpSrvCpuExecutor_.get(), spillerExecutor_.get());
+  createTaskManager();
 
   if (systemConfig->prestoNativeSidecar()) {
     registerSidecarEndpoints();
@@ -1739,4 +1738,9 @@ void PrestoServer::registerDynamicFunctions() {
   }
 }
 
+void PrestoServer::createTaskManager() {
+  taskManager_ = std::make_unique<TaskManager>(
+      driverExecutor_.get(), httpSrvCpuExecutor_.get(), spillerExecutor_.get());
+}
+
 } // namespace facebook::presto

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

@@ -227,6 +227,8 @@ class PrestoServer {
 
   void checkOverload();
 
+  virtual void createTaskManager();
+
   const std::string configDirectoryPath_;
 
   std::shared_ptr<CoordinatorDiscoverer> coordinatorDiscoverer_;

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

@@ -70,6 +70,13 @@ void updateFromSessionConfigs(
     }
   }
 
+  if (session.source) {
+    queryConfigs[velox::core::QueryConfig::kSource] = *session.source;
+  }
+  if (!session.clientTags.empty()) {
+    queryConfigs[velox::core::QueryConfig::kClientTags] = folly::join(',', session.clientTags);
+  }
+
   // If there's a timeZoneKey, convert to timezone name and add to the
   // configs. Throws if timeZoneKey can't be resolved.
   if (session.timeZoneKey != 0) {

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

@@ -59,6 +59,23 @@ void QueryContextManager::setQueryHasStartedTasks(
   queryContextCache_.wlock()->setHasStartedTasks(queryIdFromTaskId(taskId));
 }
 
+std::shared_ptr<core::QueryCtx> QueryContextManager::createAndCacheQueryCtx(
+    QueryContextCache& cache,
+    const QueryId& queryId,
+    velox::core::QueryConfig&& queryConfig,
+    std::unordered_map<std::string, std::shared_ptr<config::ConfigBase>>&& connectorConfigs,
+    std::shared_ptr<memory::MemoryPool>&& pool) {
+  auto queryCtx = core::QueryCtx::create(
+      driverExecutor_,
+      std::move(queryConfig),
+      std::move(connectorConfigs),
+      cache::AsyncDataCache::getInstance(),
+      std::move(pool),
+      spillerExecutor_,
+      queryId);
+  return cache.insert(queryId, std::move(queryCtx));
+}
+
 std::shared_ptr<core::QueryCtx> QueryContextManager::findOrCreateQueryCtx(
     const TaskId& taskId,
     velox::core::QueryConfig&& queryConfig,
@@ -90,16 +107,12 @@ std::shared_ptr<core::QueryCtx> QueryContextManager::findOrCreateQueryCtx(
       nullptr,
       poolDbgOpts);
 
-  auto queryCtx = core::QueryCtx::create(
-      driverExecutor_,
+  return createAndCacheQueryCtx(
+      *lockedCache,
+      queryId,
       std::move(queryConfig),
       std::move(connectorConfigs),
-      cache::AsyncDataCache::getInstance(),
-      std::move(pool),
-      spillerExecutor_,
-      queryId);
-
-  return lockedCache->insert(queryId, std::move(queryCtx));
+      std::move(pool));
 }
 
 void QueryContextManager::visitAllContexts(

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

@@ -123,6 +123,8 @@ class QueryContextManager {
       folly::Executor* driverExecutor,
       folly::Executor* spillerExecutor);
 
+  virtual ~QueryContextManager() = default;
+
   std::shared_ptr<velox::core::QueryCtx> findOrCreateQueryCtx(
       const protocol::TaskId& taskId,
       const protocol::TaskUpdateRequest& taskUpdateRequest);
@@ -141,17 +143,25 @@ class QueryContextManager {
   /// Test method to clear the query context cache.
   void testingClearCache();
 
+ protected:
+  folly::Executor* const driverExecutor_{nullptr};
+  folly::Executor* const spillerExecutor_{nullptr};
+
  private:
+  virtual std::shared_ptr<velox::core::QueryCtx> createAndCacheQueryCtx(
+    QueryContextCache& cache,
+    const protocol::QueryId& queryId,
+    velox::core::QueryConfig&& queryConfig,
+    std::unordered_map<std::string, std::shared_ptr<velox::config::ConfigBase>>&& connectorConfigs,
+    std::shared_ptr<velox::memory::MemoryPool>&& pool);
+
   std::shared_ptr<velox::core::QueryCtx> findOrCreateQueryCtx(
       const protocol::TaskId& taskId,
       velox::core::QueryConfig&& queryConfig,
       std::unordered_map<
           std::string,
           std::shared_ptr<velox::config::ConfigBase>>&& connectorConfigStrings);
 
-  folly::Executor* const driverExecutor_{nullptr};
-  folly::Executor* const spillerExecutor_{nullptr};
-
   folly::Synchronized<QueryContextCache> queryContextCache_;
 };
 

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

@@ -355,10 +355,10 @@ TaskManager::TaskManager(
     folly::Executor* driverExecutor,
     folly::Executor* httpSrvCpuExecutor,
     folly::Executor* spillerExecutor)
-    : bufferManager_(velox::exec::OutputBufferManager::getInstanceRef()),
-      queryContextManager_(std::make_unique<QueryContextManager>(
+    : queryContextManager_(std::make_unique<QueryContextManager>(
           driverExecutor,
           spillerExecutor)),
+      bufferManager_(velox::exec::OutputBufferManager::getInstanceRef()),
       httpSrvCpuExecutor_(httpSrvCpuExecutor) {
   VELOX_CHECK_NOT_NULL(bufferManager_, "invalid OutputBufferManager");
 }

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

@@ -34,6 +34,8 @@ class TaskManager {
       folly::Executor* httpSrvExecutor,
       folly::Executor* spillerExecutor);
 
+  virtual ~TaskManager() = default;
+
   /// Invoked by Presto server shutdown to wait for all the tasks to complete
   /// and cleanup the completed tasks.
   void shutdown();
@@ -195,6 +197,9 @@ class TaskManager {
   /// See if we have any queued tasks that can be started.
   void maybeStartNextQueuedTask();
 
+ protected:
+  std::unique_ptr<QueryContextManager> queryContextManager_;
+
  private:
   static constexpr folly::StringPiece kMaxDriversPerTask{
       "max_drivers_per_task"};
@@ -229,7 +234,6 @@ class TaskManager {
   std::shared_ptr<velox::exec::OutputBufferManager> bufferManager_;
   folly::Synchronized<TaskMap> taskMap_;
   folly::Synchronized<TaskQueue> taskQueue_;
-  std::unique_ptr<QueryContextManager> queryContextManager_;
   folly::Executor* httpSrvCpuExecutor_;
   std::atomic_bool serverOverloaded_{false};
   std::atomic_uint32_t numQueuedDrivers_{0};