Skip to content

Commit f9d15ce

Browse files
PluiesPluies
committed
Trino CLI: support --extra-header parameter
Adds an --extra-header flag to the Trino CLI to allow for passing arbitrary HTTP headers to Trino requests.
1 parent 4846bdb commit f9d15ce

File tree

9 files changed

+163
-0
lines changed

9 files changed

+163
-0
lines changed

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

Lines changed: 70 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -55,6 +55,7 @@
5555
import static io.trino.client.uri.PropertyName.EXTERNAL_AUTHENTICATION;
5656
import static io.trino.client.uri.PropertyName.EXTERNAL_AUTHENTICATION_REDIRECT_HANDLERS;
5757
import static io.trino.client.uri.PropertyName.EXTRA_CREDENTIALS;
58+
import static io.trino.client.uri.PropertyName.EXTRA_HEADERS;
5859
import static io.trino.client.uri.PropertyName.HTTP_PROXY;
5960
import static io.trino.client.uri.PropertyName.KERBEROS_CONFIG_PATH;
6061
import static io.trino.client.uri.PropertyName.KERBEROS_CREDENTIAL_CACHE_PATH;
@@ -193,6 +194,10 @@ public class ClientOptions
193194
@Option(names = "--client-tags", paramLabel = "<tags>", description = "Client tags")
194195
public Optional<String> clientTags;
195196

197+
@PropertyMapping(EXTRA_HEADERS)
198+
@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)")
199+
public final List<ExtraHeader> extraHeaders = new ArrayList<>();
200+
196201
@PropertyMapping(TRACE_TOKEN)
197202
@Option(names = "--trace-token", paramLabel = "<token>", description = "Trace token")
198203
public Optional<String> traceToken;
@@ -321,6 +326,7 @@ public ClientSession toClientSession(TrinoUri uri)
321326
.source(source.orElse("trino-cli"))
322327
.traceToken(traceToken)
323328
.clientTags(parseClientTags(clientTags.orElse("")))
329+
.extraHeaders(toExtraHeaders(extraHeaders))
324330
.clientInfo(clientInfo.orElse(null))
325331
.catalog(uri.getCatalog().orElse(catalog.orElse(null)))
326332
.schema(uri.getSchema().orElse(schema.orElse(null)))
@@ -405,6 +411,9 @@ public TrinoUri getTrinoUri(Map<PropertyName, String> restrictedProperties)
405411
source.ifPresent(builder::setSource);
406412
clientInfo.ifPresent(builder::setClientInfo);
407413
clientTags.ifPresent(builder::setClientTags);
414+
if (!extraHeaders.isEmpty()) {
415+
builder.setExtraHeaders(toExtraHeaders(extraHeaders));
416+
}
408417
traceToken.ifPresent(builder::setTraceToken);
409418
socksProxy.ifPresent(builder::setSocksProxy);
410419
httpProxy.ifPresent(builder::setHttpProxy);
@@ -476,6 +485,15 @@ public static Set<String> parseClientTags(String clientTagsString)
476485
return ImmutableSet.copyOf(splitter.split(nullToEmpty(clientTagsString)));
477486
}
478487

488+
public static Map<String, String> toExtraHeaders(List<ExtraHeader> extraHeaders)
489+
{
490+
ImmutableMap.Builder<String, String> builder = ImmutableMap.builder();
491+
for (ExtraHeader extraHeader : extraHeaders) {
492+
builder.put(extraHeader.getHeader(), extraHeader.getValue());
493+
}
494+
return builder.buildOrThrow();
495+
}
496+
479497
public static Map<String, String> toProperties(List<ClientSessionProperty> sessionProperties)
480498
{
481499
ImmutableMap.Builder<String, String> builder = ImmutableMap.builder();
@@ -570,6 +588,58 @@ public int hashCode()
570588
}
571589
}
572590

