trinodb · Pluies · Aug 11, 2023
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
@@ -55,6 +55,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_PROXY;
 import static io.trino.client.uri.PropertyName.KERBEROS_CONFIG_PATH;
 import static io.trino.client.uri.PropertyName.KERBEROS_CREDENTIAL_CACHE_PATH;
@@ -193,6 +194,10 @@ public class ClientOptions
     @Option(names = "--client-tags", paramLabel = "<tags>", description = "Client tags")
     public Optional<String> clientTags;
 
+    @PropertyMapping(EXTRA_HEADERS)
+    @Option(names = "--extra-header", paramLabel = "<header>", description = "Extra HTTP header to attach to Trino queries (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;
@@ -321,6 +326,7 @@ public ClientSession toClientSession(TrinoUri uri)
                 .source(source.orElse("trino-cli"))
                 .traceToken(traceToken)
                 .clientTags(parseClientTags(clientTags.orElse("")))
+                .extraHeaders(toExtraHeaders(extraHeaders))
                 .clientInfo(clientInfo.orElse(null))
                 .catalog(uri.getCatalog().orElse(catalog.orElse(null)))
                 .schema(uri.getSchema().orElse(schema.orElse(null)))
@@ -405,6 +411,9 @@ public TrinoUri getTrinoUri(Map<PropertyName, String> restrictedProperties)
         source.ifPresent(builder::setSource);
         clientInfo.ifPresent(builder::setClientInfo);
         clientTags.ifPresent(builder::setClientTags);
+        if (!extraHeaders.isEmpty()) {
+            builder.setExtraHeaders(toExtraHeaders(extraHeaders));
+        }
         traceToken.ifPresent(builder::setTraceToken);
         socksProxy.ifPresent(builder::setSocksProxy);
         httpProxy.ifPresent(builder::setHttpProxy);
@@ -476,6 +485,15 @@ public static Set<String> parseClientTags(String clientTagsString)
         return ImmutableSet.copyOf(splitter.split(nullToEmpty(clientTagsString)));
     }
 
+    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();
+    }
+
     public static Map<String, String> toProperties(List<ClientSessionProperty> sessionProperties)
     {
         ImmutableMap.Builder<String, String> builder = ImmutableMap.builder();
@@ -570,6 +588,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
@@ -18,6 +18,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;
@@ -54,6 +55,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)
                 .setExecutionExceptionHandler((e, cmd, parseResult) -> {

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
@@ -215,6 +215,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;
+        assertEquals(options.extraHeaders, 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
@@ -46,6 +46,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;
@@ -78,6 +79,7 @@ private ClientSession(
             String source,
             Optional<String> traceToken,
             Set<String> clientTags,
+            Map<String, String> extraHeaders,
             String clientInfo,
             Optional<String> catalog,
             Optional<String> schema,
@@ -99,6 +101,7 @@ private ClientSession(
         this.source = source;
         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;
@@ -173,6 +176,11 @@ public Set<String> getClientTags()
         return clientTags;
     }
 
+    public Map<String, String> getExtraHeaders()
+    {
+        return extraHeaders;
+    }
+
     public String getClientInfo()
     {
         return clientInfo;
@@ -259,6 +267,7 @@ public String toString()
                 .add("principal", principal)
                 .add("user", user)
                 .add("clientTags", clientTags)
+                .add("extraHeaders", extraHeaders)
                 .add("clientInfo", clientInfo)
                 .add("catalog", catalog)
                 .add("schema", schema)
@@ -280,6 +289,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;
@@ -306,6 +316,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);
@@ -358,6 +369,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;
@@ -451,6 +468,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
@@ -52,6 +52,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;
 
@@ -98,6 +99,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((k, v) -> builder.addHeader(k, v));
+            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
@@ -91,6 +91,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, 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();
@@ -131,6 +132,7 @@ enum SslVerificationMode
             .add(EXTRA_CREDENTIALS)
             .add(CLIENT_INFO)
             .add(CLIENT_TAGS)
+            .add(EXTRA_HEADERS)
             .add(TRACE_TOKEN)
             .add(SESSION_PROPERTIES)
             .add(SOURCE)
@@ -636,6 +638,22 @@ public static Map<String, String> parseExtraCredentials(String extraCredentialSt
         }
     }
 
+    private static class ExtraHeaders
+            extends AbstractConnectionProperty<String, Map<String, String>>
+    {
+        public ExtraHeaders()
+        {
+            super(PropertyName.EXTRA_HEADERS, NOT_REQUIRED, ALLOWED, ExtraHeaders::parseExtraHeaders);
+        }
+
+        // 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);
+        }
+    }
+
     private static class SessionProperties
             extends AbstractConnectionProperty<String, Map<String, String>>
     {

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
@@ -58,6 +58,7 @@ public enum PropertyName
     EXTRA_CREDENTIALS("extraCredentials"),
     CLIENT_INFO("clientInfo"),
     CLIENT_TAGS("clientTags"),
+    EXTRA_HEADERS("extraHeaders"),
     TRACE_TOKEN("traceToken"),
     SESSION_PROPERTIES("sessionProperties"),
     SOURCE("source"),