Use interceptor for Prometheus token auth

trinodb · May 17, 2022 · 7051ef1 · 7051ef1
1 parent d7777f9
commit 7051ef1
Show file tree

Hide file tree

Showing 4 changed files with 24 additions and 7 deletions.
diff --git a/plugin/trino-prometheus/src/main/java/io/trino/plugin/prometheus/PrometheusClient.java b/plugin/trino-prometheus/src/main/java/io/trino/plugin/prometheus/PrometheusClient.java
@@ -21,6 +21,7 @@
 import io.trino.spi.type.DoubleType;
 import io.trino.spi.type.Type;
 import io.trino.spi.type.TypeManager;
+import okhttp3.Interceptor;
 import okhttp3.OkHttpClient;
 import okhttp3.OkHttpClient.Builder;
 import okhttp3.Request;
@@ -39,6 +40,7 @@
 import java.util.Set;
 import java.util.function.Supplier;
 
+import static com.google.common.net.HttpHeaders.AUTHORIZATION;
 import static io.trino.plugin.prometheus.PrometheusErrorCode.PROMETHEUS_TABLES_METRICS_RETRIEVE_ERROR;
 import static io.trino.plugin.prometheus.PrometheusErrorCode.PROMETHEUS_UNKNOWN_ERROR;
 import static io.trino.spi.type.TimestampWithTimeZoneType.createTimestampWithTimeZoneType;
@@ -55,7 +57,6 @@ public class PrometheusClient
     static final String METRICS_ENDPOINT = "/api/v1/label/__name__/values";
 
     private final OkHttpClient httpClient;
-    private final Optional<File> bearerTokenFile;
     private final Supplier<Map<String, Object>> tableSupplier;
     private final Type varcharMapType;
 
@@ -66,9 +67,10 @@ public PrometheusClient(PrometheusConnectorConfig config, JsonCodec<Map<String,
         requireNonNull(metricCodec, "metricCodec is null");
         requireNonNull(typeManager, "typeManager is null");
 
-        httpClient = new Builder().readTimeout(Duration.ofMillis(config.getReadTimeout().toMillis())).build();
+        Builder clientBuilder = new Builder().readTimeout(Duration.ofMillis(config.getReadTimeout().toMillis()));
+        setupTokenAuth(clientBuilder, getBearerAuthInfoFromFile(config.getBearerTokenFile()));
+        this.httpClient = clientBuilder.build();
 
-        bearerTokenFile = config.getBearerTokenFile();
         URI prometheusMetricsUri = getPrometheusMetricsURI(config.getPrometheusURI());
         tableSupplier = Suppliers.memoizeWithExpiration(
                 () -> fetchMetrics(metricCodec, prometheusMetricsUri),
@@ -137,8 +139,6 @@ private Map<String, Object> fetchMetrics(JsonCodec<Map<String, Object>> metricsC
     public byte[] fetchUri(URI uri)
     {
         Request.Builder requestBuilder = new Request.Builder().url(uri.toString());
-        getBearerAuthInfoFromFile().ifPresent(bearerToken -> requestBuilder.header("Authorization", "Bearer " + bearerToken));
-
         Response response;
         try {
             response = httpClient.newCall(requestBuilder.build()).execute();
@@ -153,7 +153,7 @@ public byte[] fetchUri(URI uri)
         throw new TrinoException(PROMETHEUS_UNKNOWN_ERROR, "Bad response " + response.code() + " " + response.message());
     }
 
-    private Optional<String> getBearerAuthInfoFromFile()
+    private Optional<String> getBearerAuthInfoFromFile(Optional<File> bearerTokenFile)
     {
         return bearerTokenFile.map(tokenFileName -> {
             try {
@@ -164,4 +164,17 @@ private Optional<String> getBearerAuthInfoFromFile()
             }
         });
     }
+
+    private static void setupTokenAuth(OkHttpClient.Builder clientBuilder, Optional<String> accessToken)
+    {
+        accessToken.ifPresent(token -> clientBuilder.addInterceptor(tokenAuth(token)));
+    }
+
+    private static Interceptor tokenAuth(String accessToken)
+    {
+        requireNonNull(accessToken, "accessToken is null");
+        return chain -> chain.proceed(chain.request().newBuilder()
+                .addHeader(AUTHORIZATION, "Bearer " + accessToken)
+                .build());
+    }
 }
diff --git a/.../main/java/io/trino/tests/product/launcher/env/environment/EnvMultinodeAllConnectors.java b/.../main/java/io/trino/tests/product/launcher/env/environment/EnvMultinodeAllConnectors.java
@@ -87,6 +87,9 @@ public void extendEnvironment(Environment.Builder builder)
                 container.withCopyFileToContainer(
                         forHostPath(configDir.getPath("google-sheets-auth.json")),
                         CONTAINER_PRESTO_ETC + "/catalog/google-sheets-auth.json");
+                container.withCopyFileToContainer(
+                        forHostPath(configDir.getPath("prometheus-bearer.txt")),
+                        CONTAINER_PRESTO_ETC + "/catalog/prometheus-bearer.txt");
             }
         });
     }

diff --git a/...esources/docker/presto-product-tests/conf/environment/multinode-all/prometheus-bearer.txt b/...esources/docker/presto-product-tests/conf/environment/multinode-all/prometheus-bearer.txt
@@ -0,0 +1 @@
+test
diff --git a/...esources/docker/presto-product-tests/conf/environment/multinode-all/prometheus.properties b/...esources/docker/presto-product-tests/conf/environment/multinode-all/prometheus.properties
@@ -3,5 +3,5 @@ prometheus.uri=http://host1.invalid:9090
 prometheus.query.chunk.size.duration=1d
 prometheus.max.query.range.duration=21d
 prometheus.cache.ttl=30s
-prometheus.bearer.token.file=/path/to/bearer/token/file
+prometheus.bearer.token.file=/docker/presto-product-tests/conf/presto/etc/catalog/prometheus-bearer.txt
 prometheus.read-timeout=10s