Merge branch 'master' into conflict_fix

Author: sc250072

client/trino-cli/src/main/java/io/trino/cli/ClientOptions.java

@@ -47,6 +47,7 @@
 import static com.google.common.base.Strings.nullToEmpty;
 import static io.trino.cli.TerminalUtils.getTerminal;
 import static io.trino.client.KerberosUtil.defaultCredentialCachePath;
+import static io.trino.client.ProtocolHeaders.TRINO_HEADERS;
 import static io.trino.client.uri.PropertyName.ACCESS_TOKEN;
 import static io.trino.client.uri.PropertyName.CATALOG;
 import static io.trino.client.uri.PropertyName.CLIENT_INFO;
@@ -56,6 +57,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 +206,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 +428,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 +507,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 +623,59 @@ 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(!TRINO_HEADERS.isProtocolHeader(header), "Header '%s' is a protocol header and cannot be set as an extra header", header);
+            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('.');

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

client/trino-cli/src/test/java/io/trino/cli/TestClientOptions.java

@@ -246,6 +246,20 @@ 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")));
+
+        assertThatThrownBy(() -> createConsole("--extra-header", "X-Trino-User=Forbidden"))
+                .hasCauseInstanceOf(IllegalArgumentException.class)
+                .hasMessageContaining("Header 'X-Trino-User' is a protocol header and cannot be set as an extra header");
+    }
+
     @Test
     public void testSessionProperties()
     {

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

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

client/trino-client/src/main/java/io/trino/client/ProtocolHeaders.java

@@ -194,6 +194,16 @@ private ProtocolHeaders(String name)
         responseOriginalRole = RESPONSE_SET_ORIGINAL_ROLES.withProtocolName(name);
     }
 
+    public boolean isProtocolHeader(String headerName)
+    {
+        for (Headers header : Headers.values()) {
+            if (header.withProtocolName(name).equalsIgnoreCase(headerName)) {
+                return true;
+            }
+        }
+        return false;
+    }
+
     public String getProtocolName()
     {
         return name;

client/trino-client/src/main/java/io/trino/client/uri/ConnectionProperties.java

@@ -16,6 +16,7 @@
 import com.google.common.base.CharMatcher;
 import com.google.common.base.Joiner;
 import com.google.common.base.Splitter;
+import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSet;
 import com.google.common.net.HostAndPort;
 import io.airlift.units.Duration;
@@ -44,6 +45,7 @@
 import static com.google.common.collect.Streams.stream;
 import static io.trino.client.ClientSelectedRole.Type.ALL;
 import static io.trino.client.ClientSelectedRole.Type.NONE;
+import static io.trino.client.ProtocolHeaders.TRINO_HEADERS;
 import static io.trino.client.uri.AbstractConnectionProperty.Validator;
 import static io.trino.client.uri.AbstractConnectionProperty.validator;
 import static io.trino.client.uri.ConnectionProperties.SslVerificationMode.FULL;
@@ -99,6 +101,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 +142,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 +793,42 @@ public static String toString(Map<String, String> values)
         }
     }
 
+    private static class ExtraHeaders
+            extends AbstractConnectionProperty<String, Map<String, String>>
+    {
+        private static final Validator<Properties> VALIDATE_EXTRA_HEADER = validator(
+                ExtraHeaders::isNotReservedHeader,
+                format("Connection property %s cannot override any of the Trino protocol headers", PropertyName.EXTRA_HEADERS));
+
+        public ExtraHeaders()
+        {
+            super(PropertyName.EXTRA_HEADERS, NOT_REQUIRED, VALIDATE_EXTRA_HEADER, converter(ExtraHeaders::parseExtraHeaders, ExtraHeaders::toString));
+        }
+
+        // Extra headers consists of a list of header name value pairs.
+        // E.g., `jdbc:trino://example.net:8080/?extraHeaders=X-Trino-Route:foo;X-Trino-Custom:bar` will send
+        // HTTP headers `X-Trino-Route=foo` and `X-Trino-Custom=bar`.
+        // These headers must not conflict with Trino protocol headers.
+        public static Map<String, String> parseExtraHeaders(String extraHeadersString)
+        {
+            return new MapPropertyParser(PropertyName.EXTRA_HEADERS.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 boolean isNotReservedHeader(Properties properties)
+        {
+            Map<String, String> extraHeaders = EXTRA_HEADERS.getValueOrDefault(properties, ImmutableMap.of());
+            return extraHeaders.keySet().stream()
+                    .noneMatch(TRINO_HEADERS::isProtocolHeader);
+        }
+    }
+
     private static class SessionProperties
             extends AbstractConnectionProperty<String, Map<String, String>>
     {

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

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"),