Add support for --extra-header to Trino CLI and JDBC

Pluies · wendigo · commit c52a9e9c1174 · 2025-10-30T14:16:57.000+01:00
Allows passing arbitrary headers to Trino clients (JDBC/CLI/library)
diff --git a/client/trino-cli/src/main/java/io/trino/cli/ClientOptions.java b/client/trino-cli/src/main/java/io/trino/cli/ClientOptions.java
@@ -56,6 +56,7 @@
 import static io.trino.client.uri.PropertyName.EXTERNAL_AUTHENTICATION;
 import static io.trino.client.uri.PropertyName.EXTERNAL_AUTHENTICATION_REDIRECT_HANDLERS;
 import static io.trino.client.uri.PropertyName.EXTRA_CREDENTIALS;
+import static io.trino.client.uri.PropertyName.EXTRA_HEADERS;
 import static io.trino.client.uri.PropertyName.HTTP_LOGGING_LEVEL;
 import static io.trino.client.uri.PropertyName.HTTP_PROXY;
 import static io.trino.client.uri.PropertyName.KERBEROS_CONFIG_PATH;
@@ -204,6 +205,10 @@ public class ClientOptions
     @Option(names = "--client-tags", paramLabel = "<tags>", description = "Client tags", converter = ClientTagsConverter.class)
     public Optional<Set<String>> clientTags;
 
+    @PropertyMapping(EXTRA_HEADERS)
+    @Option(names = "--extra-header", paramLabel = "<header>", description = "Additional HTTP header to add to HTTP requests (property can be used multiple times; format is key=value)")
+    public final List<ExtraHeader> extraHeaders = new ArrayList<>();
+
     @PropertyMapping(TRACE_TOKEN)
     @Option(names = "--trace-token", paramLabel = "<token>", description = "Trace token")
     public Optional<String> traceToken;
@@ -422,6 +427,9 @@ public TrinoUri getTrinoUri(Map<PropertyName, String> restrictedProperties)
         if (!sessionProperties.isEmpty()) {
             builder.setSessionProperties(toProperties(sessionProperties));
         }
+        if (!extraHeaders.isEmpty()) {
+            builder.setExtraHeaders(toExtraHeaders(extraHeaders));
+        }
         if (!resourceEstimates.isEmpty()) {
             builder.setResourceEstimates(toResourceEstimates(resourceEstimates));
         }
@@ -498,6 +506,15 @@ public static URI parseServer(String server)
         }
     }
 
+    public static Map<String, String> toExtraHeaders(List<ExtraHeader> extraHeaders)
+    {
+        ImmutableMap.Builder<String, String> builder = ImmutableMap.builder();
+        for (ExtraHeader extraHeader : extraHeaders) {
+            builder.put(extraHeader.getHeader(), extraHeader.getValue());
+        }
+        return builder.buildOrThrow();
+    }
+
     private static Map<String, String> toProperties(List<ClientSessionProperty> sessionProperties)
     {
         ImmutableMap.Builder<String, String> builder = ImmutableMap.builder();
@@ -605,6 +622,58 @@ public int hashCode()
         }
     }
 
+    public static final class ExtraHeader
+    {
+        private final String header;
+        private final String value;
+
+        public ExtraHeader(String headerAndValue)
+        {
+            List<String> nameValue = NAME_VALUE_SPLITTER.splitToList(headerAndValue);
+            checkArgument(nameValue.size() == 2, "Header and value: %s", headerAndValue);
+
+            this.header = nameValue.get(0);
+            this.value = nameValue.get(1);
+            checkArgument(!header.isEmpty(), "Header name is empty");
+            checkArgument(!value.isEmpty(), "Header value is empty");
+        }
+
+        public ExtraHeader(String header, String value)
+        {
+            this.header = header;
+            this.value = value;
+        }
+
+        public String getHeader()
+        {
+            return header;
+        }
+
+        public String getValue()
+        {
+            return value;
+        }
+
+        @Override
+        public boolean equals(Object o)
+        {
+            if (this == o) {
+                return true;
+            }
+            if (o == null || getClass() != o.getClass()) {
+                return false;
+            }
+            ExtraHeader other = (ExtraHeader) o;
+            return Objects.equals(header, other.header) && Objects.equals(value, other.value);
+        }
+
+        @Override
+        public int hashCode()
+        {
+            return Objects.hash(header, value);
+        }
+    }
+
     public static final class ClientSessionProperty
     {
         private static final Splitter NAME_SPLITTER = Splitter.on('.');
diff --git a/client/trino-cli/src/main/java/io/trino/cli/Trino.java b/client/trino-cli/src/main/java/io/trino/cli/Trino.java
@@ -20,6 +20,7 @@
 import io.trino.cli.ClientOptions.ClientExtraCredential;
 import io.trino.cli.ClientOptions.ClientResourceEstimate;
 import io.trino.cli.ClientOptions.ClientSessionProperty;
+import io.trino.cli.ClientOptions.ExtraHeader;
 import org.jline.utils.AttributedStringBuilder;
 import org.jline.utils.AttributedStyle;
 import picocli.CommandLine;
@@ -62,6 +63,7 @@ public static CommandLine createCommandLine(Object command)
                 .registerConverter(ClientResourceEstimate.class, ClientResourceEstimate::new)
                 .registerConverter(ClientSessionProperty.class, ClientSessionProperty::new)
                 .registerConverter(ClientExtraCredential.class, ClientExtraCredential::new)
+                .registerConverter(ExtraHeader.class, ExtraHeader::new)
                 .registerConverter(HostAndPort.class, HostAndPort::fromString)
                 .registerConverter(Duration.class, Duration::valueOf)
                 .setResourceBundle(new TrinoResourceBundle())
diff --git a/client/trino-cli/src/test/java/io/trino/cli/TestClientOptions.java b/client/trino-cli/src/test/java/io/trino/cli/TestClientOptions.java
@@ -246,6 +246,16 @@ public void testExtraCredentials()
                 new ClientOptions.ClientExtraCredential("test.token.bar", "bar")));
     }
 
