diff --git a/src/iceberg/catalog/rest/CMakeLists.txt b/src/iceberg/catalog/rest/CMakeLists.txt index f64860ff4..0736290bc 100644 --- a/src/iceberg/catalog/rest/CMakeLists.txt +++ b/src/iceberg/catalog/rest/CMakeLists.txt @@ -34,6 +34,8 @@ set(ICEBERG_REST_SOURCES rest_catalog.cc rest_file_io.cc rest_metrics_reporter.cc + rest_table.cc + rest_table_scan.cc rest_util.cc types.cc) diff --git a/src/iceberg/catalog/rest/catalog_properties.cc b/src/iceberg/catalog/rest/catalog_properties.cc index 0e417e6c3..759c3f171 100644 --- a/src/iceberg/catalog/rest/catalog_properties.cc +++ b/src/iceberg/catalog/rest/catalog_properties.cc @@ -20,6 +20,8 @@ #include "iceberg/catalog/rest/catalog_properties.h" #include +#include +#include #include #include @@ -61,4 +63,14 @@ Result RestCatalogProperties::SnapshotLoadingMode() const { } } +Result> RestCatalogProperties::ScanPlanningModeFrom( + const std::unordered_map& config) { + auto it = config.find(kScanPlanningMode.key()); + if (it == config.end()) return std::nullopt; + std::string lower = StringUtils::ToLower(it->second); + if (lower == "client") return ScanPlanningMode::kClient; + if (lower == "server") return ScanPlanningMode::kServer; + return InvalidArgument("Invalid scan planning mode: '{}'.", it->second); +} + } // namespace iceberg::rest diff --git a/src/iceberg/catalog/rest/catalog_properties.h b/src/iceberg/catalog/rest/catalog_properties.h index d1ee0e9c4..da5bd550a 100644 --- a/src/iceberg/catalog/rest/catalog_properties.h +++ b/src/iceberg/catalog/rest/catalog_properties.h @@ -19,6 +19,7 @@ #pragma once +#include #include #include @@ -34,6 +35,9 @@ namespace iceberg::rest { /// \brief Snapshot loading mode for REST catalog. enum class SnapshotMode : uint8_t { kAll, kRefs }; +/// \brief Scan planning mode for REST catalog. +enum class ScanPlanningMode : uint8_t { kClient, kServer }; + /// \brief Configuration class for a REST Catalog. class ICEBERG_REST_EXPORT RestCatalogProperties : public ConfigBase { @@ -58,6 +62,14 @@ class ICEBERG_REST_EXPORT RestCatalogProperties /// \brief Whether to report metrics to the REST catalog server (default: true). inline static Entry kMetricsReportingEnabled{ "rest-metrics-reporting-enabled", "true"}; + /// \brief The scan planning mode (client or server). + inline static Entry kScanPlanningMode{"scan-planning-mode", "client"}; + /// \brief Maximum time to wait for a submitted scan plan, in milliseconds. + inline static Entry kScanPlanningPollTimeoutMs{ + "rest-scan-planning.poll-timeout-ms", 300000}; + /// \brief Extra polls after the first fetch of a submitted scan plan. + inline static Entry kScanPlanningPollNumRetries{ + "rest-scan-planning.poll-num-retries", 10}; /// \brief The prefix for HTTP headers. inline static constexpr std::string_view kHeaderPrefix = "header."; @@ -80,6 +92,11 @@ class ICEBERG_REST_EXPORT RestCatalogProperties /// "REFS", or an error if the value is invalid. Parsing is /// case-insensitive to match Java behavior. Result SnapshotLoadingMode() const; + + /// \brief Get the scan planning mode from the given config map, returning + /// std::nullopt if the key is absent. + static Result> ScanPlanningModeFrom( + const std::unordered_map& config); }; } // namespace iceberg::rest diff --git a/src/iceberg/catalog/rest/http_client.cc b/src/iceberg/catalog/rest/http_client.cc index 6661c5098..d48471cde 100644 --- a/src/iceberg/catalog/rest/http_client.cc +++ b/src/iceberg/catalog/rest/http_client.cc @@ -50,6 +50,14 @@ class HttpResponse::Impl { }; HttpResponse::HttpResponse() = default; + +HttpResponse::HttpResponse(int32_t status_code, std::string body) { + cpr::Response cpr_response; + cpr_response.status_code = status_code; + cpr_response.text = std::move(body); + impl_ = std::make_unique(std::move(cpr_response)); +} + HttpResponse::~HttpResponse() = default; HttpResponse::HttpResponse(HttpResponse&&) noexcept = default; HttpResponse& HttpResponse::operator=(HttpResponse&&) noexcept = default; diff --git a/src/iceberg/catalog/rest/http_client.h b/src/iceberg/catalog/rest/http_client.h index ea9c10a39..9400202dc 100644 --- a/src/iceberg/catalog/rest/http_client.h +++ b/src/iceberg/catalog/rest/http_client.h @@ -45,6 +45,7 @@ namespace iceberg::rest { class ICEBERG_REST_EXPORT HttpResponse { public: HttpResponse(); + HttpResponse(int32_t status_code, std::string body); ~HttpResponse(); HttpResponse(const HttpResponse&) = delete; @@ -71,7 +72,7 @@ class ICEBERG_REST_EXPORT HttpResponse { class ICEBERG_REST_EXPORT HttpClient { public: explicit HttpClient(std::unordered_map default_headers = {}); - ~HttpClient(); + virtual ~HttpClient(); HttpClient(const HttpClient&) = delete; HttpClient& operator=(const HttpClient&) = delete; @@ -79,36 +80,35 @@ class ICEBERG_REST_EXPORT HttpClient { HttpClient& operator=(HttpClient&&) = delete; /// \brief Sends a GET request. - Result Get(const std::string& path, - const std::unordered_map& params, - const std::unordered_map& headers, - const ErrorHandler& error_handler, auth::AuthSession& session); + virtual Result Get( + const std::string& path, const std::unordered_map& params, + const std::unordered_map& headers, + const ErrorHandler& error_handler, auth::AuthSession& session); /// \brief Sends a POST request. - Result Post(const std::string& path, const std::string& body, - const std::unordered_map& headers, - const ErrorHandler& error_handler, - auth::AuthSession& session); + virtual Result Post( + const std::string& path, const std::string& body, + const std::unordered_map& headers, + const ErrorHandler& error_handler, auth::AuthSession& session); /// \brief Sends a POST request with form data. - Result PostForm( + virtual Result PostForm( const std::string& path, const std::unordered_map& form_data, const std::unordered_map& headers, const ErrorHandler& error_handler, auth::AuthSession& session); /// \brief Sends a HEAD request. - Result Head(const std::string& path, - const std::unordered_map& headers, - const ErrorHandler& error_handler, - auth::AuthSession& session); + virtual Result Head( + const std::string& path, + const std::unordered_map& headers, + const ErrorHandler& error_handler, auth::AuthSession& session); /// \brief Sends a DELETE request. - Result Delete(const std::string& path, - const std::unordered_map& params, - const std::unordered_map& headers, - const ErrorHandler& error_handler, - auth::AuthSession& session); + virtual Result Delete( + const std::string& path, const std::unordered_map& params, + const std::unordered_map& headers, + const ErrorHandler& error_handler, auth::AuthSession& session); private: std::unordered_map default_headers_; diff --git a/src/iceberg/catalog/rest/json_serde.cc b/src/iceberg/catalog/rest/json_serde.cc index ba348dec3..c7d4662dd 100644 --- a/src/iceberg/catalog/rest/json_serde.cc +++ b/src/iceberg/catalog/rest/json_serde.cc @@ -128,8 +128,7 @@ Result StorageCredentialFromJson(const nlohmann::json& json) return credential; } -/// \brief Reads the optional `storage-credentials` array shared by the -/// LoadTable and LoadCredentials responses. +/// \brief Reads the optional `storage-credentials` array. Result> StorageCredentialsFromJson( const nlohmann::json& json) { std::vector credentials; @@ -146,6 +145,21 @@ Result> StorageCredentialsFromJson( } return credentials; } + +/// \brief Writes `storage-credentials` to json if any are present. +Status StorageCredentialsToJson(const std::vector& credentials, + nlohmann::json& json) { + if (credentials.empty()) { + return {}; + } + nlohmann::json creds_json = nlohmann::json::array(); + for (const auto& cred : credentials) { + ICEBERG_ASSIGN_OR_RAISE(auto entry, StorageCredentialToJson(cred)); + creds_json.push_back(std::move(entry)); + } + json[kStorageCredentials] = std::move(creds_json); + return {}; +} } // namespace Result>> FileScanTasksFromJson( @@ -829,6 +843,8 @@ Result ToJson(const PlanTableScanRequest& request) { if (request.snapshot_id.has_value()) { json[kSnapshotId] = request.snapshot_id.value(); } + // FIXME: Preserve explicit empty projections as select: []; omitting select + // means the default projection. Java currently has the same issue. if (!request.select.empty()) { json[kSelect] = request.select; } @@ -884,6 +900,7 @@ Result PlanTableScanResponseFromJson( } ICEBERG_RETURN_UNEXPECTED( ScanTaskFieldsFromJson(json, response, partition_specs_by_id, schema)); + ICEBERG_ASSIGN_OR_RAISE(response.storage_credentials, StorageCredentialsFromJson(json)); ICEBERG_RETURN_UNEXPECTED(response.Validate()); return response; } @@ -895,6 +912,7 @@ Result ToJson( const Schema& schema) { ICEBERG_ASSIGN_OR_RAISE(nlohmann::json json, ScanTaskFieldsToJson(response, partition_specs_by_id, schema)); + ICEBERG_RETURN_UNEXPECTED(StorageCredentialsToJson(response.storage_credentials, json)); json[kPlanStatus] = ToString(response.plan_status); if (!response.plan_id.empty()) { json[kPlanId] = response.plan_id; @@ -921,6 +939,7 @@ Result FetchPlanningResultResponseFromJson( } ICEBERG_RETURN_UNEXPECTED( ScanTaskFieldsFromJson(json, response, partition_specs_by_id, schema)); + ICEBERG_ASSIGN_OR_RAISE(response.storage_credentials, StorageCredentialsFromJson(json)); ICEBERG_RETURN_UNEXPECTED(response.Validate()); return response; } @@ -932,6 +951,7 @@ Result ToJson( const Schema& schema) { ICEBERG_ASSIGN_OR_RAISE(nlohmann::json json, ScanTaskFieldsToJson(response, partition_specs_by_id, schema)); + ICEBERG_RETURN_UNEXPECTED(StorageCredentialsToJson(response.storage_credentials, json)); json[kPlanStatus] = ToString(response.plan_status); if (response.error.has_value()) { json[kError] = ToJson(*response.error)[kError]; diff --git a/src/iceberg/catalog/rest/rest_catalog.cc b/src/iceberg/catalog/rest/rest_catalog.cc index 0ca89a725..767b6eef6 100644 --- a/src/iceberg/catalog/rest/rest_catalog.cc +++ b/src/iceberg/catalog/rest/rest_catalog.cc @@ -21,6 +21,7 @@ #include #include +#include #include #include #include @@ -41,10 +42,14 @@ #include "iceberg/catalog/rest/resource_paths.h" #include "iceberg/catalog/rest/rest_file_io.h" #include "iceberg/catalog/rest/rest_metrics_reporter_internal.h" +#include "iceberg/catalog/rest/rest_table.h" +#include "iceberg/catalog/rest/rest_table_scan.h" #include "iceberg/catalog/rest/rest_util.h" #include "iceberg/catalog/rest/types.h" #include "iceberg/json_serde_internal.h" +#include "iceberg/logging/log_level.h" #include "iceberg/logging/log_macros.h" +#include "iceberg/logging/logger.h" #include "iceberg/metrics/metrics_reporters.h" #include "iceberg/partition_spec.h" #include "iceberg/result.h" @@ -67,16 +72,18 @@ class RestStorageCredentialProvider final : public StorageCredentialProvider { public: RestStorageCredentialProvider(std::shared_ptr client, std::shared_ptr session, - std::string path) + std::string path, + std::unordered_map params) : client_(std::move(client)), session_(std::move(session)), - path_(std::move(path)) {} + path_(std::move(path)), + params_(std::move(params)) {} Result> Load() override { // The provider can be shared by several FileIOs; keep fetches serialized. std::lock_guard lock(mutex_); ICEBERG_ASSIGN_OR_RAISE(const auto response, - client_->Get(path_, /*params=*/{}, /*headers=*/{}, + client_->Get(path_, params_, /*headers=*/{}, *TableErrorHandler::Instance(), *session_)); // Parse errors can contain credential data; return a fixed message. auto json = FromJsonString(response.body()); @@ -95,6 +102,7 @@ class RestStorageCredentialProvider final : public StorageCredentialProvider { std::shared_ptr client_; std::shared_ptr session_; std::string path_; + std::unordered_map params_; std::mutex mutex_; }; @@ -121,11 +129,10 @@ std::string RestTableName(std::string_view catalog_name, } /// \brief Fetch server configuration from the REST catalog server. -Result FetchServerConfig(const ResourcePaths& paths, +Result FetchServerConfig(HttpClient& client, const ResourcePaths& paths, const RestCatalogProperties& current_config, auth::AuthSession& session) { ICEBERG_ASSIGN_OR_RAISE(auto config_path, paths.Config()); - HttpClient client(current_config.ExtractHeaders()); // Send the client's warehouse location to the service to keep in sync. // This is needed for cases where the warehouse is configured client side, but may @@ -424,7 +431,14 @@ RestCatalog::~RestCatalog() { } Result> RestCatalog::Make( - const RestCatalogProperties& config, Executor* metrics_executor) { + const RestCatalogProperties& config, Executor* metrics_executor, + ClientFactory client_factory) { + if (!client_factory) { + client_factory = [](const RestCatalogProperties& properties) + -> Result> { + return std::make_shared(properties.ExtractHeaders()); + }; + } ICEBERG_ASSIGN_OR_RAISE(auto uri, config.Uri()); std::string catalog_name = config.Get(RestCatalogProperties::kName); @@ -437,11 +451,12 @@ Result> RestCatalog::Make( config.Get(RestCatalogProperties::kNamespaceSeparator))); // Create init session for fetching server configuration - auto init_client = std::make_shared(config.ExtractHeaders()); + ICEBERG_ASSIGN_OR_RAISE(auto init_client, client_factory(config)); + ICEBERG_PRECHECK(init_client != nullptr, "HTTP client factory returned null"); ICEBERG_ASSIGN_OR_RAISE(auto init_session, auth_manager->InitSession(init_client, config.configs())); ICEBERG_ASSIGN_OR_RAISE(auto server_config, - FetchServerConfig(*paths, config, *init_session)); + FetchServerConfig(*init_client, *paths, config, *init_session)); RestCatalogProperties final_config = RestCatalogProperties::FromMap( MergeConfigs(server_config.defaults, config.configs(), server_config.overrides)); @@ -467,8 +482,12 @@ Result> RestCatalog::Make( // Get snapshot loading mode ICEBERG_ASSIGN_OR_RAISE(auto snapshot_mode, final_config.SnapshotLoadingMode()); + ICEBERG_ASSIGN_OR_RAISE( + auto client_scan_planning_mode, + RestCatalogProperties::ScanPlanningModeFrom(final_config.configs())); - auto client = std::make_shared(final_config.ExtractHeaders()); + ICEBERG_ASSIGN_OR_RAISE(auto client, client_factory(final_config)); + ICEBERG_PRECHECK(client != nullptr, "HTTP client factory returned null"); ICEBERG_ASSIGN_OR_RAISE(auto catalog_session, auth_manager->CatalogSession(client, final_config.configs())); @@ -487,16 +506,19 @@ Result> RestCatalog::Make( return std::shared_ptr(new RestCatalog( std::move(final_config), std::move(file_io), std::move(client), std::move(paths), std::move(endpoints), std::move(auth_manager), std::move(catalog_session), - snapshot_mode, std::move(default_context), std::move(reporter), metrics_executor)); + snapshot_mode, client_scan_planning_mode, std::move(default_context), + std::move(reporter), metrics_executor)); } RestCatalog::RestCatalog(RestCatalogProperties config, std::shared_ptr file_io, std::shared_ptr client, - std::unique_ptr paths, + std::shared_ptr paths, std::unordered_set endpoints, std::unique_ptr auth_manager, std::shared_ptr catalog_session, - SnapshotMode snapshot_mode, SessionContext default_context, + SnapshotMode snapshot_mode, + std::optional client_scan_planning_mode, + SessionContext default_context, std::shared_ptr reporter, Executor* metrics_executor) : config_(std::move(config)), @@ -508,6 +530,7 @@ RestCatalog::RestCatalog(RestCatalogProperties config, std::shared_ptr f auth_manager_(std::move(auth_manager)), catalog_session_(std::move(catalog_session)), snapshot_mode_(snapshot_mode), + client_scan_planning_mode_(client_scan_planning_mode), default_context_(std::move(default_context)), reporter_(std::move(reporter)), metrics_executor_(metrics_executor) { @@ -546,30 +569,28 @@ Result> RestCatalog::TableAuthSession( std::move(contextual_session)); } -std::shared_ptr RestCatalog::MakeStorageCredentialProvider( - const TableIdentifier& identifier, - std::shared_ptr table_session) const { +Result> +RestCatalog::MakeStorageCredentialProvider( + const TableIdentifier& identifier, std::shared_ptr table_session, + std::optional plan_id) const { if (!supported_endpoints_.contains(Endpoint::TableCredentials())) { // Not an error, but it surfaces much later as credentials expiring. ICEBERG_LOG_DEBUG( "Catalog does not advertise {}; vended credentials for '{}' will not be " "refreshed", Endpoint::TableCredentials().ToString(), ToString(identifier)); - return nullptr; - } - auto path = paths_->Credentials(identifier); - if (!path.has_value()) { - ICEBERG_LOG_WARN( - "Cannot build the credentials path for '{}' ({}); its vended credentials " - "will not be refreshed", - ToString(identifier), path.error().message); - return nullptr; + return std::shared_ptr{}; } + ICEBERG_ASSIGN_OR_RAISE(auto credentials_path, paths_->Credentials(identifier)); auto client = client_; - auto credentials_path = std::move(path.value()); auto session = std::move(table_session); + std::unordered_map params; + if (plan_id.has_value()) { + params.emplace("planId", std::move(*plan_id)); + } return std::make_shared( - std::move(client), std::move(session), std::move(credentials_path)); + std::move(client), std::move(session), std::move(credentials_path), + std::move(params)); } Result> RestCatalog::TableFileIO( @@ -581,7 +602,8 @@ Result> RestCatalog::TableFileIO( // Only vended credentials expire, so only they need a provider. std::shared_ptr provider; if (!storage_credentials.empty()) { - provider = MakeStorageCredentialProvider(identifier, std::move(table_session)); + ICEBERG_ASSIGN_OR_RAISE( + provider, MakeStorageCredentialProvider(identifier, std::move(table_session))); } return MakeTableFileIO(config_.configs(), table_config, storage_credentials, std::move(provider)); @@ -972,10 +994,10 @@ Result> RestCatalog::MakeTableFromLoadResult( auto table_catalog = std::make_shared( shared_from_this(), context, identifier, table_config, table_session, table_io); - return Table::Make(identifier, std::move(result.metadata), - std::move(result.metadata_location), std::move(table_io), - std::move(table_catalog), RestTableName(name_, identifier), - std::move(reporter)); + return MakeTable(identifier, std::move(result.metadata), + std::move(result.metadata_location), std::move(table_io), + std::move(table_catalog), std::move(reporter), table_config, + std::move(table_session)); } Result> RestCatalog::MakeTableFromCommitResponse( @@ -987,10 +1009,71 @@ Result> RestCatalog::MakeTableFromCommitResponse( // Reuse the bound FileIO because commit responses carry no config or credentials. auto table_catalog = std::make_shared( shared_from_this(), context, identifier, table_config, table_session, table_io); - return Table::Make(identifier, std::move(response.metadata), - std::move(response.metadata_location), std::move(table_io), - std::move(table_catalog), RestTableName(name_, identifier), - std::move(reporter)); + return MakeTable(identifier, std::move(response.metadata), + std::move(response.metadata_location), std::move(table_io), + std::move(table_catalog), std::move(reporter), table_config, + std::move(table_session)); +} + +Result> RestCatalog::MakeTable( + const TableIdentifier& identifier, std::shared_ptr metadata, + std::string metadata_location, std::shared_ptr table_io, + std::shared_ptr table_catalog, std::shared_ptr reporter, + const std::unordered_map& table_config, + std::shared_ptr table_session) { + // Table config overrides client config. + ICEBERG_ASSIGN_OR_RAISE(auto server_mode, + RestCatalogProperties::ScanPlanningModeFrom(table_config)); + if (client_scan_planning_mode_.has_value() && server_mode.has_value() && + *client_scan_planning_mode_ != *server_mode) { + Log(LogLevel::kWarn, + "Scan planning mode mismatch for table {}: client config={}, server config={}. " + "Server config will take precedence.", + identifier.ToString(), + *client_scan_planning_mode_ == ScanPlanningMode::kClient ? "client" : "server", + *server_mode == ScanPlanningMode::kClient ? "client" : "server"); + } + ScanPlanningMode effective_mode = server_mode.value_or( + client_scan_planning_mode_.value_or(ScanPlanningMode::kClient)); + + if (effective_mode == ScanPlanningMode::kServer) { + if (!supported_endpoints_.contains(Endpoint::PlanTableScan())) { + return NotSupported( + "Server requires server-side scan planning for table {} but does not support " + "the PlanTableScan endpoint.", + identifier.ToString()); + } + // Scans can outlive the table; keep its authentication sessions alive too. + auto credential_provider_factory = [catalog = shared_from_this(), identifier, + table_session](const std::string& plan_id) + -> Result> { + // Without a plan id, refresh is table-scoped (no planId query parameter). + std::optional scoped_plan_id; + if (!plan_id.empty()) { + scoped_plan_id = plan_id; + } + return catalog->MakeStorageCredentialProvider(identifier, table_session, + std::move(scoped_plan_id)); + }; + RestScanContext rest_ctx{ + .client = client_, + .paths = paths_, + .session = std::move(table_session), + .supported_endpoints = supported_endpoints_, + .identifier = identifier, + .catalog_config = config_.configs(), + .table_config = table_config, + .credential_provider_factory = std::move(credential_provider_factory), + }; + return RestTable::Make(identifier, std::move(metadata), std::move(metadata_location), + std::move(table_io), std::move(table_catalog), + RestTableName(name_, identifier), std::move(reporter), + std::move(rest_ctx)); + } + + return Table::Make(identifier, std::move(metadata), std::move(metadata_location), + std::move(table_io), std::move(table_catalog), + RestTableName(name_, identifier), std::move(reporter)); } } // namespace iceberg::rest diff --git a/src/iceberg/catalog/rest/rest_catalog.h b/src/iceberg/catalog/rest/rest_catalog.h index 6f65e6d61..967e7cc9f 100644 --- a/src/iceberg/catalog/rest/rest_catalog.h +++ b/src/iceberg/catalog/rest/rest_catalog.h @@ -19,7 +19,9 @@ #pragma once +#include #include +#include #include #include @@ -51,12 +53,18 @@ class ICEBERG_REST_EXPORT RestCatalog final RestCatalog(RestCatalog&&) = delete; RestCatalog& operator=(RestCatalog&&) = delete; + using ClientFactory = + std::function>(const RestCatalogProperties&)>; + /// \brief Create a RestCatalog instance. /// \param metrics_executor Optional non-owning executor for asynchronous REST metrics. /// It must outlive the catalog and all tables created by it. When null, REST metrics /// are reported synchronously. + /// \param client_factory Optional HTTP client factory, called with the initial config + /// and then the merged server config. When empty, uses the default HTTP client. static Result> Make(const RestCatalogProperties& config, - Executor* metrics_executor = nullptr); + Executor* metrics_executor = nullptr, + ClientFactory client_factory = {}); std::string_view name() const override; @@ -69,12 +77,14 @@ class ICEBERG_REST_EXPORT RestCatalog final class TableScopedCatalog; RestCatalog(RestCatalogProperties config, std::shared_ptr file_io, - std::shared_ptr client, std::unique_ptr paths, + std::shared_ptr client, std::shared_ptr paths, std::unordered_set endpoints, std::unique_ptr auth_manager, std::shared_ptr catalog_session, - SnapshotMode snapshot_mode, SessionContext default_context, - std::shared_ptr reporter, Executor* metrics_executor); + SnapshotMode snapshot_mode, + std::optional client_scan_planning_mode, + SessionContext default_context, std::shared_ptr reporter, + Executor* metrics_executor); Result> ContextualAuthSession( const SessionContext& context); @@ -92,10 +102,10 @@ class ICEBERG_REST_EXPORT RestCatalog final /// \brief Build a provider for this table's vended credentials. /// - /// Returns nullptr when the catalog does not serve LoadCredentials. - std::shared_ptr MakeStorageCredentialProvider( - const TableIdentifier& identifier, - std::shared_ptr table_session) const; + /// Returns a null provider when the catalog does not serve LoadCredentials. + Result> MakeStorageCredentialProvider( + const TableIdentifier& identifier, std::shared_ptr table_session, + std::optional plan_id = std::nullopt) const; Result> ListNamespaces(const Namespace& ns, auth::AuthSession& session) const; @@ -198,15 +208,24 @@ class ICEBERG_REST_EXPORT RestCatalog final const std::unordered_map& table_config, std::shared_ptr table_session, std::shared_ptr table_io); + /// \brief Returns a RestTable when scan planning mode is server, else a Table. + Result> MakeTable( + const TableIdentifier& identifier, std::shared_ptr metadata, + std::string metadata_location, std::shared_ptr table_io, + std::shared_ptr table_catalog, std::shared_ptr reporter, + const std::unordered_map& table_config, + std::shared_ptr table_session); + RestCatalogProperties config_; std::shared_ptr file_io_; std::shared_ptr client_; - std::unique_ptr paths_; + std::shared_ptr paths_; std::string name_; std::unordered_set supported_endpoints_; std::unique_ptr auth_manager_; std::shared_ptr catalog_session_; SnapshotMode snapshot_mode_; + std::optional client_scan_planning_mode_; SessionContext default_context_; std::weak_ptr default_catalog_; std::shared_ptr reporter_; diff --git a/src/iceberg/catalog/rest/rest_table.cc b/src/iceberg/catalog/rest/rest_table.cc new file mode 100644 index 000000000..dda33316f --- /dev/null +++ b/src/iceberg/catalog/rest/rest_table.cc @@ -0,0 +1,68 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you 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. + */ + +#include "iceberg/catalog/rest/rest_table.h" + +#include +#include + +#include "iceberg/catalog/rest/rest_table_scan.h" +#include "iceberg/result.h" +#include "iceberg/table_metadata.h" +#include "iceberg/util/macros.h" + +namespace iceberg::rest { + +RestTable::RestTable(TableIdentifier identifier, std::shared_ptr metadata, + std::string metadata_location, std::shared_ptr io, + std::shared_ptr catalog, std::string full_name, + std::shared_ptr reporter, + RestScanContext rest_context) + : Table(std::move(identifier), std::move(metadata), std::move(metadata_location), + std::move(io), std::move(catalog), std::move(full_name), std::move(reporter)), + rest_context_(std::make_shared(std::move(rest_context))) {} + +RestTable::~RestTable() = default; + +Result> RestTable::Make( + TableIdentifier identifier, std::shared_ptr metadata, + std::string metadata_location, std::shared_ptr io, + std::shared_ptr catalog, std::string full_name, + std::shared_ptr reporter, RestScanContext rest_context) { + if (metadata == nullptr) { + return InvalidArgument("Metadata cannot be null"); + } + return std::shared_ptr( + new RestTable(std::move(identifier), std::move(metadata), + std::move(metadata_location), std::move(io), std::move(catalog), + std::move(full_name), std::move(reporter), std::move(rest_context))); +} + +Result> RestTable::NewScan() const { + return std::make_unique(metadata_, io_, full_name_, reporter_, + rest_context_); +} + +Result> +RestTable::NewIncrementalAppendScan() const { + return std::make_unique(metadata_, io_, full_name_, + reporter_, rest_context_); +} + +} // namespace iceberg::rest diff --git a/src/iceberg/catalog/rest/rest_table.h b/src/iceberg/catalog/rest/rest_table.h new file mode 100644 index 000000000..fdfa7d952 --- /dev/null +++ b/src/iceberg/catalog/rest/rest_table.h @@ -0,0 +1,65 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you 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 +#include + +#include "iceberg/catalog/rest/iceberg_rest_export.h" +#include "iceberg/catalog/rest/type_fwd.h" +#include "iceberg/result.h" +#include "iceberg/table.h" +#include "iceberg/type_fwd.h" + +/// \file iceberg/catalog/rest/rest_table.h +/// A Table subclass that uses server-side distributed scan planning via the REST catalog. + +namespace iceberg::rest { + +/// \brief A Table whose NewScan() returns a RestTableScanBuilder, delegating +/// PlanFiles() to the REST catalog server's scan planning endpoints. +class ICEBERG_REST_EXPORT RestTable final : public Table { + public: + static Result> Make( + TableIdentifier identifier, std::shared_ptr metadata, + std::string metadata_location, std::shared_ptr io, + std::shared_ptr catalog, std::string full_name, + std::shared_ptr reporter, RestScanContext rest_context); + + ~RestTable() override; + + /// \brief Returns a RestTableScanBuilder that will delegate PlanFiles() to the + /// REST catalog server. + Result> NewScan() const override; + + /// \brief Returns a RestIncrementalAppendScanBuilder that delegates to the server. + Result> NewIncrementalAppendScan() + const override; + + private: + RestTable(TableIdentifier identifier, std::shared_ptr metadata, + std::string metadata_location, std::shared_ptr io, + std::shared_ptr catalog, std::string full_name, + std::shared_ptr reporter, RestScanContext rest_context); + + std::shared_ptr rest_context_; +}; + +} // namespace iceberg::rest diff --git a/src/iceberg/catalog/rest/rest_table_scan.cc b/src/iceberg/catalog/rest/rest_table_scan.cc new file mode 100644 index 000000000..e7dae0f49 --- /dev/null +++ b/src/iceberg/catalog/rest/rest_table_scan.cc @@ -0,0 +1,571 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you 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. + */ + +#include "iceberg/catalog/rest/rest_table_scan.h" + +#include +#include +#include +#include +#include +#include + +#include + +#include "iceberg/catalog/rest/catalog_properties.h" +#include "iceberg/catalog/rest/endpoint.h" +#include "iceberg/catalog/rest/error_handlers.h" +#include "iceberg/catalog/rest/http_client.h" +#include "iceberg/catalog/rest/json_serde_internal.h" +#include "iceberg/catalog/rest/resource_paths.h" +#include "iceberg/catalog/rest/rest_file_io.h" +#include "iceberg/catalog/rest/types.h" +#include "iceberg/expression/expression.h" +#include "iceberg/json_serde_internal.h" +#include "iceberg/logging/log_macros.h" +#include "iceberg/partition_spec.h" +#include "iceberg/result.h" +#include "iceberg/schema.h" +#include "iceberg/table_metadata.h" +#include "iceberg/util/macros.h" +#include "iceberg/util/string_util.h" +#include "iceberg/util/type_util.h" + +namespace iceberg::rest { + +namespace { + +constexpr int64_t kMinSleepMs = 1'000; +constexpr int64_t kMaxSleepMs = 60'000; + +struct PollConfig { + int64_t max_wait_ms; + int32_t max_retries; +}; + +Result ResolvePollConfig(const RestScanContext& ctx) { + const auto& timeout = RestCatalogProperties::kScanPlanningPollTimeoutMs; + const auto& retries = RestCatalogProperties::kScanPlanningPollNumRetries; + PollConfig config{ + .max_wait_ms = timeout.value(), + .max_retries = retries.value(), + }; + + if (auto it = ctx.catalog_config.find(timeout.key()); it != ctx.catalog_config.end()) { + ICEBERG_ASSIGN_OR_RAISE(config.max_wait_ms, + StringUtils::ParseNumber(it->second)); + } + ICEBERG_PRECHECK(config.max_wait_ms > 0, "Invalid value for {}: {} (must be positive)", + timeout.key(), config.max_wait_ms); + + if (auto it = ctx.catalog_config.find(retries.key()); it != ctx.catalog_config.end()) { + ICEBERG_ASSIGN_OR_RAISE(config.max_retries, + StringUtils::ParseNumber(it->second)); + } + ICEBERG_PRECHECK(config.max_retries >= 0, + "Invalid value for {}: {} (must be non-negative)", retries.key(), + config.max_retries); + + return config; +} + +using SpecsById = std::unordered_map>; + +#define ICEBERG_ENDPOINT_CHECK(endpoints, endpoint) \ + do { \ + if (!endpoints.contains(endpoint)) { \ + return NotSupported("Not supported endpoint: {}", endpoint.ToString()); \ + } \ + } while (0) + +// --------------------------------------------------------------------------- +// Shared HTTP scan planning helpers used by all REST scan implementations. +// --------------------------------------------------------------------------- + +// Column names in field ID order. +// FIXME: Sending both "s" and "s.x" asks the server for all of "s". +// Preserve partial struct projections; Java currently has the same issue. +Result> ProjectedColumnNames(const Schema& projected, + bool include_struct_ids = true) { + ICEBERG_ASSIGN_OR_RAISE(auto id_set, GetProjectedIdsVisitor::GetProjectedIds( + projected, include_struct_ids)); + std::vector field_ids(id_set.begin(), id_set.end()); + std::ranges::sort(field_ids); + + std::vector names; + names.reserve(field_ids.size()); + for (int32_t field_id : field_ids) { + ICEBERG_ASSIGN_OR_RAISE(auto name, projected.FindColumnNameById(field_id)); + ICEBERG_CHECK(name.has_value(), "Projected field {} not found in schema", field_id); + names.emplace_back(*name); + } + return names; +} + +Result> MakeScanFileIO( + const RestScanContext& ctx, const std::string& plan_id, + const std::vector& credentials) { + if (credentials.empty()) return std::shared_ptr{}; + std::shared_ptr provider; + if (ctx.credential_provider_factory) { + ICEBERG_ASSIGN_OR_RAISE(provider, ctx.credential_provider_factory(plan_id)); + } + ICEBERG_ASSIGN_OR_RAISE(std::shared_ptr io, + MakeTableFileIO(ctx.catalog_config, ctx.table_config, + credentials, std::move(provider))); + return io; +} + +Status CancelPlanning(const RestScanContext& ctx, const std::string& plan_id) { + if (plan_id.empty()) return {}; + if (!ctx.supported_endpoints.contains(Endpoint::CancelPlanning())) return {}; + + ICEBERG_ASSIGN_OR_RAISE(auto path, ctx.paths->Plan(ctx.identifier, plan_id)); + ICEBERG_RETURN_UNEXPECTED(ctx.client->Delete( + path, /*params=*/{}, /*headers=*/{}, *PlanErrorHandler::Instance(), *ctx.session)); + return {}; +} + +// Cancels the plan on scope exit unless released. +class PlanCancellationGuard { + public: + PlanCancellationGuard(const RestScanContext& ctx, const std::string& plan_id) + : ctx_(ctx), plan_id_(plan_id) {} + + PlanCancellationGuard(const PlanCancellationGuard&) = delete; + PlanCancellationGuard& operator=(const PlanCancellationGuard&) = delete; + + ~PlanCancellationGuard() { + if (!active_) return; + // Cleanup is best-effort; a missing plan has already been released. + if (auto status = CancelPlanning(ctx_, plan_id_); + !status && status.error().kind != ErrorKind::kNoSuchPlanId) { + ICEBERG_LOG_WARN("Failed to cancel scan plan {}: {}", plan_id_, + status.error().message); + } + } + + void Release() { active_ = false; } + + private: + const RestScanContext& ctx_; + const std::string& plan_id_; + bool active_ = true; +}; + +// Fetches one plan task without resolving any nested plan tasks it returns. +Result FetchScanTasksOnce(const RestScanContext& ctx, + const Schema& schema, + const std::string& plan_task, + const SpecsById& specs) { + ICEBERG_ENDPOINT_CHECK(ctx.supported_endpoints, Endpoint::FetchScanTasks()); + + ICEBERG_ASSIGN_OR_RAISE(auto path, ctx.paths->FetchScanTasks(ctx.identifier)); + FetchScanTasksRequest request{.planTask = plan_task}; + ICEBERG_ASSIGN_OR_RAISE(auto json_request, ToJsonString(ToJson(request))); + ICEBERG_ASSIGN_OR_RAISE( + const auto response, + ctx.client->Post(path, json_request, /*headers=*/{}, + *PlanTaskErrorHandler::Instance(), *ctx.session)); + ICEBERG_ASSIGN_OR_RAISE(auto json, FromJsonString(response.body())); + ICEBERG_ASSIGN_OR_RAISE(auto result, + FetchScanTasksResponseFromJson(json, specs, schema)); + ICEBERG_RETURN_UNEXPECTED(result.Validate()); + return result; +} + +// Polls until the plan completes and returns the completed response. Plan tasks +// are left unresolved. +Result FetchPlanningResult( + const RestScanContext& ctx, const Schema& schema, const std::string& plan_id, + const SpecsById& specs, const PollConfig& poll_config, + std::shared_ptr& scan_io, PlanCancellationGuard& cleanup) { + ICEBERG_ENDPOINT_CHECK(ctx.supported_endpoints, Endpoint::FetchPlanningResult()); + ICEBERG_ASSIGN_OR_RAISE(auto path, ctx.paths->Plan(ctx.identifier, plan_id)); + + auto delay_ms = kMinSleepMs; + auto start = std::chrono::steady_clock::now(); + + for (int32_t retry = 0; retry <= poll_config.max_retries; ++retry) { + ICEBERG_ASSIGN_OR_RAISE(auto response, + ctx.client->Get(path, /*params=*/{}, /*headers=*/{}, + *PlanErrorHandler::Instance(), *ctx.session)); + ICEBERG_ASSIGN_OR_RAISE(auto json, FromJsonString(response.body())); + ICEBERG_ASSIGN_OR_RAISE(auto result, + FetchPlanningResultResponseFromJson(json, specs, schema)); + ICEBERG_RETURN_UNEXPECTED(result.Validate()); + + switch (result.plan_status) { + case PlanStatus::kCompleted: { + ICEBERG_ASSIGN_OR_RAISE(scan_io, + MakeScanFileIO(ctx, plan_id, result.storage_credentials)); + return result; + } + case PlanStatus::kSubmitted: { + auto elapsed_ms = std::chrono::duration_cast( + std::chrono::steady_clock::now() - start) + .count(); + if (elapsed_ms >= poll_config.max_wait_ms) { + return IOError("Scan planning timed out after {}ms waiting for plan_id={}", + elapsed_ms, plan_id); + } + if (retry == poll_config.max_retries) break; + std::this_thread::sleep_for(std::chrono::milliseconds(delay_ms)); + delay_ms = std::min(delay_ms * 2, kMaxSleepMs); + continue; + } + case PlanStatus::kFailed: + return IOError("Scan planning failed: {}", + result.error ? result.error->message : "unknown error"); + case PlanStatus::kCancelled: + cleanup.Release(); + return IOError("Scan planning was cancelled for plan_id={}", plan_id); + } + } + + return IOError("Scan planning exceeded max retries ({}) for plan_id={}", + poll_config.max_retries, plan_id); +} + +// TODO(sandeepg): Fetch plan tasks concurrently when an executor is configured via +// PlanWith(), as Java's ScanTaskIterable does. +/// A lazy stream that drives FetchScanTasks calls on demand. +/// +/// The eager POST /plan (or poll until COMPLETED) is done before this stream is +/// constructed. The stream then yields the directly-returned file_scan_tasks first, +/// then fetches each plan_task token lazily one at a time as Next() is called, +/// queuing any nested plan tasks returned by the server. +class RestFileScanTaskStream final : public FileScanTaskStream { + public: + RestFileScanTaskStream(std::shared_ptr ctx, + std::shared_ptr schema, std::string plan_id, + std::vector> initial_tasks, + std::vector plan_task_tokens, SpecsById specs, + std::shared_ptr filter, bool ignore_residuals) + : ctx_(std::move(ctx)), + schema_(std::move(schema)), + plan_id_(std::move(plan_id)), + buffer_(std::move(initial_tasks)), + plan_task_tokens_(std::make_move_iterator(plan_task_tokens.begin()), + std::make_move_iterator(plan_task_tokens.end())), + specs_(std::move(specs)), + filter_(filter ? std::move(filter) : True::Instance()), + ignore_residuals_(ignore_residuals), + cleanup_(*ctx_, plan_id_) {} + + ~RestFileScanTaskStream() override = default; + + protected: + Result>> NextImpl() override { + while (true) { + if (buffer_pos_ < buffer_.size()) { + auto task = std::move(buffer_[buffer_pos_++]); + if (ignore_residuals_ || !task->residual_filter()) { + // The server may leave row filtering to the client. + return std::make_shared( + task->data_file(), task->delete_files(), + ignore_residuals_ ? True::Instance() : filter_); + } + return task; + } + if (plan_task_tokens_.empty()) { + cleanup_.Release(); + return std::nullopt; + } + const std::string token = std::move(plan_task_tokens_.front()); + plan_task_tokens_.pop_front(); + ICEBERG_ASSIGN_OR_RAISE(auto result, + FetchScanTasksOnce(*ctx_, *schema_, token, specs_)); + if (result.plan_tasks.has_value()) { + plan_task_tokens_.insert(plan_task_tokens_.end(), + std::make_move_iterator(result.plan_tasks->begin()), + std::make_move_iterator(result.plan_tasks->end())); + } + buffer_ = std::move(result.file_scan_tasks) + .value_or(std::vector>{}); + buffer_pos_ = 0; + } + } + + private: + std::shared_ptr ctx_; + std::shared_ptr schema_; + std::string plan_id_; + std::vector> buffer_; + size_t buffer_pos_ = 0; + std::deque plan_task_tokens_; + SpecsById specs_; + std::shared_ptr filter_; + bool ignore_residuals_; + PlanCancellationGuard cleanup_; +}; + +/// POST /plan (polling if SUBMITTED), apply credentials, return a lazy stream. +Result ExecuteScanPlanStream( + const std::shared_ptr& rest_context, const Schema& schema, + std::shared_ptr schema_ptr, PlanTableScanRequest request, + const SpecsById& specs, std::shared_ptr& scan_io, bool ignore_residuals) { + const auto& ctx = *rest_context; + ICEBERG_ENDPOINT_CHECK(ctx.supported_endpoints, Endpoint::PlanTableScan()); + ICEBERG_ASSIGN_OR_RAISE(auto poll_config, ResolvePollConfig(ctx)); + + ICEBERG_ASSIGN_OR_RAISE(auto path, ctx.paths->Plan(ctx.identifier)); + ICEBERG_ASSIGN_OR_RAISE(auto request_json, ToJson(request)); + ICEBERG_ASSIGN_OR_RAISE(auto json_request, ToJsonString(request_json)); + ICEBERG_ASSIGN_OR_RAISE(const auto response, + ctx.client->Post(path, json_request, /*headers=*/{}, + *PlanErrorHandler::Instance(), *ctx.session)); + ICEBERG_ASSIGN_OR_RAISE(auto json, FromJsonString(response.body())); + // The parser validates the response, so read the plan id first to cancel on failure. + std::string raw_plan_id; + if (auto it = json.find("plan-id"); it != json.end() && it->is_string()) { + raw_plan_id = it->get(); + } + PlanCancellationGuard cleanup(ctx, raw_plan_id); + ICEBERG_ASSIGN_OR_RAISE(auto result, + PlanTableScanResponseFromJson(json, specs, schema)); + + const std::string plan_id = result.plan_id; + + if (result.plan_status == PlanStatus::kSubmitted) { + ICEBERG_ASSIGN_OR_RAISE( + auto completed, + FetchPlanningResult(ctx, schema, plan_id, specs, poll_config, scan_io, cleanup)); + FileScanTaskStreamPtr stream = std::make_unique( + rest_context, std::move(schema_ptr), plan_id, + std::move(completed.file_scan_tasks) + .value_or(std::vector>{}), + std::move(completed.plan_tasks).value_or(std::vector{}), specs, + std::move(request.filter), ignore_residuals); + cleanup.Release(); + return stream; + } + + if (result.plan_status == PlanStatus::kFailed) { + return IOError("Scan planning failed: {}", + result.error ? result.error->message : "unknown error"); + } + if (result.plan_status == PlanStatus::kCancelled) { + cleanup.Release(); + return IOError("Scan planning was cancelled for plan_id={}", plan_id); + } + + ICEBERG_CHECK(result.plan_status == PlanStatus::kCompleted, + "Unexpected scan planning status: {}", ToString(result.plan_status)); + ICEBERG_ASSIGN_OR_RAISE(scan_io, + MakeScanFileIO(ctx, plan_id, result.storage_credentials)); + + std::vector> initial_tasks; + if (result.file_scan_tasks.has_value()) { + initial_tasks = std::move(*result.file_scan_tasks); + } + std::vector plan_task_tokens; + if (result.plan_tasks.has_value()) { + plan_task_tokens = std::move(*result.plan_tasks); + } + + FileScanTaskStreamPtr stream = std::make_unique( + rest_context, std::move(schema_ptr), plan_id, std::move(initial_tasks), + std::move(plan_task_tokens), specs, std::move(request.filter), ignore_residuals); + cleanup.Release(); + return stream; +} + +} // namespace + +// --------------------------------------------------------------------------- +// RestTableScan +// --------------------------------------------------------------------------- + +RestTableScan::RestTableScan(std::shared_ptr metadata, + std::shared_ptr schema, std::shared_ptr io, + internal::TableScanContext context, + std::shared_ptr rest_context) + : DataTableScan(std::move(metadata), std::move(schema), std::move(io), + std::move(context)), + rest_context_(std::move(rest_context)) {} + +Result> RestTableScan::Make( + std::shared_ptr metadata, std::shared_ptr schema, + std::shared_ptr io, internal::TableScanContext context, + std::shared_ptr rest_context) { + ICEBERG_PRECHECK(metadata != nullptr, "Table metadata cannot be null"); + ICEBERG_PRECHECK(schema != nullptr, "Schema cannot be null"); + ICEBERG_PRECHECK(io != nullptr, "FileIO cannot be null"); + ICEBERG_PRECHECK(rest_context != nullptr, "REST scan context cannot be null"); + ICEBERG_CHECK( + !context.from_snapshot_id.has_value() && !context.to_snapshot_id.has_value(), + "Use RestIncrementalAppendScan for incremental scans"); + return std::unique_ptr( + new RestTableScan(std::move(metadata), std::move(schema), std::move(io), + std::move(context), std::move(rest_context))); +} + +Result RestTableScan::PlanFilesStream() const { + scan_io_ = nullptr; // reset so stale credentials from a prior plan are not reused + TableMetadataCache metadata_cache(metadata_.get()); + ICEBERG_ASSIGN_OR_RAISE(auto specs, metadata_cache.GetPartitionSpecsById()); + + PlanTableScanRequest request; + ICEBERG_ASSIGN_OR_RAISE(auto projected_schema, schema()); + ICEBERG_ASSIGN_OR_RAISE(request.select, ProjectedColumnNames(*projected_schema)); + request.filter = context_.filter; + request.case_sensitive = context_.case_sensitive; + request.min_rows_requested = context_.min_rows_requested; + + if (context_.snapshot_id.has_value()) { + request.snapshot_id = context_.snapshot_id; + request.use_snapshot_schema = context_.use_snapshot_schema; + } + + if (context_.return_column_stats && context_.columns_to_keep_stats.empty()) { + ICEBERG_ASSIGN_OR_RAISE(request.stats_fields, + ProjectedColumnNames(*schema_, /*include_struct_ids=*/false)); + } else { + for (int32_t field_id : context_.columns_to_keep_stats) { + ICEBERG_ASSIGN_OR_RAISE(auto name, schema_->FindColumnNameById(field_id)); + if (name.has_value()) { + request.stats_fields.emplace_back(*name); + } + } + } + + return ExecuteScanPlanStream(rest_context_, *schema_, schema_, std::move(request), + specs, scan_io_, context_.ignore_residuals); +} + +const std::shared_ptr& RestTableScan::io() const { + return scan_io_ ? scan_io_ : io_; +} + +// --------------------------------------------------------------------------- +// RestTableScanBuilder +// --------------------------------------------------------------------------- + +RestTableScanBuilder::RestTableScanBuilder( + std::shared_ptr metadata, std::shared_ptr io, + std::string table_name, std::shared_ptr metrics_reporter, + std::shared_ptr rest_context) + : DataTableScanBuilder(std::move(metadata), std::move(io), std::move(table_name), + std::move(metrics_reporter)), + rest_context_(std::move(rest_context)) {} + +Result> RestTableScanBuilder::Build() { + ICEBERG_RETURN_UNEXPECTED(CheckErrors()); + ICEBERG_RETURN_UNEXPECTED(ResolveColumnStatsSelection()); + ICEBERG_RETURN_UNEXPECTED(context_.Validate()); + ICEBERG_ASSIGN_OR_RAISE(auto schema, ResolveSnapshotSchema()); + return RestTableScan::Make(metadata_, schema.get(), io_, std::move(context_), + rest_context_); +} + +// --------------------------------------------------------------------------- +// RestIncrementalAppendScan +// --------------------------------------------------------------------------- + +RestIncrementalAppendScan::RestIncrementalAppendScan( + std::shared_ptr metadata, std::shared_ptr schema, + std::shared_ptr io, internal::TableScanContext context, + std::shared_ptr rest_context) + : IncrementalAppendScan(std::move(metadata), std::move(schema), std::move(io), + std::move(context)), + rest_context_(std::move(rest_context)) {} + +Result> RestIncrementalAppendScan::Make( + std::shared_ptr metadata, std::shared_ptr schema, + std::shared_ptr io, internal::TableScanContext context, + std::shared_ptr rest_context) { + ICEBERG_PRECHECK(metadata != nullptr, "Table metadata cannot be null"); + ICEBERG_PRECHECK(schema != nullptr, "Schema cannot be null"); + ICEBERG_PRECHECK(io != nullptr, "FileIO cannot be null"); + ICEBERG_PRECHECK(rest_context != nullptr, "REST scan context cannot be null"); + return std::unique_ptr( + new RestIncrementalAppendScan(std::move(metadata), std::move(schema), std::move(io), + std::move(context), std::move(rest_context))); +} + +Result>> RestIncrementalAppendScan::PlanFiles( + std::optional from_snapshot_id_exclusive, + int64_t to_snapshot_id_inclusive) const { + // The REST API requires start-snapshot-id for incremental scans. + if (!from_snapshot_id_exclusive.has_value()) { + return NotSupported("REST incremental append scan requires a start snapshot"); + } + + scan_io_ = nullptr; // reset so stale credentials from a prior plan are not reused + TableMetadataCache metadata_cache(metadata_.get()); + ICEBERG_ASSIGN_OR_RAISE(auto specs, metadata_cache.GetPartitionSpecsById()); + + PlanTableScanRequest request; + ICEBERG_ASSIGN_OR_RAISE(auto projected_schema, schema()); + ICEBERG_ASSIGN_OR_RAISE(request.select, ProjectedColumnNames(*projected_schema)); + request.filter = context_.filter; + request.case_sensitive = context_.case_sensitive; + request.min_rows_requested = context_.min_rows_requested; + request.start_snapshot_id = from_snapshot_id_exclusive; + request.end_snapshot_id = to_snapshot_id_inclusive; + // Match Java's RESTTableScan: schema selection for incremental ranges is + // unspecified in the REST spec. + request.use_snapshot_schema = true; + + if (context_.return_column_stats && context_.columns_to_keep_stats.empty()) { + ICEBERG_ASSIGN_OR_RAISE(request.stats_fields, + ProjectedColumnNames(*schema_, /*include_struct_ids=*/false)); + } else { + for (int32_t field_id : context_.columns_to_keep_stats) { + ICEBERG_ASSIGN_OR_RAISE(auto name, schema_->FindColumnNameById(field_id)); + if (name.has_value()) { + request.stats_fields.emplace_back(*name); + } + } + } + + ICEBERG_ASSIGN_OR_RAISE( + auto stream, + ExecuteScanPlanStream(rest_context_, *schema_, schema_, std::move(request), specs, + scan_io_, context_.ignore_residuals)); + return stream->ToVector(); +} + +const std::shared_ptr& RestIncrementalAppendScan::io() const { + return scan_io_ ? scan_io_ : io_; +} + +// --------------------------------------------------------------------------- +// RestIncrementalAppendScanBuilder +// --------------------------------------------------------------------------- + +RestIncrementalAppendScanBuilder::RestIncrementalAppendScanBuilder( + std::shared_ptr metadata, std::shared_ptr io, + std::string table_name, std::shared_ptr metrics_reporter, + std::shared_ptr rest_context) + : IncrementalAppendScanBuilder(std::move(metadata), std::move(io), + std::move(table_name), std::move(metrics_reporter)), + rest_context_(std::move(rest_context)) {} + +Result> RestIncrementalAppendScanBuilder::Build() { + ICEBERG_RETURN_UNEXPECTED(CheckErrors()); + ICEBERG_RETURN_UNEXPECTED(ResolveColumnStatsSelection()); + ICEBERG_RETURN_UNEXPECTED(context_.Validate()); + ICEBERG_ASSIGN_OR_RAISE(auto schema, ResolveSnapshotSchema()); + return RestIncrementalAppendScan::Make(metadata_, schema.get(), io_, + std::move(context_), rest_context_); +} + +} // namespace iceberg::rest diff --git a/src/iceberg/catalog/rest/rest_table_scan.h b/src/iceberg/catalog/rest/rest_table_scan.h new file mode 100644 index 000000000..c366e10c6 --- /dev/null +++ b/src/iceberg/catalog/rest/rest_table_scan.h @@ -0,0 +1,171 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you 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 +#include +#include +#include +#include +#include +#include +#include + +#include "iceberg/catalog/rest/endpoint.h" +#include "iceberg/catalog/rest/iceberg_rest_export.h" +#include "iceberg/catalog/rest/type_fwd.h" +#include "iceberg/data_table_scan.h" +#include "iceberg/incremental_append_scan.h" +#include "iceberg/result.h" +#include "iceberg/table_identifier.h" +#include "iceberg/table_scan.h" +#include "iceberg/type_fwd.h" + +/// \file iceberg/catalog/rest/rest_table_scan.h +/// REST-specific table scans that delegate scan planning to the REST catalog server. + +namespace iceberg::rest { + +/// \brief REST table context shared by its builders, scans and task streams. +struct ICEBERG_REST_EXPORT RestScanContext { + std::shared_ptr client; + std::shared_ptr paths; + std::shared_ptr session; + std::unordered_set supported_endpoints; + TableIdentifier identifier; + /// Catalog-level config, used with table_config to build a scan-scoped FileIO + /// when the server vends storage credentials in a planning response. + std::unordered_map catalog_config; + /// Table-level config merged with catalog_config for scan-scoped FileIO creation. + std::unordered_map table_config; + /// Builds a plan-scoped provider for refreshing vended credentials; may be empty. + /// A null provider means credential refresh is unavailable. + std::function>( + const std::string& plan_id)> + credential_provider_factory; +}; + +/// \brief A DataTableScan that delegates PlanFilesStream() to the REST catalog server +/// via the scan planning endpoints (planTableScan / fetchPlanningResult / +/// cancelPlanning / fetchScanTasks). +class ICEBERG_REST_EXPORT RestTableScan : public DataTableScan { + public: + ~RestTableScan() override = default; + + /// \brief Create a point-in-time scan. Incremental ranges must use + /// RestIncrementalAppendScan. + static Result> Make( + std::shared_ptr metadata, std::shared_ptr schema, + std::shared_ptr io, internal::TableScanContext context, + std::shared_ptr rest_context); + + /// \brief Plans files lazily via the REST scan planning endpoints. + /// + /// Plan tasks are fetched serially; an executor set via PlanWith() is not used. + Result PlanFilesStream() const override; + + /// \brief Returns the effective FileIO for reading scan results. + /// + /// If the server vended storage credentials during planning, returns a FileIO + /// initialised with those credentials; otherwise returns the table's FileIO. + const std::shared_ptr& io() const override; + + private: + RestTableScan(std::shared_ptr metadata, std::shared_ptr schema, + std::shared_ptr io, internal::TableScanContext context, + std::shared_ptr rest_context); + + std::shared_ptr rest_context_; + mutable std::shared_ptr scan_io_; +}; + +/// \brief Builder that produces a RestTableScan with the REST HTTP context injected. +class ICEBERG_REST_EXPORT RestTableScanBuilder : public DataTableScanBuilder { + public: + RestTableScanBuilder(std::shared_ptr metadata, + std::shared_ptr io, std::string table_name, + std::shared_ptr metrics_reporter, + std::shared_ptr rest_context); + + /// \brief Resolves schema/context via parent logic then creates a RestTableScan. + Result> Build() override; + + private: + std::shared_ptr rest_context_; +}; + +/// \brief An IncrementalAppendScan that delegates PlanFiles() to the REST catalog server. +/// +/// Scans without an exclusive start snapshot, including an inclusive root snapshot, +/// return NotSupported. +/// +/// IncrementalChangelogScan is not delegated because the REST planTableScan response +/// only carries FileScanTask objects; reconstructing ChangelogScanTask entries +/// (AddedRowsScanTask vs DeletedDataFileScanTask) requires per-snapshot operation +/// metadata that the server does not return. Changelog scans always plan locally. +class ICEBERG_REST_EXPORT RestIncrementalAppendScan : public IncrementalAppendScan { + public: + ~RestIncrementalAppendScan() override = default; + + static Result> Make( + std::shared_ptr metadata, std::shared_ptr schema, + std::shared_ptr io, internal::TableScanContext context, + std::shared_ptr rest_context); + + using IncrementalAppendScan::PlanFiles; + + /// \brief Returns the FileIO with vended credentials if the server returned any + /// during planning; otherwise returns the table's FileIO. + const std::shared_ptr& io() const override; + + protected: + /// \brief Plans files via the REST scan planning endpoints. + /// + /// Plan tasks are fetched serially; an executor set via PlanWith() is not used. + Result>> PlanFiles( + std::optional from_snapshot_id_exclusive, + int64_t to_snapshot_id_inclusive) const override; + + private: + RestIncrementalAppendScan(std::shared_ptr metadata, + std::shared_ptr schema, std::shared_ptr io, + internal::TableScanContext context, + std::shared_ptr rest_context); + + std::shared_ptr rest_context_; + mutable std::shared_ptr scan_io_; +}; + +/// \brief Builder that produces a RestIncrementalAppendScan. +class ICEBERG_REST_EXPORT RestIncrementalAppendScanBuilder + : public IncrementalAppendScanBuilder { + public: + RestIncrementalAppendScanBuilder(std::shared_ptr metadata, + std::shared_ptr io, std::string table_name, + std::shared_ptr metrics_reporter, + std::shared_ptr rest_context); + + Result> Build() override; + + private: + std::shared_ptr rest_context_; +}; + +} // namespace iceberg::rest diff --git a/src/iceberg/catalog/rest/type_fwd.h b/src/iceberg/catalog/rest/type_fwd.h index ee684b245..ba8aa5196 100644 --- a/src/iceberg/catalog/rest/type_fwd.h +++ b/src/iceberg/catalog/rest/type_fwd.h @@ -28,6 +28,7 @@ struct ErrorResponse; struct CommitTableResponse; struct LoadTableResult; struct OAuthTokenResponse; +struct RestScanContext; class Endpoint; class ErrorHandler; diff --git a/src/iceberg/catalog/rest/types.cc b/src/iceberg/catalog/rest/types.cc index 84fba9a7c..27cd7b170 100644 --- a/src/iceberg/catalog/rest/types.cc +++ b/src/iceberg/catalog/rest/types.cc @@ -229,13 +229,14 @@ bool HasNonEmptyFileScanTasks(const Response& response) { bool PlanTableScanResponse::operator==(const PlanTableScanResponse& other) const { return ScanTaskFieldsEqual(*this, other) && plan_status == other.plan_status && - plan_id == other.plan_id && error == other.error; + plan_id == other.plan_id && error == other.error && + storage_credentials == other.storage_credentials; } bool FetchPlanningResultResponse::operator==( const FetchPlanningResultResponse& other) const { return ScanTaskFieldsEqual(*this, other) && plan_status == other.plan_status && - error == other.error; + error == other.error && storage_credentials == other.storage_credentials; } bool FetchScanTasksRequest::operator==(const FetchScanTasksRequest& other) const { diff --git a/src/iceberg/catalog/rest/types.h b/src/iceberg/catalog/rest/types.h index a47a8e851..f048d6e7d 100644 --- a/src/iceberg/catalog/rest/types.h +++ b/src/iceberg/catalog/rest/types.h @@ -352,7 +352,7 @@ struct ICEBERG_REST_EXPORT PlanTableScanResponse { PlanStatus plan_status = PlanStatus::kCompleted; std::string plan_id; std::optional error; - // TODO(sandeepg): Add storage credentials and bind scan FileIO to them. + std::vector storage_credentials; Status Validate() const; @@ -367,7 +367,7 @@ struct ICEBERG_REST_EXPORT FetchPlanningResultResponse { std::vector> delete_files; PlanStatus plan_status = PlanStatus::kCompleted; std::optional error; - // TODO(sandeepg): Add storage credentials and bind scan FileIO to them. + std::vector storage_credentials; Status Validate() const; diff --git a/src/iceberg/cpp20_compatibility_internal.h b/src/iceberg/cpp20_compatibility_internal.h index b5f46e488..3f9b5b577 100644 --- a/src/iceberg/cpp20_compatibility_internal.h +++ b/src/iceberg/cpp20_compatibility_internal.h @@ -240,6 +240,8 @@ # include "iceberg/catalog/rest/resource_paths.h" # include "iceberg/catalog/rest/rest_catalog.h" # include "iceberg/catalog/rest/rest_file_io.h" +# include "iceberg/catalog/rest/rest_table.h" +# include "iceberg/catalog/rest/rest_table_scan.h" # include "iceberg/catalog/rest/rest_util.h" # include "iceberg/catalog/rest/type_fwd.h" # include "iceberg/catalog/rest/types.h" diff --git a/src/iceberg/data_table_scan.h b/src/iceberg/data_table_scan.h index a6de0020c..96804d953 100644 --- a/src/iceberg/data_table_scan.h +++ b/src/iceberg/data_table_scan.h @@ -54,7 +54,7 @@ class ICEBERG_EXPORT DataTableScan : public TableScan { /// can outlive this scan. An executor configured through PlanWith() is borrowed and /// must remain alive until the stream is destroyed, as later Next() calls may submit /// work to it. - Result PlanFilesStream() const; + virtual Result PlanFilesStream() const; protected: using TableScan::TableScan; diff --git a/src/iceberg/expression/expression.h b/src/iceberg/expression/expression.h index 1663c69f8..43e4a2c1f 100644 --- a/src/iceberg/expression/expression.h +++ b/src/iceberg/expression/expression.h @@ -64,7 +64,7 @@ class ICEBERG_EXPORT Expression : public util::Formattable { kMin }; - virtual ~Expression() = default; + ~Expression() override = default; /// \brief Returns the operation for an expression node. virtual Operation op() const = 0; diff --git a/src/iceberg/table_scan.cc b/src/iceberg/table_scan.cc index 46f235fbe..636834d62 100644 --- a/src/iceberg/table_scan.cc +++ b/src/iceberg/table_scan.cc @@ -316,6 +316,7 @@ TableScanBuilder& TableScanBuilder::UseSnapshot(int64_t snap context_.snapshot_id.value()); ICEBERG_BUILDER_ASSIGN_OR_RETURN(std::ignore, metadata_->SnapshotById(snapshot_id)); context_.snapshot_id = snapshot_id; + context_.use_snapshot_schema = true; return *this; } @@ -323,6 +324,7 @@ template TableScanBuilder& TableScanBuilder::UseRef(const std::string& ref) { if (ref == SnapshotRef::kMainBranch) { context_.snapshot_id.reset(); + context_.use_snapshot_schema = false; return *this; } @@ -335,6 +337,7 @@ TableScanBuilder& TableScanBuilder::UseRef(const std::string const int64_t snapshot_id = iter->second->snapshot_id; ICEBERG_BUILDER_ASSIGN_OR_RETURN(std::ignore, metadata_->SnapshotById(snapshot_id)); context_.snapshot_id = snapshot_id; + context_.use_snapshot_schema = (iter->second->type() == SnapshotRefType::kTag); return *this; } @@ -434,7 +437,7 @@ template Result>> TableScanBuilder::ResolveSnapshotSchema() { if (snapshot_schema_ == nullptr) { - if (context_.snapshot_id.has_value()) { + if (context_.snapshot_id.has_value() && context_.use_snapshot_schema) { ICEBERG_ASSIGN_OR_RAISE(auto snapshot, metadata_->SnapshotById(*context_.snapshot_id)); int32_t schema_id = snapshot->schema_id.value_or(Schema::kInitialSchemaId); diff --git a/src/iceberg/table_scan.h b/src/iceberg/table_scan.h index 6ec29ad9f..abe93bbff 100644 --- a/src/iceberg/table_scan.h +++ b/src/iceberg/table_scan.h @@ -219,7 +219,7 @@ class ICEBERG_EXPORT DeletedDataFileScanTask : public ChangelogScanTask { namespace internal { // Internal table scan context used by different scan implementations. -struct TableScanContext { +struct ICEBERG_EXPORT TableScanContext { std::optional snapshot_id; std::shared_ptr filter; bool ignore_residuals{false}; @@ -234,6 +234,7 @@ struct TableScanContext { std::optional to_snapshot_id; std::string branch{}; std::optional min_rows_requested; + bool use_snapshot_schema{false}; OptionalExecutor plan_executor; std::string table_name; std::shared_ptr metrics_reporter; @@ -402,14 +403,17 @@ class ICEBERG_TEMPLATE_CLASS_EXPORT TableScanBuilder : public ErrorCollector { /// \brief Builds and returns a TableScan instance. /// \return A Result containing the TableScan or an error. - Result> Build(); + virtual Result> Build(); protected: TableScanBuilder(std::shared_ptr metadata, std::shared_ptr io, std::string table_name, std::shared_ptr metrics_reporter); - // Return the schema bound to the specified snapshot. + TableScanBuilder(TableScanBuilder&&) = default; + TableScanBuilder& operator=(TableScanBuilder&&) = default; + + // Return the snapshot schema for tags and time travel, else the table schema. Result>> ResolveSnapshotSchema(); Status ResolveColumnStatsSelection(); @@ -438,7 +442,7 @@ class ICEBERG_EXPORT TableScan { const internal::TableScanContext& context() const; /// \brief Returns the file I/O instance used for reading files. - const std::shared_ptr& io() const; + virtual const std::shared_ptr& io() const; /// \brief Returns this scan's filter expression. const std::shared_ptr& filter() const; diff --git a/src/iceberg/test/CMakeLists.txt b/src/iceberg/test/CMakeLists.txt index 058a9dd3c..659db4f61 100644 --- a/src/iceberg/test/CMakeLists.txt +++ b/src/iceberg/test/CMakeLists.txt @@ -321,11 +321,14 @@ if(ICEBERG_BUILD_REST) add_rest_iceberg_test(rest_catalog_test SOURCES auth_manager_test.cc + catalog_properties_test.cc error_handlers_test.cc endpoint_test.cc + rest_catalog_unit_test.cc rest_file_io_test.cc rest_json_serde_test.cc rest_metrics_reporter_test.cc + rest_table_scan_test.cc rest_util_test.cc) if(ICEBERG_SIGV4) diff --git a/src/iceberg/test/auth_manager_test.cc b/src/iceberg/test/auth_manager_test.cc index 61962d1c3..fd9605fc3 100644 --- a/src/iceberg/test/auth_manager_test.cc +++ b/src/iceberg/test/auth_manager_test.cc @@ -788,7 +788,7 @@ TEST(OAuth2SessionTest, InitialTokenIsUsed) { ASSERT_THAT(auth_result, IsOk()); EXPECT_EQ(auth_result.value().headers.at("Authorization"), "Bearer initial-token-123"); - session->Close(); + ASSERT_THAT(session->Close(), IsOk()); } TEST(OAuth2SessionTest, InitTokenExpirationUsesRequestStartTime) { diff --git a/src/iceberg/test/catalog_properties_test.cc b/src/iceberg/test/catalog_properties_test.cc new file mode 100644 index 000000000..0100c1488 --- /dev/null +++ b/src/iceberg/test/catalog_properties_test.cc @@ -0,0 +1,60 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you 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. + */ + +#include "iceberg/catalog/rest/catalog_properties.h" + +#include +#include + +#include + +#include "iceberg/test/matchers.h" + +namespace iceberg::rest { + +TEST(ScanPlanningModeTest, MissingKeyReturnsNullopt) { + std::unordered_map config; + auto result = RestCatalogProperties::ScanPlanningModeFrom(config); + ASSERT_THAT(result, IsOk()); + EXPECT_FALSE(result->has_value()); +} + +TEST(ScanPlanningModeTest, ParsesExplicitModes) { + const std::array, 4> cases{ + {{"client", ScanPlanningMode::kClient}, + {"CLIENT", ScanPlanningMode::kClient}, + {"server", ScanPlanningMode::kServer}, + {"SERVER", ScanPlanningMode::kServer}}}; + for (const auto& [value, expected] : cases) { + SCOPED_TRACE(value); + auto result = + RestCatalogProperties::ScanPlanningModeFrom({{"scan-planning-mode", value}}); + ASSERT_THAT(result, IsOk()); + ASSERT_TRUE(result->has_value()); + EXPECT_EQ(**result, expected); + } +} + +TEST(ScanPlanningModeTest, InvalidValueReturnsError) { + auto result = + RestCatalogProperties::ScanPlanningModeFrom({{"scan-planning-mode", "invalid"}}); + EXPECT_THAT(result, IsError(ErrorKind::kInvalidArgument)); +} + +} // namespace iceberg::rest diff --git a/src/iceberg/test/metrics_reporter_test.cc b/src/iceberg/test/metrics_reporter_test.cc index 385e20c23..ecddde07e 100644 --- a/src/iceberg/test/metrics_reporter_test.cc +++ b/src/iceberg/test/metrics_reporter_test.cc @@ -31,6 +31,7 @@ #include "iceberg/metrics/counter.h" #include "iceberg/metrics/metrics_reporters.h" +#include "iceberg/test/matchers.h" namespace iceberg { @@ -75,7 +76,7 @@ TEST(CustomMetricsReporterTest, RegisterAndLoad) { ASSERT_NE(reporter, nullptr); ScanReport scan_report{.table_name = "test.table"}; - reporter->Report(scan_report); + ASSERT_THAT(reporter->Report(scan_report), IsOk()); EXPECT_EQ(reporter->reports().size(), 1); EXPECT_EQ(GetReportType(reporter->reports()[0]), MetricsReportType::kScanReport); @@ -163,7 +164,7 @@ class CollectorTest : public ::testing::TestWithParam {}; TEST_P(CollectorTest, CollectsAndPreservesReport) { const auto& param = GetParam(); CollectingMetricsReporter reporter; - reporter.Report(param.report); + ASSERT_THAT(reporter.Report(param.report), IsOk()); ASSERT_EQ(reporter.reports().size(), 1); EXPECT_EQ(GetReportType(reporter.reports()[0]), param.expected_type); @@ -228,7 +229,7 @@ TEST(CombineTest, FlattenNestedComposite) { EXPECT_EQ(dynamic_cast(r.get()), nullptr); } - abc->Report(CommitReport{.table_name = "db.t2"}); + ASSERT_THAT(abc->Report(CommitReport{.table_name = "db.t2"}), IsOk()); EXPECT_EQ(a->reports().size(), 1u); EXPECT_EQ(b->reports().size(), 1u); EXPECT_EQ(c->reports().size(), 1u); @@ -246,7 +247,7 @@ TEST(CombineTest, DeduplicateByIdentity) { ASSERT_NE(composite, nullptr); EXPECT_EQ(composite->Reporters().size(), 2u); - result->Report(ScanReport{}); + ASSERT_THAT(result->Report(ScanReport{}), IsOk()); EXPECT_EQ(a->reports().size(), 1u); EXPECT_EQ(b->reports().size(), 1u); // delivered once, not twice } diff --git a/src/iceberg/test/rest_catalog_unit_test.cc b/src/iceberg/test/rest_catalog_unit_test.cc new file mode 100644 index 000000000..bb76a59a4 --- /dev/null +++ b/src/iceberg/test/rest_catalog_unit_test.cc @@ -0,0 +1,491 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you 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. + */ + +#include +#include +#include +#include +#include + +#include +#include +#include + +#include "iceberg/catalog.h" +#include "iceberg/catalog/rest/catalog_properties.h" +#include "iceberg/catalog/rest/endpoint.h" +#include "iceberg/catalog/rest/error_handlers.h" +#include "iceberg/catalog/rest/http_client.h" +#include "iceberg/catalog/rest/json_serde_internal.h" +#include "iceberg/catalog/rest/rest_catalog.h" +#include "iceberg/catalog/rest/rest_table.h" +#include "iceberg/catalog/rest/rest_table_scan.h" +#include "iceberg/catalog/rest/types.h" +#include "iceberg/file_io.h" +#include "iceberg/file_io_registry.h" +#include "iceberg/storage_credential.h" +#include "iceberg/table.h" +#include "iceberg/table_identifier.h" +#include "iceberg/table_requirement.h" +#include "iceberg/table_update.h" +#include "iceberg/test/matchers.h" + +namespace iceberg::rest { + +namespace { + +using ::testing::_; +using ::testing::Return; + +// -------------------------------------------------------------------------- +// Mock HTTP client (same pattern as rest_table_scan_test.cc) +// -------------------------------------------------------------------------- +class MockHttpClient : public HttpClient { + public: + MockHttpClient() : HttpClient({}) {} + + MOCK_METHOD(Result, Get, + (const std::string& path, + (const std::unordered_map&)params, + (const std::unordered_map&)headers, + const ErrorHandler& error_handler, auth::AuthSession& session), + (override)); + + MOCK_METHOD(Result, Post, + (const std::string& path, const std::string& body, + (const std::unordered_map&)headers, + const ErrorHandler& error_handler, auth::AuthSession& session), + (override)); + + MOCK_METHOD(Result, Delete, + (const std::string& path, + (const std::unordered_map&)params, + (const std::unordered_map&)headers, + const ErrorHandler& error_handler, auth::AuthSession& session), + (override)); +}; + +// -------------------------------------------------------------------------- +// Minimal FileIO stub +// -------------------------------------------------------------------------- +class NoOpFileIO : public FileIO { + public: + Result ReadFile(const std::string&, std::optional) override { + return IOError("NoOpFileIO"); + } + Status WriteFile(const std::string&, std::string_view) override { return {}; } + Status DeleteFile(const std::string&) override { return {}; } +}; + +namespace { + +std::shared_ptr installed_provider; + +class CredentialedFileIO : public NoOpFileIO, public SupportsStorageCredentials { + public: + Status InitializeStorageCredentials( + const std::vector& credentials, + std::shared_ptr provider) override { + credentials_ = credentials; + installed_provider = std::move(provider); + return {}; + } + + Status SetStorageCredentials( + const std::vector& credentials) override { + credentials_ = credentials; + return {}; + } + + std::vector credentials() const override { return credentials_; } + + SupportsStorageCredentials* AsSupportsStorageCredentials() override { return this; } + + private: + std::vector credentials_; +}; + +constexpr std::string_view kCredentialedFileIO = "rest-catalog-unit-test-credentialed"; + +void RegisterCredentialedFileIO() { + FileIORegistry::Register( + std::string(kCredentialedFileIO), + {.create = [](const std::unordered_map&) + -> Result> { + return std::make_unique(); + }}); + installed_provider = nullptr; +} + +// Plan response that vends credentials for the scan. +constexpr std::string_view kPlanWithCredentialsResponse = R"({ + "status": "completed", + "plan-id": "plan-1", + "storage-credentials": [{"prefix": "s3://bucket", "config": {"k": "v"}}] +})"; + +} // namespace + +// -------------------------------------------------------------------------- +// Helper JSON bodies for LoadTable GET responses +// -------------------------------------------------------------------------- + +// Minimal table metadata blob (no snapshots, no partitions). +constexpr std::string_view kBaseMetadataJson = + R"("metadata-location":"s3://bucket/metadata/v1.json","metadata":{"format-version":2,"table-uuid":"test-uuid","location":"s3://bucket/test","last-sequence-number":0,"last-updated-ms":0,"last-column-id":1,"schemas":[{"type":"struct","schema-id":1,"fields":[{"id":1,"name":"id","type":"int","required":true}]}],"current-schema-id":1,"partition-specs":[{"spec-id":0,"fields":[]}],"default-spec-id":0,"last-partition-id":0,"sort-orders":[{"order-id":0,"fields":[]}],"default-sort-order-id":0,"properties":{}})"; + +// LoadTable response where server config requests server-side scan planning. +const std::string kLoadTableServerScanResponse = + std::string(R"({"config":{"scan-planning-mode":"server"},)") + + std::string(kBaseMetadataJson) + "}"; + +// LoadTable response with no scan-planning-mode set (defaults to client). +const std::string kLoadTableDefaultResponse = + std::string("{") + std::string(kBaseMetadataJson) + "}"; + +// UpdateTable (commit) response. +const std::string kCommitTableResponse = + std::string("{") + std::string(kBaseMetadataJson) + "}"; + +// -------------------------------------------------------------------------- +// Use the regular catalog initialization with a mock HTTP client. +// -------------------------------------------------------------------------- +class RestCatalogLoadTableTest : public ::testing::Test { + protected: + void TearDown() override { installed_provider.reset(); } + + void SetUp() override { + mock_client_ = std::make_shared(); + FileIORegistry::Register( + "rest-catalog-unit-test-noop", + {.create = [](const std::unordered_map&) + -> Result> { + return std::make_unique(); + }}); + + identifier_ = TableIdentifier{.ns = Namespace{{"default"}}, .name = "my_table"}; + + all_plan_endpoints_ = {Endpoint::LoadTable(), Endpoint::UpdateTable(), + Endpoint::PlanTableScan(), Endpoint::FetchPlanningResult(), + Endpoint::CancelPlanning(), Endpoint::FetchScanTasks()}; + + no_plan_endpoint_set_ = {Endpoint::LoadTable()}; + } + + // Builds a RestCatalog with the given supported endpoints and an optional + // client-side scan-planning-mode setting. + Result> MakeCatalog( + const std::unordered_set& endpoints, + const std::string& client_scan_mode = "", + std::unordered_map extra_props = {}) { + std::unordered_map props = std::move(extra_props); + props["uri"] = "http://test-server"; + props.try_emplace("io-impl", "rest-catalog-unit-test-noop"); + if (!client_scan_mode.empty()) { + props["scan-planning-mode"] = client_scan_mode; + } + CatalogConfig server_config{.endpoints = {endpoints.begin(), endpoints.end()}}; + EXPECT_CALL(*mock_client_, Get(::testing::EndsWith("/v1/config"), _, _, _, _)) + .WillOnce(Return(HttpResponse(200, ToJson(server_config).dump()))); + auto config = RestCatalogProperties::FromMap(std::move(props)); + return RestCatalog::Make(config, /*metrics_executor=*/nullptr, + [client = mock_client_](const auto&) { return client; }); + } + + std::shared_ptr mock_client_; + TableIdentifier identifier_; + std::unordered_set all_plan_endpoints_; + std::unordered_set no_plan_endpoint_set_; +}; + +TEST_F(RestCatalogLoadTableTest, InitializationPropagatesClientFactoryErrors) { + auto config = RestCatalogProperties::FromMap( + {{"uri", "http://test-server"}, {"io-impl", "rest-catalog-unit-test-noop"}}); + for (int fail_on_call : {1, 2}) { + SCOPED_TRACE(fail_on_call); + if (fail_on_call == 1) { + EXPECT_CALL(*mock_client_, Get(_, _, _, _, _)).Times(0); + } else { + EXPECT_CALL(*mock_client_, Get(::testing::EndsWith("/v1/config"), _, _, _, _)) + .WillOnce(Return(HttpResponse(200, ToJson(CatalogConfig{}).dump()))); + } + int calls = 0; + auto result = RestCatalog::Make( + config, /*metrics_executor=*/nullptr, + [&](const RestCatalogProperties&) -> Result> { + if (++calls == fail_on_call) { + return IOError("HTTP client setup failed"); + } + return mock_client_; + }); + ASSERT_THAT(result, IsError(ErrorKind::kIOError)); + EXPECT_EQ(result.error().message, "HTTP client setup failed"); + EXPECT_EQ(calls, fail_on_call); + ASSERT_TRUE(::testing::Mock::VerifyAndClearExpectations(mock_client_.get())); + } +} + +TEST_F(RestCatalogLoadTableTest, InitializationUsesMergedServerConfig) { + CatalogConfig server_config{ + .defaults = {{"header.X-Default", "default"}, {"header.X-Shared", "default"}}, + .overrides = {{"prefix", "catalog"}, + {"snapshot-loading-mode", "REFS"}, + {"scan-planning-mode", "server"}, + {"header.X-Shared", "server"}}, + .endpoints = {all_plan_endpoints_.begin(), all_plan_endpoints_.end()}}; + EXPECT_CALL( + *mock_client_, + Get(::testing::EndsWith("/v1/config"), + ::testing::Contains(::testing::Pair("warehouse", "client-warehouse")), _, _, _)) + .WillOnce(Return(HttpResponse(200, ToJson(server_config).dump()))); + EXPECT_CALL(*mock_client_, + Get(::testing::EndsWith("/v1/catalog/namespaces/default/tables/my_table"), + ::testing::Contains(::testing::Pair("snapshots", "refs")), _, _, _)) + .WillOnce(Return(HttpResponse(200, kLoadTableDefaultResponse))); + EXPECT_CALL( + *mock_client_, + Post(::testing::EndsWith("/v1/catalog/namespaces/default/tables/my_table/plan"), _, + _, _, _)) + .WillOnce(Return(HttpResponse(200, R"({"status":"completed"})"))); + + auto config = + RestCatalogProperties::FromMap({{"uri", "http://test-server/"}, + {"warehouse", "client-warehouse"}, + {"io-impl", "rest-catalog-unit-test-noop"}, + {"scan-planning-mode", "client"}, + {"header.X-Shared", "client"}}); + std::vector client_configs; + ICEBERG_UNWRAP_OR_FAIL(auto root, + RestCatalog::Make(config, /*metrics_executor=*/nullptr, + [&](const RestCatalogProperties& properties) { + client_configs.push_back(properties); + return mock_client_; + })); + ASSERT_EQ(client_configs.size(), 2); + EXPECT_EQ(client_configs[0].ExtractHeaders().at("X-Shared"), "client"); + EXPECT_FALSE(client_configs[0].ExtractHeaders().contains("X-Default")); + EXPECT_EQ(client_configs[1].ExtractHeaders().at("X-Shared"), "server"); + EXPECT_EQ(client_configs[1].ExtractHeaders().at("X-Default"), "default"); + EXPECT_EQ(client_configs[1].Get(RestCatalogProperties::kPrefix), "catalog"); + ICEBERG_UNWRAP_OR_FAIL(auto catalog, root->AsCatalog()); + ICEBERG_UNWRAP_OR_FAIL(auto table, catalog->LoadTable(identifier_)); + ICEBERG_UNWRAP_OR_FAIL(auto builder, table->NewScan()); + ICEBERG_UNWRAP_OR_FAIL(auto scan, builder->Build()); + ICEBERG_UNWRAP_OR_FAIL(auto tasks, scan->PlanFiles()); + EXPECT_TRUE(tasks.empty()); +} + +// -------------------------------------------------------------------------- +// "scan-planning-mode":"server" but PlanTableScan endpoint missing → +// LoadTable returns NotSupported. +// -------------------------------------------------------------------------- +TEST_F(RestCatalogLoadTableTest, ServerScanWithMissingEndpointReturnsNotSupported) { + EXPECT_CALL(*mock_client_, Get(_, _, _, _, _)) + .WillOnce(Return(HttpResponse(200, kLoadTableServerScanResponse))); + + ICEBERG_UNWRAP_OR_FAIL(auto catalog, MakeCatalog(no_plan_endpoint_set_)); + ICEBERG_UNWRAP_OR_FAIL(auto as_catalog, catalog->AsCatalog()); + auto result = as_catalog->LoadTable(identifier_); + + EXPECT_THAT(result, IsError(ErrorKind::kNotSupported)); +} + +// -------------------------------------------------------------------------- +// A completed plan without a plan-id refreshes credentials without planId. +// -------------------------------------------------------------------------- +TEST_F(RestCatalogLoadTableTest, ScanCredentialProviderWithoutPlanIdOmitsPlanId) { + RegisterCredentialedFileIO(); + EXPECT_CALL(*mock_client_, + Get(::testing::Not(::testing::HasSubstr("/credentials")), _, _, _, _)) + .WillOnce(Return(HttpResponse(200, kLoadTableServerScanResponse))); + EXPECT_CALL(*mock_client_, Post(_, _, _, _, _)) + .WillOnce(Return(HttpResponse( + 200, + R"({"status":"completed","storage-credentials":[{"prefix":"s3://bucket","config":{"k":"v"}}]})"))); + EXPECT_CALL(*mock_client_, + Get(::testing::HasSubstr("/credentials"), ::testing::IsEmpty(), _, _, _)) + .WillOnce(Return(HttpResponse( + 200, + R"({"storage-credentials":[{"prefix":"s3://bucket","config":{"k":"v2"}}]})"))); + + auto endpoints = all_plan_endpoints_; + endpoints.insert(Endpoint::TableCredentials()); + ICEBERG_UNWRAP_OR_FAIL(auto catalog, + MakeCatalog(endpoints, /*client_scan_mode=*/"", + {{"io-impl", std::string(kCredentialedFileIO)}})); + ICEBERG_UNWRAP_OR_FAIL(auto as_catalog, catalog->AsCatalog()); + ICEBERG_UNWRAP_OR_FAIL(auto table, as_catalog->LoadTable(identifier_)); + ICEBERG_UNWRAP_OR_FAIL(auto builder, table->NewScan()); + ICEBERG_UNWRAP_OR_FAIL(auto scan, builder->Build()); + ICEBERG_UNWRAP_OR_FAIL(auto tasks, scan->PlanFiles()); + + ASSERT_NE(installed_provider, nullptr); + ICEBERG_UNWRAP_OR_FAIL(auto refreshed, installed_provider->Load()); + ASSERT_EQ(refreshed.size(), 1); +} + +// -------------------------------------------------------------------------- +// Without the TableCredentials endpoint, scan credentials are not refreshed. +// -------------------------------------------------------------------------- +TEST_F(RestCatalogLoadTableTest, + ScanCredentialsWithoutCredentialsEndpointHaveNoProvider) { + RegisterCredentialedFileIO(); + EXPECT_CALL(*mock_client_, Get(_, _, _, _, _)) + .WillOnce(Return(HttpResponse(200, kLoadTableServerScanResponse))); + EXPECT_CALL(*mock_client_, Post(_, _, _, _, _)) + .WillOnce(Return(HttpResponse(200, std::string(kPlanWithCredentialsResponse)))); + + ICEBERG_UNWRAP_OR_FAIL(auto catalog, + MakeCatalog(all_plan_endpoints_, /*client_scan_mode=*/"", + {{"io-impl", std::string(kCredentialedFileIO)}})); + ICEBERG_UNWRAP_OR_FAIL(auto as_catalog, catalog->AsCatalog()); + ICEBERG_UNWRAP_OR_FAIL(auto table, as_catalog->LoadTable(identifier_)); + ICEBERG_UNWRAP_OR_FAIL(auto builder, table->NewScan()); + ICEBERG_UNWRAP_OR_FAIL(auto scan, builder->Build()); + ICEBERG_UNWRAP_OR_FAIL(auto tasks, scan->PlanFiles()); + + EXPECT_EQ(installed_provider, nullptr); + auto* credentialed = scan->io()->AsSupportsStorageCredentials(); + ASSERT_NE(credentialed, nullptr); + ASSERT_EQ(credentialed->credentials().size(), 1); + EXPECT_EQ(credentialed->credentials()[0].config.at("k"), "v"); +} + +struct PlanningModeCase { + std::string client; + std::string table; + bool server; +}; + +class RestCatalogPlanningModeTest + : public RestCatalogLoadTableTest, + public ::testing::WithParamInterface {}; + +TEST_P(RestCatalogPlanningModeTest, LoadAndCommitKeepTheEffectivePlanningMode) { + const auto& config = GetParam(); + auto response = nlohmann::json::parse(kLoadTableDefaultResponse); + if (!config.table.empty()) response["config"]["scan-planning-mode"] = config.table; + EXPECT_CALL(*mock_client_, Get(::testing::EndsWith("/tables/my_table"), _, _, _, _)) + .WillOnce(Return(HttpResponse(200, response.dump()))); + EXPECT_CALL(*mock_client_, Post(::testing::EndsWith("/tables/my_table"), _, _, _, _)) + .WillOnce(Return(HttpResponse(200, kCommitTableResponse))); + EXPECT_CALL(*mock_client_, Post(::testing::EndsWith("/plan"), _, _, _, _)) + .Times(config.server ? 2 : 0) + .WillRepeatedly([](auto&&...) -> Result { + return HttpResponse(200, R"({"status":"completed"})"); + }); + + ICEBERG_UNWRAP_OR_FAIL(auto catalog, MakeCatalog(all_plan_endpoints_, config.client)); + ICEBERG_UNWRAP_OR_FAIL(auto as_catalog, catalog->AsCatalog()); + ICEBERG_UNWRAP_OR_FAIL(auto table, as_catalog->LoadTable(identifier_)); + ICEBERG_UNWRAP_OR_FAIL(auto committed, + table->catalog()->UpdateTable(identifier_, {}, {})); + for (const auto& current : {table, committed}) { + ICEBERG_UNWRAP_OR_FAIL(auto builder, current->NewScan()); + ICEBERG_UNWRAP_OR_FAIL(auto scan, builder->Build()); + ICEBERG_UNWRAP_OR_FAIL(auto tasks, scan->PlanFiles()); + EXPECT_TRUE(tasks.empty()); + } +} + +INSTANTIATE_TEST_SUITE_P( + PlanningModes, RestCatalogPlanningModeTest, + ::testing::Values( + PlanningModeCase{.client = "", .table = "", .server = false}, + PlanningModeCase{.client = "server", .table = "", .server = true}, + PlanningModeCase{.client = "", .table = "server", .server = true}, + PlanningModeCase{.client = "server", .table = "client", .server = false}, + PlanningModeCase{.client = "client", .table = "server", .server = true})); + +class RestCatalogScanLifetimeTest : public RestCatalogLoadTableTest, + public ::testing::WithParamInterface {}; + +TEST_P(RestCatalogScanLifetimeTest, ScanKeepsCredentialRefreshAlive) { + RegisterCredentialedFileIO(); + auto response = nlohmann::json::parse(kLoadTableServerScanResponse); + auto& metadata = response["metadata"]; + metadata["current-snapshot-id"] = 200; + metadata["last-sequence-number"] = 2; + metadata["snapshots"] = + nlohmann::json::array({{{"snapshot-id", 100}, + {"sequence-number", 1}, + {"timestamp-ms", 1}, + {"manifest-list", "s3://bucket/manifests-1.avro"}, + {"schema-id", 1}}, + {{"snapshot-id", 200}, + {"parent-snapshot-id", 100}, + {"sequence-number", 2}, + {"timestamp-ms", 2}, + {"manifest-list", "s3://bucket/manifests-2.avro"}, + {"schema-id", 1}}}); + EXPECT_CALL(*mock_client_, Get(::testing::EndsWith("/tables/my_table"), _, _, _, _)) + .WillOnce(Return(HttpResponse(200, response.dump()))); + EXPECT_CALL(*mock_client_, Post(::testing::EndsWith("/plan"), _, _, _, _)) + .WillOnce(Return(HttpResponse(200, std::string(kPlanWithCredentialsResponse)))); + EXPECT_CALL(*mock_client_, + Get(::testing::EndsWith("/credentials"), + ::testing::Contains(::testing::Pair("planId", "plan-1")), _, _, _)) + .WillOnce(Return(HttpResponse( + 200, + R"({"storage-credentials":[{"prefix":"s3://bucket","config":{"k":"v2"}}]})"))); + + auto endpoints = all_plan_endpoints_; + endpoints.insert(Endpoint::TableCredentials()); + ICEBERG_UNWRAP_OR_FAIL( + auto catalog, + MakeCatalog(endpoints, "", {{"io-impl", std::string(kCredentialedFileIO)}})); + std::weak_ptr weak_catalog = catalog; + ICEBERG_UNWRAP_OR_FAIL(auto as_catalog, catalog->AsCatalog()); + ICEBERG_UNWRAP_OR_FAIL(auto table, as_catalog->LoadTable(identifier_)); + auto exercise = [&](auto& builder) { + ICEBERG_UNWRAP_OR_FAIL(auto scan, builder->Build()); + builder.reset(); + table.reset(); + as_catalog.reset(); + catalog.reset(); + EXPECT_FALSE(weak_catalog.expired()); + + ICEBERG_UNWRAP_OR_FAIL(auto tasks, scan->PlanFiles()); + ASSERT_NE(installed_provider, nullptr); + auto* credentialed = scan->io()->AsSupportsStorageCredentials(); + ASSERT_NE(credentialed, nullptr); + ASSERT_EQ(credentialed->credentials().size(), 1); + EXPECT_EQ(credentialed->credentials()[0].config.at("k"), "v"); + ICEBERG_UNWRAP_OR_FAIL(auto refreshed, installed_provider->Load()); + ASSERT_EQ(refreshed.size(), 1); + EXPECT_EQ(refreshed[0].config.at("k"), "v2"); + + scan.reset(); + EXPECT_TRUE(weak_catalog.expired()); + }; + if (GetParam()) { + ICEBERG_UNWRAP_OR_FAIL(auto builder, table->NewIncrementalAppendScan()); + builder->FromSnapshot(100); + exercise(builder); + } else { + ICEBERG_UNWRAP_OR_FAIL(auto builder, table->NewScan()); + exercise(builder); + } +} + +INSTANTIATE_TEST_SUITE_P(ScanModes, RestCatalogScanLifetimeTest, ::testing::Bool()); + +} // namespace + +} // namespace iceberg::rest diff --git a/src/iceberg/test/rest_table_scan_test.cc b/src/iceberg/test/rest_table_scan_test.cc new file mode 100644 index 000000000..a1b8e6e86 --- /dev/null +++ b/src/iceberg/test/rest_table_scan_test.cc @@ -0,0 +1,1395 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you 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. + */ + +#include "iceberg/catalog/rest/rest_table_scan.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include +#include +#include + +#include "iceberg/catalog/rest/auth/auth_session.h" +#include "iceberg/catalog/rest/endpoint.h" +#include "iceberg/catalog/rest/error_handlers.h" +#include "iceberg/catalog/rest/http_client.h" +#include "iceberg/catalog/rest/resource_paths.h" +#include "iceberg/catalog/rest/rest_table.h" +#include "iceberg/catalog/rest/rest_table_scan.h" +#include "iceberg/constants.h" +#include "iceberg/expression/expressions.h" +#include "iceberg/expression/json_serde_internal.h" +#include "iceberg/expression/literal.h" +#include "iceberg/file_io.h" +#include "iceberg/file_io_registry.h" +#include "iceberg/manifest/manifest_entry.h" +#include "iceberg/partition_spec.h" +#include "iceberg/schema.h" +#include "iceberg/snapshot.h" +#include "iceberg/storage_credential.h" +#include "iceberg/table_identifier.h" +#include "iceberg/table_metadata.h" +#include "iceberg/table_scan.h" +#include "iceberg/test/executor.h" +#include "iceberg/test/matchers.h" +#include "iceberg/type.h" + +namespace iceberg::rest { + +namespace { + +using ::testing::_; +using ::testing::Return; + +// Matches a JSON string body where `key` has exactly `expected_value`. +MATCHER_P2(JsonBodyHas, key, expected_value, "") { + try { + auto json = nlohmann::json::parse(arg); + if (!json.contains(key)) { + *result_listener << "JSON body missing key \"" << key << "\""; + return false; + } + nlohmann::json expected = expected_value; + if (json.at(key) != expected) { + *result_listener << "JSON[\"" << key << "\"] = " << json.at(key) << ", expected " + << expected; + return false; + } + return true; + } catch (...) { + *result_listener << "failed to parse JSON body"; + return false; + } +} + +// Matches a JSON string body that does NOT contain `key`. +MATCHER_P(JsonBodyLacks, key, "") { + try { + auto json = nlohmann::json::parse(arg); + if (json.contains(key)) { + *result_listener << "JSON body unexpectedly contains key \"" << key << "\""; + return false; + } + return true; + } catch (...) { + *result_listener << "failed to parse JSON body"; + return false; + } +} + +// -------------------------------------------------------------------------- +// Mock HTTP client that overrides the virtual methods of HttpClient. +// The base class constructor creates a cpr::ConnectionPool, which is a +// lightweight allocation (no network connections are opened at construction). +// -------------------------------------------------------------------------- +class MockHttpClient : public HttpClient { + public: + MockHttpClient() : HttpClient({}) {} + + MOCK_METHOD(Result, Get, + (const std::string& path, + (const std::unordered_map&)params, + (const std::unordered_map&)headers, + const ErrorHandler& error_handler, auth::AuthSession& session), + (override)); + + MOCK_METHOD(Result, Post, + (const std::string& path, const std::string& body, + (const std::unordered_map&)headers, + const ErrorHandler& error_handler, auth::AuthSession& session), + (override)); + + MOCK_METHOD(Result, Delete, + (const std::string& path, + (const std::unordered_map&)params, + (const std::unordered_map&)headers, + const ErrorHandler& error_handler, auth::AuthSession& session), + (override)); +}; + +// -------------------------------------------------------------------------- +// Minimal FileIO stub (no real I/O needed for server-side scan planning tests) +// -------------------------------------------------------------------------- +class NoOpFileIO : public FileIO { + public: + Result ReadFile(const std::string&, std::optional) override { + return IOError("NoOpFileIO"); + } + Status WriteFile(const std::string&, std::string_view) override { return {}; } + Status DeleteFile(const std::string&) override { return {}; } +}; + +namespace { + +// Records the credentials and provider installed on the scan FileIO. +std::vector installed_credentials; +std::shared_ptr installed_provider; + +class CredentialedFileIO : public NoOpFileIO, public SupportsStorageCredentials { + public: + Status InitializeStorageCredentials( + const std::vector& credentials, + std::shared_ptr provider) override { + installed_credentials = credentials; + installed_provider = std::move(provider); + return {}; + } + + Status SetStorageCredentials( + const std::vector& credentials) override { + installed_credentials = credentials; + return {}; + } + + std::vector credentials() const override { + return installed_credentials; + } + + SupportsStorageCredentials* AsSupportsStorageCredentials() override { return this; } +}; + +constexpr std::string_view kCredentialedFileIO = "rest-table-scan-test-credentialed"; + +nlohmann::json FileTaskJson(const std::string& path) { + return {{"data-file", + {{"content", "data"}, + {"file-path", path}, + {"file-format", "PARQUET"}, + {"spec-id", 0}, + {"partition", nlohmann::json::array()}, + {"file-size-in-bytes", 100}, + {"record-count", 10}}}}; +} + +} // namespace + +// -------------------------------------------------------------------------- +// Test fixture shared by RestTableScan tests. +// -------------------------------------------------------------------------- +class RestTableScanTest : public ::testing::Test { + protected: + void TearDown() override { installed_provider.reset(); } + + void SetUp() override { + schema_ = std::make_shared( + std::vector{SchemaField::MakeRequired(1, "id", int32()), + SchemaField::MakeRequired(2, "data", string())}); + + auto spec = PartitionSpec::Unpartitioned(); + + constexpr int64_t kSnapshotId = 1000L; + auto snapshot = std::make_shared( + Snapshot{.snapshot_id = kSnapshotId, + .sequence_number = 1L, + .timestamp_ms = TimePointMsFromUnixMs(1609459200000L), + .manifest_list = "/tmp/manifest-list.avro", + .schema_id = schema_->schema_id()}); + + metadata_ = std::make_shared( + TableMetadata{.format_version = 2, + .table_uuid = "test-uuid", + .location = "/tmp/table", + .last_sequence_number = 1L, + .last_updated_ms = TimePointMsFromUnixMs(1609459200000L), + .last_column_id = 2, + .schemas = {schema_}, + .current_schema_id = schema_->schema_id(), + .partition_specs = {spec}, + .default_spec_id = spec->spec_id(), + .last_partition_id = 999, + .current_snapshot_id = kSnapshotId, + .snapshots = {snapshot}, + .refs = {{"main", std::make_shared(SnapshotRef{ + .snapshot_id = kSnapshotId, + .retention = SnapshotRef::Branch{}})}}}); + + file_io_ = std::make_shared(); + + mock_client_ = std::make_shared(); + + ICEBERG_UNWRAP_OR_FAIL(paths_, + ResourcePaths::Make("http://test-server", /*prefix=*/"", + /*namespace_separator=*/"%1F")); + + session_ = auth::AuthSession::MakeDefault(/*headers=*/{}); + + identifier_ = TableIdentifier{.ns = Namespace{{"default"}}, .name = "my_table"}; + + all_plan_endpoints_ = {Endpoint::PlanTableScan(), Endpoint::FetchPlanningResult(), + Endpoint::CancelPlanning(), Endpoint::FetchScanTasks()}; + } + + // Pass std::nullopt to get the full set of plan endpoints (default). + // Pass an explicit set (including empty) to use exactly that set. + RestScanContext MakeContext( + std::optional> endpoints = std::nullopt) { + auto effective = endpoints.has_value() ? std::move(*endpoints) : all_plan_endpoints_; + return RestScanContext{ + .client = mock_client_, + .paths = paths_, + .session = session_, + .supported_endpoints = std::move(effective), + .identifier = identifier_, + }; + } + + static void RegisterCredentialedFileIO() { + FileIORegistry::Register( + std::string(kCredentialedFileIO), + {.create = [](const std::unordered_map&) + -> Result> { + return std::make_unique(); + }}); + installed_credentials.clear(); + installed_provider = nullptr; + } + + Result> MakeScan(RestScanContext ctx) { + return RestTableScan::Make(metadata_, schema_, file_io_, internal::TableScanContext{}, + std::make_shared(std::move(ctx))); + } + + std::shared_ptr schema_; + std::shared_ptr metadata_; + std::shared_ptr file_io_; + std::shared_ptr mock_client_; + std::shared_ptr paths_; + std::shared_ptr session_; + TableIdentifier identifier_; + std::unordered_set all_plan_endpoints_; +}; + +// -------------------------------------------------------------------------- +// Cancel is a no-op when plan_id is empty (server returned no plan-id). +// -------------------------------------------------------------------------- +TEST_F(RestTableScanTest, CancelIsNoOpWithEmptyPlanId) { + constexpr std::string_view kPlanResponse = + R"({"status":"completed","plan-tasks":["tok-b"]})"; + + EXPECT_CALL(*mock_client_, Post(_, _, _, _, _)) + .WillOnce(Return(HttpResponse(200, std::string(kPlanResponse)))) + .WillOnce(Return(IOError("FetchScanTasks failed"))); + EXPECT_CALL(*mock_client_, Delete(_, _, _, _, _)).Times(0); + + ICEBERG_UNWRAP_OR_FAIL(auto scan, MakeScan(MakeContext())); + auto result = scan->PlanFiles(); + EXPECT_THAT(result, IsError(ErrorKind::kIOError)); +} + +// -------------------------------------------------------------------------- +// Cancel is a no-op when CancelPlanning endpoint is not advertised. +// -------------------------------------------------------------------------- +TEST_F(RestTableScanTest, CancelIsNoOpWhenEndpointNotAdvertised) { + constexpr std::string_view kPlanResponse = + R"({"status":"completed","plan-id":"plan-2","plan-tasks":["tok-c"]})"; + + std::unordered_set endpoints_without_cancel = { + Endpoint::PlanTableScan(), Endpoint::FetchPlanningResult(), + Endpoint::FetchScanTasks()}; + + EXPECT_CALL(*mock_client_, Post(_, _, _, _, _)) + .WillOnce(Return(HttpResponse(200, std::string(kPlanResponse)))) + .WillOnce(Return(IOError("FetchScanTasks failed"))); + EXPECT_CALL(*mock_client_, Delete(_, _, _, _, _)).Times(0); + + ICEBERG_UNWRAP_OR_FAIL(auto scan, MakeScan(MakeContext(endpoints_without_cancel))); + auto result = scan->PlanFiles(); + EXPECT_THAT(result, IsError(ErrorKind::kIOError)); +} + +// -------------------------------------------------------------------------- +// UseSnapshot(): the POST body sent to the server must contain both +// "snapshot-id" and "use-snapshot-schema": true. +// -------------------------------------------------------------------------- +TEST_F(RestTableScanTest, UseSnapshotPropagatesUseSnapshotSchemaInContext) { + constexpr int64_t kSnapshotId = 1000L; + constexpr std::string_view kResponseBody = R"({"status":"completed"})"; + + EXPECT_CALL(*mock_client_, + Post(_, + testing::AllOf(JsonBodyHas("snapshot-id", kSnapshotId), + JsonBodyHas("use-snapshot-schema", true)), + _, _, _)) + .WillOnce(Return(HttpResponse(200, std::string(kResponseBody)))); + + RestTableScanBuilder builder( + metadata_, file_io_, "test.my_table", nullptr, + std::make_shared(MakeContext(std::nullopt))); + builder.UseSnapshot(kSnapshotId); + ICEBERG_UNWRAP_OR_FAIL(auto scan, builder.Build()); + ICEBERG_UNWRAP_OR_FAIL(auto tasks, scan->PlanFiles()); + EXPECT_TRUE(tasks.empty()); +} + +// -------------------------------------------------------------------------- +// Select("*") sends every column, not a literal "*". +// -------------------------------------------------------------------------- +TEST_F(RestTableScanTest, SelectAllColumnsSendsAllColumns) { + constexpr std::string_view kResponseBody = R"({"status":"completed"})"; + EXPECT_CALL( + *mock_client_, + Post(_, JsonBodyHas("select", std::vector{"id", "data"}), _, _, _)) + .WillOnce(Return(HttpResponse(200, std::string(kResponseBody)))); + + RestTableScanBuilder builder(metadata_, file_io_, "test.my_table", nullptr, + std::make_shared(MakeContext())); + builder.Select({std::string(Schema::kAllColumns)}); + ICEBERG_UNWRAP_OR_FAIL(auto scan, builder.Build()); + ICEBERG_UNWRAP_OR_FAIL(auto tasks, scan->PlanFiles()); + EXPECT_TRUE(tasks.empty()); +} + +// -------------------------------------------------------------------------- +// UseRef(): branches use the table schema, tags use the snapshot schema. +// -------------------------------------------------------------------------- +class RestTableScanRefSchemaTest : public RestTableScanTest { + protected: + static constexpr int64_t kEvolvedSnapshotId = 2000L; + + void SetUp() override { + RestTableScanTest::SetUp(); + evolved_schema_ = std::make_shared( + std::vector{SchemaField::MakeRequired(10, "id", int32()), + SchemaField::MakeRequired(11, "data", string())}, + /*schema_id=*/1); + metadata_->schemas.push_back(evolved_schema_); + metadata_->last_column_id = 11; + metadata_->snapshots.push_back(std::make_shared( + Snapshot{.snapshot_id = kEvolvedSnapshotId, + .parent_snapshot_id = 1000L, + .sequence_number = 2L, + .timestamp_ms = TimePointMsFromUnixMs(1609459201000L), + .manifest_list = "/tmp/manifest-list-2.avro", + .schema_id = evolved_schema_->schema_id()})); + metadata_->refs["evolved-branch"] = std::make_shared(SnapshotRef{ + .snapshot_id = kEvolvedSnapshotId, .retention = SnapshotRef::Branch{}}); + metadata_->refs["evolved-tag"] = std::make_shared( + SnapshotRef{.snapshot_id = kEvolvedSnapshotId, .retention = SnapshotRef::Tag{}}); + } + + std::shared_ptr evolved_schema_; +}; + +TEST_F(RestTableScanRefSchemaTest, UseRefBranchUsesTableSchema) { + constexpr std::string_view kResponseBody = R"({"status":"completed"})"; + EXPECT_CALL(*mock_client_, + Post(_, + testing::AllOf(JsonBodyHas("snapshot-id", kEvolvedSnapshotId), + JsonBodyHas("use-snapshot-schema", false)), + _, _, _)) + .WillOnce(Return(HttpResponse(200, std::string(kResponseBody)))); + + RestTableScanBuilder builder(metadata_, file_io_, "test.my_table", nullptr, + std::make_shared(MakeContext())); + builder.UseRef("evolved-branch"); + ICEBERG_UNWRAP_OR_FAIL(auto scan, builder.Build()); + ICEBERG_UNWRAP_OR_FAIL(auto scan_schema, scan->schema()); + EXPECT_EQ(scan_schema->schema_id(), schema_->schema_id()); + ICEBERG_UNWRAP_OR_FAIL(auto tasks, scan->PlanFiles()); + EXPECT_TRUE(tasks.empty()); +} + +TEST_F(RestTableScanRefSchemaTest, UseRefTagUsesSnapshotSchema) { + constexpr std::string_view kResponseBody = R"({"status":"completed"})"; + EXPECT_CALL(*mock_client_, + Post(_, + testing::AllOf(JsonBodyHas("snapshot-id", kEvolvedSnapshotId), + JsonBodyHas("use-snapshot-schema", true)), + _, _, _)) + .WillOnce(Return(HttpResponse(200, std::string(kResponseBody)))); + + RestTableScanBuilder builder(metadata_, file_io_, "test.my_table", nullptr, + std::make_shared(MakeContext())); + builder.UseRef("evolved-tag"); + ICEBERG_UNWRAP_OR_FAIL(auto scan, builder.Build()); + ICEBERG_UNWRAP_OR_FAIL(auto scan_schema, scan->schema()); + EXPECT_EQ(scan_schema->schema_id(), evolved_schema_->schema_id()); + ICEBERG_UNWRAP_OR_FAIL(auto tasks, scan->PlanFiles()); + EXPECT_TRUE(tasks.empty()); +} + +// -------------------------------------------------------------------------- +// Storage credentials in a FetchScanTasks response are ignored. +// -------------------------------------------------------------------------- +TEST_F(RestTableScanTest, StorageCredentialsInFetchScanTasksResponseAreIgnored) { + constexpr std::string_view kPlanResponse = + R"({"status":"completed","plan-id":"plan-cred","plan-tasks":["tok-cred"]})"; + constexpr std::string_view kTasksResponse = R"({ + "file-scan-tasks": [], + "storage-credentials": [ + {"prefix": "s3://bucket/prefix", "config": {"key": "value"}} + ] + })"; + + EXPECT_CALL(*mock_client_, Post(_, _, _, _, _)) + .WillOnce(Return(HttpResponse(200, std::string(kPlanResponse)))) + .WillOnce(Return(HttpResponse(200, std::string(kTasksResponse)))); + + ICEBERG_UNWRAP_OR_FAIL(auto scan, MakeScan(MakeContext())); + ICEBERG_UNWRAP_OR_FAIL(auto tasks, scan->PlanFiles()); + EXPECT_TRUE(tasks.empty()); + + EXPECT_EQ(scan->io().get(), file_io_.get()); +} + +// -------------------------------------------------------------------------- +// PlanFilesStream: Next() propagates a FetchScanTasks error and DELETE /plan +// is called via the stream destructor since consumed_ is never set. +// -------------------------------------------------------------------------- +TEST_F(RestTableScanTest, PlanFilesStreamNextReturnsErrorOnFetchFailure) { + constexpr std::string_view kPlanResponse = + R"({"status":"completed","plan-id":"plan-next-err","plan-tasks":["tok-err"]})"; + + EXPECT_CALL(*mock_client_, Post(_, _, _, _, _)) + .WillOnce(Return(HttpResponse(200, std::string(kPlanResponse)))) + .WillOnce(Return(IOError("FetchScanTasks network error"))); + EXPECT_CALL(*mock_client_, Delete(_, _, _, _, _)) + .WillOnce(Return(HttpResponse(200, "{}"))); + + ICEBERG_UNWRAP_OR_FAIL(auto scan, MakeScan(MakeContext())); + ICEBERG_UNWRAP_OR_FAIL(auto stream, scan->PlanFilesStream()); + auto result = stream->Next(); + EXPECT_THAT(result, IsError(ErrorKind::kIOError)); +} + +// -------------------------------------------------------------------------- +// PlanFilesStream: stream destroyed after consuming the first task but before +// the second token is fetched → DELETE /plan called by the destructor. +// -------------------------------------------------------------------------- +TEST_F(RestTableScanTest, PlanFilesStreamCancelAfterPartialConsumption) { + constexpr std::string_view kPlanResponse = + R"({"status":"completed","plan-id":"plan-partial","plan-tasks":["tok-p1","tok-p2"]})"; + constexpr std::string_view kTaskResponse = R"({ + "file-scan-tasks": [ + {"data-file":{"content":"data","file-path":"s3://b/fp1.parquet", + "file-format":"PARQUET","spec-id":0,"partition":[],"file-size-in-bytes":1,"record-count":1}} + ] + })"; + + EXPECT_CALL(*mock_client_, Post(_, _, _, _, _)) + .WillOnce(Return(HttpResponse(200, std::string(kPlanResponse)))) + .WillOnce(Return(HttpResponse(200, std::string(kTaskResponse)))); + EXPECT_CALL(*mock_client_, Delete(_, _, _, _, _)) + .WillOnce(Return(HttpResponse(200, "{}"))); + + ICEBERG_UNWRAP_OR_FAIL(auto scan, MakeScan(MakeContext())); + { + ICEBERG_UNWRAP_OR_FAIL(auto stream, scan->PlanFilesStream()); + // Consume the first task from tok-p1; tok-p2 has never been fetched. + ICEBERG_UNWRAP_OR_FAIL(auto task, stream->Next()); + ASSERT_TRUE(task.has_value()); + // Destroy stream here — tok-p2 is still pending, so destructor calls DELETE. + } +} + +// -------------------------------------------------------------------------- +// PlanFilesStream: nested plan tasks returned by FetchScanTasks are fetched +// only when the stream reaches them. +// -------------------------------------------------------------------------- +TEST_F(RestTableScanTest, PlanFilesStreamFetchesNestedPlanTasksLazily) { + constexpr std::string_view kPlanResponse = + R"({"status":"completed","plan-id":"plan-nested","plan-tasks":["tok-1"]})"; + constexpr std::string_view kTask1Response = R"({ + "plan-tasks": ["tok-nested"], + "file-scan-tasks": [ + {"data-file":{"content":"data","file-path":"s3://b/f1.parquet", + "file-format":"PARQUET","spec-id":0,"partition":[],"file-size-in-bytes":1,"record-count":1}} + ] + })"; + constexpr std::string_view kNestedResponse = R"({ + "file-scan-tasks": [ + {"data-file":{"content":"data","file-path":"s3://b/f2.parquet", + "file-format":"PARQUET","spec-id":0,"partition":[],"file-size-in-bytes":1,"record-count":1}} + ] + })"; + + int fetch_count = 0; + EXPECT_CALL(*mock_client_, Post(::testing::EndsWith("/plan"), _, _, _, _)) + .WillOnce(Return(HttpResponse(200, std::string(kPlanResponse)))); + EXPECT_CALL(*mock_client_, Post(::testing::EndsWith("/tasks"), + JsonBodyHas("plan-task", "tok-1"), _, _, _)) + .WillOnce([&](auto&&...) -> Result { + ++fetch_count; + return HttpResponse(200, std::string(kTask1Response)); + }); + EXPECT_CALL(*mock_client_, Post(::testing::EndsWith("/tasks"), + JsonBodyHas("plan-task", "tok-nested"), _, _, _)) + .WillOnce([&](auto&&...) -> Result { + ++fetch_count; + return HttpResponse(200, std::string(kNestedResponse)); + }); + + ICEBERG_UNWRAP_OR_FAIL(auto scan, MakeScan(MakeContext())); + ICEBERG_UNWRAP_OR_FAIL(auto stream, scan->PlanFilesStream()); + + ICEBERG_UNWRAP_OR_FAIL(auto t1, stream->Next()); + ASSERT_TRUE(t1.has_value()); + EXPECT_EQ(fetch_count, 1); + EXPECT_EQ((*t1)->data_file()->file_path, "s3://b/f1.parquet"); + + ICEBERG_UNWRAP_OR_FAIL(auto t2, stream->Next()); + ASSERT_TRUE(t2.has_value()); + EXPECT_EQ(fetch_count, 2); + EXPECT_EQ((*t2)->data_file()->file_path, "s3://b/f2.parquet"); + + ICEBERG_UNWRAP_OR_FAIL(auto end, stream->Next()); + EXPECT_FALSE(end.has_value()); +} + +// -------------------------------------------------------------------------- +// PlanWith(): plan tasks are fetched serially and the executor is not used. +// -------------------------------------------------------------------------- +TEST_F(RestTableScanTest, PlanFilesStreamIgnoresPlanExecutor) { + constexpr std::string_view kPlanResponse = + R"({"status":"completed","plan-id":"plan-exec","plan-tasks":["tok-1","tok-2"]})"; + constexpr std::string_view kTask1Response = R"({ + "file-scan-tasks": [ + {"data-file":{"content":"data","file-path":"s3://b/f1.parquet", + "file-format":"PARQUET","spec-id":0,"partition":[],"file-size-in-bytes":1,"record-count":1}} + ] + })"; + constexpr std::string_view kTask2Response = R"({ + "file-scan-tasks": [ + {"data-file":{"content":"data","file-path":"s3://b/f2.parquet", + "file-format":"PARQUET","spec-id":0,"partition":[],"file-size-in-bytes":1,"record-count":1}} + ] + })"; + EXPECT_CALL(*mock_client_, Post(_, _, _, _, _)) + .WillOnce(Return(HttpResponse(200, std::string(kPlanResponse)))) + .WillOnce(Return(HttpResponse(200, std::string(kTask1Response)))) + .WillOnce(Return(HttpResponse(200, std::string(kTask2Response)))); + + test::ThreadExecutor executor; + RestTableScanBuilder builder(metadata_, file_io_, "test.my_table", nullptr, + std::make_shared(MakeContext())); + builder.PlanWith(executor); + ICEBERG_UNWRAP_OR_FAIL(auto scan, builder.Build()); + ICEBERG_UNWRAP_OR_FAIL(auto tasks, scan->PlanFiles()); + + ASSERT_EQ(tasks.size(), 2); + EXPECT_EQ(tasks[0]->data_file()->file_path, "s3://b/f1.parquet"); + EXPECT_EQ(tasks[1]->data_file()->file_path, "s3://b/f2.parquet"); + EXPECT_EQ(executor.submit_count(), 0); +} + +// -------------------------------------------------------------------------- +// poll-num-retries=0: a single GET for a SUBMITTED plan, then cancel. +// -------------------------------------------------------------------------- +TEST_F(RestTableScanTest, PlanFilesStreamPollNumRetriesZero) { + constexpr std::string_view kSubmittedBody = + R"({"status":"submitted","plan-id":"plan-no-retry"})"; + EXPECT_CALL(*mock_client_, Post(_, _, _, _, _)) + .WillOnce(Return(HttpResponse(200, std::string(kSubmittedBody)))); + EXPECT_CALL(*mock_client_, Get(_, _, _, _, _)) + .WillOnce(Return(HttpResponse(200, std::string(kSubmittedBody)))); + EXPECT_CALL(*mock_client_, Delete(_, _, _, _, _)) + .WillOnce(Return(HttpResponse(200, "{}"))); + + auto ctx = MakeContext(); + ctx.catalog_config = {{"rest-scan-planning.poll-num-retries", "0"}}; + ICEBERG_UNWRAP_OR_FAIL(auto scan, MakeScan(std::move(ctx))); + EXPECT_THAT(scan->PlanFilesStream(), + ::testing::AllOf(IsError(ErrorKind::kIOError), + HasErrorMessage("exceeded max retries (0)"))); +} + +// -------------------------------------------------------------------------- +// Without a provider factory, vended credentials are installed as static. +// -------------------------------------------------------------------------- +TEST_F(RestTableScanTest, StorageCredentialsWithoutFactoryInstallNoProvider) { + constexpr std::string_view kPlanResponse = R"({ + "status": "completed", + "plan-id": "plan-static", + "storage-credentials": [{"prefix": "s3://bucket", "config": {"k": "v"}}] + })"; + EXPECT_CALL(*mock_client_, Post(_, _, _, _, _)) + .WillOnce(Return(HttpResponse(200, std::string(kPlanResponse)))); + RegisterCredentialedFileIO(); + + auto ctx = MakeContext(); + ctx.catalog_config = {{"io-impl", std::string(kCredentialedFileIO)}}; + ICEBERG_UNWRAP_OR_FAIL(auto scan, MakeScan(std::move(ctx))); + ICEBERG_UNWRAP_OR_FAIL(auto tasks, scan->PlanFiles()); + + EXPECT_EQ(installed_provider, nullptr); + EXPECT_EQ(installed_credentials.size(), 1); +} + +// -------------------------------------------------------------------------- +// FetchPlanningResult: GET fails after SUBMITTED → DELETE /plan called before +// returning the error (tests the FetchPlanningResult error-path cancel fix). +// -------------------------------------------------------------------------- +TEST_F(RestTableScanTest, CancelCalledOnFetchPlanningResultGetError) { + constexpr std::string_view kSubmittedBody = + R"({"status":"submitted","plan-id":"plan-fetch-err"})"; + + EXPECT_CALL(*mock_client_, Post(_, _, _, _, _)) + .WillOnce(Return(HttpResponse(200, std::string(kSubmittedBody)))); + EXPECT_CALL(*mock_client_, Get(_, _, _, _, _)) + .WillOnce(Return(IOError("network failure"))); + EXPECT_CALL(*mock_client_, Delete(_, _, _, _, _)) + .WillOnce(Return(IOError("cancel failure"))); + + ICEBERG_UNWRAP_OR_FAIL(auto scan, MakeScan(MakeContext())); + auto result = scan->PlanFiles(); + EXPECT_THAT(result, ::testing::AllOf(IsError(ErrorKind::kIOError), + HasErrorMessage("network failure"))); +} + +// ========================================================================== +// RestIncrementalAppendScan tests +// ========================================================================== +// Lineage: 100 <- 200 <- 300 (main), and 200 <- 400 (branch "audit"). +class RestIncrementalAppendScanTest : public RestTableScanTest { + protected: + static constexpr int64_t kRootId = 100L; + static constexpr int64_t kMiddleId = 200L; + static constexpr int64_t kMainHeadId = 300L; + static constexpr int64_t kBranchHeadId = 400L; + static constexpr std::string_view kBranch = "audit"; + + void SetUp() override { + RestTableScanTest::SetUp(); + + auto make_snapshot = [&](int64_t id, std::optional parent_id, + int64_t sequence_number) { + return std::make_shared( + Snapshot{.snapshot_id = id, + .parent_snapshot_id = parent_id, + .sequence_number = sequence_number, + .timestamp_ms = TimePointMsFromUnixMs(1609459200000L + id), + .manifest_list = "/tmp/manifest-list.avro", + .schema_id = schema_->schema_id()}); + }; + auto make_branch = [](int64_t id) { + return std::make_shared( + SnapshotRef{.snapshot_id = id, .retention = SnapshotRef::Branch{}}); + }; + + metadata_->snapshots = {make_snapshot(kRootId, std::nullopt, 1L), + make_snapshot(kMiddleId, kRootId, 2L), + make_snapshot(kMainHeadId, kMiddleId, 3L), + make_snapshot(kBranchHeadId, kMiddleId, 4L)}; + metadata_->last_sequence_number = 4L; + metadata_->current_snapshot_id = kMainHeadId; + metadata_->refs = {{"main", make_branch(kMainHeadId)}, + {std::string(kBranch), make_branch(kBranchHeadId)}}; + } + + RestIncrementalAppendScanBuilder MakeBuilder(RestScanContext ctx) { + return {metadata_, file_io_, "test.my_table", nullptr, + std::make_shared(std::move(ctx))}; + } + + Result> MakeIncrementalScan( + RestScanContext ctx, std::optional from_snapshot_id = kRootId, + bool from_inclusive = false, std::optional to_snapshot_id = std::nullopt) { + auto builder = MakeBuilder(std::move(ctx)); + if (from_snapshot_id.has_value()) { + builder.FromSnapshot(*from_snapshot_id, from_inclusive); + } + if (to_snapshot_id.has_value()) { + builder.ToSnapshot(*to_snapshot_id); + } + return builder.Build(); + } + + template + Result>> PlanScan(bool incremental, + RestScanContext ctx, + Configure configure) { + auto plan = [&](auto& builder) -> Result>> { + configure(builder); + ICEBERG_ASSIGN_OR_RAISE(auto scan, builder.Build()); + return scan->PlanFiles(); + }; + if (incremental) { + auto builder = MakeBuilder(std::move(ctx)); + builder.FromSnapshot(kRootId); + return plan(builder); + } + RestTableScanBuilder builder(metadata_, file_io_, "test.my_table", nullptr, + std::make_shared(std::move(ctx))); + return plan(builder); + } +}; + +// -------------------------------------------------------------------------- +// No current snapshot → PlanFiles returns empty without calling the server. +// -------------------------------------------------------------------------- +TEST_F(RestIncrementalAppendScanTest, PlanFilesEmptyWhenNoCurrentSnapshot) { + // Build metadata with no current snapshot. + auto spec = PartitionSpec::Unpartitioned(); + auto empty_metadata = std::make_shared(TableMetadata{ + .format_version = 2, + .table_uuid = "no-snap-uuid", + .location = "/tmp/table", + .last_sequence_number = 0L, + .last_updated_ms = TimePointMsFromUnixMs(1609459200000L), + .last_column_id = 2, + .schemas = {schema_}, + .current_schema_id = schema_->schema_id(), + .partition_specs = {spec}, + .default_spec_id = spec->spec_id(), + .last_partition_id = 999, + .current_snapshot_id = kInvalidSnapshotId, + }); + + EXPECT_CALL(*mock_client_, Post(_, _, _, _, _)).Times(0); + + RestIncrementalAppendScanBuilder builder( + empty_metadata, file_io_, "test.my_table", nullptr, + std::make_shared(MakeContext())); + ICEBERG_UNWRAP_OR_FAIL(auto scan, builder.Build()); + ICEBERG_UNWRAP_OR_FAIL(auto tasks, scan->PlanFiles()); + EXPECT_TRUE(tasks.empty()); +} + +// -------------------------------------------------------------------------- +// POST body carries the resolved snapshot range. +// -------------------------------------------------------------------------- +TEST_F(RestIncrementalAppendScanTest, PlanFilesSendsResolvedRange) { + constexpr std::string_view kResponseBody = R"({"status":"completed"})"; + EXPECT_CALL(*mock_client_, + Post(_, + testing::AllOf(JsonBodyHas("start-snapshot-id", kRootId), + JsonBodyHas("end-snapshot-id", kMainHeadId), + JsonBodyHas("use-snapshot-schema", true)), + _, _, _)) + .WillOnce(Return(HttpResponse(200, std::string(kResponseBody)))); + + ICEBERG_UNWRAP_OR_FAIL(auto scan, MakeIncrementalScan(MakeContext())); + ICEBERG_UNWRAP_OR_FAIL(auto tasks, scan->PlanFiles()); + EXPECT_TRUE(tasks.empty()); +} + +// -------------------------------------------------------------------------- +// Explicit to_snapshot_id is sent as end-snapshot-id. +// -------------------------------------------------------------------------- +TEST_F(RestIncrementalAppendScanTest, PlanFilesWithExplicitToSnapshotId) { + constexpr std::string_view kResponseBody = R"({"status":"completed"})"; + EXPECT_CALL(*mock_client_, + Post(_, + testing::AllOf(JsonBodyHas("start-snapshot-id", kRootId), + JsonBodyHas("end-snapshot-id", kMiddleId)), + _, _, _)) + .WillOnce(Return(HttpResponse(200, std::string(kResponseBody)))); + + ICEBERG_UNWRAP_OR_FAIL(auto scan, + MakeIncrementalScan(MakeContext(), kRootId, + /*from_inclusive=*/false, kMiddleId)); + ICEBERG_UNWRAP_OR_FAIL(auto tasks, scan->PlanFiles()); + EXPECT_TRUE(tasks.empty()); +} + +// -------------------------------------------------------------------------- +// Exclusive from_snapshot_id is sent as start-snapshot-id. +// -------------------------------------------------------------------------- +TEST_F(RestIncrementalAppendScanTest, PlanFilesWithFromSnapshotIdExclusive) { + constexpr std::string_view kResponseBody = R"({"status":"completed"})"; + EXPECT_CALL(*mock_client_, + Post(_, + testing::AllOf(JsonBodyHas("start-snapshot-id", kMiddleId), + JsonBodyHas("end-snapshot-id", kMainHeadId)), + _, _, _)) + .WillOnce(Return(HttpResponse(200, std::string(kResponseBody)))); + + ICEBERG_UNWRAP_OR_FAIL(auto scan, MakeIncrementalScan(MakeContext(), kMiddleId, + /*from_inclusive=*/false)); + ICEBERG_UNWRAP_OR_FAIL(auto tasks, scan->PlanFiles()); + EXPECT_TRUE(tasks.empty()); +} + +// -------------------------------------------------------------------------- +// Inclusive from_snapshot_id sends its parent as start-snapshot-id. +// -------------------------------------------------------------------------- +TEST_F(RestIncrementalAppendScanTest, PlanFilesWithFromSnapshotIdInclusiveWithParent) { + constexpr std::string_view kResponseBody = R"({"status":"completed"})"; + EXPECT_CALL(*mock_client_, + Post(_, + testing::AllOf(JsonBodyHas("start-snapshot-id", kRootId), + JsonBodyHas("end-snapshot-id", kMainHeadId)), + _, _, _)) + .WillOnce(Return(HttpResponse(200, std::string(kResponseBody)))); + + ICEBERG_UNWRAP_OR_FAIL(auto scan, MakeIncrementalScan(MakeContext(), kMiddleId, + /*from_inclusive=*/true)); + ICEBERG_UNWRAP_OR_FAIL(auto tasks, scan->PlanFiles()); + EXPECT_TRUE(tasks.empty()); +} + +// -------------------------------------------------------------------------- +// Inclusive from_snapshot_id on the root snapshot → NotSupported. +// -------------------------------------------------------------------------- +TEST_F(RestIncrementalAppendScanTest, PlanFilesWithFromSnapshotIdInclusiveNoParent) { + EXPECT_CALL(*mock_client_, Post(_, _, _, _, _)).Times(0); + + ICEBERG_UNWRAP_OR_FAIL( + auto scan, MakeIncrementalScan(MakeContext(), kRootId, /*from_inclusive=*/true)); + auto result = scan->PlanFiles(); + EXPECT_THAT(result, IsError(ErrorKind::kNotSupported)); +} + +// -------------------------------------------------------------------------- +// No from_snapshot_id → NotSupported. +// -------------------------------------------------------------------------- +TEST_F(RestIncrementalAppendScanTest, PlanFilesWithoutFromSnapshotIdNotSupported) { + EXPECT_CALL(*mock_client_, Post(_, _, _, _, _)).Times(0); + + ICEBERG_UNWRAP_OR_FAIL( + auto scan, MakeIncrementalScan(MakeContext(), /*from_snapshot_id=*/std::nullopt)); + auto result = scan->PlanFiles(); + EXPECT_THAT(result, IsError(ErrorKind::kNotSupported)); +} + +// -------------------------------------------------------------------------- +// UseBranch() without to_snapshot_id sends the branch head as end-snapshot-id. +// -------------------------------------------------------------------------- +TEST_F(RestIncrementalAppendScanTest, UseBranchWithoutToSnapshotUsesBranchHead) { + constexpr std::string_view kResponseBody = R"({"status":"completed"})"; + EXPECT_CALL(*mock_client_, + Post(_, + testing::AllOf(JsonBodyHas("start-snapshot-id", kRootId), + JsonBodyHas("end-snapshot-id", kBranchHeadId)), + _, _, _)) + .WillOnce(Return(HttpResponse(200, std::string(kResponseBody)))); + + auto builder = MakeBuilder(MakeContext()); + builder.FromSnapshot(kRootId).UseBranch(std::string(kBranch)); + ICEBERG_UNWRAP_OR_FAIL(auto scan, builder.Build()); + ICEBERG_UNWRAP_OR_FAIL(auto tasks, scan->PlanFiles()); + EXPECT_TRUE(tasks.empty()); +} + +// -------------------------------------------------------------------------- +// UseBranch() with to_snapshot_id on the branch. +// -------------------------------------------------------------------------- +TEST_F(RestIncrementalAppendScanTest, UseBranchWithToSnapshotOnBranch) { + constexpr std::string_view kResponseBody = R"({"status":"completed"})"; + EXPECT_CALL(*mock_client_, + Post(_, + testing::AllOf(JsonBodyHas("start-snapshot-id", kRootId), + JsonBodyHas("end-snapshot-id", kMiddleId)), + _, _, _)) + .WillOnce(Return(HttpResponse(200, std::string(kResponseBody)))); + + auto builder = MakeBuilder(MakeContext()); + builder.FromSnapshot(kRootId).ToSnapshot(kMiddleId).UseBranch(std::string(kBranch)); + ICEBERG_UNWRAP_OR_FAIL(auto scan, builder.Build()); + ICEBERG_UNWRAP_OR_FAIL(auto tasks, scan->PlanFiles()); + EXPECT_TRUE(tasks.empty()); +} + +// -------------------------------------------------------------------------- +// UseBranch() with to_snapshot_id not on the branch → ValidationFailed. +// -------------------------------------------------------------------------- +TEST_F(RestIncrementalAppendScanTest, UseBranchRejectsToSnapshotNotOnBranch) { + EXPECT_CALL(*mock_client_, Post(_, _, _, _, _)).Times(0); + + auto builder = MakeBuilder(MakeContext()); + builder.FromSnapshot(kRootId).ToSnapshot(kMainHeadId).UseBranch(std::string(kBranch)); + ICEBERG_UNWRAP_OR_FAIL(auto scan, builder.Build()); + auto result = scan->PlanFiles(); + EXPECT_THAT(result, IsError(ErrorKind::kValidationFailed)); + EXPECT_THAT(result, HasErrorMessage("not a valid snapshot on the current branch")); +} + +// -------------------------------------------------------------------------- +// from_snapshot_id not an ancestor of the end snapshot → ValidationFailed. +// -------------------------------------------------------------------------- +TEST_F(RestIncrementalAppendScanTest, RejectsFromSnapshotNotAncestorOfEnd) { + EXPECT_CALL(*mock_client_, Post(_, _, _, _, _)).Times(0); + + ICEBERG_UNWRAP_OR_FAIL(auto scan, MakeIncrementalScan(MakeContext(), kBranchHeadId, + /*from_inclusive=*/false)); + auto result = scan->PlanFiles(); + EXPECT_THAT(result, IsError(ErrorKind::kValidationFailed)); +} + +// Exercise residual handling for both scan types and each response path. +class RestScanResidualTest + : public RestIncrementalAppendScanTest, + public ::testing::WithParamInterface> {}; + +TEST_P(RestScanResidualTest, PreservesFilteringAndDeleteFiles) { + const auto [incremental, submitted, fetch_token, server_residual, ignore_residuals] = + GetParam(); + auto filter = Expressions::GreaterThan("id", Literal::Int(10)); + auto residual = Expressions::GreaterThan("id", Literal::Int(20)); + ICEBERG_UNWRAP_OR_FAIL(auto filter_json, iceberg::ToJson(*filter)); + ICEBERG_UNWRAP_OR_FAIL(auto residual_json, iceberg::ToJson(*residual)); + + auto task = FileTaskJson("s3://bucket/data.parquet"); + task["delete-file-references"] = {0}; + if (server_residual) task["residual-filter"] = residual_json; + auto delete_file = FileTaskJson("s3://bucket/delete.parquet")["data-file"]; + delete_file["content"] = "position-deletes"; + nlohmann::json tasks_response = { + {"file-scan-tasks", nlohmann::json::array({task})}, + {"delete-files", nlohmann::json::array({delete_file})}}; + auto completed = + fetch_token ? nlohmann::json{{"plan-tasks", {"residual-token"}}} : tasks_response; + completed["status"] = "completed"; + completed["plan-id"] = "residual-plan"; + auto initial = + submitted ? nlohmann::json{{"status", "submitted"}, {"plan-id", "residual-plan"}} + : completed; + EXPECT_CALL(*mock_client_, + Post(::testing::EndsWith("/plan"), + ::testing::AllOf( + JsonBodyHas("filter", filter_json), + JsonBodyHas("select", std::vector{"id", "data"})), + _, _, _)) + .WillOnce(Return(HttpResponse(200, initial.dump()))); + if (submitted) { + completed.erase("plan-id"); + EXPECT_CALL(*mock_client_, + Get(::testing::EndsWith("/plan/residual-plan"), _, _, _, _)) + .WillOnce(Return(HttpResponse(200, completed.dump()))); + } + if (fetch_token) { + EXPECT_CALL(*mock_client_, Post(::testing::EndsWith("/tasks"), + JsonBodyHas("plan-task", "residual-token"), _, _, _)) + .WillOnce(Return(HttpResponse(200, tasks_response.dump()))); + } + EXPECT_CALL(*mock_client_, Delete(_, _, _, _, _)).Times(0); + + ICEBERG_UNWRAP_OR_FAIL(auto tasks, + PlanScan(incremental, MakeContext(), [&](auto& builder) { + builder.Select({"data"}).Filter(filter); + if (ignore_residuals) builder.IgnoreResiduals(); + })); + ASSERT_EQ(tasks.size(), 1); + EXPECT_EQ(tasks[0]->data_file()->file_path, "s3://bucket/data.parquet"); + ASSERT_EQ(tasks[0]->delete_files().size(), 1); + EXPECT_EQ(tasks[0]->delete_files()[0]->file_path, "s3://bucket/delete.parquet"); + ASSERT_NE(tasks[0]->residual_filter(), nullptr); + if (ignore_residuals) { + EXPECT_EQ(tasks[0]->residual_filter()->op(), Expression::Operation::kTrue); + } else { + ICEBERG_UNWRAP_OR_FAIL(auto actual, iceberg::ToJson(*tasks[0]->residual_filter())); + EXPECT_EQ(actual, server_residual ? residual_json : filter_json); + } +} + +INSTANTIATE_TEST_SUITE_P(ScanModes, RestScanResidualTest, + ::testing::Combine(::testing::Bool(), ::testing::Bool(), + ::testing::Bool(), ::testing::Bool(), + ::testing::Bool())); + +class RestScanRegressionTest : public RestIncrementalAppendScanTest, + public ::testing::WithParamInterface {}; + +TEST_P(RestScanRegressionTest, NoFilterProducesTrueResidual) { + nlohmann::json response = { + {"status", "completed"}, + {"file-scan-tasks", + nlohmann::json::array({FileTaskJson("s3://bucket/data.parquet")})}}; + EXPECT_CALL( + *mock_client_, + Post(::testing::EndsWith("/plan"), + ::testing::AllOf(JsonBodyHas("select", std::vector{"id", "data"}), + JsonBodyHas("use-snapshot-schema", GetParam()), + JsonBodyLacks("snapshot-id")), + _, _, _)) + .WillOnce(Return(HttpResponse(200, response.dump()))); + ICEBERG_UNWRAP_OR_FAIL(auto tasks, PlanScan(GetParam(), MakeContext(), [](auto&) {})); + ASSERT_EQ(tasks.size(), 1); + ASSERT_NE(tasks[0]->residual_filter(), nullptr); + EXPECT_EQ(tasks[0]->residual_filter()->op(), Expression::Operation::kTrue); +} + +TEST_P(RestScanRegressionTest, AllColumnStatsIncludeUnprojectedColumns) { + EXPECT_CALL( + *mock_client_, + Post(::testing::EndsWith("/plan"), + ::testing::AllOf( + JsonBodyHas("select", std::vector{"data"}), + JsonBodyHas("stats-fields", std::vector{"id", "data"})), + _, _, _)) + .WillOnce(Return(HttpResponse(200, R"({"status":"completed"})"))); + ICEBERG_UNWRAP_OR_FAIL(auto tasks, + PlanScan(GetParam(), MakeContext(), [](auto& builder) { + builder.Select({"data"}).IncludeColumnStats(); + })); + EXPECT_TRUE(tasks.empty()); +} + +TEST_P(RestScanRegressionTest, InvalidPollingConfigFailsBeforePlanning) { + struct PropertyCase { + std::string_view key; + std::string_view value; + std::string_view error; + }; + const std::array cases{{{.key = "rest-scan-planning.poll-timeout-ms", + .value = "0", + .error = "must be positive"}, + {.key = "rest-scan-planning.poll-timeout-ms", + .value = "-1", + .error = "must be positive"}, + {.key = "rest-scan-planning.poll-timeout-ms", + .value = "abc", + .error = "invalid argument"}, + {.key = "rest-scan-planning.poll-timeout-ms", + .value = "9223372036854775808", + .error = "out of range"}, + {.key = "rest-scan-planning.poll-num-retries", + .value = "-1", + .error = "must be non-negative"}, + {.key = "rest-scan-planning.poll-num-retries", + .value = "abc", + .error = "invalid argument"}, + {.key = "rest-scan-planning.poll-num-retries", + .value = "2147483648", + .error = "out of range"}}}; + for (const auto& [key, value, error] : cases) { + SCOPED_TRACE(std::string(key) + "=" + std::string(value)); + auto ctx = MakeContext(); + ctx.catalog_config = {{std::string(key), std::string(value)}}; + EXPECT_CALL(*mock_client_, Post(_, _, _, _, _)).Times(0); + EXPECT_THAT( + PlanScan(GetParam(), std::move(ctx), [](auto&) {}), + ::testing::AllOf(IsError(ErrorKind::kInvalidArgument), HasErrorMessage(error))); + ASSERT_TRUE(::testing::Mock::VerifyAndClearExpectations(mock_client_.get())); + } +} + +TEST_P(RestScanRegressionTest, MissingPollingEndpointCancelsSubmittedPlan) { + auto ctx = MakeContext(); + ctx.supported_endpoints.erase(Endpoint::FetchPlanningResult()); + EXPECT_CALL(*mock_client_, Post(::testing::EndsWith("/plan"), _, _, _, _)) + .WillOnce(Return( + HttpResponse(200, R"({"status":"submitted","plan-id":"missing-poll"})"))); + EXPECT_CALL(*mock_client_, Get(_, _, _, _, _)).Times(0); + EXPECT_CALL(*mock_client_, + Delete(::testing::EndsWith("/plan/missing-poll"), _, _, _, _)) + .WillOnce(Return(HttpResponse(204, ""))); + EXPECT_THAT(PlanScan(GetParam(), std::move(ctx), [](auto&) {}), + IsError(ErrorKind::kNotSupported)); +} + +TEST_P(RestScanRegressionTest, LongContinuationChain) { + constexpr int kPages = 2048; + EXPECT_CALL(*mock_client_, Post(::testing::EndsWith("/plan"), _, _, _, _)) + .WillOnce(Return(HttpResponse( + 200, R"({"status":"completed","plan-id":"pages","plan-tasks":["0"]})"))); + int next_page = 0; + EXPECT_CALL(*mock_client_, Post(::testing::EndsWith("/tasks"), _, _, _, _)) + .Times(kPages) + .WillRepeatedly([&](const std::string&, const std::string& body, + auto&&...) -> Result { + EXPECT_EQ(nlohmann::json::parse(body).at("plan-task"), std::to_string(next_page)); + auto task = FileTaskJson("s3://bucket/" + std::to_string(next_page) + ".parquet"); + nlohmann::json response = { + {"file-scan-tasks", nlohmann::json::array({std::move(task)})}}; + if (++next_page < kPages) { + response["plan-tasks"] = {std::to_string(next_page)}; + } + return HttpResponse(200, response.dump()); + }); + EXPECT_CALL(*mock_client_, Delete(_, _, _, _, _)).Times(0); + + ICEBERG_UNWRAP_OR_FAIL( + auto table, RestTable::Make(identifier_, metadata_, "/tmp/metadata.json", file_io_, + nullptr, "test.my_table", nullptr, MakeContext())); + auto plan = [&]() -> Result>> { + if (GetParam()) { + ICEBERG_ASSIGN_OR_RAISE(auto builder, table->NewIncrementalAppendScan()); + builder->FromSnapshot(kRootId); + ICEBERG_ASSIGN_OR_RAISE(auto scan, builder->Build()); + return scan->PlanFiles(); + } + ICEBERG_ASSIGN_OR_RAISE(auto builder, table->NewScan()); + ICEBERG_ASSIGN_OR_RAISE(auto scan, builder->Build()); + ICEBERG_ASSIGN_OR_RAISE(auto stream, scan->PlanFilesStream()); + scan.reset(); + table.reset(); + return stream->ToVector(); + }; + ICEBERG_UNWRAP_OR_FAIL(auto tasks, plan()); + ASSERT_EQ(tasks.size(), kPages); + for (int i = 0; i < kPages; ++i) { + EXPECT_EQ(tasks[i]->data_file()->file_path, + "s3://bucket/" + std::to_string(i) + ".parquet"); + } +} + +TEST_P(RestScanRegressionTest, TaskFetchFailureCancelsPlan) { + EXPECT_CALL(*mock_client_, Post(::testing::EndsWith("/plan"), _, _, _, _)) + .WillOnce(Return(HttpResponse( + 200, R"({"status":"completed","plan-id":"failed-page","plan-tasks":["0"]})"))); + EXPECT_CALL(*mock_client_, + Post(::testing::EndsWith("/tasks"), JsonBodyHas("plan-task", "0"), _, _, _)) + .WillOnce(Return(IOError("task fetch failed"))); + EXPECT_CALL(*mock_client_, Delete(::testing::EndsWith("/plan/failed-page"), _, _, _, _)) + .WillOnce(Return(HttpResponse(204, ""))); + EXPECT_THAT(PlanScan(GetParam(), MakeContext(), [](auto&) {}), + IsError(ErrorKind::kIOError)); +} + +class RestScanStreamTest : public RestTableScanTest, + public ::testing::WithParamInterface { + protected: + void ExpectPlan(nlohmann::json completed) { + completed["status"] = "completed"; + completed["plan-id"] = "stream"; + auto initial = GetParam() + ? nlohmann::json{{"status", "submitted"}, {"plan-id", "stream"}} + : completed; + EXPECT_CALL(*mock_client_, Post(::testing::EndsWith("/plan"), _, _, _, _)) + .WillOnce(Return(HttpResponse(200, initial.dump()))); + if (GetParam()) { + completed.erase("plan-id"); + EXPECT_CALL(*mock_client_, Get(::testing::EndsWith("/plan/stream"), _, _, _, _)) + .WillOnce(Return(HttpResponse(200, completed.dump()))); + } else { + EXPECT_CALL(*mock_client_, Get(_, _, _, _, _)).Times(0); + } + } +}; + +TEST_P(RestScanStreamTest, FetchesOnlyWhenBufferIsExhausted) { + ExpectPlan( + {{"file-scan-tasks", nlohmann::json::array({FileTaskJson("s3://b/0.parquet")})}, + {"plan-tasks", {"first", "second"}}}); + int fetch_count = 0; + EXPECT_CALL(*mock_client_, Post(::testing::EndsWith("/tasks"), + JsonBodyHas("plan-task", "first"), _, _, _)) + .WillOnce([&](auto&&...) -> Result { + EXPECT_EQ(fetch_count++, 0); + return HttpResponse(200, R"({"plan-tasks":["third"]})"); + }); + EXPECT_CALL(*mock_client_, Post(::testing::EndsWith("/tasks"), + JsonBodyHas("plan-task", "second"), _, _, _)) + .WillOnce([&](auto&&...) -> Result { + EXPECT_EQ(fetch_count++, 1); + nlohmann::json response = { + {"file-scan-tasks", + nlohmann::json::array( + {FileTaskJson("s3://b/1.parquet"), FileTaskJson("s3://b/2.parquet")})}}; + return HttpResponse(200, response.dump()); + }); + EXPECT_CALL(*mock_client_, Post(::testing::EndsWith("/tasks"), + JsonBodyHas("plan-task", "third"), _, _, _)) + .WillOnce([&](auto&&...) -> Result { + EXPECT_EQ(fetch_count++, 2); + nlohmann::json response = { + {"file-scan-tasks", + nlohmann::json::array({FileTaskJson("s3://b/3.parquet")})}}; + return HttpResponse(200, response.dump()); + }); + EXPECT_CALL(*mock_client_, Delete(_, _, _, _, _)).Times(0); + ICEBERG_UNWRAP_OR_FAIL(auto scan, MakeScan(MakeContext())); + ICEBERG_UNWRAP_OR_FAIL(auto stream, scan->PlanFilesStream()); + EXPECT_EQ(fetch_count, 0); + for (int i = 0; i <= 3; ++i) { + ICEBERG_UNWRAP_OR_FAIL(auto task, stream->Next()); + ASSERT_TRUE(task.has_value()); + EXPECT_EQ((*task)->data_file()->file_path, + "s3://b/" + std::to_string(i) + ".parquet"); + EXPECT_EQ(fetch_count, i == 0 ? 0 : (i <= 2 ? 2 : 3)); + } + ICEBERG_UNWRAP_OR_FAIL(auto end, stream->Next()); + EXPECT_FALSE(end.has_value()); +} + +TEST_P(RestScanStreamTest, UnconsumedPlanIsCancelled) { + ExpectPlan({{"plan-tasks", {"pending"}}}); + EXPECT_CALL(*mock_client_, Post(::testing::EndsWith("/tasks"), _, _, _, _)).Times(0); + EXPECT_CALL(*mock_client_, Delete(::testing::EndsWith("/plan/stream"), _, _, _, _)) + .WillOnce(Return(HttpResponse(204, ""))); + ICEBERG_UNWRAP_OR_FAIL(auto scan, MakeScan(MakeContext())); + ICEBERG_UNWRAP_OR_FAIL(auto stream, scan->PlanFilesStream()); + stream.reset(); +} + +INSTANTIATE_TEST_SUITE_P(PlanningResponses, RestScanStreamTest, ::testing::Bool()); + +TEST_P(RestScanRegressionTest, EmptyCompletedPlan) { + for (bool submitted : {false, true}) { + SCOPED_TRACE(submitted); + std::string initial = submitted ? R"({"status":"submitted","plan-id":"empty"})" + : R"({"status":"completed"})"; + EXPECT_CALL(*mock_client_, Post(::testing::EndsWith("/plan"), _, _, _, _)) + .WillOnce(Return(HttpResponse(200, initial))); + if (submitted) { + EXPECT_CALL(*mock_client_, Get(::testing::EndsWith("/plan/empty"), _, _, _, _)) + .WillOnce(Return(HttpResponse(200, R"({"status":"completed"})"))); + } else { + EXPECT_CALL(*mock_client_, Get(_, _, _, _, _)).Times(0); + } + EXPECT_CALL(*mock_client_, Delete(_, _, _, _, _)).Times(0); + ICEBERG_UNWRAP_OR_FAIL(auto tasks, PlanScan(GetParam(), MakeContext(), [](auto&) {})); + EXPECT_TRUE(tasks.empty()); + ASSERT_TRUE(::testing::Mock::VerifyAndClearExpectations(mock_client_.get())); + } +} + +TEST_P(RestScanRegressionTest, FailedPlanReturnsError) { + EXPECT_CALL(*mock_client_, Post(::testing::EndsWith("/plan"), _, _, _, _)) + .WillOnce(Return(HttpResponse( + 200, + R"({"status":"failed","error":{"message":"server error","type":"ServerError","code":500}})"))); + EXPECT_CALL(*mock_client_, Get(_, _, _, _, _)).Times(0); + EXPECT_CALL(*mock_client_, Delete(_, _, _, _, _)).Times(0); + EXPECT_THAT( + PlanScan(GetParam(), MakeContext(), [](auto&) {}), + ::testing::AllOf(IsError(ErrorKind::kIOError), HasErrorMessage("server error"))); +} + +TEST_P(RestScanRegressionTest, MissingPlanningEndpointDoesNotSendRequest) { + auto ctx = MakeContext(); + ctx.supported_endpoints.erase(Endpoint::PlanTableScan()); + EXPECT_CALL(*mock_client_, Post(_, _, _, _, _)).Times(0); + EXPECT_THAT(PlanScan(GetParam(), std::move(ctx), [](auto&) {}), + IsError(ErrorKind::kNotSupported)); +} + +TEST_P(RestScanRegressionTest, MissingTaskEndpointCancelsPlan) { + auto ctx = MakeContext(); + ctx.supported_endpoints.erase(Endpoint::FetchScanTasks()); + EXPECT_CALL(*mock_client_, Post(::testing::EndsWith("/plan"), _, _, _, _)) + .WillOnce(Return(HttpResponse( + 200, + R"({"status":"completed","plan-id":"no-tasks","plan-tasks":["pending"]})"))); + EXPECT_CALL(*mock_client_, Delete(::testing::EndsWith("/plan/no-tasks"), _, _, _, _)) + .WillOnce(Return(HttpResponse(204, ""))); + EXPECT_THAT(PlanScan(GetParam(), std::move(ctx), [](auto&) {}), + IsError(ErrorKind::kNotSupported)); +} + +TEST_P(RestScanRegressionTest, InvalidPlanResponseCancelsKnownId) { + struct Case { + std::string body; + ErrorKind error; + bool cancel; + }; + const std::array cases{ + {{.body = R"({"status":"submitted","plan-id":"invalid","plan-tasks":["token"]})", + .error = ErrorKind::kValidationFailed, + .cancel = true}, + {.body = + R"({"status":"failed","plan-id":"invalid","error":{"message":"error","type":"ServerError","code":500}})", + .error = ErrorKind::kValidationFailed, + .cancel = true}, + {.body = R"({"status":"failed"})", + .error = ErrorKind::kJsonParseError, + .cancel = false}}}; + for (const auto& test_case : cases) { + SCOPED_TRACE(test_case.body); + EXPECT_CALL(*mock_client_, Post(::testing::EndsWith("/plan"), _, _, _, _)) + .WillOnce(Return(HttpResponse(200, test_case.body))); + EXPECT_CALL(*mock_client_, Get(_, _, _, _, _)).Times(0); + if (test_case.cancel) { + EXPECT_CALL(*mock_client_, Delete(::testing::EndsWith("/plan/invalid"), _, _, _, _)) + .WillOnce(Return(HttpResponse(204, ""))); + } else { + EXPECT_CALL(*mock_client_, Delete(_, _, _, _, _)).Times(0); + } + EXPECT_THAT(PlanScan(GetParam(), MakeContext(), [](auto&) {}), + IsError(test_case.error)); + ASSERT_TRUE(::testing::Mock::VerifyAndClearExpectations(mock_client_.get())); + } +} + +TEST_P(RestScanRegressionTest, RequestedColumnStatsAreSent) { + EXPECT_CALL(*mock_client_, + Post(::testing::EndsWith("/plan"), + JsonBodyHas("stats-fields", std::vector{"id"}), _, _, _)) + .WillOnce(Return(HttpResponse(200, R"({"status":"completed"})"))); + ICEBERG_UNWRAP_OR_FAIL( + auto tasks, PlanScan(GetParam(), MakeContext(), + [](auto& builder) { builder.IncludeColumnStats({"id"}); })); + EXPECT_TRUE(tasks.empty()); +} + +TEST_P(RestScanRegressionTest, UnknownStatsColumnFailsBeforePlanning) { + EXPECT_CALL(*mock_client_, Post(_, _, _, _, _)).Times(0); + EXPECT_THAT(PlanScan(GetParam(), MakeContext(), + [](auto& builder) { builder.IncludeColumnStats({"missing"}); }), + ::testing::AllOf(IsError(ErrorKind::kValidationFailed), + HasErrorMessage("Cannot find stats column: missing"))); +} + +TEST_P(RestScanRegressionTest, ProjectUsesTheProjectedSchema) { + EXPECT_CALL(*mock_client_, + Post(::testing::EndsWith("/plan"), + JsonBodyHas("select", std::vector{"id"}), _, _, _)) + .WillOnce(Return(HttpResponse(200, R"({"status":"completed"})"))); + ICEBERG_UNWRAP_OR_FAIL( + auto tasks, PlanScan(GetParam(), MakeContext(), [](auto& builder) { + builder.Project(std::make_shared( + std::vector{SchemaField::MakeRequired(1, "id", int32())})); + })); + EXPECT_TRUE(tasks.empty()); +} + +TEST_P(RestScanRegressionTest, SelectResolvesCaseInsensitiveNames) { + EXPECT_CALL(*mock_client_, + Post(::testing::EndsWith("/plan"), + JsonBodyHas("select", std::vector{"id"}), _, _, _)) + .WillOnce(Return(HttpResponse(200, R"({"status":"completed"})"))); + ICEBERG_UNWRAP_OR_FAIL(auto tasks, + PlanScan(GetParam(), MakeContext(), [](auto& builder) { + builder.CaseSensitive(false).Select({"ID"}); + })); + EXPECT_TRUE(tasks.empty()); +} + +TEST_P(RestScanRegressionTest, ReplanningDropsOldCredentials) { + EXPECT_CALL(*mock_client_, Post(::testing::EndsWith("/plan"), _, _, _, _)) + .WillOnce(Return(HttpResponse( + 200, + R"({"status":"completed","storage-credentials":[{"prefix":"s3://bucket","config":{"key":"value"}}]})"))) + .WillOnce(Return(HttpResponse(200, R"({"status":"completed"})"))); + auto exercise = [&](auto& builder) { + ICEBERG_UNWRAP_OR_FAIL(auto scan, builder.Build()); + EXPECT_EQ(scan->io(), file_io_); + ICEBERG_UNWRAP_OR_FAIL(auto first, scan->PlanFiles()); + auto* credentialed = scan->io()->AsSupportsStorageCredentials(); + ASSERT_NE(credentialed, nullptr); + ASSERT_EQ(credentialed->credentials().size(), 1); + EXPECT_EQ(credentialed->credentials()[0].prefix, "s3://bucket"); + ICEBERG_UNWRAP_OR_FAIL(auto second, scan->PlanFiles()); + EXPECT_EQ(scan->io(), file_io_); + }; + if (GetParam()) { + auto builder = MakeBuilder(MakeContext()); + builder.FromSnapshot(kRootId); + exercise(builder); + } else { + RestTableScanBuilder builder(metadata_, file_io_, "test.my_table", nullptr, + std::make_shared(MakeContext())); + exercise(builder); + } +} + +INSTANTIATE_TEST_SUITE_P(ScanModes, RestScanRegressionTest, ::testing::Bool()); + +} // namespace + +} // namespace iceberg::rest diff --git a/src/iceberg/test/table_scan_test.cc b/src/iceberg/test/table_scan_test.cc index f08440a10..40b51bd5e 100644 --- a/src/iceberg/test/table_scan_test.cc +++ b/src/iceberg/test/table_scan_test.cc @@ -252,6 +252,8 @@ TEST_P(TableScanTest, IncludeColumnStatsUsesFinalSnapshotSchema) { .schema_id = evolved_schema->schema_id()})); table_metadata_->refs["evolved-branch"] = std::make_shared( SnapshotRef{.snapshot_id = kEvolvedSnapshotId, .retention = SnapshotRef::Branch{}}); + table_metadata_->refs["evolved-tag"] = std::make_shared( + SnapshotRef{.snapshot_id = kEvolvedSnapshotId, .retention = SnapshotRef::Tag{}}); { ICEBERG_UNWRAP_OR_FAIL(auto builder, MakeScanBuilder(table_metadata_)); @@ -268,7 +270,7 @@ TEST_P(TableScanTest, IncludeColumnStatsUsesFinalSnapshotSchema) { { ICEBERG_UNWRAP_OR_FAIL(auto builder, MakeScanBuilder(table_metadata_)); - builder->IncludeColumnStats({"id"}).UseRef("evolved-branch"); + builder->IncludeColumnStats({"id"}).UseRef("evolved-tag"); ICEBERG_UNWRAP_OR_FAIL(auto scan, builder->Build()); ICEBERG_UNWRAP_OR_FAIL(auto scan_schema, scan->schema()); @@ -278,6 +280,20 @@ TEST_P(TableScanTest, IncludeColumnStatsUsesFinalSnapshotSchema) { EXPECT_TRUE(stats_fields.contains(kEvolvedIdFieldId)); EXPECT_FALSE(stats_fields.contains(kBaseIdFieldId)); } + + // Branches use the table schema. + { + ICEBERG_UNWRAP_OR_FAIL(auto builder, MakeScanBuilder(table_metadata_)); + builder->IncludeColumnStats({"id"}).UseRef("evolved-branch"); + ICEBERG_UNWRAP_OR_FAIL(auto scan, builder->Build()); + ICEBERG_UNWRAP_OR_FAIL(auto scan_schema, scan->schema()); + + EXPECT_EQ(scan_schema->schema_id(), schema_->schema_id()); + const auto& stats_fields = scan->context().columns_to_keep_stats; + EXPECT_EQ(stats_fields.size(), 1); + EXPECT_TRUE(stats_fields.contains(kBaseIdFieldId)); + EXPECT_FALSE(stats_fields.contains(kEvolvedIdFieldId)); + } } TEST_P(TableScanTest, IncludeColumnStatsRejectsMissingColumn) { diff --git a/src/iceberg/type_fwd.h b/src/iceberg/type_fwd.h index 84a991909..511b14eb0 100644 --- a/src/iceberg/type_fwd.h +++ b/src/iceberg/type_fwd.h @@ -204,6 +204,7 @@ struct ReaderOptions; struct WriterOptions; struct StorageCredential; class FileIO; +class StorageCredentialProvider; class Reader; class Writer; diff --git a/src/iceberg/util/type_util.h b/src/iceberg/util/type_util.h index 8623ad254..72e6bacee 100644 --- a/src/iceberg/util/type_util.h +++ b/src/iceberg/util/type_util.h @@ -127,7 +127,7 @@ class PruneColumnVisitor { }; /// \brief Visitor for getting projected field IDs. -class GetProjectedIdsVisitor { +class ICEBERG_EXPORT GetProjectedIdsVisitor { public: explicit GetProjectedIdsVisitor(bool include_struct_ids = false);