591+
public static final class ExtraHeader
592+
{
593+
private final String header;
594+
private final String value;
595+
596+
public ExtraHeader(String headerAndValue)
597+
{
598+
List<String> nameValue = NAME_VALUE_SPLITTER.splitToList(headerAndValue);
599+
checkArgument(nameValue.size() == 2, "Header and value: %s", headerAndValue);
600+
601+
this.header = nameValue.get(0);
602+
this.value = nameValue.get(1);
603+
checkArgument(!header.isEmpty(), "Header name is empty");
604+
checkArgument(!value.isEmpty(), "Header value is empty");
605+
}
606+
607+
public ExtraHeader(String header, String value)
608+
{
609+
this.header = header;
610+
this.value = value;
611+
}
612+
613+
public String getHeader()
614+
{
615+
return header;
616+
}
617+
618+
public String getValue()
619+
{
620+
return value;
621+
}
622+
623+
@Override
624+
public boolean equals(Object o)
625+
{
626+
if (this == o) {
627+
return true;
628+
}
629+
if (o == null || getClass() != o.getClass()) {
630+
return false;
631+
}
632+
ExtraHeader other = (ExtraHeader) o;
633+
return Objects.equals(header, other.header) && Objects.equals(value, other.value);
634+
}
635+
636+
@Override
637+
public int hashCode()
638+
{
639+
return Objects.hash(header, value);
640+
}
641+
}
642+
573643
public static final class ClientSessionProperty
574644
{
575645
private static final Splitter NAME_SPLITTER = Splitter.on('.');

‎client/trino-cli/src/main/java/io/trino/cli/Trino.java‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@
1818
import io.trino.cli.ClientOptions.ClientExtraCredential;
1919
import io.trino.cli.ClientOptions.ClientResourceEstimate;
2020
import io.trino.cli.ClientOptions.ClientSessionProperty;
21+
import io.trino.cli.ClientOptions.ExtraHeader;
2122
import org.jline.utils.AttributedStringBuilder;
2223
import org.jline.utils.AttributedStyle;
2324
import picocli.CommandLine;
@@ -54,6 +55,7 @@ public static CommandLine createCommandLine(Object command)
5455
.registerConverter(ClientResourceEstimate.class, ClientResourceEstimate::new)
5556
.registerConverter(ClientSessionProperty.class, ClientSessionProperty::new)
5657
.registerConverter(ClientExtraCredential.class, ClientExtraCredential::new)
58+
.registerConverter(ExtraHeader.class, ExtraHeader::new)
5759
.registerConverter(HostAndPort.class, HostAndPort::fromString)
5860
.registerConverter(Duration.class, Duration::valueOf)
5961
.setExecutionExceptionHandler((e, cmd, parseResult) -> {

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

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -215,6 +215,16 @@ public void testExtraCredentials()
215215
new ClientOptions.ClientExtraCredential("test.token.bar", "bar")));
216216
}
217217

