From 118eab27222a86b035261fe329d3dd055a04ad55 Mon Sep 17 00:00:00 2001 From: Dmitri Bourlatchkov Date: Thu, 4 Sep 2025 20:08:40 -0400 Subject: [PATCH] Allow `PolarisServerManager` implementations to define custom client headers `PolarisServerManager` is a plugin point for downstream builds to define runtime env. for tests under `integration-tests`. This change allows more flexibility for test runtime environments by allowing injecting extra headers into test clients. --- .../polaris/service/it/env/IcebergHelper.java | 2 +- .../service/it/env/PolarisApiEndpoints.java | 20 ++++++++++--- .../service/it/env/PolarisRestApi.java | 3 +- .../apache/polaris/service/it/env/Server.java | 5 ++++ .../ext/PolarisIntegrationTestExtension.java | 2 +- .../PolarisApplicationIntegrationTest.java | 30 ++++++++----------- .../service/it/RestCatalogMinIOSpecialIT.java | 2 +- .../it/ApplicationIntegrationTest.java | 6 ++-- 8 files changed, 40 insertions(+), 30 deletions(-) diff --git a/integration-tests/src/main/java/org/apache/polaris/service/it/env/IcebergHelper.java b/integration-tests/src/main/java/org/apache/polaris/service/it/env/IcebergHelper.java index 027c68905a..480995727f 100644 --- a/integration-tests/src/main/java/org/apache/polaris/service/it/env/IcebergHelper.java +++ b/integration-tests/src/main/java/org/apache/polaris/service/it/env/IcebergHelper.java @@ -39,7 +39,7 @@ public static RESTCatalog restCatalog( .put(CatalogProperties.URI, endpoints.catalogApiEndpoint().toString()) .put(OAuth2Properties.TOKEN, authToken) .put("warehouse", catalog) - .put("header." + endpoints.realmHeaderName(), endpoints.realmId()) + .putAll(endpoints.extraHeaders("header.")) .putAll(extraProperties); restCatalog.initialize("polaris", propertiesBuilder.buildKeepingLast()); diff --git a/integration-tests/src/main/java/org/apache/polaris/service/it/env/PolarisApiEndpoints.java b/integration-tests/src/main/java/org/apache/polaris/service/it/env/PolarisApiEndpoints.java index 45fe7809b4..78e62bcd29 100644 --- a/integration-tests/src/main/java/org/apache/polaris/service/it/env/PolarisApiEndpoints.java +++ b/integration-tests/src/main/java/org/apache/polaris/service/it/env/PolarisApiEndpoints.java @@ -20,6 +20,8 @@ import java.io.Serializable; import java.net.URI; +import java.util.Map; +import java.util.stream.Collectors; /** * This class contains the most fundamental information for accessing Polaris APIs, such as the base @@ -30,12 +32,16 @@ public final class PolarisApiEndpoints implements Serializable { private final URI baseUri; private final String realmId; - private final String realmHeaderName; + private final Map headers; public PolarisApiEndpoints(URI baseUri, String realmId, String realmHeaderName) { + this(baseUri, realmId, Map.of(realmHeaderName, realmId)); + } + + public PolarisApiEndpoints(URI baseUri, String realmId, Map headers) { this.baseUri = baseUri; this.realmId = realmId; - this.realmHeaderName = realmHeaderName; + this.headers = headers; } public URI catalogApiEndpoint() { @@ -50,7 +56,13 @@ public String realmId() { return realmId; } - public String realmHeaderName() { - return realmHeaderName; + public Map extraHeaders() { + return headers; + } + + public Map extraHeaders(String keyPrefix) { + return headers.entrySet().stream() + .map(e -> Map.entry(keyPrefix + e.getKey(), e.getValue())) + .collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue)); } } diff --git a/integration-tests/src/main/java/org/apache/polaris/service/it/env/PolarisRestApi.java b/integration-tests/src/main/java/org/apache/polaris/service/it/env/PolarisRestApi.java index ddf9e0c50b..40222c2b3f 100644 --- a/integration-tests/src/main/java/org/apache/polaris/service/it/env/PolarisRestApi.java +++ b/integration-tests/src/main/java/org/apache/polaris/service/it/env/PolarisRestApi.java @@ -39,8 +39,7 @@ public class PolarisRestApi extends RestApi { } protected Map defaultHeaders() { - Map headers = new HashMap<>(); - headers.put(endpoints.realmHeaderName(), endpoints.realmId()); + Map headers = new HashMap<>(endpoints.extraHeaders()); if (authToken != null) { headers.put("Authorization", "Bearer " + authToken); } diff --git a/integration-tests/src/main/java/org/apache/polaris/service/it/env/Server.java b/integration-tests/src/main/java/org/apache/polaris/service/it/env/Server.java index 4fea80070e..fea53d8917 100644 --- a/integration-tests/src/main/java/org/apache/polaris/service/it/env/Server.java +++ b/integration-tests/src/main/java/org/apache/polaris/service/it/env/Server.java @@ -19,6 +19,7 @@ package org.apache.polaris.service.it.env; import java.net.URI; +import java.util.Map; /** * This is a holder for access information to a particular Polaris Server. Test cases may use only @@ -37,6 +38,10 @@ default String realmHeaderName() { return DEFAULT_REALM_HEADER; } + default Map headers() { + return Map.of(realmHeaderName(), realmId()); + } + /** * The base URI to all Polaris APIs (e.g. the common base of the Iceberg REST API endpoints and * Polaris Management API endpoints). diff --git a/integration-tests/src/main/java/org/apache/polaris/service/it/ext/PolarisIntegrationTestExtension.java b/integration-tests/src/main/java/org/apache/polaris/service/it/ext/PolarisIntegrationTestExtension.java index 9031247fc5..fd64273e11 100644 --- a/integration-tests/src/main/java/org/apache/polaris/service/it/ext/PolarisIntegrationTestExtension.java +++ b/integration-tests/src/main/java/org/apache/polaris/service/it/ext/PolarisIntegrationTestExtension.java @@ -92,7 +92,7 @@ private static class Env implements AutoCloseable { private Env(Server server) { this.server = server; this.endpoints = - new PolarisApiEndpoints(server.baseUri(), server.realmId(), server.realmHeaderName()); + new PolarisApiEndpoints(server.baseUri(), server.realmId(), server.headers()); } PolarisApiEndpoints endpoints() { diff --git a/integration-tests/src/main/java/org/apache/polaris/service/it/test/PolarisApplicationIntegrationTest.java b/integration-tests/src/main/java/org/apache/polaris/service/it/test/PolarisApplicationIntegrationTest.java index a7e21427ea..87f2c66e55 100644 --- a/integration-tests/src/main/java/org/apache/polaris/service/it/test/PolarisApplicationIntegrationTest.java +++ b/integration-tests/src/main/java/org/apache/polaris/service/it/test/PolarisApplicationIntegrationTest.java @@ -259,15 +259,12 @@ private static RESTSessionCatalog newSessionCatalog(String catalog) { RESTSessionCatalog sessionCatalog = new RESTSessionCatalog(); sessionCatalog.initialize( "polaris_catalog_test", - Map.of( - "uri", - endpoints.catalogApiEndpoint().toString(), - OAuth2Properties.TOKEN, - authToken, - "warehouse", - catalog, - "header." + endpoints.realmHeaderName(), - realm)); + ImmutableMap.builder() + .put("uri", endpoints.catalogApiEndpoint().toString()) + .put(OAuth2Properties.TOKEN, authToken) + .put("warehouse", catalog) + .putAll(endpoints.extraHeaders("header.")) + .build()); return sessionCatalog; } @@ -590,15 +587,12 @@ public void testWarehouseNotSpecified() throws IOException { () -> sessionCatalog.initialize( "polaris_catalog_test", - Map.of( - "uri", - endpoints.catalogApiEndpoint().toString(), - OAuth2Properties.TOKEN, - authToken, - "warehouse", - emptyEnvironmentVariable, - "header." + endpoints.realmHeaderName(), - realm))) + ImmutableMap.builder() + .put("uri", endpoints.catalogApiEndpoint().toString()) + .put(OAuth2Properties.TOKEN, authToken) + .put("warehouse", emptyEnvironmentVariable) + .putAll(endpoints.extraHeaders("header.")) + .build())) .isInstanceOf(BadRequestException.class) .hasMessage("Malformed request: Please specify a warehouse"); } diff --git a/runtime/service/src/intTest/java/org/apache/polaris/service/it/RestCatalogMinIOSpecialIT.java b/runtime/service/src/intTest/java/org/apache/polaris/service/it/RestCatalogMinIOSpecialIT.java index 83ddd0c570..80845608da 100644 --- a/runtime/service/src/intTest/java/org/apache/polaris/service/it/RestCatalogMinIOSpecialIT.java +++ b/runtime/service/src/intTest/java/org/apache/polaris/service/it/RestCatalogMinIOSpecialIT.java @@ -198,7 +198,7 @@ private RESTCatalog createCatalog( org.apache.iceberg.CatalogProperties.URI, endpoints.catalogApiEndpoint().toString()) .put(OAuth2Properties.TOKEN, authToken) .put("warehouse", catalogName) - .put("header." + endpoints.realmHeaderName(), endpoints.realmId()) + .putAll(endpoints.extraHeaders("header.")) .put("header.X-Iceberg-Access-Delegation", "vended-credentials"); restCatalog.initialize("polaris", propertiesBuilder.buildKeepingLast()); diff --git a/runtime/service/src/test/java/org/apache/polaris/service/it/ApplicationIntegrationTest.java b/runtime/service/src/test/java/org/apache/polaris/service/it/ApplicationIntegrationTest.java index ef3b808d40..b0f27fba59 100644 --- a/runtime/service/src/test/java/org/apache/polaris/service/it/ApplicationIntegrationTest.java +++ b/runtime/service/src/test/java/org/apache/polaris/service/it/ApplicationIntegrationTest.java @@ -75,7 +75,7 @@ public void testIcebergRestApiRefreshExpiredToken( String path = endpoints.catalogApiEndpoint() + "/v1/oauth/tokens"; try (RESTClient client = HTTPClient.builder(Map.of()) - .withHeader(endpoints.realmHeaderName(), endpoints.realmId()) + .withHeaders(endpoints.extraHeaders()) .uri(path) .withAuthSession(AuthSession.EMPTY) .build()) { @@ -106,7 +106,7 @@ public void testIcebergRestApiRefreshValidToken( String path = endpoints.catalogApiEndpoint() + "/v1/oauth/tokens"; try (RESTClient client = HTTPClient.builder(Map.of()) - .withHeader(endpoints.realmHeaderName(), endpoints.realmId()) + .withHeaders(endpoints.extraHeaders()) .uri(path) .withAuthSession(AuthSession.EMPTY) .build()) { @@ -147,7 +147,7 @@ public void testIcebergRestApiInvalidToken( String path = endpoints.catalogApiEndpoint() + "/v1/oauth/tokens"; try (RESTClient client = HTTPClient.builder(Map.of()) - .withHeader(endpoints.realmHeaderName(), endpoints.realmId()) + .withHeaders(endpoints.extraHeaders()) .uri(path) .withAuthSession(AuthSession.EMPTY) .build()) {