This is an automated email from the ASF dual-hosted git repository. jerryshao pushed a commit to branch branch-1.3 in repository https://gitbox.apache.org/repos/asf/gravitino.git
commit af13eb94298610f723237b3bc391fa414e666547 Author: MaSai <[email protected]> AuthorDate: Wed Sep 9 19:59:28 2026 +0800 [#12949] fix(iceberg-rest): Forward access-delegation header on federated loadTable (#12950) Forward `X-Iceberg-Access-Delegation: vended-credentials` when a federated Iceberg REST catalog (`catalog-backend: rest`) loads a table, and rewrite remote credential refresh endpoints to this IRC catalog. `FederatedCatalogWrapper.loadTable` previously ignored `requestCredential` and used Iceberg's `RESTCatalog.loadTable`, which does not send the access-delegation header. Scan-plan federation already forwarded the header; load table now uses the same authenticated REST GET path. Fix: #12949 A REST-backend Iceberg catalog dropped the client's credential-vending request. The near-end IRC logged `credential vending: true`, but the forwarded load arrived at the remote catalog with `access delegation: null`. Direct loads against the remote returned `storage-credentials`; federated loads returned only metadata, so engines failed on the first data read. Yes. Clients that send `X-Iceberg-Access-Delegation: vended-credentials` through a federated Iceberg REST catalog now receive remote `storage-credentials` on `loadTable`, matching a direct load against the remote catalog. No new APIs or property keys. - `./gradlew :iceberg:iceberg-rest-server:test --tests org.apache.gravitino.iceberg.service.TestCatalogWrapperForREST --tests org.apache.gravitino.iceberg.service.TestIcebergRESTUtils -PskipITs` - New unit tests cover header forwarding on vended federated load, omitting the header when vending is not requested, and rewriting upstream refresh endpoints. Made with [Cursor](https://cursor.com) --------- Co-authored-by: Cursor <[email protected]> --- .../iceberg/service/FederatedCatalogWrapper.java | 403 ++++---- .../iceberg/service/IcebergRESTUtils.java | 57 ++ .../service/rest/IcebergTableOperations.java | 12 +- .../iceberg/service/TestCatalogWrapperForREST.java | 1016 +++++++++++++++----- .../iceberg/service/TestIcebergRESTUtils.java | 99 ++ .../service/rest/TestIcebergTableOperations.java | 79 +- 6 files changed, 1240 insertions(+), 426 deletions(-) diff --git a/iceberg/iceberg-rest-server/src/main/java/org/apache/gravitino/iceberg/service/FederatedCatalogWrapper.java b/iceberg/iceberg-rest-server/src/main/java/org/apache/gravitino/iceberg/service/FederatedCatalogWrapper.java index cbe5724921..eb3892004a 100644 --- a/iceberg/iceberg-rest-server/src/main/java/org/apache/gravitino/iceberg/service/FederatedCatalogWrapper.java +++ b/iceberg/iceberg-rest-server/src/main/java/org/apache/gravitino/iceberg/service/FederatedCatalogWrapper.java @@ -21,28 +21,20 @@ package org.apache.gravitino.iceberg.service; import com.google.common.annotations.VisibleForTesting; import com.google.common.base.Preconditions; +import com.google.common.collect.ImmutableMap; import com.google.common.collect.Maps; -import java.time.OffsetDateTime; -import java.time.ZoneOffset; import java.util.Collections; -import java.util.HashMap; import java.util.List; import java.util.Map; import java.util.Optional; +import java.util.function.Function; import java.util.stream.Collectors; import org.apache.gravitino.credential.CredentialPrivilege; -import org.apache.gravitino.credential.CredentialPropertyUtils; import org.apache.gravitino.iceberg.common.IcebergConfig; -import org.apache.gravitino.utils.MapUtils; -import org.apache.iceberg.BaseMetadataTable; -import org.apache.iceberg.BaseTable; import org.apache.iceberg.BaseTransaction; import org.apache.iceberg.CatalogProperties; import org.apache.iceberg.MetadataUpdate; -import org.apache.iceberg.PartitionSpec; import org.apache.iceberg.Schema; -import org.apache.iceberg.SortOrder; -import org.apache.iceberg.Table; import org.apache.iceberg.TableMetadata; import org.apache.iceberg.TableOperations; import org.apache.iceberg.Transaction; @@ -50,10 +42,6 @@ import org.apache.iceberg.UpdateRequirement; import org.apache.iceberg.catalog.Catalog; import org.apache.iceberg.catalog.Namespace; import org.apache.iceberg.catalog.TableIdentifier; -import org.apache.iceberg.exceptions.AlreadyExistsException; -import org.apache.iceberg.exceptions.NoSuchTableException; -import org.apache.iceberg.inmemory.InMemoryFileIO; -import org.apache.iceberg.io.FileIO; import org.apache.iceberg.rest.CatalogHandlers; import org.apache.iceberg.rest.ErrorHandlers; import org.apache.iceberg.rest.HTTPClient; @@ -64,7 +52,6 @@ import org.apache.iceberg.rest.ResourcePaths; import org.apache.iceberg.rest.auth.AuthManager; import org.apache.iceberg.rest.auth.AuthManagers; import org.apache.iceberg.rest.auth.AuthSession; -import org.apache.iceberg.rest.credentials.Credential; import org.apache.iceberg.rest.requests.CreateTableRequest; import org.apache.iceberg.rest.requests.RegisterTableRequest; import org.apache.iceberg.rest.requests.UpdateTableRequest; @@ -76,12 +63,12 @@ import org.apache.iceberg.rest.responses.LoadTableResponse; * {@link RESTCatalog}). * * <p>Federation-specific behavior is expressed through polymorphic overrides instead of {@code - * instanceof RESTCatalog} checks scattered across the base class. Table operations are routed to - * federation-aware {@code *Internal} methods so client-facing FileIO and credential properties are - * extracted from the remote catalog's {@code table.io()}. Credentials are vended by the remote - * catalog, so this wrapper never injects Gravitino-generated credentials. + * instanceof RESTCatalog} checks scattered across the base class. Table load, create and register + * use authenticated REST calls so {@code X-Iceberg-Access-Delegation} can be forwarded when the + * client requested credential vending. Update still uses the Iceberg Catalog API. This wrapper + * never injects Gravitino-generated credentials. * - * <p>Portions of the table create and update handling are derived from Apache Iceberg's {@code + * <p>Portions of the table update handling are derived from Apache Iceberg's {@code * org.apache.iceberg.rest.CatalogHandlers}: * https://github.com/apache/iceberg/blob/2abac79fcae94b5ad039bd09f7235be191b0761e/core/src/main/java/org/apache/iceberg/rest/CatalogHandlers.java */ @@ -89,6 +76,8 @@ public class FederatedCatalogWrapper extends CatalogWrapperForREST { private static final String FORMAT_VERSION = "format-version"; private static final Schema EMPTY_SCHEMA = new Schema(); + private static final String X_ICEBERG_ACCESS_DELEGATION = "X-Iceberg-Access-Delegation"; + private static final String VENDED_CREDENTIALS = "vended-credentials"; /** * Creates a federated wrapper. @@ -100,24 +89,63 @@ public class FederatedCatalogWrapper extends CatalogWrapperForREST { super(catalogName, config); } + /** + * Creates a table on the remote REST catalog. + * + * <p>Always uses a dedicated REST POST, including staged create, rather than Iceberg's Catalog + * API. When credential vending is requested the {@code X-Iceberg-Access-Delegation: + * vended-credentials} header is forwarded so the remote catalog returns {@code + * storage-credentials} inline. Upstream credential refresh endpoints are rewritten to this IRC + * catalog. + * + * @param namespace the namespace that will own the table. + * @param request the create-table request. + * @param requestCredential whether the client requested vended credentials. + * @return the create response, including rewritten remote credentials when requested. + */ @Override public LoadTableResponse createTable( Namespace namespace, CreateTableRequest request, boolean requestCredential) { - // The remote REST catalog vends its own credentials, so the requestCredential flag is not used - // here; FileIO-derived client config is extracted by createTableInternal. - return createTableInternal(namespace, request); + return createTableViaREST(namespace, request, requestCredential); } + /** + * Loads a table from the remote REST catalog. + * + * <p>Always uses a dedicated REST GET rather than Iceberg's {@link RESTCatalog#loadTable}, which + * cannot send {@code X-Iceberg-Access-Delegation}. When credential vending is requested the + * header is forwarded so the remote catalog returns {@code storage-credentials} inline. Upstream + * credential refresh endpoints are rewritten to this IRC catalog. The {@code privilege} is + * ignored because the remote catalog decides what to vend. + * + * @param identifier the table identifier. + * @param requestCredential whether the client requested vended credentials. + * @param privilege ignored; the remote REST catalog vends its own credentials. + * @return the load-table response, including rewritten remote credentials when requested. + */ @Override public LoadTableResponse loadTable( TableIdentifier identifier, boolean requestCredential, CredentialPrivilege privilege) { - return loadTableInternal(identifier); + return loadTableViaREST(identifier, requestCredential); } + /** + * Registers a table on the remote REST catalog. + * + * <p>Always uses a dedicated REST POST rather than Iceberg's Catalog API. When credential vending + * is requested the {@code X-Iceberg-Access-Delegation: vended-credentials} header is forwarded so + * the remote catalog returns {@code storage-credentials} inline. Upstream credential refresh + * endpoints are rewritten to this IRC catalog. + * + * @param namespace the namespace that will own the table. + * @param request the register-table request. + * @param requestCredential whether the client requested vended credentials. + * @return the register response, including rewritten remote credentials when requested. + */ @Override public LoadTableResponse registerTable( Namespace namespace, RegisterTableRequest request, boolean requestCredential) { - return registerTableInternal(namespace, request); + return registerTableViaREST(namespace, request, requestCredential); } @Override @@ -159,6 +187,34 @@ public class FederatedCatalogWrapper extends CatalogWrapperForREST { String credentialsPath = ResourcePaths.forCatalogProperties(properties).table(identifier) + "/credentials"; + return callRemoteCatalog( + restCatalog, + String.format("loading credentials for table: %s", identifier), + client -> + client.get( + credentialsPath, + LoadCredentialsResponse.class, + Collections.emptyMap(), + ErrorHandlers.tableErrorHandler())); + } + + /** + * Runs an action against the remote REST catalog through a short-lived authenticated client. + * + * <p>Centralizes the auth manager, HTTP client and auth session lifecycle shared by the federated + * credential and load/create/register requests. Resources are closed in reverse order of + * creation, and a close failure on one does not prevent the others from being closed. + * + * @param restCatalog the underlying REST catalog whose properties supply the URI and auth config. + * @param description what the action is doing, used in close-failure log messages. + * @param action invoked with a client bound to an authenticated session. + * @param <T> the action's result type. + * @return the action's result. + */ + private static <T> T callRemoteCatalog( + RESTCatalog restCatalog, String description, Function<RESTClient, T> action) { + Map<String, String> properties = Maps.newHashMap(restCatalog.properties()); + AuthManager authManager = null; RESTClient client = null; AuthSession authSession = null; @@ -170,145 +226,146 @@ public class FederatedCatalogWrapper extends CatalogWrapperForREST { .withHeaders(RESTUtil.configHeaders(properties)) .build(); authSession = authManager.catalogSession(client, properties); - return client - .withAuthSession(authSession) - .get( - credentialsPath, - LoadCredentialsResponse.class, - Collections.emptyMap(), - ErrorHandlers.tableErrorHandler()); + return action.apply(client.withAuthSession(authSession)); } finally { - if (authSession != null) { - try { - authSession.close(); - } catch (Exception e) { - LOG.warn( - "Failed to close auth session when loading credentials for table: {}", identifier, e); - } - } + closeQuietly(authSession, "auth session", description); + closeQuietly(client, "REST client", description); + closeQuietly(authManager, "auth manager", description); + } + } - if (client != null) { - try { - client.close(); - } catch (Exception e) { - LOG.warn( - "Failed to close REST client when loading credentials for table: {}", identifier, e); - } - } + private static void closeQuietly( + AutoCloseable closeable, String resourceName, String description) { + if (closeable == null) { + return; + } - if (authManager != null) { - try { - authManager.close(); - } catch (Exception e) { - LOG.warn( - "Failed to close auth manager when loading credentials for table: {}", identifier, e); - } - } + try { + closeable.close(); + } catch (Exception e) { + LOG.warn("Failed to close {} when {}", resourceName, description, e); } } /** - * Federation-aware {@code createTable}: creates the table on the underlying (remote) catalog and - * extracts client-facing FileIO/credential properties from {@code table.io()}. + * Sends a {@code GET {table}} request to the remote REST catalog. + * + * <p>Follows the same HTTP client lifecycle as {@link #getRESTTableCredentials}. When credential + * vending is requested, the {@code X-Iceberg-Access-Delegation: vended-credentials} header is + * included so the remote catalog returns credentials inline in the load-table response. + * + * @param restCatalog the underlying REST catalog whose properties supply the URI and auth config. + * @param identifier the table to load. + * @param requestCredentialVending whether to include the access-delegation header. + * @return the load-table response from the remote catalog. */ - private LoadTableResponse createTableInternal(Namespace namespace, CreateTableRequest request) { - Catalog loadedCatalog = getCatalog(); - - request.validate(); - - if (request.stageCreate()) { - return stageTableCreateInternal(namespace, request); - } - - TableIdentifier ident = TableIdentifier.of(namespace, request.name()); - Table table = - loadedCatalog - .buildTable(ident, request.schema()) - .withLocation(request.location()) - .withPartitionSpec(request.spec()) - .withSortOrder(request.writeOrder()) - .withProperties(request.properties()) - .create(); - - if (table instanceof BaseTable) { - return buildLoadTableResponseFromFileIo(ident, (BaseTable) table); - } + private static LoadTableResponse getRESTLoadTable( + RESTCatalog restCatalog, TableIdentifier identifier, boolean requestCredentialVending) { + Map<String, String> properties = Maps.newHashMap(restCatalog.properties()); + String tablePath = ResourcePaths.forCatalogProperties(properties).table(identifier); + Map<String, String> queryParams = ImmutableMap.of("snapshots", IcebergRESTUtils.SNAPSHOT_ALL); + + return callRemoteCatalog( + restCatalog, + String.format("loading table: %s", identifier), + client -> + client.get( + tablePath, + queryParams, + LoadTableResponse.class, + accessDelegationHeaders(requestCredentialVending), + ErrorHandlers.tableErrorHandler())); + } - throw new IllegalStateException("Cannot wrap catalog that does not produce BaseTable"); + /** + * Sends a {@code POST {namespace}/tables} request to the remote REST catalog. + * + * @param restCatalog the underlying REST catalog whose properties supply the URI and auth config. + * @param namespace the namespace that will own the table. + * @param request the create-table request (including staged create). + * @param requestCredentialVending whether to include the access-delegation header. + * @return the create response from the remote catalog. + */ + private static LoadTableResponse getRESTCreateTable( + RESTCatalog restCatalog, + Namespace namespace, + CreateTableRequest request, + boolean requestCredentialVending) { + Map<String, String> properties = Maps.newHashMap(restCatalog.properties()); + String tablesPath = ResourcePaths.forCatalogProperties(properties).tables(namespace); + + return callRemoteCatalog( + restCatalog, + String.format("creating table: %s.%s", namespace, request.name()), + client -> + client.post( + tablesPath, + request, + LoadTableResponse.class, + accessDelegationHeaders(requestCredentialVending), + ErrorHandlers.createTableErrorHandler())); } - private LoadTableResponse stageTableCreateInternal( - Namespace namespace, CreateTableRequest request) { - Catalog loadedCatalog = getCatalog(); - TableIdentifier ident = TableIdentifier.of(namespace, request.name()); - if (loadedCatalog.tableExists(ident)) { - throw new AlreadyExistsException("Table already exists: %s", ident); - } + /** + * Sends a {@code POST {namespace}/register} request to the remote REST catalog. + * + * @param restCatalog the underlying REST catalog whose properties supply the URI and auth config. + * @param namespace the namespace that will own the table. + * @param request the register-table request. + * @param requestCredentialVending whether to include the access-delegation header. + * @return the register response from the remote catalog. + */ + private static LoadTableResponse getRESTRegisterTable( + RESTCatalog restCatalog, + Namespace namespace, + RegisterTableRequest request, + boolean requestCredentialVending) { + Map<String, String> properties = Maps.newHashMap(restCatalog.properties()); + String registerPath = ResourcePaths.forCatalogProperties(properties).register(namespace); + + return callRemoteCatalog( + restCatalog, + String.format("registering table: %s.%s", namespace, request.name()), + client -> + client.post( + registerPath, + request, + LoadTableResponse.class, + accessDelegationHeaders(requestCredentialVending), + ErrorHandlers.tableErrorHandler())); + } - Map<String, String> properties = Maps.newHashMap(); - properties.put("created-at", OffsetDateTime.now(ZoneOffset.UTC).toString()); - properties.putAll(request.properties()); - - Map<String, String> config = Maps.newHashMap(); - Catalog.TableBuilder tableBuilder = - loadedCatalog - .buildTable(ident, request.schema()) - .withPartitionSpec(request.spec()) - .withSortOrder(request.writeOrder()) - .withProperties(properties); - - Table table; - if (request.location() != null) { - table = tableBuilder.withLocation(request.location()).createTransaction().table(); - } else { - table = tableBuilder.createTransaction().table(); - } + private static Map<String, String> accessDelegationHeaders(boolean requestCredentialVending) { + return requestCredentialVending + ? ImmutableMap.of(X_ICEBERG_ACCESS_DELEGATION, VENDED_CREDENTIALS) + : Collections.emptyMap(); + } - Map<String, String> tableProperties = retrieveFileIOProperties(table.io()); - Map<String, String> filteredCredentialProperties = - CredentialPropertyUtils.filterCredentialProperties(tableProperties); - config.putAll( - MapUtils.getFilteredMap( - tableProperties, key -> catalogPropertiesToClientKeys.contains(key))); - config.putAll(filteredCredentialProperties); - config.putAll( - IcebergRESTUtils.buildRefreshProps( - catalogCredentialManager.catalogName(), ident, filteredCredentialProperties)); - - List<Credential> credentials = - IcebergRESTUtils.buildStorageCreds( - catalogCredentialManager.catalogName(), ident, table.io()); - - TableMetadata metadata = - TableMetadata.newTableMetadata( - request.schema(), - request.spec() != null ? request.spec() : PartitionSpec.unpartitioned(), - request.writeOrder() != null ? request.writeOrder() : SortOrder.unsorted(), - table.location(), - properties); - - return LoadTableResponse.builder() - .withTableMetadata(metadata) - .addAllConfig(config) - .addAllCredentials(credentials) - .build(); + private LoadTableResponse createTableViaREST( + Namespace namespace, CreateTableRequest request, boolean requestCredential) { + LoadTableResponse upstream = + getRESTCreateTable((RESTCatalog) getCatalog(), namespace, request, requestCredential); + return rewriteRemoteLoadTable(TableIdentifier.of(namespace, request.name()), upstream); } - /** - * Federation-aware {@code registerTable}: registers the existing table metadata on the underlying - * (remote) catalog and extracts client-facing FileIO/credential properties from {@code - * table.io()}, mirroring {@link #loadTableInternal(TableIdentifier)}. - */ - private LoadTableResponse registerTableInternal( - Namespace namespace, RegisterTableRequest request) { - TableIdentifier ident = TableIdentifier.of(namespace, request.name()); - Table table = getCatalog().registerTable(ident, request.metadataLocation()); + private LoadTableResponse loadTableViaREST( + TableIdentifier identifier, boolean requestCredential) { + LoadTableResponse upstream = + getRESTLoadTable((RESTCatalog) getCatalog(), identifier, requestCredential); + return rewriteRemoteLoadTable(identifier, upstream); + } - if (table instanceof BaseTable) { - return buildLoadTableResponseFromFileIo(ident, (BaseTable) table); - } + private LoadTableResponse registerTableViaREST( + Namespace namespace, RegisterTableRequest request, boolean requestCredential) { + LoadTableResponse upstream = + getRESTRegisterTable((RESTCatalog) getCatalog(), namespace, request, requestCredential); + return rewriteRemoteLoadTable(TableIdentifier.of(namespace, request.name()), upstream); + } - throw new IllegalStateException("Cannot wrap catalog that does not produce BaseTable"); + private LoadTableResponse rewriteRemoteLoadTable( + TableIdentifier identifier, LoadTableResponse upstream) { + return IcebergRESTUtils.rewriteLoadTableCredentials( + catalogCredentialManager.catalogName(), identifier, upstream); } /** @@ -361,52 +418,6 @@ public class FederatedCatalogWrapper extends CatalogWrapperForREST { } } - /** - * Federation-aware {@code loadTable}: loads the table from the underlying (remote) catalog and - * extracts client-facing FileIO/credential properties from {@code table.io()}. - */ - private LoadTableResponse loadTableInternal(TableIdentifier ident) { - Table table = getCatalog().loadTable(ident); - - if (table instanceof BaseTable) { - return buildLoadTableResponseFromFileIo(ident, (BaseTable) table); - } else if (table instanceof BaseMetadataTable) { - // metadata tables are loaded on the client side, return NoSuchTableException for now - throw new NoSuchTableException("Table does not exist: %s", ident.toString()); - } - - throw new IllegalStateException("Cannot wrap catalog that does not produce BaseTable"); - } - - /** - * Builds a {@link LoadTableResponse} from a remote {@link BaseTable}, exposing the client-facing - * FileIO and credential properties extracted from {@code table.io()}, including the refreshable - * vended credentials and refresh properties for the remote storage. - * - * @param ident the table identifier, used to build the credential-refresh endpoint. - * @param table the remote base table whose {@code io()} carries the storage credentials. - * @return the load-table response including FileIO-derived client config and vended credentials. - */ - private LoadTableResponse buildLoadTableResponseFromFileIo( - TableIdentifier ident, BaseTable table) { - Map<String, String> properties = retrieveFileIOProperties(table.io()); - Map<String, String> filteredCredentialProperties = - CredentialPropertyUtils.filterCredentialProperties(properties); - return LoadTableResponse.builder() - .withTableMetadata(table.operations().current()) - .addAllConfig( - MapUtils.getFilteredMap(properties, key -> catalogPropertiesToClientKeys.contains(key))) - // Keep only credential fields from FileIO properties before returning them to the client. - .addAllConfig(filteredCredentialProperties) - .addAllConfig( - IcebergRESTUtils.buildRefreshProps( - catalogCredentialManager.catalogName(), ident, filteredCredentialProperties)) - .addAllCredentials( - IcebergRESTUtils.buildStorageCreds( - catalogCredentialManager.catalogName(), ident, table.io())) - .build(); - } - private static boolean isCreate(UpdateTableRequest request) { boolean isCreate = request.requirements().stream() @@ -501,10 +512,4 @@ public class FederatedCatalogWrapper extends CatalogWrapperForREST { return true; } - - private static Map<String, String> retrieveFileIOProperties(FileIO fileIO) { - return fileIO instanceof InMemoryFileIO - ? Maps.newHashMap() - : new HashMap<>(fileIO.properties()); - } } diff --git a/iceberg/iceberg-rest-server/src/main/java/org/apache/gravitino/iceberg/service/IcebergRESTUtils.java b/iceberg/iceberg-rest-server/src/main/java/org/apache/gravitino/iceberg/service/IcebergRESTUtils.java index 9a4d74091d..bedd609d37 100644 --- a/iceberg/iceberg-rest-server/src/main/java/org/apache/gravitino/iceberg/service/IcebergRESTUtils.java +++ b/iceberg/iceberg-rest-server/src/main/java/org/apache/gravitino/iceberg/service/IcebergRESTUtils.java @@ -21,6 +21,7 @@ package org.apache.gravitino.iceberg.service; import com.fasterxml.jackson.databind.ObjectMapper; import com.google.common.base.Preconditions; import com.google.common.collect.ImmutableMap; +import com.google.common.collect.ImmutableSet; import java.io.IOException; import java.nio.charset.StandardCharsets; import java.security.MessageDigest; @@ -36,6 +37,7 @@ import java.util.HashMap; import java.util.List; import java.util.Map; import java.util.Optional; +import java.util.Set; import java.util.stream.Stream; import javax.servlet.http.HttpServletRequest; import javax.ws.rs.core.EntityTag; @@ -77,6 +79,18 @@ public class IcebergRESTUtils { public static final String SNAPSHOT_REFS = "refs"; + /** + * Iceberg refresh-endpoint keys that may appear in {@link LoadTableResponse#config()}. Kept in + * sync with {@link CredentialPropertyUtils#buildRefreshProps}; they are not retained by {@link + * CredentialPropertyUtils#filterCredentialProperties}, so top-level config must drop them + * explicitly before IRC-local endpoints are re-applied. + */ + private static final Set<String> REFRESH_CREDENTIALS_ENDPOINT_KEYS = + ImmutableSet.of( + "client.refresh-credentials-endpoint", + "gcs.oauth2.refresh-credentials-endpoint", + "adls.refresh-credentials-endpoint"); + /** Snapshot modes for the Iceberg loadTable endpoint. */ public enum SnapshotMode { ALL(SNAPSHOT_ALL), @@ -232,6 +246,49 @@ public class IcebergRESTUtils { return toRESTCredential(prefix, ImmutableMap.copyOf(filteredConfig)); } + /** + * Rewrites credentials in a federated {@link LoadTableResponse} so their {@code + * refresh-credentials-endpoint} entries, including any flattened into {@code config}, point at + * this IRC instance instead of the upstream catalog. + * + * <p>Upstream refresh endpoints are removed from top-level {@code config} before IRC-local ones + * are re-applied, matching {@link #rewriteCredential}. Without that step a refresh URL that lived + * only in {@code config} (with tokens only in {@code storage-credentials}) would leak. + * + * @param catalogName IRC catalog name used to build refresh paths + * @param tableIdentifier table receiving the credentials + * @param upstream the load-table response returned by the upstream REST catalog + * @return a load-table response with IRC-local refresh endpoints + */ + public static LoadTableResponse rewriteLoadTableCredentials( + String catalogName, TableIdentifier tableIdentifier, LoadTableResponse upstream) { + Map<String, String> config = new HashMap<>(); + if (upstream.config() != null) { + config.putAll(upstream.config()); + } + // filterCredentialProperties drops refresh endpoints from its return value but putAll does not + // remove keys already copied from upstream.config(). Drop them first so a refresh URL that + // lived only in config (tokens only in storage-credentials) cannot leak to clients. + config.keySet().removeAll(REFRESH_CREDENTIALS_ENDPOINT_KEYS); + Map<String, String> filteredCredentialProperties = + CredentialPropertyUtils.filterCredentialProperties(config); + config.putAll(filteredCredentialProperties); + config.putAll(buildRefreshProps(catalogName, tableIdentifier, filteredCredentialProperties)); + + LoadTableResponse.Builder builder = + LoadTableResponse.builder() + .withTableMetadata(upstream.tableMetadata()) + .addAllConfig(config); + if (upstream.credentials() != null) { + for (org.apache.iceberg.rest.credentials.Credential credential : upstream.credentials()) { + builder.addCredential( + rewriteCredential( + catalogName, tableIdentifier, credential.prefix(), credential.config())); + } + } + return builder.build(); + } + public static <T> Response ok(T t) { return Response.status(Response.Status.OK).entity(t).type(MediaType.APPLICATION_JSON).build(); } diff --git a/iceberg/iceberg-rest-server/src/main/java/org/apache/gravitino/iceberg/service/rest/IcebergTableOperations.java b/iceberg/iceberg-rest-server/src/main/java/org/apache/gravitino/iceberg/service/rest/IcebergTableOperations.java index 24d567e6c7..1782fe8ec9 100644 --- a/iceberg/iceberg-rest-server/src/main/java/org/apache/gravitino/iceberg/service/rest/IcebergTableOperations.java +++ b/iceberg/iceberg-rest-server/src/main/java/org/apache/gravitino/iceberg/service/rest/IcebergTableOperations.java @@ -567,10 +567,14 @@ public class IcebergTableOperations { } TableMetadata filteredMetadata = TableMetadata.buildFrom(metadata).suppressHistoricalSnapshots().build(); - return LoadTableResponse.builder() - .withTableMetadata(filteredMetadata) - .addAllConfig(loadTableResponse.config()) - .build(); + LoadTableResponse.Builder builder = + LoadTableResponse.builder() + .withTableMetadata(filteredMetadata) + .addAllConfig(loadTableResponse.config()); + if (loadTableResponse.credentials() != null) { + builder.addAllCredentials(loadTableResponse.credentials()); + } + return builder.build(); } private static Response buildResponseWithETag(LoadTableResponse loadTableResponse) { diff --git a/iceberg/iceberg-rest-server/src/test/java/org/apache/gravitino/iceberg/service/TestCatalogWrapperForREST.java b/iceberg/iceberg-rest-server/src/test/java/org/apache/gravitino/iceberg/service/TestCatalogWrapperForREST.java index e372170585..612a618524 100644 --- a/iceberg/iceberg-rest-server/src/test/java/org/apache/gravitino/iceberg/service/TestCatalogWrapperForREST.java +++ b/iceberg/iceberg-rest-server/src/test/java/org/apache/gravitino/iceberg/service/TestCatalogWrapperForREST.java @@ -20,14 +20,12 @@ package org.apache.gravitino.iceberg.service; import static org.mockito.Mockito.any; -import static org.mockito.Mockito.anyMap; import static org.mockito.Mockito.anyString; import static org.mockito.Mockito.doAnswer; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.never; import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; -import static org.mockito.Mockito.withSettings; import com.google.common.collect.ImmutableMap; import com.sun.net.httpserver.HttpServer; @@ -47,17 +45,15 @@ import org.apache.gravitino.credential.CredentialConstants; import org.apache.gravitino.credential.CredentialPrivilege; import org.apache.gravitino.iceberg.common.IcebergConfig; import org.apache.gravitino.iceberg.service.extension.DummyCredentialProvider; -import org.apache.iceberg.BaseTable; import org.apache.iceberg.BaseTransaction; import org.apache.iceberg.CatalogProperties; import org.apache.iceberg.MetadataUpdate; import org.apache.iceberg.PartitionSpec; import org.apache.iceberg.Schema; import org.apache.iceberg.SortOrder; -import org.apache.iceberg.Table; import org.apache.iceberg.TableMetadata; +import org.apache.iceberg.TableMetadataParser; import org.apache.iceberg.TableOperations; -import org.apache.iceberg.Transaction; import org.apache.iceberg.UpdateRequirement; import org.apache.iceberg.catalog.Catalog; import org.apache.iceberg.catalog.Namespace; @@ -67,10 +63,7 @@ import org.apache.iceberg.exceptions.ForbiddenException; import org.apache.iceberg.exceptions.NoSuchTableException; import org.apache.iceberg.exceptions.NotAuthorizedException; import org.apache.iceberg.exceptions.ServiceFailureException; -import org.apache.iceberg.io.FileIO; import org.apache.iceberg.io.ResolvingFileIO; -import org.apache.iceberg.io.StorageCredential; -import org.apache.iceberg.io.SupportsStorageCredentials; import org.apache.iceberg.rest.RESTCatalog; import org.apache.iceberg.rest.auth.AuthProperties; import org.apache.iceberg.rest.credentials.Credential; @@ -80,6 +73,7 @@ import org.apache.iceberg.rest.requests.RegisterTableRequest; import org.apache.iceberg.rest.requests.UpdateTableRequest; import org.apache.iceberg.rest.responses.LoadCredentialsResponse; import org.apache.iceberg.rest.responses.LoadTableResponse; +import org.apache.iceberg.rest.responses.LoadTableResponseParser; import org.apache.iceberg.types.Types; import org.junit.jupiter.api.Assertions; import org.junit.jupiter.api.Test; @@ -431,70 +425,93 @@ public class TestCatalogWrapperForREST { } @Test - void testLoadTableRefreshEndpoint() { + void testLoadTableRefreshEndpoint() throws Exception { TableIdentifier ident = TableIdentifier.of(Namespace.of("db"), "tbl"); - RESTCatalog catalog = mock(RESTCatalog.class); - BaseTable baseTable = mock(BaseTable.class); - TableOperations ops = mock(TableOperations.class); - FileIO fileIO = mock(FileIO.class); TableMetadata metadata = - TableMetadata.newTableMetadata( - new Schema(Types.NestedField.required(1, "id", Types.IntegerType.get())), - PartitionSpec.unpartitioned(), - SortOrder.unsorted(), - "s3://bucket/db/tbl", - Collections.emptyMap()); - - when(catalog.loadTable(ident)).thenReturn(baseTable); - when(baseTable.operations()).thenReturn(ops); - when(ops.current()).thenReturn(metadata); - when(baseTable.io()).thenReturn(fileIO); - when(fileIO.properties()) - .thenReturn( - ImmutableMap.of( - "s3.session-token", - "token", - "s3.session-token-expires-at-ms", - "123", - "client.refresh-credentials-endpoint", - "v1/upstream/namespaces/db/tables/tbl/credentials")); + TableMetadataParser.fromJson( + "s3://bucket/db/tbl/metadata/v1.metadata.json", + TableMetadataParser.toJson( + TableMetadata.newTableMetadata( + new Schema(Types.NestedField.required(1, "id", Types.IntegerType.get())), + PartitionSpec.unpartitioned(), + SortOrder.unsorted(), + "s3://bucket/db/tbl", + Collections.emptyMap()))); + LoadTableResponse upstreamResponse = + LoadTableResponse.builder() + .withTableMetadata(metadata) + .addAllConfig( + ImmutableMap.of( + "s3.session-token", + "token", + "s3.session-token-expires-at-ms", + "123", + "client.refresh-credentials-endpoint", + "v1/upstream/namespaces/db/tables/tbl/credentials")) + .build(); + String upstreamJson = LoadTableResponseParser.toJson(upstreamResponse); - IcebergConfig config = - new IcebergConfig( - ImmutableMap.of( - IcebergConstants.CATALOG_BACKEND, - "memory", - IcebergConstants.WAREHOUSE, - "/tmp/warehouse")); - CatalogWrapperForREST wrapper = new StaticCatalogWrapperForREST("irc1", config, catalog); + HttpServer server = HttpServer.create(new InetSocketAddress(0), 0); + server.createContext( + "/", + exchange -> { + byte[] body = upstreamJson.getBytes(StandardCharsets.UTF_8); + exchange.getResponseHeaders().add("Content-Type", "application/json"); + exchange.sendResponseHeaders(200, body.length); + try (OutputStream os = exchange.getResponseBody()) { + os.write(body); + } + }); + server.start(); + try { + String uri = "http://127.0.0.1:"; + server.getAddress().getPort(); + RESTCatalog catalog = mock(RESTCatalog.class); + when(catalog.name()).thenReturn("upstream"); + when(catalog.properties()) + .thenReturn( + ImmutableMap.of( + CatalogProperties.URI, + uri, + AuthProperties.AUTH_TYPE, + AuthProperties.AUTH_TYPE_NONE, + "prefix", + "upstream")); + + IcebergConfig config = + new IcebergConfig( + ImmutableMap.of( + IcebergConstants.CATALOG_BACKEND, + "memory", + IcebergConstants.WAREHOUSE, + "/tmp/warehouse")); + CatalogWrapperForREST wrapper = new StaticCatalogWrapperForREST("irc1", config, catalog); - LoadTableResponse response = wrapper.loadTable(ident, false, CredentialPrivilege.READ); + LoadTableResponse response = wrapper.loadTable(ident, true, CredentialPrivilege.READ); - Assertions.assertEquals( - "v1/irc1/namespaces/db/tables/tbl/credentials", - response.config().get("client.refresh-credentials-endpoint")); - Assertions.assertEquals("token", response.config().get("s3.session-token")); + Assertions.assertEquals( + "v1/irc1/namespaces/db/tables/tbl/credentials", + response.config().get("client.refresh-credentials-endpoint")); + Assertions.assertEquals("token", response.config().get("s3.session-token")); + } finally { + server.stop(0); + } } @Test - void testLoadTableStorageCreds() { + void testLoadTableStorageCreds() throws Exception { TableIdentifier ident = TableIdentifier.of(Namespace.of("db"), "tbl"); - RESTCatalog catalog = mock(RESTCatalog.class); - BaseTable baseTable = mock(BaseTable.class); - TableOperations ops = mock(TableOperations.class); - FileIO fileIO = - mock(FileIO.class, withSettings().extraInterfaces(SupportsStorageCredentials.class)); - SupportsStorageCredentials storageCredentialsFileIO = (SupportsStorageCredentials) fileIO; TableMetadata metadata = - TableMetadata.newTableMetadata( - new Schema(Types.NestedField.required(1, "id", Types.IntegerType.get())), - PartitionSpec.unpartitioned(), - SortOrder.unsorted(), - "s3://bucket/db/tbl", - Collections.emptyMap()); - - StorageCredential upstreamCredential = - StorageCredential.create( + TableMetadataParser.fromJson( + "s3://bucket/db/tbl/metadata/v1.metadata.json", + TableMetadataParser.toJson( + TableMetadata.newTableMetadata( + new Schema(Types.NestedField.required(1, "id", Types.IntegerType.get())), + PartitionSpec.unpartitioned(), + SortOrder.unsorted(), + "s3://bucket/db/tbl", + Collections.emptyMap()))); + Credential upstreamCredential = + IcebergRESTUtils.toRESTCredential( "s3://bucket/db/tbl/", ImmutableMap.of( "s3.access-key-id", @@ -507,33 +524,61 @@ public class TestCatalogWrapperForREST { "123", "client.refresh-credentials-endpoint", "v1/upstream/namespaces/db/tables/tbl/credentials")); + LoadTableResponse upstreamResponse = + LoadTableResponse.builder() + .withTableMetadata(metadata) + .addCredential(upstreamCredential) + .build(); + String upstreamJson = LoadTableResponseParser.toJson(upstreamResponse); - when(catalog.loadTable(ident)).thenReturn(baseTable); - when(baseTable.operations()).thenReturn(ops); - when(ops.current()).thenReturn(metadata); - when(baseTable.io()).thenReturn(fileIO); - when(fileIO.properties()).thenReturn(Collections.emptyMap()); - when(storageCredentialsFileIO.credentials()).thenReturn(List.of(upstreamCredential)); + HttpServer server = HttpServer.create(new InetSocketAddress(0), 0); + server.createContext( + "/", + exchange -> { + byte[] body = upstreamJson.getBytes(StandardCharsets.UTF_8); + exchange.getResponseHeaders().add("Content-Type", "application/json"); + exchange.sendResponseHeaders(200, body.length); + try (OutputStream os = exchange.getResponseBody()) { + os.write(body); + } + }); + server.start(); + try { + String uri = "http://127.0.0.1:"; + server.getAddress().getPort(); + RESTCatalog catalog = mock(RESTCatalog.class); + when(catalog.name()).thenReturn("upstream"); + when(catalog.properties()) + .thenReturn( + ImmutableMap.of( + CatalogProperties.URI, + uri, + AuthProperties.AUTH_TYPE, + AuthProperties.AUTH_TYPE_NONE, + "prefix", + "upstream")); - IcebergConfig config = - new IcebergConfig( - ImmutableMap.of( - IcebergConstants.CATALOG_BACKEND, - "memory", - IcebergConstants.WAREHOUSE, - "/tmp/warehouse")); - CatalogWrapperForREST wrapper = new StaticCatalogWrapperForREST("irc1", config, catalog); + IcebergConfig config = + new IcebergConfig( + ImmutableMap.of( + IcebergConstants.CATALOG_BACKEND, + "memory", + IcebergConstants.WAREHOUSE, + "/tmp/warehouse")); + CatalogWrapperForREST wrapper = new StaticCatalogWrapperForREST("irc1", config, catalog); - LoadTableResponse response = wrapper.loadTable(ident, false, CredentialPrivilege.READ); + LoadTableResponse response = wrapper.loadTable(ident, true, CredentialPrivilege.READ); - Assertions.assertEquals(1, response.credentials().size()); - Credential credential = response.credentials().get(0); - Assertions.assertEquals("s3://bucket/db/tbl/", credential.prefix()); - Assertions.assertEquals("upstream-token", credential.config().get("s3.session-token")); - Assertions.assertEquals( - "v1/irc1/namespaces/db/tables/tbl/credentials", - credential.config().get("client.refresh-credentials-endpoint")); - Assertions.assertFalse(response.config().containsKey("client.refresh-credentials-endpoint")); + Assertions.assertEquals(1, response.credentials().size()); + Credential credential = response.credentials().get(0); + Assertions.assertEquals("s3://bucket/db/tbl/", credential.prefix()); + Assertions.assertEquals("upstream-token", credential.config().get("s3.session-token")); + Assertions.assertEquals( + "v1/irc1/namespaces/db/tables/tbl/credentials", + credential.config().get("client.refresh-credentials-endpoint")); + Assertions.assertFalse(response.config().containsKey("client.refresh-credentials-endpoint")); + } finally { + server.stop(0); + } } @Test @@ -621,47 +666,6 @@ public class TestCatalogWrapperForREST { () -> CatalogWrapperForREST.filterCatalogConfigForClients(source)); } - @Test - void testFederatedRegisterTableIncludesFileIo() { - RESTCatalog catalog = mock(RESTCatalog.class); - BaseTable table = mock(BaseTable.class); - TableOperations ops = mock(TableOperations.class); - FileIO fileIO = mock(FileIO.class); - when(catalog.registerTable(any(TableIdentifier.class), anyString())).thenReturn(table); - when(table.operations()).thenReturn(ops); - when(ops.current()).thenReturn(minimalTableMetadataForStagedCreateTest()); - when(table.io()).thenReturn(fileIO); - when(fileIO.properties()) - .thenReturn( - ImmutableMap.of( - IcebergConstants.IO_IMPL, - "org.apache.iceberg.aws.s3.S3FileIO", - IcebergConstants.ICEBERG_S3_ENDPOINT, - "http://localhost:9000";)); - - IcebergConfig config = - new IcebergConfig( - ImmutableMap.of( - IcebergConstants.CATALOG_BACKEND, - "memory", - IcebergConstants.WAREHOUSE, - "/tmp/warehouse")); - CatalogWrapperForREST wrapper = new StaticCatalogWrapperForREST("test", config, catalog); - - RegisterTableRequest request = - ImmutableRegisterTableRequest.builder() - .name("tbl") - .metadataLocation("s3://bucket/warehouse/tbl/metadata/v1.metadata.json") - .build(); - - LoadTableResponse response = wrapper.registerTable(Namespace.of("db"), request, false); - - Assertions.assertEquals( - "org.apache.iceberg.aws.s3.S3FileIO", response.config().get(IcebergConstants.IO_IMPL)); - Assertions.assertEquals( - "http://localhost:9000";, response.config().get(IcebergConstants.ICEBERG_S3_ENDPOINT)); - } - @Test void testWrapperLazyLoadsCatalog() { IcebergConfig config = @@ -680,102 +684,6 @@ public class TestCatalogWrapperForREST { } } - @Test - void testStageCreateWithLocationIncludesFileIo() throws Exception { - RESTCatalog catalog = mock(RESTCatalog.class); - Catalog.TableBuilder tableBuilder = mock(Catalog.TableBuilder.class); - Transaction transaction = mock(Transaction.class); - Table table = mock(Table.class); - FileIO fileIO = mock(FileIO.class); - when(catalog.buildTable(any(TableIdentifier.class), any())).thenReturn(tableBuilder); - when(tableBuilder.withPartitionSpec(any())).thenReturn(tableBuilder); - when(tableBuilder.withSortOrder(any())).thenReturn(tableBuilder); - when(tableBuilder.withProperties(anyMap())).thenReturn(tableBuilder); - when(tableBuilder.withLocation("s3://bucket/warehouse/table")).thenReturn(tableBuilder); - when(tableBuilder.createTransaction()).thenReturn(transaction); - when(transaction.table()).thenReturn(table); - when(table.io()).thenReturn(fileIO); - when(table.location()).thenReturn("s3://bucket/warehouse/table"); - when(fileIO.properties()) - .thenReturn( - ImmutableMap.of( - IcebergConstants.IO_IMPL, - "org.apache.iceberg.aws.s3.S3FileIO", - IcebergConstants.ICEBERG_S3_ENDPOINT, - "http://localhost:9000";)); - - IcebergConfig config = - new IcebergConfig( - ImmutableMap.of( - IcebergConstants.CATALOG_BACKEND, - "memory", - IcebergConstants.WAREHOUSE, - "/tmp/warehouse")); - CatalogWrapperForREST wrapper = new StaticCatalogWrapperForREST("test", config, catalog); - - Schema schema = new Schema(Types.NestedField.required(1, "id", Types.IntegerType.get())); - CreateTableRequest request = - CreateTableRequest.builder() - .withName("tbl") - .withSchema(schema) - .withLocation("s3://bucket/warehouse/table") - .stageCreate() - .build(); - - LoadTableResponse response = wrapper.createTable(Namespace.of("db"), request, false); - - Assertions.assertEquals( - "org.apache.iceberg.aws.s3.S3FileIO", response.config().get(IcebergConstants.IO_IMPL)); - Assertions.assertEquals( - "http://localhost:9000";, response.config().get(IcebergConstants.ICEBERG_S3_ENDPOINT)); - verify(tableBuilder).withLocation("s3://bucket/warehouse/table"); - } - - @Test - void testStageCreateNullLocationSkipsWithLocation() { - RESTCatalog catalog = mock(RESTCatalog.class); - Catalog.TableBuilder tableBuilder = mock(Catalog.TableBuilder.class); - Transaction transaction = mock(Transaction.class); - Table table = mock(Table.class); - FileIO fileIO = mock(FileIO.class); - when(catalog.buildTable(any(TableIdentifier.class), any())).thenReturn(tableBuilder); - when(tableBuilder.withPartitionSpec(any())).thenReturn(tableBuilder); - when(tableBuilder.withSortOrder(any())).thenReturn(tableBuilder); - when(tableBuilder.withProperties(anyMap())).thenReturn(tableBuilder); - when(tableBuilder.createTransaction()).thenReturn(transaction); - when(transaction.table()).thenReturn(table); - when(table.io()).thenReturn(fileIO); - when(table.location()).thenReturn("s3://bucket/warehouse/default-location"); - when(fileIO.properties()) - .thenReturn( - ImmutableMap.of( - IcebergConstants.IO_IMPL, - "org.apache.iceberg.aws.s3.S3FileIO", - IcebergConstants.ICEBERG_S3_ENDPOINT, - "http://localhost:9000";)); - - IcebergConfig config = - new IcebergConfig( - ImmutableMap.of( - IcebergConstants.CATALOG_BACKEND, - "memory", - IcebergConstants.WAREHOUSE, - "/tmp/warehouse")); - CatalogWrapperForREST wrapper = new StaticCatalogWrapperForREST("test", config, catalog); - - Schema schema = new Schema(Types.NestedField.required(1, "id", Types.IntegerType.get())); - CreateTableRequest request = - CreateTableRequest.builder().withName("tbl").withSchema(schema).stageCreate().build(); - - LoadTableResponse response = wrapper.createTable(Namespace.of("db"), request, false); - - Assertions.assertEquals( - "org.apache.iceberg.aws.s3.S3FileIO", response.config().get(IcebergConstants.IO_IMPL)); - Assertions.assertEquals( - "http://localhost:9000";, response.config().get(IcebergConstants.ICEBERG_S3_ENDPOINT)); - verify(tableBuilder, never()).withLocation(any()); - } - @Test void testStagedCreateRejectsExtraRequirements() { RESTCatalog catalog = mock(RESTCatalog.class); @@ -1027,4 +935,668 @@ public class TestCatalogWrapperForREST { return catalog; } } + + @Test + void testFederatedLoadTableDelegatesToRemote() throws Exception { + TableIdentifier table = TableIdentifier.of(Namespace.of("db"), "tbl"); + String expectedPath = "/v1/upstream/namespaces/db/tables/tbl"; + + TableMetadata metadata = + TableMetadataParser.fromJson( + "s3://bucket/db/tbl/metadata/v1.metadata.json", + TableMetadataParser.toJson( + TableMetadata.newTableMetadata( + new Schema(Types.NestedField.required(1, "id", Types.IntegerType.get())), + PartitionSpec.unpartitioned(), + SortOrder.unsorted(), + "s3://bucket/db/tbl", + Collections.emptyMap()))); + org.apache.iceberg.rest.credentials.Credential cred = + IcebergRESTUtils.toRESTCredential( + "s3://bucket/db/tbl/", + ImmutableMap.of( + "s3.access-key-id", "upstream-key", + "s3.secret-access-key", "upstream-secret", + "s3.session-token", "upstream-token", + "client.refresh-credentials-endpoint", + "v1/upstream/namespaces/db/tables/tbl/credentials")); + LoadTableResponse upstreamResponse = + LoadTableResponse.builder() + .withTableMetadata(metadata) + .addAllConfig(ImmutableMap.of("io-impl", "org.apache.iceberg.aws.s3.S3FileIO")) + .addCredential(cred) + .build(); + String upstreamJson = LoadTableResponseParser.toJson(upstreamResponse); + + AtomicReference<String> requestPath = new AtomicReference<>(); + AtomicReference<String> requestMethod = new AtomicReference<>(); + AtomicReference<String> requestQuery = new AtomicReference<>(); + AtomicReference<String> accessDelegationHeader = new AtomicReference<>(); + HttpServer server = HttpServer.create(new InetSocketAddress(0), 0); + server.createContext( + "/", + exchange -> { + requestPath.set(exchange.getRequestURI().getPath()); + requestMethod.set(exchange.getRequestMethod()); + requestQuery.set(exchange.getRequestURI().getQuery()); + accessDelegationHeader.set( + exchange.getRequestHeaders().getFirst("X-Iceberg-Access-Delegation")); + byte[] body = upstreamJson.getBytes(StandardCharsets.UTF_8); + exchange.getResponseHeaders().add("Content-Type", "application/json"); + exchange.sendResponseHeaders(200, body.length); + try (OutputStream os = exchange.getResponseBody()) { + os.write(body); + } + }); + server.start(); + try { + String uri = "http://127.0.0.1:"; + server.getAddress().getPort(); + RESTCatalog restCatalog = mock(RESTCatalog.class); + when(restCatalog.name()).thenReturn("upstream"); + when(restCatalog.properties()) + .thenReturn( + ImmutableMap.of( + CatalogProperties.URI, + uri, + AuthProperties.AUTH_TYPE, + AuthProperties.AUTH_TYPE_NONE, + "prefix", + "upstream")); + + IcebergConfig config = + new IcebergConfig( + ImmutableMap.of( + IcebergConstants.CATALOG_BACKEND, + "memory", + IcebergConstants.WAREHOUSE, + "/tmp/warehouse")); + CatalogWrapperForREST wrapper = new StaticCatalogWrapperForREST("local", config, restCatalog); + + LoadTableResponse response = wrapper.loadTable(table, true, CredentialPrivilege.READ); + + Assertions.assertEquals(expectedPath, requestPath.get()); + Assertions.assertEquals("GET", requestMethod.get()); + Assertions.assertEquals("snapshots=all", requestQuery.get()); + Assertions.assertEquals("vended-credentials", accessDelegationHeader.get()); + verify(restCatalog, never()).loadTable(table); + Assertions.assertEquals(1, response.credentials().size()); + Credential credential = response.credentials().get(0); + Assertions.assertEquals("s3://bucket/db/tbl/", credential.prefix()); + Assertions.assertEquals("upstream-key", credential.config().get("s3.access-key-id")); + Assertions.assertEquals("upstream-token", credential.config().get("s3.session-token")); + Assertions.assertEquals( + "v1/local/namespaces/db/tables/tbl/credentials", + credential.config().get("client.refresh-credentials-endpoint")); + Assertions.assertEquals( + "org.apache.iceberg.aws.s3.S3FileIO", response.config().get("io-impl")); + } finally { + server.stop(0); + } + } + + @Test + void testFederatedLoadTableNoCredentials() throws Exception { + TableIdentifier table = TableIdentifier.of(Namespace.of("db"), "tbl"); + TableMetadata metadata = + TableMetadataParser.fromJson( + "s3://bucket/db/tbl/metadata/v1.metadata.json", + TableMetadataParser.toJson( + TableMetadata.newTableMetadata( + new Schema(Types.NestedField.required(1, "id", Types.IntegerType.get())), + PartitionSpec.unpartitioned(), + SortOrder.unsorted(), + "s3://bucket/db/tbl", + Collections.emptyMap()))); + LoadTableResponse upstreamResponse = + LoadTableResponse.builder() + .withTableMetadata(metadata) + .addAllConfig(ImmutableMap.of("io-impl", "org.apache.iceberg.aws.s3.S3FileIO")) + .build(); + String upstreamJson = LoadTableResponseParser.toJson(upstreamResponse); + + AtomicReference<String> accessDelegationHeader = new AtomicReference<>(); + HttpServer server = HttpServer.create(new InetSocketAddress(0), 0); + server.createContext( + "/", + exchange -> { + accessDelegationHeader.set( + exchange.getRequestHeaders().getFirst("X-Iceberg-Access-Delegation")); + byte[] body = upstreamJson.getBytes(StandardCharsets.UTF_8); + exchange.getResponseHeaders().add("Content-Type", "application/json"); + exchange.sendResponseHeaders(200, body.length); + try (OutputStream os = exchange.getResponseBody()) { + os.write(body); + } + }); + server.start(); + try { + String uri = "http://127.0.0.1:"; + server.getAddress().getPort(); + RESTCatalog restCatalog = mock(RESTCatalog.class); + when(restCatalog.name()).thenReturn("upstream"); + when(restCatalog.properties()) + .thenReturn( + ImmutableMap.of( + CatalogProperties.URI, + uri, + AuthProperties.AUTH_TYPE, + AuthProperties.AUTH_TYPE_NONE, + "prefix", + "upstream")); + + IcebergConfig config = + new IcebergConfig( + ImmutableMap.of( + IcebergConstants.CATALOG_BACKEND, + "memory", + IcebergConstants.WAREHOUSE, + "/tmp/warehouse")); + CatalogWrapperForREST wrapper = new StaticCatalogWrapperForREST("local", config, restCatalog); + + LoadTableResponse response = wrapper.loadTable(table, false, CredentialPrivilege.READ); + + Assertions.assertNull( + accessDelegationHeader.get(), + "X-Iceberg-Access-Delegation header should not be sent without credential vending"); + verify(restCatalog, never()).loadTable(table); + Assertions.assertTrue( + response.credentials() == null || response.credentials().isEmpty(), + "Non-vended request should not return remote storage-credentials"); + Assertions.assertEquals( + "org.apache.iceberg.aws.s3.S3FileIO", response.config().get("io-impl")); + } finally { + server.stop(0); + } + } + + @Test + void testFederatedCreateTableWithCredentials() throws Exception { + Namespace namespace = Namespace.of("db"); + TableMetadata metadata = + TableMetadataParser.fromJson( + "s3://bucket/db/tbl/metadata/v1.metadata.json", + TableMetadataParser.toJson( + TableMetadata.newTableMetadata( + new Schema(Types.NestedField.required(1, "id", Types.IntegerType.get())), + PartitionSpec.unpartitioned(), + SortOrder.unsorted(), + "s3://bucket/db/tbl", + Collections.emptyMap()))); + Credential cred = + IcebergRESTUtils.toRESTCredential( + "s3://bucket/db/tbl/", + ImmutableMap.of( + "s3.session-token", + "upstream-token", + "client.refresh-credentials-endpoint", + "v1/upstream/namespaces/db/tables/tbl/credentials")); + LoadTableResponse upstreamResponse = + LoadTableResponse.builder().withTableMetadata(metadata).addCredential(cred).build(); + String upstreamJson = LoadTableResponseParser.toJson(upstreamResponse); + + AtomicReference<String> requestPath = new AtomicReference<>(); + AtomicReference<String> requestMethod = new AtomicReference<>(); + AtomicReference<String> accessDelegationHeader = new AtomicReference<>(); + HttpServer server = HttpServer.create(new InetSocketAddress(0), 0); + server.createContext( + "/", + exchange -> { + requestPath.set(exchange.getRequestURI().getPath()); + requestMethod.set(exchange.getRequestMethod()); + accessDelegationHeader.set( + exchange.getRequestHeaders().getFirst("X-Iceberg-Access-Delegation")); + byte[] body = upstreamJson.getBytes(StandardCharsets.UTF_8); + exchange.getResponseHeaders().add("Content-Type", "application/json"); + exchange.sendResponseHeaders(200, body.length); + try (OutputStream os = exchange.getResponseBody()) { + os.write(body); + } + }); + server.start(); + try { + String uri = "http://127.0.0.1:"; + server.getAddress().getPort(); + RESTCatalog restCatalog = mock(RESTCatalog.class); + when(restCatalog.name()).thenReturn("upstream"); + when(restCatalog.properties()) + .thenReturn( + ImmutableMap.of( + CatalogProperties.URI, + uri, + AuthProperties.AUTH_TYPE, + AuthProperties.AUTH_TYPE_NONE, + "prefix", + "upstream")); + + IcebergConfig config = + new IcebergConfig( + ImmutableMap.of( + IcebergConstants.CATALOG_BACKEND, + "memory", + IcebergConstants.WAREHOUSE, + "/tmp/warehouse")); + CatalogWrapperForREST wrapper = new StaticCatalogWrapperForREST("local", config, restCatalog); + CreateTableRequest request = + CreateTableRequest.builder() + .withName("tbl") + .withSchema(new Schema(Types.NestedField.required(1, "id", Types.IntegerType.get()))) + .build(); + + LoadTableResponse response = wrapper.createTable(namespace, request, true); + + Assertions.assertEquals("/v1/upstream/namespaces/db/tables", requestPath.get()); + Assertions.assertEquals("POST", requestMethod.get()); + Assertions.assertEquals("vended-credentials", accessDelegationHeader.get()); + Assertions.assertEquals(1, response.credentials().size()); + Assertions.assertEquals( + "v1/local/namespaces/db/tables/tbl/credentials", + response.credentials().get(0).config().get("client.refresh-credentials-endpoint")); + } finally { + server.stop(0); + } + } + + @Test + void testFederatedCreateTableNoCredentials() throws Exception { + Namespace namespace = Namespace.of("db"); + TableMetadata metadata = + TableMetadataParser.fromJson( + "s3://bucket/db/tbl/metadata/v1.metadata.json", + TableMetadataParser.toJson( + TableMetadata.newTableMetadata( + new Schema(Types.NestedField.required(1, "id", Types.IntegerType.get())), + PartitionSpec.unpartitioned(), + SortOrder.unsorted(), + "s3://bucket/db/tbl", + Collections.emptyMap()))); + LoadTableResponse upstreamResponse = + LoadTableResponse.builder() + .withTableMetadata(metadata) + .addAllConfig(ImmutableMap.of("io-impl", "org.apache.iceberg.aws.s3.S3FileIO")) + .build(); + String upstreamJson = LoadTableResponseParser.toJson(upstreamResponse); + + AtomicReference<String> requestPath = new AtomicReference<>(); + AtomicReference<String> requestMethod = new AtomicReference<>(); + AtomicReference<String> accessDelegationHeader = new AtomicReference<>(); + HttpServer server = HttpServer.create(new InetSocketAddress(0), 0); + server.createContext( + "/", + exchange -> { + requestPath.set(exchange.getRequestURI().getPath()); + requestMethod.set(exchange.getRequestMethod()); + accessDelegationHeader.set( + exchange.getRequestHeaders().getFirst("X-Iceberg-Access-Delegation")); + byte[] body = upstreamJson.getBytes(StandardCharsets.UTF_8); + exchange.getResponseHeaders().add("Content-Type", "application/json"); + exchange.sendResponseHeaders(200, body.length); + try (OutputStream os = exchange.getResponseBody()) { + os.write(body); + } + }); + server.start(); + try { + String uri = "http://127.0.0.1:"; + server.getAddress().getPort(); + RESTCatalog restCatalog = mock(RESTCatalog.class); + when(restCatalog.name()).thenReturn("upstream"); + when(restCatalog.properties()) + .thenReturn( + ImmutableMap.of( + CatalogProperties.URI, + uri, + AuthProperties.AUTH_TYPE, + AuthProperties.AUTH_TYPE_NONE, + "prefix", + "upstream")); + + IcebergConfig config = + new IcebergConfig( + ImmutableMap.of( + IcebergConstants.CATALOG_BACKEND, + "memory", + IcebergConstants.WAREHOUSE, + "/tmp/warehouse")); + CatalogWrapperForREST wrapper = new StaticCatalogWrapperForREST("local", config, restCatalog); + CreateTableRequest request = + CreateTableRequest.builder() + .withName("tbl") + .withSchema(new Schema(Types.NestedField.required(1, "id", Types.IntegerType.get()))) + .build(); + + LoadTableResponse response = wrapper.createTable(namespace, request, false); + + Assertions.assertEquals("/v1/upstream/namespaces/db/tables", requestPath.get()); + Assertions.assertEquals("POST", requestMethod.get()); + Assertions.assertNull( + accessDelegationHeader.get(), + "X-Iceberg-Access-Delegation header should not be sent without credential vending"); + Assertions.assertTrue( + response.credentials() == null || response.credentials().isEmpty(), + "Non-vended request should not return remote storage-credentials"); + Assertions.assertEquals( + "org.apache.iceberg.aws.s3.S3FileIO", response.config().get("io-impl")); + } finally { + server.stop(0); + } + } + + @Test + void testFederatedCreateTableForwardsStageCreate() throws Exception { + Namespace namespace = Namespace.of("db"); + TableMetadata metadata = + TableMetadataParser.fromJson( + "s3://bucket/db/tbl/metadata/v1.metadata.json", + TableMetadataParser.toJson( + TableMetadata.newTableMetadata( + new Schema(Types.NestedField.required(1, "id", Types.IntegerType.get())), + PartitionSpec.unpartitioned(), + SortOrder.unsorted(), + "s3://bucket/db/tbl", + Collections.emptyMap()))); + LoadTableResponse upstreamResponse = + LoadTableResponse.builder().withTableMetadata(metadata).build(); + String upstreamJson = LoadTableResponseParser.toJson(upstreamResponse); + + AtomicReference<String> requestBody = new AtomicReference<>(); + HttpServer server = HttpServer.create(new InetSocketAddress(0), 0); + server.createContext( + "/", + exchange -> { + requestBody.set( + new String(exchange.getRequestBody().readAllBytes(), StandardCharsets.UTF_8)); + byte[] body = upstreamJson.getBytes(StandardCharsets.UTF_8); + exchange.getResponseHeaders().add("Content-Type", "application/json"); + exchange.sendResponseHeaders(200, body.length); + try (OutputStream os = exchange.getResponseBody()) { + os.write(body); + } + }); + server.start(); + try { + String uri = "http://127.0.0.1:"; + server.getAddress().getPort(); + RESTCatalog restCatalog = mock(RESTCatalog.class); + when(restCatalog.name()).thenReturn("upstream"); + when(restCatalog.properties()) + .thenReturn( + ImmutableMap.of( + CatalogProperties.URI, + uri, + AuthProperties.AUTH_TYPE, + AuthProperties.AUTH_TYPE_NONE, + "prefix", + "upstream")); + + IcebergConfig config = + new IcebergConfig( + ImmutableMap.of( + IcebergConstants.CATALOG_BACKEND, + "memory", + IcebergConstants.WAREHOUSE, + "/tmp/warehouse")); + CatalogWrapperForREST wrapper = new StaticCatalogWrapperForREST("local", config, restCatalog); + Schema schema = new Schema(Types.NestedField.required(1, "id", Types.IntegerType.get())); + CreateTableRequest request = + CreateTableRequest.builder() + .withName("tbl") + .withSchema(schema) + .withLocation("s3://bucket/warehouse/table") + .stageCreate() + .build(); + + wrapper.createTable(namespace, request, false); + + CreateTableRequest forwarded = + IcebergObjectMapper.getInstance().readValue(requestBody.get(), CreateTableRequest.class); + Assertions.assertTrue(forwarded.stageCreate()); + Assertions.assertEquals("s3://bucket/warehouse/table", forwarded.location()); + } finally { + server.stop(0); + } + } + + @Test + void testFederatedRegisterTableWithCredentials() throws Exception { + Namespace namespace = Namespace.of("db"); + TableMetadata metadata = + TableMetadataParser.fromJson( + "s3://bucket/db/tbl/metadata/v1.metadata.json", + TableMetadataParser.toJson( + TableMetadata.newTableMetadata( + new Schema(Types.NestedField.required(1, "id", Types.IntegerType.get())), + PartitionSpec.unpartitioned(), + SortOrder.unsorted(), + "s3://bucket/db/tbl", + Collections.emptyMap()))); + Credential cred = + IcebergRESTUtils.toRESTCredential( + "s3://bucket/db/tbl/", + ImmutableMap.of( + "s3.session-token", + "upstream-token", + "client.refresh-credentials-endpoint", + "v1/upstream/namespaces/db/tables/tbl/credentials")); + LoadTableResponse upstreamResponse = + LoadTableResponse.builder().withTableMetadata(metadata).addCredential(cred).build(); + String upstreamJson = LoadTableResponseParser.toJson(upstreamResponse); + + AtomicReference<String> requestPath = new AtomicReference<>(); + AtomicReference<String> requestMethod = new AtomicReference<>(); + AtomicReference<String> accessDelegationHeader = new AtomicReference<>(); + HttpServer server = HttpServer.create(new InetSocketAddress(0), 0); + server.createContext( + "/", + exchange -> { + requestPath.set(exchange.getRequestURI().getPath()); + requestMethod.set(exchange.getRequestMethod()); + accessDelegationHeader.set( + exchange.getRequestHeaders().getFirst("X-Iceberg-Access-Delegation")); + byte[] body = upstreamJson.getBytes(StandardCharsets.UTF_8); + exchange.getResponseHeaders().add("Content-Type", "application/json"); + exchange.sendResponseHeaders(200, body.length); + try (OutputStream os = exchange.getResponseBody()) { + os.write(body); + } + }); + server.start(); + try { + String uri = "http://127.0.0.1:"; + server.getAddress().getPort(); + RESTCatalog restCatalog = mock(RESTCatalog.class); + when(restCatalog.name()).thenReturn("upstream"); + when(restCatalog.properties()) + .thenReturn( + ImmutableMap.of( + CatalogProperties.URI, + uri, + AuthProperties.AUTH_TYPE, + AuthProperties.AUTH_TYPE_NONE, + "prefix", + "upstream")); + + IcebergConfig config = + new IcebergConfig( + ImmutableMap.of( + IcebergConstants.CATALOG_BACKEND, + "memory", + IcebergConstants.WAREHOUSE, + "/tmp/warehouse")); + CatalogWrapperForREST wrapper = new StaticCatalogWrapperForREST("local", config, restCatalog); + RegisterTableRequest request = + ImmutableRegisterTableRequest.builder() + .name("tbl") + .metadataLocation("s3://bucket/db/tbl/metadata/v1.metadata.json") + .build(); + + LoadTableResponse response = wrapper.registerTable(namespace, request, true); + + Assertions.assertEquals("/v1/upstream/namespaces/db/register", requestPath.get()); + Assertions.assertEquals("POST", requestMethod.get()); + Assertions.assertEquals("vended-credentials", accessDelegationHeader.get()); + Assertions.assertEquals(1, response.credentials().size()); + Assertions.assertEquals( + "v1/local/namespaces/db/tables/tbl/credentials", + response.credentials().get(0).config().get("client.refresh-credentials-endpoint")); + } finally { + server.stop(0); + } + } + + @Test + void testFederatedRegisterTableIncludesRemoteConfig() throws Exception { + Namespace namespace = Namespace.of("db"); + TableMetadata metadata = + TableMetadataParser.fromJson( + "s3://bucket/db/tbl/metadata/v1.metadata.json", + TableMetadataParser.toJson( + TableMetadata.newTableMetadata( + new Schema(Types.NestedField.required(1, "id", Types.IntegerType.get())), + PartitionSpec.unpartitioned(), + SortOrder.unsorted(), + "s3://bucket/db/tbl", + Collections.emptyMap()))); + LoadTableResponse upstreamResponse = + LoadTableResponse.builder() + .withTableMetadata(metadata) + .addAllConfig( + ImmutableMap.of( + IcebergConstants.IO_IMPL, + "org.apache.iceberg.aws.s3.S3FileIO", + IcebergConstants.ICEBERG_S3_ENDPOINT, + "http://localhost:9000";)) + .build(); + String upstreamJson = LoadTableResponseParser.toJson(upstreamResponse); + + AtomicReference<String> requestPath = new AtomicReference<>(); + AtomicReference<String> requestMethod = new AtomicReference<>(); + AtomicReference<String> accessDelegationHeader = new AtomicReference<>(); + HttpServer server = HttpServer.create(new InetSocketAddress(0), 0); + server.createContext( + "/", + exchange -> { + requestPath.set(exchange.getRequestURI().getPath()); + requestMethod.set(exchange.getRequestMethod()); + accessDelegationHeader.set( + exchange.getRequestHeaders().getFirst("X-Iceberg-Access-Delegation")); + byte[] body = upstreamJson.getBytes(StandardCharsets.UTF_8); + exchange.getResponseHeaders().add("Content-Type", "application/json"); + exchange.sendResponseHeaders(200, body.length); + try (OutputStream os = exchange.getResponseBody()) { + os.write(body); + } + }); + server.start(); + try { + String uri = "http://127.0.0.1:"; + server.getAddress().getPort(); + RESTCatalog catalog = mock(RESTCatalog.class); + when(catalog.name()).thenReturn("upstream"); + when(catalog.properties()) + .thenReturn( + ImmutableMap.of( + CatalogProperties.URI, + uri, + AuthProperties.AUTH_TYPE, + AuthProperties.AUTH_TYPE_NONE, + "prefix", + "upstream")); + + IcebergConfig config = + new IcebergConfig( + ImmutableMap.of( + IcebergConstants.CATALOG_BACKEND, + "memory", + IcebergConstants.WAREHOUSE, + "/tmp/warehouse")); + CatalogWrapperForREST wrapper = new StaticCatalogWrapperForREST("test", config, catalog); + + RegisterTableRequest request = + ImmutableRegisterTableRequest.builder() + .name("tbl") + .metadataLocation("s3://bucket/warehouse/tbl/metadata/v1.metadata.json") + .build(); + + LoadTableResponse response = wrapper.registerTable(namespace, request, false); + + Assertions.assertEquals("/v1/upstream/namespaces/db/register", requestPath.get()); + Assertions.assertEquals("POST", requestMethod.get()); + Assertions.assertNull(accessDelegationHeader.get()); + Assertions.assertEquals( + "org.apache.iceberg.aws.s3.S3FileIO", response.config().get(IcebergConstants.IO_IMPL)); + Assertions.assertEquals( + "http://localhost:9000";, response.config().get(IcebergConstants.ICEBERG_S3_ENDPOINT)); + } finally { + server.stop(0); + } + } + + @Test + void testFederatedRegisterTableOverwrite() throws Exception { + Namespace namespace = Namespace.of("db"); + TableMetadata metadata = + TableMetadataParser.fromJson( + "s3://bucket/db/tbl/metadata/v1.metadata.json", + TableMetadataParser.toJson( + TableMetadata.newTableMetadata( + new Schema(Types.NestedField.required(1, "id", Types.IntegerType.get())), + PartitionSpec.unpartitioned(), + SortOrder.unsorted(), + "s3://bucket/db/tbl", + Collections.emptyMap()))); + LoadTableResponse upstreamResponse = + LoadTableResponse.builder().withTableMetadata(metadata).build(); + String upstreamJson = LoadTableResponseParser.toJson(upstreamResponse); + + AtomicReference<String> requestBody = new AtomicReference<>(); + HttpServer server = HttpServer.create(new InetSocketAddress(0), 0); + server.createContext( + "/", + exchange -> { + requestBody.set( + new String(exchange.getRequestBody().readAllBytes(), StandardCharsets.UTF_8)); + byte[] body = upstreamJson.getBytes(StandardCharsets.UTF_8); + exchange.getResponseHeaders().add("Content-Type", "application/json"); + exchange.sendResponseHeaders(200, body.length); + try (OutputStream os = exchange.getResponseBody()) { + os.write(body); + } + }); + server.start(); + try { + String uri = "http://127.0.0.1:"; + server.getAddress().getPort(); + RESTCatalog catalog = mock(RESTCatalog.class); + when(catalog.name()).thenReturn("upstream"); + when(catalog.properties()) + .thenReturn( + ImmutableMap.of( + CatalogProperties.URI, + uri, + AuthProperties.AUTH_TYPE, + AuthProperties.AUTH_TYPE_NONE, + "prefix", + "upstream")); + + IcebergConfig config = + new IcebergConfig( + ImmutableMap.of( + IcebergConstants.CATALOG_BACKEND, + "memory", + IcebergConstants.WAREHOUSE, + "/tmp/warehouse")); + CatalogWrapperForREST wrapper = new StaticCatalogWrapperForREST("test", config, catalog); + + RegisterTableRequest request = + ImmutableRegisterTableRequest.builder() + .name("tbl") + .metadataLocation("s3://bucket/warehouse/tbl/metadata/v2.metadata.json") + .overwrite(true) + .build(); + + wrapper.registerTable(namespace, request, false); + + RegisterTableRequest forwarded = + IcebergObjectMapper.getInstance() + .readValue(requestBody.get(), RegisterTableRequest.class); + Assertions.assertTrue(forwarded.overwrite()); + Assertions.assertEquals( + "s3://bucket/warehouse/tbl/metadata/v2.metadata.json", forwarded.metadataLocation()); + } finally { + server.stop(0); + } + } } diff --git a/iceberg/iceberg-rest-server/src/test/java/org/apache/gravitino/iceberg/service/TestIcebergRESTUtils.java b/iceberg/iceberg-rest-server/src/test/java/org/apache/gravitino/iceberg/service/TestIcebergRESTUtils.java index 2ad1f85c3d..e3e4f3225d 100644 --- a/iceberg/iceberg-rest-server/src/test/java/org/apache/gravitino/iceberg/service/TestIcebergRESTUtils.java +++ b/iceberg/iceberg-rest-server/src/test/java/org/apache/gravitino/iceberg/service/TestIcebergRESTUtils.java @@ -24,6 +24,7 @@ import static org.mockito.Mockito.when; import static org.mockito.Mockito.withSettings; import com.google.common.collect.ImmutableMap; +import java.util.Collections; import java.util.List; import java.util.Map; import org.apache.gravitino.NameIdentifier; @@ -36,7 +37,9 @@ import org.apache.gravitino.credential.S3SecretKeyCredential; import org.apache.gravitino.credential.S3TokenCredential; import org.apache.gravitino.iceberg.service.authorization.IcebergRESTServerContext; import org.apache.gravitino.iceberg.service.provider.IcebergConfigProvider; +import org.apache.iceberg.PartitionSpec; import org.apache.iceberg.Schema; +import org.apache.iceberg.SortOrder; import org.apache.iceberg.TableMetadata; import org.apache.iceberg.catalog.Namespace; import org.apache.iceberg.catalog.TableIdentifier; @@ -47,6 +50,7 @@ import org.apache.iceberg.rest.credentials.Credential; import org.apache.iceberg.rest.requests.CreateTableRequest; import org.apache.iceberg.rest.responses.ImmutableLoadCredentialsResponse; import org.apache.iceberg.rest.responses.LoadCredentialsResponse; +import org.apache.iceberg.rest.responses.LoadTableResponse; import org.apache.iceberg.types.Types.IntegerType; import org.apache.iceberg.types.Types.NestedField; import org.apache.iceberg.types.Types.StringType; @@ -350,4 +354,99 @@ public class TestIcebergRESTUtils { "v1/irc1/namespaces/db/tables/tbl/credentials", credential.config().get("client.refresh-credentials-endpoint")); } + + @Test + void testRewriteLoadTableCredentials() { + TableIdentifier table = TableIdentifier.of(Namespace.of("db"), "tbl"); + TableMetadata metadata = + TableMetadata.newTableMetadata( + new Schema(NestedField.required(1, "id", IntegerType.get())), + PartitionSpec.unpartitioned(), + SortOrder.unsorted(), + "s3://bucket/db/tbl", + Collections.emptyMap()); + LoadTableResponse upstream = + LoadTableResponse.builder() + .withTableMetadata(metadata) + .addAllConfig( + ImmutableMap.of( + "io-impl", + "org.apache.iceberg.aws.s3.S3FileIO", + "s3.session-token", + "upstream-token", + "client.refresh-credentials-endpoint", + "v1/upstream/namespaces/db/tables/tbl/credentials")) + .addCredential( + IcebergRESTUtils.toRESTCredential( + "s3://bucket/db/tbl/", + ImmutableMap.of( + "s3.session-token", + "upstream-token", + "client.refresh-credentials-endpoint", + "v1/upstream/namespaces/db/tables/tbl/credentials"))) + .build(); + + LoadTableResponse rewritten = + IcebergRESTUtils.rewriteLoadTableCredentials("irc1", table, upstream); + + Assertions.assertEquals("s3://bucket/db/tbl", rewritten.tableMetadata().location()); + Assertions.assertEquals( + "org.apache.iceberg.aws.s3.S3FileIO", rewritten.config().get("io-impl")); + Assertions.assertEquals("upstream-token", rewritten.config().get("s3.session-token")); + Assertions.assertEquals( + "v1/irc1/namespaces/db/tables/tbl/credentials", + rewritten.config().get("client.refresh-credentials-endpoint")); + Assertions.assertEquals(1, rewritten.credentials().size()); + Credential credential = rewritten.credentials().get(0); + Assertions.assertEquals("s3://bucket/db/tbl/", credential.prefix()); + Assertions.assertEquals("upstream-token", credential.config().get("s3.session-token")); + Assertions.assertEquals( + "v1/irc1/namespaces/db/tables/tbl/credentials", + credential.config().get("client.refresh-credentials-endpoint")); + } + + @Test + void testRewriteLoadTableCredentialsDropsConfigRefreshWithoutTokens() { + TableIdentifier table = TableIdentifier.of(Namespace.of("db"), "tbl"); + TableMetadata metadata = + TableMetadata.newTableMetadata( + new Schema(NestedField.required(1, "id", IntegerType.get())), + PartitionSpec.unpartitioned(), + SortOrder.unsorted(), + "s3://bucket/db/tbl", + Collections.emptyMap()); + // Upstream put the refresh URL in top-level config while the session token lives only in + // storage-credentials. The config refresh key must be dropped, not left pointing upstream. + LoadTableResponse upstream = + LoadTableResponse.builder() + .withTableMetadata(metadata) + .addAllConfig( + ImmutableMap.of( + "io-impl", + "org.apache.iceberg.aws.s3.S3FileIO", + "client.refresh-credentials-endpoint", + "v1/upstream/namespaces/db/tables/tbl/credentials")) + .addCredential( + IcebergRESTUtils.toRESTCredential( + "s3://bucket/db/tbl/", + ImmutableMap.of( + "s3.session-token", + "upstream-token", + "client.refresh-credentials-endpoint", + "v1/upstream/namespaces/db/tables/tbl/credentials"))) + .build(); + + LoadTableResponse rewritten = + IcebergRESTUtils.rewriteLoadTableCredentials("irc1", table, upstream); + + Assertions.assertEquals( + "org.apache.iceberg.aws.s3.S3FileIO", rewritten.config().get("io-impl")); + Assertions.assertFalse( + rewritten.config().containsKey("client.refresh-credentials-endpoint"), + "Top-level config must not keep an upstream refresh endpoint when tokens are absent"); + Assertions.assertEquals(1, rewritten.credentials().size()); + Assertions.assertEquals( + "v1/irc1/namespaces/db/tables/tbl/credentials", + rewritten.credentials().get(0).config().get("client.refresh-credentials-endpoint")); + } } diff --git a/iceberg/iceberg-rest-server/src/test/java/org/apache/gravitino/iceberg/service/rest/TestIcebergTableOperations.java b/iceberg/iceberg-rest-server/src/test/java/org/apache/gravitino/iceberg/service/rest/TestIcebergTableOperations.java index 14d2782b4f..8cb0557837 100644 --- a/iceberg/iceberg-rest-server/src/test/java/org/apache/gravitino/iceberg/service/rest/TestIcebergTableOperations.java +++ b/iceberg/iceberg-rest-server/src/test/java/org/apache/gravitino/iceberg/service/rest/TestIcebergTableOperations.java @@ -542,6 +542,11 @@ public class TestIcebergTableOperations extends IcebergNamespaceTestBase { } private Response doLoadTableWithSnapshots(Namespace ns, String name, String snapshots) { + return doLoadTableWithSnapshots(ns, name, snapshots, false); + } + + private Response doLoadTableWithSnapshots( + Namespace ns, String name, String snapshots, boolean credentialVending) { String path = IcebergRestTestUtil.NAMESPACE_PATH + "/" @@ -549,7 +554,12 @@ public class TestIcebergTableOperations extends IcebergNamespaceTestBase { + "/tables/" + name; Map<String, String> queryParams = ImmutableMap.of("snapshots", snapshots); - return getIcebergClientBuilder(path, Optional.of(queryParams)).get(); + Invocation.Builder builder = getIcebergClientBuilder(path, Optional.of(queryParams)); + if (credentialVending) { + builder = + builder.header(IcebergTableOperations.X_ICEBERG_ACCESS_DELEGATION, "vended-credentials"); + } + return builder.get(); } private Response doPlanTableScan(Namespace ns, String tableName, PlanTableScanRequest request) { @@ -1083,6 +1093,73 @@ public class TestIcebergTableOperations extends IcebergNamespaceTestBase { "Refs should be preserved in filtered response"); } + @ParameterizedTest + @MethodSource("org.apache.gravitino.iceberg.service.rest.IcebergRestTestUtil#testNamespaces") + void testLoadTableSnapshotsRefsKeepsCredentials(Namespace namespace) { + verifyCreateNamespaceSucc(namespace); + String tableName = "snapshots_refs_creds"; + CreateTableRequest createTableRequest = + CreateTableRequest.builder() + .withName(tableName) + .withSchema(tableSchema) + .withLocation("s3://bucket/" + tableName) + .setProperties( + ImmutableMap.of( + CatalogWrapperForTest.GENERATE_PLAN_TASKS_DATA_PROP, Boolean.TRUE.toString())) + .build(); + Response createResponse = + getTableClientBuilder(namespace, Optional.empty()) + .header(IcebergTableOperations.X_ICEBERG_ACCESS_DELEGATION, "vended-credentials") + .post(Entity.entity(createTableRequest, MediaType.APPLICATION_JSON_TYPE)); + Assertions.assertEquals(Status.OK.getStatusCode(), createResponse.getStatus()); + + Response refsResponse = doLoadTableWithSnapshots(namespace, tableName, "refs", true); + Assertions.assertEquals(Status.OK.getStatusCode(), refsResponse.getStatus()); + LoadTableResponse refsTableResponse = refsResponse.readEntity(LoadTableResponse.class); + + Assertions.assertTrue( + refsTableResponse.tableMetadata().snapshots().size() >= 1, + "Filtered response should keep at least the current ref snapshot"); + Assertions.assertEquals( + DummyCredentialProvider.DUMMY_CREDENTIAL_TYPE, + refsTableResponse.config().get(Credential.CREDENTIAL_TYPE), + "snapshots=refs must keep vended credentials after filtering"); + Assertions.assertFalse( + refsTableResponse.credentials().isEmpty(), + "snapshots=refs must keep storage-credentials after filtering"); + } + + @Test + void testFilterSnapshotsByRefsKeepsCredentials() { + TableMetadata metadata = + TableMetadata.newTableMetadata( + tableSchema, + org.apache.iceberg.PartitionSpec.unpartitioned(), + "s3://bucket/db/tbl", + ImmutableMap.of()); + org.apache.iceberg.rest.credentials.Credential credential = + IcebergRESTUtils.toRESTCredential( + "s3://bucket/db/tbl/", + ImmutableMap.of( + "s3.session-token", + "token", + "client.refresh-credentials-endpoint", + "v1/c/ns/t/credentials")); + LoadTableResponse original = + LoadTableResponse.builder() + .withTableMetadata(metadata) + .addAllConfig(ImmutableMap.of("io-impl", "org.apache.iceberg.aws.s3.S3FileIO")) + .addCredential(credential) + .build(); + + LoadTableResponse filtered = IcebergTableOperations.filterSnapshotsByRefs(original); + + Assertions.assertEquals(1, filtered.credentials().size()); + Assertions.assertEquals( + "token", filtered.credentials().get(0).config().get("s3.session-token")); + Assertions.assertEquals("org.apache.iceberg.aws.s3.S3FileIO", filtered.config().get("io-impl")); + } + @ParameterizedTest @MethodSource("org.apache.gravitino.iceberg.service.rest.IcebergRestTestUtil#testNamespaces") void testLoadTableSnapshotsAllReturnsAllSnapshots(Namespace namespace) {
