Allow PolarisServerManager implementations to define custom client headers (#2510)

`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.

Author: dimas-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());

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<String, String> headers;
 
   public PolarisApiEndpoints(URI baseUri, String realmId, String realmHeaderName) {
+    this(baseUri, realmId, Map.of(realmHeaderName, realmId));
+  }
+
+  public PolarisApiEndpoints(URI baseUri, String realmId, Map<String, String> 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<String, String> extraHeaders() {
+    return headers;
+  }
+
+  public Map<String, String> 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));
   }
 }

integration-tests/src/main/java/org/apache/polaris/service/it/env/PolarisRestApi.java

@@ -39,8 +39,7 @@ public class PolarisRestApi extends RestApi {
   }
 
   protected Map<String, String> defaultHeaders() {
-    Map<String, String> headers = new HashMap<>();
-    headers.put(endpoints.realmHeaderName(), endpoints.realmId());
+    Map<String, String> headers = new HashMap<>(endpoints.extraHeaders());
     if (authToken != null) {
       headers.put("Authorization", "Bearer " + authToken);
     }

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<String, String> 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).

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() {

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.<String, String>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.<String, String>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");
     }

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());

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()) {