Allow a custom header name for authorization in Prometheus

docs/src/main/sphinx/connector/prometheus.md

@@ -49,6 +49,7 @@ The following configuration properties are available:
 | `prometheus.cache.ttl`                      | How long values from this config file are cached                                             |
 | `prometheus.auth.user`                      | Username for basic authentication                                                            |
 | `prometheus.auth.password`                  | Password for basic authentication                                                            |
+| `prometheus.auth.http.header.name`          | Name of the header to use for authorization                                                  |
 | `prometheus.bearer.token.file`              | File holding bearer token if needed for access to Prometheus                                 |
 | `prometheus.read-timeout`                   | How much time a query to Prometheus has before timing out                                    |
 | `prometheus.case-insensitive-name-matching` | Match Prometheus metric names case insensitively. Defaults to `false`                        |
@@ -79,7 +80,8 @@ settings are meant to protect against an unlimited query.
 
 Prometheus can be setup to require a Authorization header with every query. The value in
 `prometheus.bearer.token.file` allows for a bearer token to be read from the configured file. This file
-is optional and not required unless your Prometheus setup requires it.
+is optional and not required unless your Prometheus setup requires it.  
+`prometheus.auth.http.header.name` allows you to use a custom header name for bearer token. Default value is `Authorization`.
 
 (prometheus-type-mapping)=
 

plugin/trino-prometheus/src/main/java/io/trino/plugin/prometheus/PrometheusClient.java

@@ -73,7 +73,7 @@ public PrometheusClient(PrometheusConnectorConfig config, JsonCodec<Map<String,
 
         Builder clientBuilder = new Builder().readTimeout(Duration.ofMillis(config.getReadTimeout().toMillis()));
         setupBasicAuth(clientBuilder, config.getUser(), config.getPassword());
-        setupTokenAuth(clientBuilder, getBearerAuthInfoFromFile(config.getBearerTokenFile()));
+        setupTokenAuth(clientBuilder, getBearerAuthInfoFromFile(config.getBearerTokenFile()), config.getHttpAuthHeaderName());
         this.httpClient = clientBuilder.build();
 
         URI prometheusMetricsUri = getPrometheusMetricsURI(config.getPrometheusURI());
@@ -193,9 +193,9 @@ private static void setupBasicAuth(OkHttpClient.Builder clientBuilder, Optional<
         }
     }
 
-    private static void setupTokenAuth(OkHttpClient.Builder clientBuilder, Optional<String> accessToken)
+    private static void setupTokenAuth(OkHttpClient.Builder clientBuilder, Optional<String> accessToken, String httpAuthHeaderName)
     {
-        accessToken.ifPresent(token -> clientBuilder.addInterceptor(tokenAuth(token)));
+        accessToken.ifPresent(token -> clientBuilder.addInterceptor(tokenAuth(token, httpAuthHeaderName)));
     }
 
     private static Interceptor basicAuth(String user, String password)
@@ -212,11 +212,13 @@ private static Interceptor basicAuth(String user, String password)
                 .build());
     }
 
-    private static Interceptor tokenAuth(String accessToken)
+    private static Interceptor tokenAuth(String accessToken, String httpAuthHeaderName)
     {
         requireNonNull(accessToken, "accessToken is null");
+        boolean useBearer = httpAuthHeaderName.equals(AUTHORIZATION);
+        String token = useBearer ? "Bearer " + accessToken : accessToken;
         return chain -> chain.proceed(chain.request().newBuilder()
-                .addHeader(AUTHORIZATION, "Bearer " + accessToken)
+                .addHeader(httpAuthHeaderName, token)
                 .build());
     }
 }

plugin/trino-prometheus/src/main/java/io/trino/plugin/prometheus/PrometheusConnectorConfig.java

@@ -14,6 +14,7 @@
 package io.trino.plugin.prometheus;
 
 import com.google.common.collect.ImmutableList;
+import com.google.common.net.HttpHeaders;
 import com.google.inject.ConfigurationException;
 import com.google.inject.spi.Message;
 import io.airlift.configuration.Config;
@@ -36,6 +37,7 @@ public class PrometheusConnectorConfig
     private Duration maxQueryRangeDuration = new Duration(21, TimeUnit.DAYS);
     private Duration cacheDuration = new Duration(30, TimeUnit.SECONDS);
     private Duration readTimeout = new Duration(10, TimeUnit.SECONDS);
+    private String httpAuthHeaderName = HttpHeaders.AUTHORIZATION;
     private File bearerTokenFile;
     private String user;
     private String password;
@@ -97,6 +99,19 @@ public PrometheusConnectorConfig setCacheDuration(Duration cacheConfigDuration)
         return this;
     }
 
+    public String getHttpAuthHeaderName()
+    {
+        return httpAuthHeaderName;
+    }
+
+    @Config("prometheus.auth.http.header.name")
+    @ConfigDescription("Name of the HTTP header to use for authorization")
+    public PrometheusConnectorConfig setHttpAuthHeaderName(String httpHeaderName)
+    {
+        this.httpAuthHeaderName = httpHeaderName;
+        return this;
+    }
+
     public Optional<File> getBearerTokenFile()
     {
         return Optional.ofNullable(bearerTokenFile);

plugin/trino-prometheus/src/test/java/io/trino/plugin/prometheus/TestPrometheusConnectorConfig.java

@@ -14,6 +14,7 @@
 package io.trino.plugin.prometheus;
 
 import com.google.common.collect.ImmutableMap;
+import com.google.common.net.HttpHeaders;
 import com.google.inject.ConfigurationException;
 import io.airlift.units.Duration;
 import org.junit.jupiter.api.Test;
@@ -40,6 +41,7 @@ public void testDefaults()
                 .setMaxQueryRangeDuration(new Duration(21, DAYS))
                 .setCacheDuration(new Duration(30, SECONDS))
                 .setBearerTokenFile(null)
+                .setHttpAuthHeaderName(HttpHeaders.AUTHORIZATION)
                 .setUser(null)
                 .setPassword(null)
                 .setReadTimeout(new Duration(10, SECONDS))
@@ -54,6 +56,7 @@ public void testExplicitPropertyMappings()
                 .put("prometheus.query.chunk.size.duration", "365d")
                 .put("prometheus.max.query.range.duration", "1095d")
                 .put("prometheus.cache.ttl", "60s")
+                .put("prometheus.auth.http.header.name", "X-team-auth")
                 .put("prometheus.bearer.token.file", "/tmp/bearer_token.txt")
                 .put("prometheus.auth.user", "admin")
                 .put("prometheus.auth.password", "password")
@@ -67,6 +70,7 @@ public void testExplicitPropertyMappings()
         expected.setQueryChunkSizeDuration(new Duration(365, DAYS));
         expected.setMaxQueryRangeDuration(new Duration(1095, DAYS));
         expected.setCacheDuration(new Duration(60, SECONDS));
+        expected.setHttpAuthHeaderName("X-team-auth");
         expected.setBearerTokenFile(new File("/tmp/bearer_token.txt"));
         expected.setUser("admin");
         expected.setPassword("password");