218+
@Test
219+
public void testExtraHeaders()
220+
{
221+
Console console = createConsole("--extra-header", "X-Trino-Routing-Group=foo", "--extra-header", "x-foo=bar");
222+
ClientOptions options = console.clientOptions;
223+
assertEquals(options.extraHeaders, ImmutableList.of(
224+
new ClientOptions.ExtraHeader("X-Trino-Routing-Group", "foo"),
225+
new ClientOptions.ExtraHeader("x-foo", "bar")));
226+
}
227+
218228
@Test
219229
public void testSessionProperties()
220230
{

‎client/trino-client/src/main/java/io/trino/client/ClientSession.java‎

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -46,6 +46,7 @@ public class ClientSession
4646
private final ZoneId timeZone;
4747
private final Locale locale;
4848
private final Map<String, String> resourceEstimates;
49+
private final Map<String, String> extraHeaders;
4950
private final Map<String, String> properties;
5051
private final Map<String, String> preparedStatements;
5152
private final Map<String, ClientSelectedRole> roles;
@@ -78,6 +79,7 @@ private ClientSession(
7879
String source,
7980
Optional<String> traceToken,
8081
Set<String> clientTags,
82+
Map<String, String> extraHeaders,
8183
String clientInfo,
8284
Optional<String> catalog,
8385
Optional<String> schema,
@@ -99,6 +101,7 @@ private ClientSession(
99101
this.source = source;
100102
this.traceToken = requireNonNull(traceToken, "traceToken is null");
101103
this.clientTags = ImmutableSet.copyOf(requireNonNull(clientTags, "clientTags is null"));
104+
this.extraHeaders = ImmutableMap.copyOf(requireNonNull(extraHeaders, "extraHeaders is null"));
102105
this.clientInfo = clientInfo;
103106
this.catalog = catalog;
104107
this.schema = schema;
@@ -173,6 +176,11 @@ public Set<String> getClientTags()
173176
return clientTags;
174177
}
175178

179+
public Map<String, String> getExtraHeaders()
180+
{
181+
return extraHeaders;
182+
}
183+
176184
public String getClientInfo()
177185
{
178186
return clientInfo;
@@ -259,6 +267,7 @@ public String toString()
259267
.add("principal", principal)
260268
.add("user", user)
261269
.add("clientTags", clientTags)
270+
.add("extraHeaders", extraHeaders)
262271
.add("clientInfo", clientInfo)
263272
.add("catalog", catalog)
264273
.add("schema", schema)
@@ -280,6 +289,7 @@ public static final class Builder
280289
private String source;
281290
private Optional<String> traceToken = Optional.empty();
282291
private Set<String> clientTags = ImmutableSet.of();
292+
private Map<String, String> extraHeaders = ImmutableMap.of();
283293
private String clientInfo;
284294
private String catalog;
285295
private String schema;
@@ -306,6 +316,7 @@ private Builder(ClientSession clientSession)
306316
source = clientSession.getSource();
307317
traceToken = clientSession.getTraceToken();
308318
clientTags = clientSession.getClientTags();
319+
extraHeaders = clientSession.getExtraHeaders();
309320
clientInfo = clientSession.getClientInfo();
310321
catalog = clientSession.getCatalog().orElse(null);
311322
schema = clientSession.getSchema().orElse(null);
@@ -358,6 +369,12 @@ public Builder clientTags(Set<String> clientTags)
358369
return this;
359370
}
360371

372+
public Builder extraHeaders(Map<String, String> extraHeaders)
373+
{
374+
this.extraHeaders = extraHeaders;
375+
return this;
376+
}
377+
361378
public Builder clientInfo(String clientInfo)
362379
{
363380
this.clientInfo = clientInfo;
@@ -451,6 +468,7 @@ public ClientSession build()
451468
source,
452469
traceToken,
453470
clientTags,
471+
extraHeaders,
454472
clientInfo,
455473
Optional.ofNullable(catalog),
456474
Optional.ofNullable(schema),

‎client/trino-client/src/main/java/io/trino/client/OkHttpUtil.java‎

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -52,6 +52,7 @@
5252
import java.security.cert.X509Certificate;
5353
import java.util.Arrays;
5454
import java.util.List;
55+
import java.util.Map;
5556
import java.util.Optional;
5657
import java.util.concurrent.TimeUnit;
5758

@@ -98,6 +99,17 @@ public static Interceptor tokenAuth(String accessToken)
9899
.build());
99100
}
100101

102+
public static Interceptor extraHeaders(Map<String, String> extraHeaders)
103+
{
104+
requireNonNull(extraHeaders, "extraHeaders is null");
105+
106+
return chain -> {
107+
okhttp3.Request.Builder builder = chain.request().newBuilder();
108+
extraHeaders.forEach((k, v) -> builder.addHeader(k, v));
109+
return chain.proceed(builder.build());
110+
};
111+
}
112+
101113
public static void setupTimeouts(OkHttpClient.Builder clientBuilder, int timeout, TimeUnit unit)
102114
{
103115
clientBuilder

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

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -91,6 +91,7 @@ enum SslVerificationMode
9191
public static final ConnectionProperty<String, Map<String, String>> EXTRA_CREDENTIALS = new ExtraCredentials();
9292
public static final ConnectionProperty<String, String> CLIENT_INFO = new ClientInfo();
9393
public static final ConnectionProperty<String, String> CLIENT_TAGS = new ClientTags();
94+
public static final ConnectionProperty<String, Map<String, String>> EXTRA_HEADERS = new ExtraHeaders();
9495
public static final ConnectionProperty<String, String> TRACE_TOKEN = new TraceToken();
9596
public static final ConnectionProperty<String, Map<String, String>> SESSION_PROPERTIES = new SessionProperties();
9697
public static final ConnectionProperty<String, String> SOURCE = new Source();
@@ -131,6 +132,7 @@ enum SslVerificationMode
131132
.add(EXTRA_CREDENTIALS)
132133
.add(CLIENT_INFO)
133134
.add(CLIENT_TAGS)
135+
.add(EXTRA_HEADERS)
134136
.add(TRACE_TOKEN)
135137
.add(SESSION_PROPERTIES)
136138
.add(SOURCE)
@@ -636,6 +638,22 @@ public static Map<String, String> parseExtraCredentials(String extraCredentialSt
636638
}
637639
}
638640

641+
private static class ExtraHeaders
642+
extends AbstractConnectionProperty<String, Map<String, String>>
643+
{
644+
public ExtraHeaders()
645+
{
646+
super(PropertyName.EXTRA_HEADERS, NOT_REQUIRED, ALLOWED, ExtraHeaders::parseExtraHeaders);
647+
}
648+
649+
// Extra credentials consists of a list of credential name value pairs.
650+
// E.g., `jdbc:trino://example.net:8080/?extraHeaders=abc:xyz;foo:bar` will create credentials `abc=xyz` and `foo=bar`
651+
public static Map<String, String> parseExtraHeaders(String extraHeadersString)
652+
{
653+
return new MapPropertyParser(PropertyName.EXTRA_CREDENTIALS.toString()).parse(extraHeadersString);
654+
}
655+
}
656+
639657
private static class SessionProperties
640658
extends AbstractConnectionProperty<String, Map<String, String>>
641659
{

‎client/trino-client/src/main/java/io/trino/client/uri/PropertyName.java‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -58,6 +58,7 @@ public enum PropertyName
5858
EXTRA_CREDENTIALS("extraCredentials"),
5959
CLIENT_INFO("clientInfo"),
6060
CLIENT_TAGS("clientTags"),
61+
EXTRA_HEADERS("extraHeaders"),
6162
TRACE_TOKEN("traceToken"),
6263
SESSION_PROPERTIES("sessionProperties"),
6364
SOURCE("source"),

0 commit comments

Comments
 (0)