feat: Add TokenProvider to ConnectorQueryCtx and include it in the cache key for FileHandle (facebookincubator#13919)

Summary:
Summary
As discussed in facebookincubator#13875 and facebookincubator#13914, we need to provide user information to the filesystem to support authentication and impersonation. This requirement can be generalized with a TokenProvider abstraction.

This PR introduces an abstract `TokenProvide`r class and integrates it into both `ConnectorQueryCtx` and `FileOptions`. It also becomes part of the `FileHandle` cache key.

To use it, users simply need to implement the TokenProvider interface with hash, equals, and getToken methods. The filesystem can then access relevant user or token information when reading from or writing to files.

Pull Request resolved: facebookincubator#13919

Reviewed By: Yuhta

Differential Revision: D78280984

Pulled By: xiaoxmeng

fbshipit-source-id: fa26c566aef45a66863e0328ad633b27a974f218

Author: wraymo

velox/common/file/FileSystems.h

@@ -17,6 +17,7 @@
 
 #include "velox/common/base/Exceptions.h"
 #include "velox/common/base/RuntimeMetrics.h"
+#include "velox/common/file/TokenProvider.h"
 #include "velox/common/memory/MemoryPool.h"
 
 #include <functional>
@@ -84,6 +85,9 @@ struct FileOptions {
   /// A hint to the file system for which region size of the file should be
   /// read. Specifically, the read length.
   std::optional<int64_t> readRangeHint{std::nullopt};
+
+  /// A token provider that can be used to get tokens for accessing the file.
+  std::shared_ptr<TokenProvider> tokenProvider{nullptr};
 };
 
 /// Defines directory options

velox/common/file/TokenProvider.h

@@ -0,0 +1,48 @@
+/*
+ * Copyright (c) Facebook, Inc. and its affiliates.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#pragma once
+
+#include <string>
+
+namespace facebook::velox::filesystems {
+
+/// Identifier for the file systems to implement to differentiate different
+/// tokens needed in the same query (user).  Such information usually needs to
+/// be passed down and stored in the ReadFile/WriteFile object of the specific
+/// file system.
+class AccessTokenKey {
+  virtual ~AccessTokenKey() = default;
+};
+
+/// Abstract token each file system can implement and cast.
+class AccessToken {
+  virtual ~AccessToken() = default;
+};
+
+/// Interface for providing access tokens to file systems.
+class TokenProvider {
+ public:
+  virtual ~TokenProvider() = default;
+
+  virtual bool equals(const TokenProvider& other) const = 0;
+  virtual size_t hash() const = 0;
+
+  virtual std::shared_ptr<AccessToken> getToken(
+      const AccessTokenKey& key) const = 0;
+};
+
+} // namespace facebook::velox::filesystems

velox/connectors/Connector.h

@@ -23,6 +23,7 @@
 #include "velox/common/base/SpillStats.h"
 #include "velox/common/caching/AsyncDataCache.h"
 #include "velox/common/caching/ScanTracker.h"
+#include "velox/common/file/TokenProvider.h"
 #include "velox/common/future/VeloxPromise.h"
 #include "velox/core/ExpressionEvaluator.h"
 #include "velox/type/Filter.h"
@@ -413,7 +414,8 @@ class ConnectorQueryCtx {
       int driverId,
       const std::string& sessionTimezone,
       bool adjustTimestampToTimezone = false,
-      folly::CancellationToken cancellationToken = {})
+      folly::CancellationToken cancellationToken = {},
+      std::shared_ptr<filesystems::TokenProvider> tokenProvider = {})
       : operatorPool_(operatorPool),
         connectorPool_(connectorPool),
         sessionProperties_(sessionProperties),
@@ -428,7 +430,8 @@ class ConnectorQueryCtx {
         planNodeId_(planNodeId),
         sessionTimezone_(sessionTimezone),
         adjustTimestampToTimezone_(adjustTimestampToTimezone),
-        cancellationToken_(std::move(cancellationToken)) {
+        cancellationToken_(std::move(cancellationToken)),
+        fsTokenProvider_(std::move(tokenProvider)) {
     VELOX_CHECK_NOT_NULL(sessionProperties);
   }
 
@@ -516,6 +519,10 @@ class ConnectorQueryCtx {
     selectiveNimbleReaderEnabled_ = value;
   }
 
+  std::shared_ptr<filesystems::TokenProvider> fsTokenProvider() const {
+    return fsTokenProvider_;
+  }
+
  private:
   memory::MemoryPool* const operatorPool_;
   memory::MemoryPool* const connectorPool_;
@@ -532,6 +539,7 @@ class ConnectorQueryCtx {
   const std::string sessionTimezone_;
   const bool adjustTimestampToTimezone_;
   const folly::CancellationToken cancellationToken_;
+  const std::shared_ptr<filesystems::TokenProvider> fsTokenProvider_;
   bool selectiveNimbleReaderEnabled_{false};
 };
 

velox/connectors/hive/FileHandle.cpp

@@ -40,7 +40,7 @@ std::string groupName(const std::string& filename) {
 } // namespace
 
 std::unique_ptr<FileHandle> FileHandleGenerator::operator()(
-    const std::string& filename,
+    const FileHandleKey& key,
     const FileProperties* properties,
     filesystems::File::IoStats* stats) {
   // We have seen cases where drivers are stuck when creating file handles.
@@ -53,11 +53,13 @@ std::unique_ptr<FileHandle> FileHandleGenerator::operator()(
     fileHandle = std::make_unique<FileHandle>();
     filesystems::FileOptions options;
     options.stats = stats;
+    options.tokenProvider = key.tokenProvider;
     if (properties) {
       options.fileSize = properties->fileSize;
       options.readRangeHint = properties->readRangeHint;
       options.extraFileInfo = properties->extraFileInfo;
     }
+    const auto& filename = key.filename;
     fileHandle->file = filesystems::getFileSystem(filename, properties_)
                            ->openFileForRead(filename, options);
     fileHandle->uuid = StringIdLease(fileIds(), filename);

velox/connectors/hive/FileHandle.h

@@ -25,10 +25,12 @@
 
 #pragma once
 
+#include "velox/common/base/BitUtil.h"
 #include "velox/common/caching/CachedFactory.h"
 #include "velox/common/caching/FileIds.h"
 #include "velox/common/config/Config.h"
 #include "velox/common/file/File.h"
+#include "velox/common/file/TokenProvider.h"
 #include "velox/connectors/hive/FileProperties.h"
 
 namespace facebook::velox {
@@ -59,7 +61,44 @@ struct FileHandleSizer {
   uint64_t operator()(const FileHandle& a);
 };
 
-using FileHandleCache = SimpleLRUCache<std::string, FileHandle>;
+struct FileHandleKey {
+  std::string filename;
+  std::shared_ptr<filesystems::TokenProvider> tokenProvider{nullptr};
+
+  bool operator==(const FileHandleKey& other) const {
+    if (filename != other.filename) {
+      return false;
+    }
+
+    if (tokenProvider == other.tokenProvider) {
+      return true;
+    }
+
+    if (!tokenProvider || !other.tokenProvider) {
+      return false;
+    }
+
+    return tokenProvider->equals(*other.tokenProvider);
+  }
+};
+
+} // namespace facebook::velox
+
+namespace std {
+template <>
+struct hash<facebook::velox::FileHandleKey> {
+  size_t operator()(const facebook::velox::FileHandleKey& key) const noexcept {
+    size_t filenameHash = std::hash<std::string>()(key.filename);
+    return key.tokenProvider ? facebook::velox::bits::hashMix(
+                                   filenameHash, key.tokenProvider->hash())
+                             : filenameHash;
+  }
+};
+} // namespace std
+
+namespace facebook::velox {
+using FileHandleCache =
+    SimpleLRUCache<facebook::velox::FileHandleKey, FileHandle>;
 
 // Creates FileHandles via the Generator interface the CachedFactory requires.
 class FileHandleGenerator {
@@ -68,7 +107,7 @@ class FileHandleGenerator {
   FileHandleGenerator(std::shared_ptr<const config::ConfigBase> properties)
       : properties_(std::move(properties)) {}
   std::unique_ptr<FileHandle> operator()(
-      const std::string& filename,
+      const FileHandleKey& filename,
       const FileProperties* properties,
       filesystems::File::IoStats* stats);
 
@@ -77,14 +116,14 @@ class FileHandleGenerator {
 };
 
 using FileHandleFactory = CachedFactory<
-    std::string,
+    FileHandleKey,
     FileHandle,
     FileHandleGenerator,
     FileProperties,
     filesystems::File::IoStats,
     FileHandleSizer>;
 
-using FileHandleCachedPtr = CachedPtr<std::string, FileHandle>;
+using FileHandleCachedPtr = CachedPtr<FileHandleKey, FileHandle>;
 
 using FileHandleCacheStats = SimpleLRUCacheStats;
 

velox/connectors/hive/HiveConnector.cpp

@@ -44,7 +44,7 @@ HiveConnector::HiveConnector(
       hiveConfig_(std::make_shared<HiveConfig>(config)),
       fileHandleFactory_(
           hiveConfig_->isFileHandleCacheEnabled()
-              ? std::make_unique<SimpleLRUCache<std::string, FileHandle>>(
+              ? std::make_unique<SimpleLRUCache<FileHandleKey, FileHandle>>(
                     hiveConfig_->numCacheFileHandles())
               : nullptr,
           std::make_unique<FileHandleGenerator>(config)),

velox/connectors/hive/SplitReader.cpp

@@ -286,9 +286,12 @@ void SplitReader::createReader() {
       baseReaderOpts_.fileFormat(), dwio::common::FileFormat::UNKNOWN);
 
   FileHandleCachedPtr fileHandleCachePtr;
+  FileHandleKey fileHandleKey{
+      .filename = hiveSplit_->filePath,
+      .tokenProvider = connectorQueryCtx_->fsTokenProvider()};
   try {
     fileHandleCachePtr = fileHandleFactory_->generate(
-        hiveSplit_->filePath,
+        fileHandleKey,
         hiveSplit_->properties.has_value() ? &*hiveSplit_->properties : nullptr,
         fsStats_ ? fsStats_.get() : nullptr);
     VELOX_CHECK_NOT_NULL(fileHandleCachePtr.get());

velox/connectors/hive/iceberg/PositionalDeleteFileReader.cpp

@@ -92,8 +92,10 @@ PositionalDeleteFileReader::PositionalDeleteFileReader(
       /*tableParameters=*/{},
       deleteReaderOpts);
 
