feat(rest): route each object path to its vended S3 credential (per-path FileIO)

Author: plusplusjiajia

src/iceberg/arrow/arrow_io.cc

@@ -22,6 +22,7 @@
 #include <limits>
 #include <mutex>
 #include <optional>
+#include <unordered_map>
 #include <vector>
 
 #include <arrow/buffer.h>
@@ -35,6 +36,8 @@
 #include "iceberg/arrow/arrow_io_internal.h"
 #include "iceberg/arrow/arrow_io_util.h"
 #include "iceberg/arrow/arrow_status_internal.h"
+#include "iceberg/arrow/s3/arrow_s3_internal.h"
+#include "iceberg/util/location_util.h"
 #include "iceberg/util/macros.h"
 
 namespace iceberg::arrow {
@@ -483,13 +486,53 @@ class ArrowOutputFile : public OutputFile {
 
 }  // namespace
 
-Result<std::string> ArrowFileSystemFileIO::ResolvePath(const std::string& file_location) {
+const std::shared_ptr<::arrow::fs::FileSystem>& ArrowFileSystemFileIO::FileSystemForPath(
+    std::string_view location) const {
+  if (fs_by_prefix_.empty()) {
+    return arrow_fs_;
+  }
+  // Longest matching prefix wins; fall back to the default file system.
+  const std::string canonical = LocationUtil::CanonicalizeS3Scheme(location);
+  const std::shared_ptr<::arrow::fs::FileSystem>* best = &arrow_fs_;
+  size_t best_len = 0;
+  for (const auto& [prefix, fs] : fs_by_prefix_) {
+    if (prefix.size() > best_len && LocationUtil::PathHasPrefix(canonical, prefix)) {
+      best = &fs;
+      best_len = prefix.size();
+    }
+  }
+  return *best;
+}
+
+Status ArrowFileSystemFileIO::SetStorageCredentials(
+    const std::vector<
+        std::pair<std::string, std::unordered_map<std::string, std::string>>>&
+        properties_by_prefix) {
+#if ICEBERG_S3_ENABLED
+  std::vector<std::pair<std::string, std::shared_ptr<::arrow::fs::FileSystem>>>
+      fs_by_prefix;
+  fs_by_prefix.reserve(properties_by_prefix.size());
+  for (const auto& [prefix, properties] : properties_by_prefix) {
+    ICEBERG_ASSIGN_OR_RAISE(auto fs, BuildArrowS3FileSystem(properties));
+    fs_by_prefix.emplace_back(prefix, std::move(fs));
+  }
+  fs_by_prefix_ = std::move(fs_by_prefix);
+  return {};
+#else
+  (void)properties_by_prefix;
+  return NotSupported("S3 storage credentials require Arrow S3 support");
+#endif
+}
+
+Result<std::string> ArrowFileSystemFileIO::ResolvePath(
+    const std::shared_ptr<::arrow::fs::FileSystem>& fs,
+    const std::string& file_location) {
   const auto pos = file_location.find("://");
   if (pos == std::string::npos) {
     return file_location;
   }
 
-  auto path = arrow_fs_->PathFromUri(file_location);
+  auto path = fs->PathFromUri(file_location);
   if (path.ok()) {
     return std::move(path).ValueOrDie();
   }
@@ -512,14 +555,14 @@ Result<std::shared_ptr<::arrow::io::RandomAccessFile>> OpenArrowInputStream(
   ICEBERG_PRECHECK(io != nullptr, "FileIO cannot be null");
 
   if (auto arrow_io = std::dynamic_pointer_cast<ArrowFileSystemFileIO>(io)) {
-    ICEBERG_ASSIGN_OR_RAISE(auto resolved_path, arrow_io->ResolvePath(path));
+    const auto& fs = arrow_io->FileSystemForPath(path);
+    ICEBERG_ASSIGN_OR_RAISE(auto resolved_path, arrow_io->ResolvePath(fs, path));
     ::arrow::fs::FileInfo file_info(resolved_path, ::arrow::fs::FileType::File);
     if (length.has_value()) {
       ICEBERG_ASSIGN_OR_RAISE(auto size, ToInt64Length(*length));
       file_info.set_size(size);
     }
-    ICEBERG_ARROW_ASSIGN_OR_RETURN(auto input,
-                                   arrow_io->arrow_fs_->OpenInputFile(file_info));
+    ICEBERG_ARROW_ASSIGN_OR_RETURN(auto input, fs->OpenInputFile(file_info));
     return input;
   }
 
@@ -543,16 +586,15 @@ Result<std::shared_ptr<::arrow::io::OutputStream>> OpenArrowOutputStream(
   ICEBERG_PRECHECK(io != nullptr, "FileIO cannot be null");
 
   if (auto arrow_io = std::dynamic_pointer_cast<ArrowFileSystemFileIO>(io)) {
-    ICEBERG_ASSIGN_OR_RAISE(auto resolved_path, arrow_io->ResolvePath(path));
+    const auto& fs = arrow_io->FileSystemForPath(path);
+    ICEBERG_ASSIGN_OR_RAISE(auto resolved_path, arrow_io->ResolvePath(fs, path));
     if (!overwrite) {
-      ICEBERG_ARROW_ASSIGN_OR_RETURN(auto info,
-                                     arrow_io->arrow_fs_->GetFileInfo(resolved_path));
+      ICEBERG_ARROW_ASSIGN_OR_RETURN(auto info, fs->GetFileInfo(resolved_path));
       if (info.type() != ::arrow::fs::FileType::NotFound) {
         return AlreadyExists("File already exists: {}", path);
       }
     }
-    ICEBERG_ARROW_ASSIGN_OR_RETURN(auto output,
-                                   arrow_io->arrow_fs_->OpenOutputStream(resolved_path));
+    ICEBERG_ARROW_ASSIGN_OR_RETURN(auto output, fs->OpenOutputStream(resolved_path));
     return output;
   }
 
@@ -568,42 +610,60 @@ Result<std::shared_ptr<::arrow::io::OutputStream>> OpenArrowOutputStream(
 
 Result<std::unique_ptr<InputFile>> ArrowFileSystemFileIO::NewInputFile(
     std::string file_location) {
-  ICEBERG_ASSIGN_OR_RAISE(auto path, ResolvePath(file_location));
-  return std::make_unique<ArrowInputFile>(arrow_fs_, std::move(file_location),
-                                          std::move(path), std::nullopt);
+  const auto& fs = FileSystemForPath(file_location);
+  ICEBERG_ASSIGN_OR_RAISE(auto path, ResolvePath(fs, file_location));
+  return std::make_unique<ArrowInputFile>(fs, std::move(file_location), std::move(path),
+                                          std::nullopt);
 }
 
 Result<std::unique_ptr<InputFile>> ArrowFileSystemFileIO::NewInputFile(
     std::string file_location, size_t length) {
   ICEBERG_ASSIGN_OR_RAISE(auto size, ToInt64Length(length));
-  ICEBERG_ASSIGN_OR_RAISE(auto path, ResolvePath(file_location));
-  return std::make_unique<ArrowInputFile>(arrow_fs_, std::move(file_location),
-                                          std::move(path), size);
+  const auto& fs = FileSystemForPath(file_location);
+  ICEBERG_ASSIGN_OR_RAISE(auto path, ResolvePath(fs, file_location));
+  return std::make_unique<ArrowInputFile>(fs, std::move(file_location), std::move(path),
+                                          size);
 }
 
 Result<std::unique_ptr<OutputFile>> ArrowFileSystemFileIO::NewOutputFile(
     std::string file_location) {
-  ICEBERG_ASSIGN_OR_RAISE(auto path, ResolvePath(file_location));
-  return std::make_unique<ArrowOutputFile>(arrow_fs_, std::move(file_location),
-                                           std::move(path));
+  const auto& fs = FileSystemForPath(file_location);
+  ICEBERG_ASSIGN_OR_RAISE(auto path, ResolvePath(fs, file_location));
+  return std::make_unique<ArrowOutputFile>(fs, std::move(file_location), std::move(path));
 }
 
 /// \brief Delete a file at the given location.
 Status ArrowFileSystemFileIO::DeleteFile(const std::string& file_location) {
-  ICEBERG_ASSIGN_OR_RAISE(auto path, ResolvePath(file_location));
-  ICEBERG_ARROW_RETURN_NOT_OK(arrow_fs_->DeleteFile(path));
+  const auto& fs = FileSystemForPath(file_location);
+  ICEBERG_ASSIGN_OR_RAISE(auto path, ResolvePath(fs, file_location));
+  ICEBERG_ARROW_RETURN_NOT_OK(fs->DeleteFile(path));
   return {};
 }
 
 Status ArrowFileSystemFileIO::DeleteFiles(
     const std::vector<std::string>& file_locations) {
-  std::vector<std::string> paths;
-  paths.reserve(file_locations.size());
+  if (fs_by_prefix_.empty()) {
+    // No per-prefix routing: one batched delete on the default file system.
+    std::vector<std::string> paths;
+    paths.reserve(file_locations.size());
+    for (const auto& file_location : file_locations) {
+      ICEBERG_ASSIGN_OR_RAISE(auto path, ResolvePath(arrow_fs_, file_location));
+      paths.push_back(std::move(path));
+    }
+    ICEBERG_ARROW_RETURN_NOT_OK(arrow_fs_->DeleteFiles(paths));
+    return {};
+  }
+
+  // Paths may route to different file systems: group by fs, then batch per fs.
+  std::unordered_map<::arrow::fs::FileSystem*, std::vector<std::string>> paths_by_fs;
   for (const auto& file_location : file_locations) {
-    ICEBERG_ASSIGN_OR_RAISE(auto path, ResolvePath(file_location));
-    paths.push_back(std::move(path));
+    const auto& fs = FileSystemForPath(file_location);
+    ICEBERG_ASSIGN_OR_RAISE(auto path, ResolvePath(fs, file_location));
+    paths_by_fs[fs.get()].push_back(std::move(path));
+  }
+  for (auto& [fs, paths] : paths_by_fs) {
+    ICEBERG_ARROW_RETURN_NOT_OK(fs->DeleteFiles(paths));
   }
-  ICEBERG_ARROW_RETURN_NOT_OK(arrow_fs_->DeleteFiles(paths));
   return {};
 }
 

src/iceberg/arrow/arrow_io_internal.h

@@ -23,6 +23,9 @@
 #include <memory>
 #include <optional>
 #include <string>
+#include <string_view>
+#include <unordered_map>
+#include <utility>
 #include <vector>
 
 #include <arrow/filesystem/type_fwd.h>
@@ -52,7 +55,8 @@ OpenArrowOutputStream(const std::shared_ptr<FileIO>& io, const std::string& path
                       bool overwrite = true);
 
 /// \brief A concrete implementation of FileIO for Arrow file system.
-class ICEBERG_BUNDLE_EXPORT ArrowFileSystemFileIO : public FileIO {
+class ICEBERG_BUNDLE_EXPORT ArrowFileSystemFileIO : public FileIO,
+                                                    public SupportsStorageCredentials {
  public:
   explicit ArrowFileSystemFileIO(std::shared_ptr<::arrow::fs::FileSystem> arrow_fs)
       : arrow_fs_(std::move(arrow_fs)) {}
@@ -81,6 +85,12 @@ class ICEBERG_BUNDLE_EXPORT ArrowFileSystemFileIO : public FileIO {
   /// \brief Delete files at the given locations.
   Status DeleteFiles(const std::vector<std::string>& file_locations) override;
 
+  /// \brief Build one S3 file system per credential prefix for per-path routing.
+  Status SetStorageCredentials(
+      const std::vector<
+          std::pair<std::string, std::unordered_map<std::string, std::string>>>&
+          properties_by_prefix) override;
+
   /// \brief Get the Arrow file system.
   const std::shared_ptr<::arrow::fs::FileSystem>& fs() const { return arrow_fs_; }
 
@@ -92,10 +102,18 @@ class ICEBERG_BUNDLE_EXPORT ArrowFileSystemFileIO : public FileIO {
   friend Result<std::shared_ptr<::arrow::io::OutputStream>> OpenArrowOutputStream(
       const std::shared_ptr<FileIO>& io, const std::string& path, bool overwrite);
 
-  /// \brief Resolve a file location to a filesystem path.
-  Result<std::string> ResolvePath(const std::string& file_location);
+  /// \brief Pick the file system for `location` (longest matching prefix, else
+  /// the default).
+  const std::shared_ptr<::arrow::fs::FileSystem>& FileSystemForPath(
+      std::string_view location) const;
+
+  /// \brief Resolve a file location to a filesystem path using `fs`.
+  Result<std::string> ResolvePath(const std::shared_ptr<::arrow::fs::FileSystem>& fs,
+                                  const std::string& file_location);
 
   std::shared_ptr<::arrow::fs::FileSystem> arrow_fs_;
+  std::vector<std::pair<std::string, std::shared_ptr<::arrow::fs::FileSystem>>>
+      fs_by_prefix_;
 };
 
 }  // namespace iceberg::arrow

src/iceberg/arrow/s3/arrow_s3_file_io.cc

@@ -21,6 +21,7 @@
 #include <optional>
 #include <string>
 #include <string_view>
+#include <utility>
 
 #include <arrow/filesystem/filesystem.h>
 #if ICEBERG_S3_ENABLED
@@ -172,15 +173,21 @@ Result<::arrow::fs::S3Options> ConfigureS3Options(
 }
 #endif
 
-Result<std::unique_ptr<FileIO>> MakeS3FileIO(
-    const std::unordered_map<std::string, std::string>& properties) {
 #if ICEBERG_S3_ENABLED
+Result<std::shared_ptr<::arrow::fs::FileSystem>> BuildArrowS3FileSystem(
+    const std::unordered_map<std::string, std::string>& properties) {
   ICEBERG_RETURN_UNEXPECTED(EnsureS3Initialized());
-
-  // Configure S3 options from properties (uses default credentials if empty)
   ICEBERG_ASSIGN_OR_RAISE(auto options, ConfigureS3Options(properties));
   ICEBERG_ARROW_ASSIGN_OR_RETURN(auto fs, ::arrow::fs::S3FileSystem::Make(options));
+  return std::shared_ptr<::arrow::fs::FileSystem>(std::move(fs));
+}
+#endif
 
+Result<std::unique_ptr<FileIO>> MakeS3FileIO(
+    const std::unordered_map<std::string, std::string>& properties) {
+#if ICEBERG_S3_ENABLED
+  // Uses default credentials if properties are empty.
+  ICEBERG_ASSIGN_OR_RAISE(auto fs, BuildArrowS3FileSystem(properties));
   return std::make_unique<ArrowFileSystemFileIO>(std::move(fs));
 #else
   return NotSupported("Arrow S3 support is not enabled");

src/iceberg/arrow/s3/arrow_s3_internal.h

@@ -19,6 +19,7 @@
 
 #pragma once
 
+#include <memory>
 #include <string>
 #include <unordered_map>
 
@@ -39,6 +40,11 @@ namespace iceberg::arrow {
 /// addressing style) can be unit tested without a live S3 endpoint.
 ICEBERG_BUNDLE_EXPORT Result<::arrow::fs::S3Options> ConfigureS3Options(
     const std::unordered_map<std::string, std::string>& properties);
+
+/// \brief Build an Arrow S3 file system from a properties map (initializes S3 if
+/// needed). Exposed so the credential-aware FileIO can build one fs per prefix.
+ICEBERG_BUNDLE_EXPORT Result<std::shared_ptr<::arrow::fs::FileSystem>>
+BuildArrowS3FileSystem(const std::unordered_map<std::string, std::string>& properties);
 #endif
 
 }  // namespace iceberg::arrow

src/iceberg/catalog/rest/rest_catalog.cc

@@ -525,31 +525,25 @@ Result<std::shared_ptr<auth::AuthSession>> RestCatalog::TableAuthSession(
 Result<std::shared_ptr<FileIO>> RestCatalog::TableFileIO(
     const SessionContext& /*context*/,
     const std::unordered_map<std::string, std::string>& table_config,
-    const std::vector<StorageCredential>& storage_credentials,
-    std::string_view storage_location) const {
-  // Bind one S3 FileIO from the vended credential whose prefix best matches the
-  // table location (the common one-credential-per-table case).
-  // TODO: support per-path credential routing, STS refresh via credentials.uri,
-  // and non-S3 (GCS/ADLS) credentials, like Java's S3FileIO.
-  if (const StorageCredential* s3_cred =
-          SelectS3StorageCredential(storage_credentials, storage_location);
-      s3_cred != nullptr) {
-    ICEBERG_ASSIGN_OR_RAISE(
-        auto io, MakeS3FileIOFromCredential(config_.configs(), table_config, *s3_cred));
+    const std::vector<StorageCredential>& storage_credentials) const {
+  if (!storage_credentials.empty()) {
+    // Only non-S3 (GCS/ADLS) credentials vended, which we can't honor; fail fast.
+    if (HasOnlyNonS3StorageCredentials(storage_credentials)) {
+      auto prefixes =
+          storage_credentials | std::views::transform(&StorageCredential::prefix);
+      return NotSupported(
+          "Vended storage credentials {} are unsupported (only S3-family "
+          "credentials are supported)",
+          FormatRange(prefixes, ", ", "[", "]"));
+    }
+    // Build a FileIO that routes each path to its vended credential. TODO: STS
+    // refresh via credentials.uri is not yet supported.
+    ICEBERG_ASSIGN_OR_RAISE(auto io,
+                            MakeS3FileIOFromCredentials(config_.configs(), table_config,
+                                                        storage_credentials));
     return std::shared_ptr<FileIO>(std::move(io));
   }
 
-  // Only non-S3 (e.g. GCS/ADLS) credentials were vended and we can't honor them;
-  // fail fast. (S3 credentials that simply don't match fall through below.)
-  if (HasOnlyNonS3StorageCredentials(storage_credentials)) {
-    auto prefixes =
-        storage_credentials | std::views::transform(&StorageCredential::prefix);
-    return NotSupported(
-        "Vended storage credentials {} are unsupported (only S3-family "
-        "credentials are supported)",
-        FormatRange(prefixes, ", ", "[", "]"));
-  }
-
   ICEBERG_RETURN_UNEXPECTED(ValidateNoFileIOConfig(table_config));
   return file_io_;
 }
@@ -789,9 +783,8 @@ Result<std::shared_ptr<Transaction>> RestCatalog::StageCreateTable(
                           /*stage_create=*/true, *contextual_session));
   auto table_config = std::move(result.config);
   auto storage_credentials = std::move(result.storage_credentials);
-  ICEBERG_ASSIGN_OR_RAISE(
-      auto table_io,
-      TableFileIO(context, table_config, storage_credentials, result.metadata->location));
+  ICEBERG_ASSIGN_OR_RAISE(auto table_io,
+                          TableFileIO(context, table_config, storage_credentials));
   ICEBERG_ASSIGN_OR_RAISE(
       auto table_session,
       TableAuthSession(identifier, table_config, std::move(contextual_session)));
@@ -906,9 +899,8 @@ Result<std::shared_ptr<Table>> RestCatalog::MakeTableFromLoadResult(
     std::shared_ptr<auth::AuthSession> contextual_session) {
   auto table_config = std::move(result.config);
   auto storage_credentials = std::move(result.storage_credentials);
-  ICEBERG_ASSIGN_OR_RAISE(
-      auto table_io,
-      TableFileIO(context, table_config, storage_credentials, result.metadata->location));
+  ICEBERG_ASSIGN_OR_RAISE(auto table_io,
+                          TableFileIO(context, table_config, storage_credentials));
   ICEBERG_ASSIGN_OR_RAISE(
       auto table_session,
       TableAuthSession(identifier, table_config, std::move(contextual_session)));

src/iceberg/catalog/rest/rest_catalog.h

@@ -80,8 +80,7 @@ class ICEBERG_REST_EXPORT RestCatalog final
   Result<std::shared_ptr<FileIO>> TableFileIO(
       const SessionContext& context,
       const std::unordered_map<std::string, std::string>& table_config,
-      const std::vector<StorageCredential>& storage_credentials,
-      std::string_view storage_location) const;
+      const std::vector<StorageCredential>& storage_credentials) const;
 
   Result<std::vector<Namespace>> ListNamespaces(const Namespace& ns,
                                                 auth::AuthSession& session) const;