GH-49078: [FS][Azure] Fix lossy pickling of SubTreeFileSystem(base_path, AzureFileSystem(...))

cpp/src/arrow/filesystem/azurefs.cc

@@ -248,13 +248,17 @@ Result<AzureOptions> AzureOptions::FromUri(const std::string& uri_string,
 }
 
 bool AzureOptions::Equals(const AzureOptions& other) const {
-  const bool equals = blob_storage_authority == other.blob_storage_authority &&
-                      dfs_storage_authority == other.dfs_storage_authority &&
-                      blob_storage_scheme == other.blob_storage_scheme &&
-                      dfs_storage_scheme == other.dfs_storage_scheme &&
-                      default_metadata == other.default_metadata &&
-                      account_name == other.account_name &&
-                      credential_kind_ == other.credential_kind_;
+  const bool equals =
+      account_name == other.account_name &&
+      blob_storage_authority == other.blob_storage_authority &&
+      dfs_storage_authority == other.dfs_storage_authority &&
+      blob_storage_scheme == other.blob_storage_scheme &&
+      dfs_storage_scheme == other.dfs_storage_scheme &&
+      default_metadata == other.default_metadata &&
+      background_writes == other.background_writes &&
+      credential_kind_ == other.credential_kind_ && account_key_ == other.account_key_ &&
+      sas_token_ == other.sas_token_ && tenant_id_ == other.tenant_id_ &&
+      client_id_ == other.client_id_ && client_secret_ == other.client_secret_;
   if (!equals) {
     return false;
   }
@@ -318,65 +322,90 @@ std::string AzureOptions::AccountDfsUrl(const std::string& account_name) const {
   return BuildBaseUrl(dfs_storage_scheme, dfs_storage_authority, account_name);
 }
 
+void AzureOptions::ClearCredentials() {
+  credential_kind_ = CredentialKind::kDefault;
+  storage_shared_key_credential_ = nullptr;
+  account_key_.clear();
+  sas_token_.clear();
+  tenant_id_.clear();
+  client_id_.clear();
+  client_secret_.clear();
+  token_credential_ = nullptr;
+}
+
 Status AzureOptions::ConfigureDefaultCredential() {
+  ClearCredentials();
   credential_kind_ = CredentialKind::kDefault;
   token_credential_ = std::make_shared<Azure::Identity::DefaultAzureCredential>();
   return Status::OK();
 }
 
 Status AzureOptions::ConfigureAnonymousCredential() {
+  ClearCredentials();
   credential_kind_ = CredentialKind::kAnonymous;
   return Status::OK();
 }
 
 Status AzureOptions::ConfigureAccountKeyCredential(const std::string& account_key) {
+  ClearCredentials();
   credential_kind_ = CredentialKind::kStorageSharedKey;
   if (account_name.empty()) {
     return Status::Invalid("AzureOptions doesn't contain a valid account name");
   }
+  account_key_ = account_key;
   storage_shared_key_credential_ =
       std::make_shared<Storage::StorageSharedKeyCredential>(account_name, account_key);
   return Status::OK();
 }
 
 Status AzureOptions::ConfigureSASCredential(const std::string& sas_token) {
-  credential_kind_ = CredentialKind::kSASToken;
+  ClearCredentials();
   if (account_name.empty()) {
     return Status::Invalid("AzureOptions doesn't contain a valid account name");
   }
   sas_token_ = sas_token;
+  credential_kind_ = CredentialKind::kSASToken;
   return Status::OK();
 }
 
 Status AzureOptions::ConfigureClientSecretCredential(const std::string& tenant_id,
                                                      const std::string& client_id,
                                                      const std::string& client_secret) {
+  ClearCredentials();
+  tenant_id_ = tenant_id;
+  client_id_ = client_id;
+  client_secret_ = client_secret;
   credential_kind_ = CredentialKind::kClientSecret;
   token_credential_ = std::make_shared<Azure::Identity::ClientSecretCredential>(
       tenant_id, client_id, client_secret);
   return Status::OK();
 }
 
 Status AzureOptions::ConfigureManagedIdentityCredential(const std::string& client_id) {
+  ClearCredentials();
+  client_id_ = client_id;
   credential_kind_ = CredentialKind::kManagedIdentity;
   token_credential_ =
       std::make_shared<Azure::Identity::ManagedIdentityCredential>(client_id);
   return Status::OK();
 }
 
 Status AzureOptions::ConfigureCLICredential() {
+  ClearCredentials();
   credential_kind_ = CredentialKind::kCLI;
   token_credential_ = std::make_shared<Azure::Identity::AzureCliCredential>();
   return Status::OK();
 }
 
 Status AzureOptions::ConfigureWorkloadIdentityCredential() {
+  ClearCredentials();
   credential_kind_ = CredentialKind::kWorkloadIdentity;
   token_credential_ = std::make_shared<Azure::Identity::WorkloadIdentityCredential>();
   return Status::OK();
 }
 
 Status AzureOptions::ConfigureEnvironmentCredential() {
+  ClearCredentials();
   credential_kind_ = CredentialKind::kEnvironment;
   token_credential_ = std::make_shared<Azure::Identity::EnvironmentCredential>();
   return Status::OK();

cpp/src/arrow/filesystem/azurefs.h

@@ -99,10 +99,6 @@ struct ARROW_EXPORT AzureOptions {
   /// Default: "https"
   std::string dfs_storage_scheme = "https";
 
-  // TODO(GH-38598): Add support for more auth methods.
-  // std::string connection_string;
-  // std::string sas_token;
-
   /// \brief Default metadata for OpenOutputStream.
   ///
   /// This will be ignored if non-empty metadata is passed to OpenOutputStream.
@@ -126,7 +122,11 @@ struct ARROW_EXPORT AzureOptions {
 
   std::shared_ptr<Azure::Storage::StorageSharedKeyCredential>
       storage_shared_key_credential_;
+  std::string account_key_;
   std::string sas_token_;
+  std::string tenant_id_;
+  std::string client_id_;
+  std::string client_secret_;
   mutable std::shared_ptr<Azure::Core::Credentials::TokenCredential> token_credential_;
 
  public:
@@ -136,6 +136,7 @@ struct ARROW_EXPORT AzureOptions {
  private:
   void ExtractFromUriSchemeAndHierPart(const Uri& uri, std::string* out_path);
   Status ExtractFromUriQuery(const Uri& uri);
+  void ClearCredentials();
 
  public:
   /// \brief Construct a new AzureOptions from an URI.
@@ -204,6 +205,12 @@ struct ARROW_EXPORT AzureOptions {
   std::string AccountBlobUrl(const std::string& account_name) const;
   std::string AccountDfsUrl(const std::string& account_name) const;
 
+  std::string AccountKey() const { return account_key_; }
+  std::string SasToken() const { return sas_token_; }
+  std::string TenantId() const { return tenant_id_; }
+  std::string ClientId() const { return client_id_; }
+  std::string ClientSecret() const { return client_secret_; }
+
   Result<std::unique_ptr<Azure::Storage::Blobs::BlobServiceClient>>
   MakeBlobServiceClient() const;
 

cpp/src/arrow/filesystem/azurefs_test.cc

@@ -493,6 +493,11 @@ TEST(AzureFileSystem, InitializeWithDefaultCredential) {
   AzureOptions options;
   options.account_name = "dummy-account-name";
   ARROW_EXPECT_OK(options.ConfigureDefaultCredential());
+  ASSERT_EQ(options.AccountKey(), "");
+  ASSERT_EQ(options.SasToken(), "");
+  ASSERT_EQ(options.TenantId(), "");
+  ASSERT_EQ(options.ClientId(), "");
+  ASSERT_EQ(options.ClientSecret(), "");
   EXPECT_OK_AND_ASSIGN(auto fs, AzureFileSystem::Make(options));
 }
 
@@ -509,6 +514,23 @@ TEST(AzureFileSystem, InitializeWithAnonymousCredential) {
   AzureOptions options;
   options.account_name = "dummy-account-name";
   ARROW_EXPECT_OK(options.ConfigureAnonymousCredential());
+  ASSERT_EQ(options.AccountKey(), "");
+  ASSERT_EQ(options.SasToken(), "");
+  ASSERT_EQ(options.TenantId(), "");
+  ASSERT_EQ(options.ClientId(), "");
+  ASSERT_EQ(options.ClientSecret(), "");
+  EXPECT_OK_AND_ASSIGN(auto fs, AzureFileSystem::Make(options));
+}
+
+TEST(AzureFileSystem, InitializeWithAccountKeyCredential) {
+  AzureOptions options;
+  options.account_name = "dummy-account-name";
+  ARROW_EXPECT_OK(options.ConfigureAccountKeyCredential("account_key"));
+  ASSERT_EQ(options.AccountKey(), "account_key");
+  ASSERT_EQ(options.SasToken(), "");
+  ASSERT_EQ(options.TenantId(), "");
+  ASSERT_EQ(options.ClientId(), "");
+  ASSERT_EQ(options.ClientSecret(), "");
   EXPECT_OK_AND_ASSIGN(auto fs, AzureFileSystem::Make(options));
 }
 
@@ -517,43 +539,98 @@ TEST(AzureFileSystem, InitializeWithClientSecretCredential) {
   options.account_name = "dummy-account-name";
   ARROW_EXPECT_OK(
       options.ConfigureClientSecretCredential("tenant_id", "client_id", "client_secret"));
+  ASSERT_EQ(options.AccountKey(), "");
+  ASSERT_EQ(options.SasToken(), "");
+  ASSERT_EQ(options.TenantId(), "tenant_id");
+  ASSERT_EQ(options.ClientId(), "client_id");
+  ASSERT_EQ(options.ClientSecret(), "client_secret");
   EXPECT_OK_AND_ASSIGN(auto fs, AzureFileSystem::Make(options));
 }
 
 TEST(AzureFileSystem, InitializeWithManagedIdentityCredential) {
   AzureOptions options;
   options.account_name = "dummy-account-name";
   ARROW_EXPECT_OK(options.ConfigureManagedIdentityCredential());
+  ASSERT_EQ(options.AccountKey(), "");
+  ASSERT_EQ(options.SasToken(), "");
+  ASSERT_EQ(options.TenantId(), "");
+  ASSERT_EQ(options.ClientId(), "");
+  ASSERT_EQ(options.ClientSecret(), "");
   EXPECT_OK_AND_ASSIGN(auto fs, AzureFileSystem::Make(options));
 
   ARROW_EXPECT_OK(options.ConfigureManagedIdentityCredential("specific-client-id"));
+  ASSERT_EQ(options.AccountKey(), "");
+  ASSERT_EQ(options.SasToken(), "");
+  ASSERT_EQ(options.TenantId(), "");
+  ASSERT_EQ(options.ClientId(), "specific-client-id");
+  ASSERT_EQ(options.ClientSecret(), "");
   EXPECT_OK_AND_ASSIGN(fs, AzureFileSystem::Make(options));
 }
 
 TEST(AzureFileSystem, InitializeWithCLICredential) {
   AzureOptions options;
   options.account_name = "dummy-account-name";
   ARROW_EXPECT_OK(options.ConfigureCLICredential());
+  ASSERT_EQ(options.AccountKey(), "");
+  ASSERT_EQ(options.SasToken(), "");
+  ASSERT_EQ(options.TenantId(), "");
+  ASSERT_EQ(options.ClientId(), "");
+  ASSERT_EQ(options.ClientSecret(), "");
   EXPECT_OK_AND_ASSIGN(auto fs, AzureFileSystem::Make(options));
 }
 
 TEST(AzureFileSystem, InitializeWithWorkloadIdentityCredential) {
   AzureOptions options;
   options.account_name = "dummy-account-name";
   ARROW_EXPECT_OK(options.ConfigureWorkloadIdentityCredential());
+  ASSERT_EQ(options.AccountKey(), "");
+  ASSERT_EQ(options.SasToken(), "");
+  ASSERT_EQ(options.TenantId(), "");
+  ASSERT_EQ(options.ClientId(), "");
+  ASSERT_EQ(options.ClientSecret(), "");
   EXPECT_OK_AND_ASSIGN(auto fs, AzureFileSystem::Make(options));
 }
 
 TEST(AzureFileSystem, InitializeWithEnvironmentCredential) {
   AzureOptions options;
   options.account_name = "dummy-account-name";
   ARROW_EXPECT_OK(options.ConfigureEnvironmentCredential());
+  ASSERT_EQ(options.AccountKey(), "");
+  ASSERT_EQ(options.SasToken(), "");
+  ASSERT_EQ(options.TenantId(), "");
+  ASSERT_EQ(options.ClientId(), "");
+  ASSERT_EQ(options.ClientSecret(), "");
   EXPECT_OK_AND_ASSIGN(auto fs, AzureFileSystem::Make(options));
 }
 
 TEST(AzureFileSystem, OptionsCompare) {
-  AzureOptions options;
-  EXPECT_TRUE(options.Equals(options));
+  AzureOptions options0;
+  EXPECT_TRUE(options0.Equals(options0));
+
+  AzureOptions options1;
+  options1.account_name = "account_name";
+  EXPECT_FALSE(options1.Equals(options0));
+
+  AzureOptions options2;
+  options2.account_name = "account_name";
+  ASSERT_OK(options2.ConfigureAccountKeyCredential("fake_account_key"));
+  EXPECT_FALSE(options2.Equals(options1));
+
+  AzureOptions options3;
+  options3.account_name = "account_name";
+  ASSERT_OK(options3.ConfigureAccountKeyCredential("different_fake_account_key"));
+  EXPECT_FALSE(options3.Equals(options2));
+
+  AzureOptions options4;
+  options4.account_name = "account_name";
+  ASSERT_OK(options4.ConfigureSASCredential("fake_sas_token"));
+  EXPECT_FALSE(options4.Equals(options3));
+
+  AzureOptions options5;
+  options5.account_name = "account_name";
+  ASSERT_OK(options5.ConfigureClientSecretCredential("fake_tenant_id", "fake_client_id",
+                                                     "fake_client_secret"));
+  EXPECT_FALSE(options5.Equals(options4));
 }
 
 class TestAzureOptions : public ::testing::Test {
@@ -1679,9 +1756,14 @@ class TestAzureFileSystem : public ::testing::Test {
                                  env->account_name(), env->account_key())));
     // AzureOptions::FromUri will not cut off extra query parameters that it consumes, so
     // make sure these don't cause problems.
-    ARROW_EXPECT_OK(options.ConfigureSASCredential(
-        "?blob_storage_authority=dummy_value0&" + sas_token.substr(1) +
-        "&credential_kind=dummy-value1"));
+    auto polluted_sas_token = "?blob_storage_authority=dummy_value0&" +
+                              sas_token.substr(1) + "&credential_kind=dummy-value1";
+    ARROW_EXPECT_OK(options.ConfigureSASCredential(polluted_sas_token));
+    ASSERT_EQ(options.AccountKey(), "");
+    ASSERT_EQ(options.SasToken(), polluted_sas_token);
+    ASSERT_EQ(options.TenantId(), "");
+    ASSERT_EQ(options.ClientId(), "");
+    ASSERT_EQ(options.ClientSecret(), "");
     EXPECT_OK_AND_ASSIGN(auto fs, AzureFileSystem::Make(options));
 
     AssertFileInfo(fs.get(), data.ObjectPath(), FileType::File);