-  auto deleteFileHandleCachePtr =
-      fileHandleFactory_->generate(deleteFile_.filePath);
+  const FileHandleKey fileHandleKey{
+      .filename = deleteFile_.filePath,
+      .tokenProvider = connectorQueryCtx_->fsTokenProvider()};
+  auto deleteFileHandleCachePtr = fileHandleFactory_->generate(fileHandleKey);
   auto deleteFileInput = createBufferedInput(
       *deleteFileHandleCachePtr,
       deleteReaderOpts,

velox/connectors/hive/iceberg/tests/IcebergSplitReaderBenchmark.cpp

@@ -324,7 +324,7 @@ void IcebergSplitReaderBenchmark::readSingleColumn(
           "");
 
   FileHandleFactory fileHandleFactory(
-      std::make_unique<SimpleLRUCache<std::string, FileHandle>>(
+      std::make_unique<SimpleLRUCache<FileHandleKey, FileHandle>>(
           hiveConfig->numCacheFileHandles()),
       std::make_unique<FileHandleGenerator>(connectorSessionProperties_));
 

velox/connectors/hive/storage_adapters/abfs/tests/AbfsFileSystemTest.cpp

@@ -163,15 +163,14 @@ TEST_F(AbfsFileSystemTest, openFileForReadWithInvalidOptions) {
 
 TEST_F(AbfsFileSystemTest, fileHandleWithProperties) {
   FileHandleFactory factory(
-      std::make_unique<SimpleLRUCache<std::string, FileHandle>>(1),
+      std::make_unique<SimpleLRUCache<FileHandleKey, FileHandle>>(1),
       std::make_unique<FileHandleGenerator>(azuriteServer_->hiveConfig()));
   FileProperties properties = {15 + kOneMB, 1};
-  auto fileHandleProperties =
-      factory.generate(azuriteServer_->fileURI(), &properties);
+  FileHandleKey key{azuriteServer_->fileURI()};
+  auto fileHandleProperties = factory.generate(key, &properties);
   readData(fileHandleProperties->file.get());
 
-  auto fileHandleWithoutProperties =
-      factory.generate(azuriteServer_->fileURI());
+  auto fileHandleWithoutProperties = factory.generate(key);
   readData(fileHandleWithoutProperties->file.get());
 }