feat(ABFS): Support SAS and OAuth config

CMakeLists.txt

@@ -254,6 +254,7 @@ if(VELOX_ENABLE_ABFS)
   endif()
   # files-datalake is built on blobs
   find_package(azure-storage-files-datalake-cpp CONFIG REQUIRED)
+  find_package(azure-identity-cpp CONFIG REQUIRED)
   add_definitions(-DVELOX_ENABLE_ABFS)
 endif()
 

velox/connectors/hive/storage_adapters/abfs/AbfsConfig.cpp

@@ -19,18 +19,20 @@
 #include "velox/common/config/Config.h"
 #include "velox/connectors/hive/storage_adapters/abfs/AbfsUtil.h"
 
+#include <azure/identity/client_secret_credential.hpp>
+
 namespace facebook::velox::filesystems {
 
 AbfsConfig::AbfsConfig(
     std::string_view path,
     const config::ConfigBase& config) {
   std::string_view file;
-  bool isHttps = true;
+  isHttps_ = true;
   if (path.find(kAbfssScheme) == 0) {
     file = path.substr(kAbfssScheme.size());
   } else if (path.find(kAbfsScheme) == 0) {
     file = path.substr(kAbfsScheme.size());
-    isHttps = false;
+    isHttps_ = false;
   } else {
     VELOX_FAIL("Invalid ABFS Path {}", path);
   }
@@ -39,30 +41,118 @@ AbfsConfig::AbfsConfig(
   fileSystem_ = file.substr(0, firstAt);
   auto firstSep = file.find_first_of("/");
   filePath_ = file.substr(firstSep + 1);
+  accountNameWithSuffix_ = file.substr(firstAt + 1, firstSep - firstAt - 1);
 
-  auto accountNameWithSuffix = file.substr(firstAt + 1, firstSep - firstAt - 1);
-  auto firstDot = accountNameWithSuffix.find_first_of(".");
-  auto accountName = accountNameWithSuffix.substr(0, firstDot);
-  auto endpointSuffix = accountNameWithSuffix.substr(firstDot + 5);
-  auto credKey = fmt::format("fs.azure.account.key.{}", accountNameWithSuffix);
-  std::stringstream ss;
-  ss << "DefaultEndpointsProtocol=" << (isHttps ? "https" : "http");
-  ss << ";AccountName=" << accountName;
-
-  if (config.valueExists(credKey)) {
+  auto authTypeKey =
+      fmt::format("{}.{}", kAzureAccountAuthType, accountNameWithSuffix_);
+  authType_ = kAzureSharedKeyAuthType;
+  if (config.valueExists(authTypeKey)) {
+    authType_ = config.get<std::string>(authTypeKey).value();
+  }
+  if (authType_ == kAzureSharedKeyAuthType) {
+    auto credKey =
+        fmt::format("{}.{}", kAzureAccountKey, accountNameWithSuffix_);
+    VELOX_USER_CHECK(
+        config.valueExists(credKey), "Config {} not found", credKey);
+    auto firstDot = accountNameWithSuffix_.find_first_of(".");
+    auto accountName = accountNameWithSuffix_.substr(0, firstDot);
+    auto endpointSuffix = accountNameWithSuffix_.substr(firstDot + 5);
+    std::stringstream ss;
+    ss << "DefaultEndpointsProtocol=" << (isHttps_ ? "https" : "http");
+    ss << ";AccountName=" << accountName;
     ss << ";AccountKey=" << config.get<std::string>(credKey).value();
+    ss << ";EndpointSuffix=" << endpointSuffix;
+
+    if (config.valueExists(kAzureBlobEndpoint)) {
+      ss << ";BlobEndpoint="
+         << config.get<std::string>(kAzureBlobEndpoint).value();
+    }
+    ss << ";";
+    connectionString_ = ss.str();
+  } else if (authType_ == kAzureOAuthAuthType) {
+    auto clientIdKey = fmt::format(
+        "{}.{}", kAzureAccountOAuth2ClientId, accountNameWithSuffix_);
+    auto clientSecretKey = fmt::format(
+        "{}.{}", kAzureAccountOAuth2ClientSecret, accountNameWithSuffix_);
+    auto clientEndpointKey = fmt::format(
+        "{}.{}", kAzureAccountOAuth2ClientEndpoint, accountNameWithSuffix_);
+    VELOX_USER_CHECK(
+        config.valueExists(clientIdKey), "Config {} not found", clientIdKey);
+    VELOX_USER_CHECK(
+        config.valueExists(clientSecretKey),
+        "Config {} not found",
+        clientSecretKey);
+    VELOX_USER_CHECK(
+        config.valueExists(clientEndpointKey),
+        "Config {} not found",
+        clientEndpointKey);
+    auto clientEndpoint = config.get<std::string>(clientEndpointKey).value();
+    auto firstSep = clientEndpoint.find_first_of("/", /* https:// */ 8);
+    authorityHost_ = clientEndpoint.substr(0, firstSep + 1);
+    auto sedondSep = clientEndpoint.find_first_of("/", firstSep + 1);
+    tenentId_ = clientEndpoint.substr(firstSep + 1, sedondSep - firstSep - 1);
+    Azure::Identity::ClientSecretCredentialOptions options;
+    options.AuthorityHost = authorityHost_;
+    tokenCredential_ =
+        std::make_shared<Azure::Identity::ClientSecretCredential>(
+            tenentId_,
+            config.get<std::string>(clientIdKey).value(),
+            config.get<std::string>(clientSecretKey).value(),
+            options);
+  } else if (authType_ == kAzureSASAuthType) {
+    auto sasKey = fmt::format("{}.{}", kAzureSASKey, accountNameWithSuffix_);
+    VELOX_USER_CHECK(config.valueExists(sasKey), "Config {} not found", sasKey);
+    sas_ = config.get<std::string>(sasKey).value();
+  } else {
+    VELOX_USER_FAIL(
+        "Unsupported auth type {}, supported auth types are SharedKey, OAuth and SAS.",
+        authType_);
+  }
+}
+
+std::unique_ptr<BlobClient> AbfsConfig::getReadFileClient() {
+  if (authType_ == kAzureSASAuthType) {
+    auto url = getUrl(true);
+    return std::make_unique<BlobClient>(fmt::format("{}?{}", url, sas_));
+  } else if (authType_ == kAzureOAuthAuthType) {
+    auto url = getUrl(true);
+    return std::make_unique<BlobClient>(url, tokenCredential_);
   } else {
-    VELOX_USER_FAIL("Config {} not found", credKey);
+    return std::make_unique<BlobClient>(BlobClient::CreateFromConnectionString(
+        connectionString_, fileSystem_, filePath_));
   }
+}
 
-  ss << ";EndpointSuffix=" << endpointSuffix;
+std::unique_ptr<DataLakeFileClient> AbfsConfig::getWriteFileClient() {
+  if (authType_ == kAzureSASAuthType) {
+    auto url = getUrl(false);
+    return std::make_unique<DataLakeFileClient>(
+        fmt::format("{}?{}", url, sas_));
+  } else if (authType_ == kAzureOAuthAuthType) {
+    auto url = getUrl(false);
+    return std::make_unique<DataLakeFileClient>(url, tokenCredential_);
+  } else {
+    return std::make_unique<DataLakeFileClient>(
+        DataLakeFileClient::CreateFromConnectionString(
+            connectionString_, fileSystem_, filePath_));
+  }
+}
 
-  if (config.valueExists(kAzureBlobEndpoint)) {
-    ss << ";BlobEndpoint="
-       << config.get<std::string>(kAzureBlobEndpoint).value();
+std::string AbfsConfig::getUrl(bool withblobSuffix) {
+  std::string accountNameWithSuffixForUrl(accountNameWithSuffix_);
+  if (withblobSuffix) {
+    // We should use correct suffix for blob client.
+    size_t start_pos = accountNameWithSuffixForUrl.find("dfs");
+    if (start_pos != std::string::npos) {
+      accountNameWithSuffixForUrl.replace(start_pos, 3, "blob");
+    }
   }
-  ss << ";";
-  connectionString_ = ss.str();
+  return fmt::format(
+      "{}{}/{}/{}",
+      isHttps_ ? "https://" : "http://",
+      accountNameWithSuffixForUrl,
+      fileSystem_,
+      filePath_);
 }
 
 } // namespace facebook::velox::filesystems

velox/connectors/hive/storage_adapters/abfs/AbfsConfig.h

@@ -16,44 +16,100 @@
 
 #pragma once
 
+#include <azure/core/credentials/credentials.hpp>
+#include <azure/storage/blobs/blob_client.hpp>
+#include <azure/storage/files/datalake.hpp>
 #include <folly/hash/Hash.h>
 #include <string>
 
+using namespace Azure::Storage::Blobs;
+using namespace Azure::Storage::Files::DataLake;
+
 namespace facebook::velox::config {
 class ConfigBase;
 }
 
 namespace facebook::velox::filesystems {
 
 // This is used to specify the Azurite endpoint in testing.
-static std::string kAzureBlobEndpoint{"fs.azure.blob-endpoint"};
+static constexpr const char* kAzureBlobEndpoint{"fs.azure.blob-endpoint"};
+
+// The authentication mechanism is set in `fs.azure.account.auth.type` (or the
+// account specific variant). The supported values are SharedKey, OAuth and SAS.
+static constexpr const char* kAzureAccountAuthType =
+    "fs.azure.account.auth.type";
+
+static constexpr const char* kAzureAccountKey = "fs.azure.account.key";
+
+static constexpr const char* kAzureSASKey = "fs.azure.sas.fixed.token";
+
+static constexpr const char* kAzureAccountOAuth2ClientId =
+    "fs.azure.account.oauth2.client.id";
+
+static constexpr const char* kAzureAccountOAuth2ClientSecret =
+    "fs.azure.account.oauth2.client.secret";
+
+// Token end point, this can be found through Azure portal. For example:
+// https://login.microsoftonline.com/{TENANTID}/oauth2/token
+static constexpr const char* kAzureAccountOAuth2ClientEndpoint =
+    "fs.azure.account.oauth2.client.endpoint";
+
+static constexpr const char* kAzureSharedKeyAuthType = "SharedKey";
+
+static constexpr const char* kAzureOAuthAuthType = "OAuth";
+
+static constexpr const char* kAzureSASAuthType = "SAS";
 
 class AbfsConfig {
  public:
   explicit AbfsConfig(std::string_view path, const config::ConfigBase& config);
 
-  std::string identity() const {
-    const auto hash = folly::Hash();
-    return std::to_string(hash(connectionString_));
+  std::unique_ptr<BlobClient> getReadFileClient();
+
+  std::unique_ptr<DataLakeFileClient> getWriteFileClient();
+
+  std::string filePath() const {
+    return filePath_;
   }
 
+  /// Test only.
+  std::string fileSystem() const {
+    return fileSystem_;
+  }
+
+  /// Test only.
   std::string connectionString() const {
     return connectionString_;
   }
 
-  std::string fileSystem() const {
-    return fileSystem_;
+  /// Test only.
+  std::string tenentId() const {
+    return tenentId_;
   }
 
-  std::string filePath() const {
-    return filePath_;
+  /// Test only.
+  std::string authorityHost() const {
+    return authorityHost_;
   }
 
  private:
+  std::string getUrl(bool withblobSuffix);
+
+  std::string authType_;
+
   // Container name is called FileSystem in some Azure API.
   std::string fileSystem_;
   std::string filePath_;
   std::string connectionString_;
+
+  bool isHttps_;
+  std::string accountNameWithSuffix_;
+
+  std::string sas_;
+
+  std::string tenentId_;
+  std::string authorityHost_;
+  std::shared_ptr<Azure::Core::Credentials::TokenCredential> tokenCredential_;
 };
 
 } // namespace facebook::velox::filesystems

velox/connectors/hive/storage_adapters/abfs/AbfsFileSystem.cpp

@@ -16,7 +16,6 @@
 
 #include "velox/connectors/hive/storage_adapters/abfs/AbfsFileSystem.h"
 
-#include <azure/storage/blobs/blob_client.hpp>
 #include <fmt/format.h>
 #include <folly/executors/IOThreadPoolExecutor.h>
 #include <glog/logging.h>
@@ -27,19 +26,16 @@
 #include "velox/connectors/hive/storage_adapters/abfs/AbfsWriteFile.h"
 
 namespace facebook::velox::filesystems {
-using namespace Azure::Storage::Blobs;
 
 class AbfsReadFile::Impl {
   constexpr static uint64_t kNaturalReadSize = 4 << 20; // 4M
   constexpr static uint64_t kReadConcurrency = 8;
 
  public:
   explicit Impl(std::string_view path, const config::ConfigBase& config) {
-    auto account = AbfsConfig(path, config);
-    filePath_ = account.filePath();
-    fileClient_ =
-        std::make_unique<BlobClient>(BlobClient::CreateFromConnectionString(
-            account.connectionString(), account.fileSystem(), filePath_));
+    auto abfsConfig = AbfsConfig(path, config);
+    filePath_ = abfsConfig.filePath();
+    fileClient_ = abfsConfig.getReadFileClient();
   }
 
   void initialize(const FileOptions& options) {
@@ -59,7 +55,6 @@ class AbfsReadFile::Impl {
     } catch (Azure::Storage::StorageException& e) {
       throwStorageExceptionWithOperationDetails("GetProperties", filePath_, e);
     }
-
     VELOX_CHECK_GE(length_, 0);
   }
 
@@ -141,7 +136,6 @@ class AbfsReadFile::Impl {
 
     Azure::Storage::Blobs::DownloadBlobOptions blob;
     blob.Range = range;
-
     auto response = fileClient_->Download(blob);
     response.Value.BodyStream->ReadToCount(
         reinterpret_cast<uint8_t*>(position), length);

velox/connectors/hive/storage_adapters/abfs/AbfsWriteFile.cpp

@@ -15,12 +15,9 @@
  */
 
 #include "velox/connectors/hive/storage_adapters/abfs/AbfsWriteFile.h"
-#include <azure/storage/files/datalake.hpp>
 #include "velox/connectors/hive/storage_adapters/abfs/AbfsConfig.h"
 #include "velox/connectors/hive/storage_adapters/abfs/AbfsUtil.h"
 
-using namespace Azure::Storage::Files::DataLake;
-
 namespace facebook::velox::filesystems {
 class DataLakeFileClientWrapper final : public AzureDataLakeFileClient {
  public:
@@ -31,7 +28,8 @@ class DataLakeFileClientWrapper final : public AzureDataLakeFileClient {
     client_->Create();
   }
 
-  Models::PathProperties getProperties() override {
+  Azure::Storage::Files::DataLake::Models::PathProperties getProperties()
+      override {
     return client_->GetProperties().Value;
   }
 
@@ -120,16 +118,11 @@ class AbfsWriteFile::Impl {
 AbfsWriteFile::AbfsWriteFile(
     std::string_view path,
     const config::ConfigBase& config) {
-  auto abfsAccount = AbfsConfig(path, config);
-  std::unique_ptr<AzureDataLakeFileClient> client =
+  auto abfsConfig = AbfsConfig(path, config);
+  std::unique_ptr<AzureDataLakeFileClient> clientWrapper =
       std::make_unique<DataLakeFileClientWrapper>(
-          std::make_unique<DataLakeFileClient>(
-              DataLakeFileClient::CreateFromConnectionString(
-                  abfsAccount.connectionString(),
-                  abfsAccount.fileSystem(),
-                  abfsAccount.filePath())));
-
-  impl_ = std::make_unique<Impl>(path, client);
+          abfsConfig.getWriteFileClient());
+  impl_ = std::make_unique<Impl>(path, clientWrapper);
 }
 
 AbfsWriteFile::AbfsWriteFile(

velox/connectors/hive/storage_adapters/abfs/CMakeLists.txt

@@ -30,6 +30,7 @@ if(VELOX_ENABLE_ABFS)
            velox_core
            velox_hive_config
            velox_dwio_common_exception
+           Azure::azure-identity
            Azure::azure-storage-blobs
            Azure::azure-storage-files-datalake
            Folly::folly