+    @Test
+    public void testExtraHeaders()
+    {
+        Console console = createConsole("--extra-header", "X-Trino-Routing-Group=foo", "--extra-header", "x-foo=bar");
+        ClientOptions options = console.clientOptions;
+        assertThat(options.extraHeaders).isEqualTo(ImmutableList.of(
+                new ClientOptions.ExtraHeader("X-Trino-Routing-Group", "foo"),
+                new ClientOptions.ExtraHeader("x-foo", "bar")));
+    }
+
     @Test
     public void testSessionProperties()
     {
diff --git a/client/trino-client/src/main/java/io/trino/client/ClientSession.java b/client/trino-client/src/main/java/io/trino/client/ClientSession.java
@@ -51,6 +51,7 @@ public class ClientSession
     private final ZoneId timeZone;
     private final Locale locale;
     private final Map<String, String> resourceEstimates;
+    private final Map<String, String> extraHeaders;
     private final Map<String, String> properties;
     private final Map<String, String> preparedStatements;
     private final Map<String, ClientSelectedRole> roles;
@@ -87,6 +88,7 @@ private ClientSession(
             String source,
             Optional<String> traceToken,
             Set<String> clientTags,
+            Map<String, String> extraHeaders,
             String clientInfo,
             Optional<String> catalog,
             Optional<String> schema,
@@ -112,6 +114,7 @@ private ClientSession(
         this.source = requireNonNull(source, "source is null");
         this.traceToken = requireNonNull(traceToken, "traceToken is null");
         this.clientTags = ImmutableSet.copyOf(requireNonNull(clientTags, "clientTags is null"));
+        this.extraHeaders = ImmutableMap.copyOf(requireNonNull(extraHeaders, "extraHeaders is null"));
         this.clientInfo = clientInfo;
         this.catalog = catalog;
         this.schema = schema;
@@ -198,6 +201,11 @@ public Set<String> getClientTags()
         return clientTags;
     }
 
+    public Map<String, String> getExtraHeaders()
+    {
+        return extraHeaders;
+    }
+
     public String getClientInfo()
     {
         return clientInfo;
@@ -295,6 +303,7 @@ public String toString()
                 .add("sessionUser", sessionUser)
                 .add("authorizationUser", authorizationUser)
                 .add("clientTags", clientTags)
+                .add("extraHeaders", extraHeaders)
                 .add("clientInfo", clientInfo)
                 .add("catalog", catalog)
                 .add("schema", schema)
@@ -324,6 +333,7 @@ public static final class Builder
         private String source;
         private Optional<String> traceToken = Optional.empty();
         private Set<String> clientTags = ImmutableSet.of();
+        private Map<String, String> extraHeaders = ImmutableMap.of();
         private String clientInfo;
         private String catalog;
         private String schema;
@@ -354,6 +364,7 @@ private Builder(ClientSession clientSession)
             source = clientSession.getSource();
             traceToken = clientSession.getTraceToken();
             clientTags = clientSession.getClientTags();
+            extraHeaders = clientSession.getExtraHeaders();
             clientInfo = clientSession.getClientInfo();
             catalog = clientSession.getCatalog().orElse(null);
             schema = clientSession.getSchema().orElse(null);
@@ -419,6 +430,12 @@ public Builder clientTags(Set<String> clientTags)
             return this;
         }
 
+        public Builder extraHeaders(Map<String, String> extraHeaders)
+        {
+            this.extraHeaders = extraHeaders;
+            return this;
+        }
+
         public Builder clientInfo(String clientInfo)
         {
             this.clientInfo = clientInfo;
@@ -526,6 +543,7 @@ public ClientSession build()
                     source,
                     traceToken,
                     clientTags,
+                    extraHeaders,
                     clientInfo,
                     Optional.ofNullable(catalog),
                     Optional.ofNullable(schema),
diff --git a/client/trino-client/src/main/java/io/trino/client/OkHttpUtil.java b/client/trino-client/src/main/java/io/trino/client/OkHttpUtil.java
@@ -55,6 +55,7 @@
 import java.security.cert.X509Certificate;
 import java.util.Arrays;
 import java.util.List;
+import java.util.Map;
 import java.util.Optional;
 import java.util.concurrent.TimeUnit;
 
@@ -109,6 +110,17 @@ public static Interceptor tokenAuth(String accessToken)
                 .build());
     }
 
+    public static Interceptor extraHeaders(Map<String, String> extraHeaders)
+    {
+        requireNonNull(extraHeaders, "extraHeaders is null");
+
+        return chain -> {
+            okhttp3.Request.Builder builder = chain.request().newBuilder();
+            extraHeaders.forEach(builder::addHeader);
+            return chain.proceed(builder.build());
+        };
+    }
+
     public static void setupTimeouts(OkHttpClient.Builder clientBuilder, int timeout, TimeUnit unit)
     {
         clientBuilder
diff --git a/client/trino-client/src/main/java/io/trino/client/uri/ConnectionProperties.java b/client/trino-client/src/main/java/io/trino/client/uri/ConnectionProperties.java
@@ -99,6 +99,7 @@ enum SslVerificationMode
     public static final ConnectionProperty<String, Map<String, String>> EXTRA_CREDENTIALS = new ExtraCredentials();
     public static final ConnectionProperty<String, String> CLIENT_INFO = new ClientInfo();
     public static final ConnectionProperty<String, Set<String>> CLIENT_TAGS = new ClientTags();
+    public static final ConnectionProperty<String, Map<String, String>> EXTRA_HEADERS = new ExtraHeaders();
     public static final ConnectionProperty<String, String> TRACE_TOKEN = new TraceToken();
     public static final ConnectionProperty<String, Map<String, String>> SESSION_PROPERTIES = new SessionProperties();
     public static final ConnectionProperty<String, String> SOURCE = new Source();
@@ -139,6 +140,7 @@ enum SslVerificationMode
             .add(EXTERNAL_AUTHENTICATION_TIMEOUT)
             .add(EXTERNAL_AUTHENTICATION_TOKEN_CACHE)
             .add(EXTRA_CREDENTIALS)
+            .add(EXTRA_HEADERS)
             .add(HOSTNAME_IN_CERTIFICATE)
             .add(HTTP_LOGGING_LEVEL)
             .add(HTTP_PROXY)
@@ -789,6 +791,29 @@ public static String toString(Map<String, String> values)
         }
     }
 
