Add Azure CLI auth

Tom-Newton · May 18, 2024 · 99c1f81 · 99c1f81
1 parent 82ebdb2
commit 99c1f81
Show file tree

Hide file tree

Showing 3 changed files with 33 additions and 0 deletions.
diff --git a/cpp/src/arrow/filesystem/azurefs.cc b/cpp/src/arrow/filesystem/azurefs.cc
@@ -117,6 +117,8 @@ Status AzureOptions::ExtractFromUriQuery(const Uri& uri) {
         credential_kind = CredentialKind::kDefault;
       } else if (kv.second == "anonymous") {
         credential_kind = CredentialKind::kAnonymous;
+      } else if (kv.second == "cli") {
+        credential_kind = CredentialKind::kCLI;
       } else if (kv.second == "workload_identity") {
         credential_kind = CredentialKind::kWorkloadIdentity;
       } else if (kv.second == "environment") {
@@ -170,6 +172,9 @@ Status AzureOptions::ExtractFromUriQuery(const Uri& uri) {
       case CredentialKind::kAnonymous:
         RETURN_NOT_OK(ConfigureAnonymousCredential());
         break;
+      case CredentialKind::kCLI:
+        RETURN_NOT_OK(ConfigureCLICredential());
+        break;
       case CredentialKind::kWorkloadIdentity:
         RETURN_NOT_OK(ConfigureWorkloadIdentityCredential());
         break;
@@ -255,6 +260,7 @@ bool AzureOptions::Equals(const AzureOptions& other) const {
       return storage_shared_key_credential_->AccountName ==
              other.storage_shared_key_credential_->AccountName;
     case CredentialKind::kClientSecret:
+    case CredentialKind::kCLI:
     case CredentialKind::kManagedIdentity:
     case CredentialKind::kWorkloadIdentity:
     case CredentialKind::kEnvironment:
@@ -337,6 +343,12 @@ Status AzureOptions::ConfigureManagedIdentityCredential(const std::string& clien
   return Status::OK();
 }
 
+Status AzureOptions::ConfigureCLICredential() {
+  credential_kind_ = CredentialKind::kCLI;
+  token_credential_ = std::make_shared<Azure::Identity::AzureCliCredential>();
+  return Status::OK();
+}
+
 Status AzureOptions::ConfigureWorkloadIdentityCredential() {
   credential_kind_ = CredentialKind::kWorkloadIdentity;
   token_credential_ = std::make_shared<Azure::Identity::WorkloadIdentityCredential>();
@@ -364,6 +376,7 @@ Result<std::unique_ptr<Blobs::BlobServiceClient>> AzureOptions::MakeBlobServiceC
       [[fallthrough]];
     case CredentialKind::kClientSecret:
     case CredentialKind::kManagedIdentity:
+    case CredentialKind::kCLI:
     case CredentialKind::kWorkloadIdentity:
     case CredentialKind::kEnvironment:
       return std::make_unique<Blobs::BlobServiceClient>(AccountBlobUrl(account_name),
@@ -391,6 +404,7 @@ AzureOptions::MakeDataLakeServiceClient() const {
       [[fallthrough]];
     case CredentialKind::kClientSecret:
     case CredentialKind::kManagedIdentity:
+    case CredentialKind::kCLI:
     case CredentialKind::kWorkloadIdentity:
     case CredentialKind::kEnvironment:
       return std::make_unique<DataLake::DataLakeServiceClient>(

diff --git a/cpp/src/arrow/filesystem/azurefs.h b/cpp/src/arrow/filesystem/azurefs.h
@@ -119,6 +119,7 @@ struct ARROW_EXPORT AzureOptions {
     kStorageSharedKey,
     kClientSecret,
     kManagedIdentity,
+    kCLI,
     kWorkloadIdentity,
     kEnvironment,
   } credential_kind_ = CredentialKind::kDefault;
@@ -188,6 +189,7 @@ struct ARROW_EXPORT AzureOptions {
                                          const std::string& client_id,
                                          const std::string& client_secret);
   Status ConfigureManagedIdentityCredential(const std::string& client_id = std::string());
+  Status ConfigureCLICredential();
   Status ConfigureWorkloadIdentityCredential();
   Status ConfigureEnvironmentCredential();
 

diff --git a/cpp/src/arrow/filesystem/azurefs_test.cc b/cpp/src/arrow/filesystem/azurefs_test.cc
@@ -521,6 +521,13 @@ TEST(AzureFileSystem, InitializeWithManagedIdentityCredential) {
   EXPECT_OK_AND_ASSIGN(fs, AzureFileSystem::Make(options));
 }
 
+TEST(AzureFileSystem, InitializeWithCLICredential) {
+  AzureOptions options;
+  options.account_name = "dummy-account-name";
+  ARROW_EXPECT_OK(options.ConfigureCLICredential());
+  EXPECT_OK_AND_ASSIGN(auto fs, AzureFileSystem::Make(options));
+}
+
 TEST(AzureFileSystem, InitializeWithWorkloadIdentityCredential) {
   AzureOptions options;
   options.account_name = "dummy-account-name";
@@ -667,6 +674,15 @@ class TestAzureOptions : public ::testing::Test {
     ASSERT_EQ(options.credential_kind_, AzureOptions::CredentialKind::kManagedIdentity);
   }
 
+  void TestFromUriCredentialCLI() {
+    ASSERT_OK_AND_ASSIGN(
+        auto options,
+        AzureOptions::FromUri("abfs://account.blob.core.windows.net/container/dir/blob?"
+                              "credential_kind=cli",
+                              nullptr));
+    ASSERT_EQ(options.credential_kind_, AzureOptions::CredentialKind::kCLI);
+  }
+
   void TestFromUriCredentialWorkloadIdentity() {
     ASSERT_OK_AND_ASSIGN(
         auto options,
@@ -733,6 +749,7 @@ TEST_F(TestAzureOptions, FromUriCredentialClientSecret) {
 TEST_F(TestAzureOptions, FromUriCredentialManagedIdentity) {
   TestFromUriCredentialManagedIdentity();
 }
+TEST_F(TestAzureOptions, FromUriCredentialCLI) { TestFromUriCredentialCLI(); }
 TEST_F(TestAzureOptions, FromUriCredentialWorkloadIdentity) {
   TestFromUriCredentialWorkloadIdentity();
 }