diff --git a/core/src/main/java/org/apache/calcite/avatica/BuiltInConnectionProperty.java b/core/src/main/java/org/apache/calcite/avatica/BuiltInConnectionProperty.java index 62c570298d..4cb18d10b0 100644 --- a/core/src/main/java/org/apache/calcite/avatica/BuiltInConnectionProperty.java +++ b/core/src/main/java/org/apache/calcite/avatica/BuiltInConnectionProperty.java @@ -91,7 +91,10 @@ public enum BuiltInConnectionProperty implements ConnectionProperty { TRANSPARENT_RECONNECTION("transparent_reconnection", Type.BOOLEAN, Boolean.FALSE, false), /** Number of rows to fetch per call. */ - FETCH_SIZE("fetch_size", Type.NUMBER, AvaticaStatement.DEFAULT_FETCH_SIZE, false); + FETCH_SIZE("fetch_size", Type.NUMBER, AvaticaStatement.DEFAULT_FETCH_SIZE, false), + + /** User-Agent header to identify requests performed by a particular application **/ + USER_AGENT("userAgent", Type.STRING, "", false); private final String camelName; private final Type type; diff --git a/core/src/main/java/org/apache/calcite/avatica/ConnectionConfig.java b/core/src/main/java/org/apache/calcite/avatica/ConnectionConfig.java index bc06f8af52..e82379b8b8 100644 --- a/core/src/main/java/org/apache/calcite/avatica/ConnectionConfig.java +++ b/core/src/main/java/org/apache/calcite/avatica/ConnectionConfig.java @@ -66,6 +66,8 @@ public interface ConnectionConfig { boolean transparentReconnectionEnabled(); /** @see BuiltInConnectionProperty#FETCH_SIZE */ int fetchSize(); + /** @see BuiltInConnectionProperty#USER_AGENT */ + String userAgent(); } // End ConnectionConfig.java diff --git a/core/src/main/java/org/apache/calcite/avatica/ConnectionConfigImpl.java b/core/src/main/java/org/apache/calcite/avatica/ConnectionConfigImpl.java index 0bb367732a..41e4f8aead 100644 --- a/core/src/main/java/org/apache/calcite/avatica/ConnectionConfigImpl.java +++ b/core/src/main/java/org/apache/calcite/avatica/ConnectionConfigImpl.java @@ -137,6 +137,10 @@ public int fetchSize() { return BuiltInConnectionProperty.FETCH_SIZE.wrap(properties).getInt(); } + public String userAgent() { + return BuiltInConnectionProperty.USER_AGENT.wrap(properties).getString(); + } + /** Converts a {@link Properties} object containing (name, value) * pairs into a map whose keys are * {@link org.apache.calcite.avatica.InternalProperty} objects. diff --git a/core/src/main/java/org/apache/calcite/avatica/remote/AvaticaCommonsHttpClientImpl.java b/core/src/main/java/org/apache/calcite/avatica/remote/AvaticaCommonsHttpClientImpl.java index 66b7ca8fcc..61bf4b9a33 100644 --- a/core/src/main/java/org/apache/calcite/avatica/remote/AvaticaCommonsHttpClientImpl.java +++ b/core/src/main/java/org/apache/calcite/avatica/remote/AvaticaCommonsHttpClientImpl.java @@ -86,6 +86,7 @@ public class AvaticaCommonsHttpClientImpl implements AvaticaHttpClient, HttpClie protected CredentialsProvider credentialsProvider = null; protected Lookup authRegistry = null; protected Object userToken; + protected String userAgent; public AvaticaCommonsHttpClientImpl(URL url) { this.uri = toURI(Objects.requireNonNull(url)); @@ -119,6 +120,10 @@ protected void initializeClient(PoolingHttpClientConnectionManager pool) { HttpPost post = new HttpPost(uri); post.setEntity(entity); + if (null != userAgent) { + post.addHeader("User-Agent", userAgent); + } + try (CloseableHttpResponse response = execute(post, context)) { final int statusCode = response.getCode(); if (HttpURLConnection.HTTP_OK == statusCode @@ -145,6 +150,11 @@ protected void initializeClient(PoolingHttpClientConnectionManager pool) { } } + @Override + public void setUserAgent(String userAgent) { + this.userAgent = userAgent; + } + // Visible for testing CloseableHttpResponse execute(HttpPost post, HttpClientContext context) throws IOException, ClientProtocolException { diff --git a/core/src/main/java/org/apache/calcite/avatica/remote/AvaticaHttpClient.java b/core/src/main/java/org/apache/calcite/avatica/remote/AvaticaHttpClient.java index eac1b74f15..984c1c22c1 100644 --- a/core/src/main/java/org/apache/calcite/avatica/remote/AvaticaHttpClient.java +++ b/core/src/main/java/org/apache/calcite/avatica/remote/AvaticaHttpClient.java @@ -29,6 +29,14 @@ public interface AvaticaHttpClient { */ byte[] send(byte[] request); + /** + * Sets the "User-Agent" HTTP header for all request to the Avatica server. + * + * @param userAgent The "User-Agent" to use. Configurable via + * {@link org.apache.calcite.avatica.BuiltInConnectionProperty#USER_AGENT} + */ + void setUserAgent(String userAgent); + } // End AvaticaHttpClient.java diff --git a/core/src/main/java/org/apache/calcite/avatica/remote/AvaticaHttpClientFactoryImpl.java b/core/src/main/java/org/apache/calcite/avatica/remote/AvaticaHttpClientFactoryImpl.java index e0f3446ae3..ea248cf70e 100644 --- a/core/src/main/java/org/apache/calcite/avatica/remote/AvaticaHttpClientFactoryImpl.java +++ b/core/src/main/java/org/apache/calcite/avatica/remote/AvaticaHttpClientFactoryImpl.java @@ -134,6 +134,10 @@ public static AvaticaHttpClientFactoryImpl getInstance() { client = new DoAsAvaticaHttpClient(client, kerberosUtil); } + if (!config.userAgent().isEmpty()) { + client.setUserAgent(config.userAgent()); + } + return client; } diff --git a/core/src/main/java/org/apache/calcite/avatica/remote/AvaticaHttpClientImpl.java b/core/src/main/java/org/apache/calcite/avatica/remote/AvaticaHttpClientImpl.java index 6fda17954e..1783c47538 100644 --- a/core/src/main/java/org/apache/calcite/avatica/remote/AvaticaHttpClientImpl.java +++ b/core/src/main/java/org/apache/calcite/avatica/remote/AvaticaHttpClientImpl.java @@ -30,6 +30,7 @@ */ public class AvaticaHttpClientImpl implements AvaticaHttpClient { protected final URL url; + protected String userAgent; public AvaticaHttpClientImpl(URL url) { this.url = url; @@ -43,6 +44,11 @@ public byte[] send(byte[] request) { connection.setRequestMethod("POST"); connection.setDoInput(true); connection.setDoOutput(true); + + if (null != userAgent) { + connection.setRequestProperty("User-Agent", userAgent); + } + try (DataOutputStream wr = new DataOutputStream(connection.getOutputStream())) { wr.write(request); wr.flush(); @@ -69,6 +75,11 @@ public byte[] send(byte[] request) { } } + @Override + public void setUserAgent(String userAgent) { + this.userAgent = userAgent; + } + HttpURLConnection openConnection() throws IOException { return (HttpURLConnection) url.openConnection(); } diff --git a/core/src/main/java/org/apache/calcite/avatica/remote/DoAsAvaticaHttpClient.java b/core/src/main/java/org/apache/calcite/avatica/remote/DoAsAvaticaHttpClient.java index 123f821703..8b516b265f 100644 --- a/core/src/main/java/org/apache/calcite/avatica/remote/DoAsAvaticaHttpClient.java +++ b/core/src/main/java/org/apache/calcite/avatica/remote/DoAsAvaticaHttpClient.java @@ -40,6 +40,11 @@ public DoAsAvaticaHttpClient(AvaticaHttpClient wrapped, KerberosConnection kerbe } }); } + + @Override + public void setUserAgent(String userAgent) { + wrapped.setUserAgent(userAgent); + } } // End DoAsAvaticaHttpClient.java diff --git a/server/src/test/java/org/apache/calcite/avatica/remote/AlternatingRemoteMetaTest.java b/server/src/test/java/org/apache/calcite/avatica/remote/AlternatingRemoteMetaTest.java index 41c84fe3d1..58534d7203 100644 --- a/server/src/test/java/org/apache/calcite/avatica/remote/AlternatingRemoteMetaTest.java +++ b/server/src/test/java/org/apache/calcite/avatica/remote/AlternatingRemoteMetaTest.java @@ -124,6 +124,9 @@ public byte[] send(byte[] request) { return client.send(request); } + @Override + public void setUserAgent(String userAgent) {} + } /**