+    private static class ExtraHeaders
+            extends AbstractConnectionProperty<String, Map<String, String>>
+    {
+        public ExtraHeaders()
+        {
+            super(PropertyName.EXTRA_HEADERS, NOT_REQUIRED, ALLOWED, converter(ExtraHeaders::parseExtraHeaders, ExtraHeaders::toString));
+        }
+
+        // Extra credentials consists of a list of credential name value pairs.
+        // E.g., `jdbc:trino://example.net:8080/?extraHeaders=abc:xyz;foo:bar` will create credentials `abc=xyz` and `foo=bar`
+        public static Map<String, String> parseExtraHeaders(String extraHeadersString)
+        {
+            return new MapPropertyParser(PropertyName.EXTRA_CREDENTIALS.toString()).parse(extraHeadersString);
+        }
+
+        public static String toString(Map<String, String> values)
+        {
+            return values.entrySet().stream()
+                    .map(entry -> entry.getKey() + ":" + entry.getValue())
+                    .collect(Collectors.joining(";"));
+        }
+    }
+
     private static class SessionProperties
             extends AbstractConnectionProperty<String, Map<String, String>>
     {
diff --git a/client/trino-client/src/main/java/io/trino/client/uri/HttpClientFactory.java b/client/trino-client/src/main/java/io/trino/client/uri/HttpClientFactory.java
@@ -30,6 +30,7 @@
 
 import static io.trino.client.KerberosUtil.defaultCredentialCachePath;
 import static io.trino.client.OkHttpUtil.basicAuth;
+import static io.trino.client.OkHttpUtil.extraHeaders;
 import static io.trino.client.OkHttpUtil.setupAlternateHostnameVerification;
 import static io.trino.client.OkHttpUtil.setupCookieJar;
 import static io.trino.client.OkHttpUtil.setupHttpLogging;
@@ -87,6 +88,10 @@ public static OkHttpClient.Builder toHttpClientBuilder(TrinoUri uri, String user
             builder.addNetworkInterceptor(tokenAuth(uri.getAccessToken().get()));
         }
 
+        if (!uri.getExtraHeaders().isEmpty()) {
+            builder.addNetworkInterceptor(extraHeaders(uri.getExtraHeaders()));
+        }
+
         if (uri.isExternalAuthenticationEnabled()) {
             if (!uri.isUseSecureConnection()) {
                 throw new RuntimeException("TLS/SSL required for authentication using external authorization");
diff --git a/client/trino-client/src/main/java/io/trino/client/uri/PropertyName.java b/client/trino-client/src/main/java/io/trino/client/uri/PropertyName.java
@@ -31,6 +31,7 @@ public enum PropertyName
     CLIENT_INFO("clientInfo"),
     CLIENT_TAGS("clientTags"),
     DISABLE_COMPRESSION("disableCompression"),
+    EXTRA_HEADERS("extraHeaders"),
     DNS_RESOLVER("dnsResolver"),
     DNS_RESOLVER_CONTEXT("dnsResolverContext"),
     ENCODING("encoding"),
diff --git a/client/trino-client/src/main/java/io/trino/client/uri/TrinoUri.java b/client/trino-client/src/main/java/io/trino/client/uri/TrinoUri.java
@@ -61,6 +61,7 @@
 import static io.trino.client.uri.ConnectionProperties.EXTERNAL_AUTHENTICATION_TIMEOUT;
 import static io.trino.client.uri.ConnectionProperties.EXTERNAL_AUTHENTICATION_TOKEN_CACHE;
 import static io.trino.client.uri.ConnectionProperties.EXTRA_CREDENTIALS;
+import static io.trino.client.uri.ConnectionProperties.EXTRA_HEADERS;
 import static io.trino.client.uri.ConnectionProperties.HOSTNAME_IN_CERTIFICATE;
 import static io.trino.client.uri.ConnectionProperties.HTTP_LOGGING_LEVEL;
 import static io.trino.client.uri.ConnectionProperties.HTTP_PROXY;
@@ -244,6 +245,11 @@ public Map<String, String> getSessionProperties()
         return resolveWithDefault(SESSION_PROPERTIES, ImmutableMap.of());
     }
 
+    public Map<String, String> getExtraHeaders()
+    {
+        return resolveWithDefault(EXTRA_HEADERS, ImmutableMap.of());
+    }
+
     public Optional<String> getSource()
     {
         return resolveOptional(SOURCE);
@@ -521,6 +527,7 @@ public ClientSession.Builder toClientSessionBuilder()
                 .timeZone(getTimeZone())
                 .locale(getLocale())
                 .properties(getSessionProperties())
+                .extraHeaders(getExtraHeaders())
                 .credentials(getExtraCredentials())
                 .transactionId(null)
                 .resourceEstimates(getResourceEstimates())
@@ -1009,6 +1016,11 @@ public Builder setSessionProperties(Map<String, String> sessionProperties)
             return setProperty(SESSION_PROPERTIES, requireNonNull(sessionProperties, "sessionProperties is null"));
         }
 
+        public Builder setExtraHeaders(Map<String, String> extraHeaders)
+        {
+            return setProperty(EXTRA_HEADERS, requireNonNull(extraHeaders, "extraHeaders is null"));
+        }
+
         public Builder setSource(String source)
         {
             return setProperty(SOURCE, requireNonNull(source, "source is null"));
diff --git a/docs/src/main/sphinx/client/cli.md b/docs/src/main/sphinx/client/cli.md
@@ -161,6 +161,9 @@ mode:
 * - `--editing-mode`
   - Sets key bindings in the CLI to be compatible with VI or
     EMACS editors. Defaults to `EMACS`.
+* - ``--extra-header``
+  - Arbitrary headers to attach to authenticated HTTP requests made by the CLI. Property can be
+    used multiple times with the format ``header=value``.
 * - `--extra-credential`
   - Extra credentials (property can be used multiple times; format is key=value)
 * - `--http-proxy`
