make first part of user agent header configurable (#9471)

apache · Oct 19, 2022 · 3e1aba3 · 3e1aba3
1 parent 5209000
commit 3e1aba3
Show file tree

Hide file tree

Showing 8 changed files with 107 additions and 56 deletions.
diff --git a/pinot-clients/pinot-java-client/src/main/java/org/apache/pinot/client/BrokerCache.java b/pinot-clients/pinot-java-client/src/main/java/org/apache/pinot/client/BrokerCache.java
@@ -104,6 +104,7 @@ public BrokerCache(Properties properties, String controllerUrl) {
         DEFAULT_CONTROLLER_CONNECT_TIMEOUT_MS));
     int handshakeTimeoutMs = Integer.parseInt(properties.getProperty("controllerHandshakeTimeoutMs",
         DEFAULT_CONTROLLER_HANDSHAKE_TIMEOUT_MS));
+    String appId = properties.getProperty("appId");
     boolean tlsV10Enabled = Boolean.parseBoolean(properties.getProperty("controllerTlsV10Enabled",
         DEFAULT_CONTROLLER_TLS_V10_ENABLED))
         || Boolean.parseBoolean(System.getProperties().getProperty("controller.tlsV10Enabled",
@@ -113,7 +114,7 @@ public BrokerCache(Properties properties, String controllerUrl) {
     builder.setReadTimeout(readTimeoutMs)
         .setConnectTimeout(connectTimeoutMs)
         .setHandshakeTimeout(handshakeTimeoutMs)
-        .setUserAgent(ConnectionUtils.getUserAgentVersionFromClassPath("ua_broker_cache"))
+        .setUserAgent(ConnectionUtils.getUserAgentVersionFromClassPath("ua_broker_cache", appId))
         .setEnabledProtocols(tlsProtocols.getEnabledProtocols().toArray(new String[0]));
 
     _client = Dsl.asyncHttpClient(builder.build());

diff --git a/...-java-client/src/main/java/org/apache/pinot/client/JsonAsyncHttpPinotClientTransport.java b/...-java-client/src/main/java/org/apache/pinot/client/JsonAsyncHttpPinotClientTransport.java
@@ -70,7 +70,8 @@ public JsonAsyncHttpPinotClientTransport() {
   }
 
   public JsonAsyncHttpPinotClientTransport(Map<String, String> headers, String scheme, String extraOptionString,
-    @Nullable SSLContext sslContext, ConnectionTimeouts connectionTimeouts, TlsProtocols tlsProtocols) {
+      @Nullable SSLContext sslContext, ConnectionTimeouts connectionTimeouts, TlsProtocols tlsProtocols,
+      @Nullable String appId) {
     _brokerReadTimeout = connectionTimeouts.getReadTimeoutMs();
     _headers = headers;
     _scheme = scheme;
@@ -82,15 +83,16 @@ public JsonAsyncHttpPinotClientTransport(Map<String, String> headers, String sch
     }
 
     builder.setReadTimeout(connectionTimeouts.getReadTimeoutMs())
-            .setConnectTimeout(connectionTimeouts.getConnectTimeoutMs())
-            .setHandshakeTimeout(connectionTimeouts.getHandshakeTimeoutMs())
-            .setUserAgent(ConnectionUtils.getUserAgentVersionFromClassPath("ua"))
-            .setEnabledProtocols(tlsProtocols.getEnabledProtocols().toArray(new String[0]));
+        .setConnectTimeout(connectionTimeouts.getConnectTimeoutMs())
+        .setHandshakeTimeout(connectionTimeouts.getHandshakeTimeoutMs())
+        .setUserAgent(ConnectionUtils.getUserAgentVersionFromClassPath("ua", appId))
+        .setEnabledProtocols(tlsProtocols.getEnabledProtocols().toArray(new String[0]));
     _httpClient = Dsl.asyncHttpClient(builder.build());
   }
 
   public JsonAsyncHttpPinotClientTransport(Map<String, String> headers, String scheme, String extraOptionStr,
-    @Nullable SslContext sslContext, ConnectionTimeouts connectionTimeouts, TlsProtocols tlsProtocols) {
+      @Nullable SslContext sslContext, ConnectionTimeouts connectionTimeouts, TlsProtocols tlsProtocols,
+      @Nullable String appId) {
     _brokerReadTimeout = connectionTimeouts.getReadTimeoutMs();
     _headers = headers;
     _scheme = scheme;
@@ -102,16 +104,16 @@ public JsonAsyncHttpPinotClientTransport(Map<String, String> headers, String sch
     }
 
     builder.setReadTimeout(connectionTimeouts.getReadTimeoutMs())
-            .setConnectTimeout(connectionTimeouts.getConnectTimeoutMs())
-            .setHandshakeTimeout(connectionTimeouts.getHandshakeTimeoutMs())
-            .setUserAgent(ConnectionUtils.getUserAgentVersionFromClassPath("ua"))
-            .setEnabledProtocols(tlsProtocols.getEnabledProtocols().toArray(new String[0]));
+        .setConnectTimeout(connectionTimeouts.getConnectTimeoutMs())
+        .setHandshakeTimeout(connectionTimeouts.getHandshakeTimeoutMs())
+        .setUserAgent(ConnectionUtils.getUserAgentVersionFromClassPath("ua", appId))
+        .setEnabledProtocols(tlsProtocols.getEnabledProtocols().toArray(new String[0]));
     _httpClient = Dsl.asyncHttpClient(builder.build());
   }
 
   @Override
   public BrokerResponse executeQuery(String brokerAddress, String query)
-    throws PinotClientException {
+      throws PinotClientException {
     try {
       return executeQueryAsync(brokerAddress, query).get(_brokerReadTimeout, TimeUnit.MILLISECONDS);
     } catch (Exception e) {

diff --git a/...lient/src/main/java/org/apache/pinot/client/JsonAsyncHttpPinotClientTransportFactory.java b/...lient/src/main/java/org/apache/pinot/client/JsonAsyncHttpPinotClientTransportFactory.java
@@ -43,16 +43,18 @@ public class JsonAsyncHttpPinotClientTransportFactory implements PinotClientTran
   private int _readTimeoutMs = Integer.parseInt(DEFAULT_BROKER_READ_TIMEOUT_MS);
   private int _connectTimeoutMs = Integer.parseInt(DEFAULT_BROKER_READ_TIMEOUT_MS);
   private int _handshakeTimeoutMs = Integer.parseInt(DEFAULT_BROKER_HANDSHAKE_TIMEOUT_MS);
+  private String _appId = null;
   private String _extraOptionString;
 
   @Override
   public PinotClientTransport buildTransport() {
-    ConnectionTimeouts connectionTimeouts = ConnectionTimeouts.create(_readTimeoutMs, _connectTimeoutMs,
-            _handshakeTimeoutMs);
+    ConnectionTimeouts connectionTimeouts =
+        ConnectionTimeouts.create(_readTimeoutMs, _connectTimeoutMs, _handshakeTimeoutMs);
     TlsProtocols tlsProtocols = TlsProtocols.defaultProtocols(_tlsV10Enabled);
     return new JsonAsyncHttpPinotClientTransport(_headers, _scheme, _extraOptionString, _sslContext, connectionTimeouts,
-        tlsProtocols);
+        tlsProtocols, _appId);
   }
+
   public Map<String, String> getHeaders() {
     return _headers;
   }
@@ -90,16 +92,15 @@ public JsonAsyncHttpPinotClientTransportFactory withConnectionProperties(Propert
       _sslContext = ConnectionUtils.getSSLContextFromProperties(properties);
     }
 
-    _readTimeoutMs = Integer.parseInt(properties.getProperty("brokerReadTimeoutMs",
-            DEFAULT_BROKER_READ_TIMEOUT_MS));
-    _connectTimeoutMs = Integer.parseInt(properties.getProperty("brokerConnectTimeoutMs",
-            DEFAULT_BROKER_CONNECT_TIMEOUT_MS));
-    _handshakeTimeoutMs = Integer.parseInt(properties.getProperty("brokerHandshakeTimeoutMs",
-            DEFAULT_BROKER_HANDSHAKE_TIMEOUT_MS));
-    _tlsV10Enabled = Boolean.parseBoolean(properties.getProperty("brokerTlsV10Enabled",
-            DEFAULT_BROKER_TLS_V10_ENABLED))
-            || Boolean.parseBoolean(System.getProperties().getProperty("broker.tlsV10Enabled",
-            DEFAULT_BROKER_TLS_V10_ENABLED));
+    _readTimeoutMs = Integer.parseInt(properties.getProperty("brokerReadTimeoutMs", DEFAULT_BROKER_READ_TIMEOUT_MS));
+    _connectTimeoutMs =
+        Integer.parseInt(properties.getProperty("brokerConnectTimeoutMs", DEFAULT_BROKER_CONNECT_TIMEOUT_MS));
+    _handshakeTimeoutMs =
+        Integer.parseInt(properties.getProperty("brokerHandshakeTimeoutMs", DEFAULT_BROKER_HANDSHAKE_TIMEOUT_MS));
+    _appId = properties.getProperty("appId");
+    _tlsV10Enabled = Boolean.parseBoolean(properties.getProperty("brokerTlsV10Enabled", DEFAULT_BROKER_TLS_V10_ENABLED))
+        || Boolean.parseBoolean(
+        System.getProperties().getProperty("broker.tlsV10Enabled", DEFAULT_BROKER_TLS_V10_ENABLED));
 
     _extraOptionString = properties.getProperty("queryOptions", "");
     return this;

diff --git a/...lients/pinot-java-client/src/main/java/org/apache/pinot/client/utils/ConnectionUtils.java b/...lients/pinot-java-client/src/main/java/org/apache/pinot/client/utils/ConnectionUtils.java
@@ -22,8 +22,10 @@
 import java.util.Map;
 import java.util.Properties;
 import java.util.stream.Collectors;
+import javax.annotation.Nullable;
 import javax.net.ssl.SSLContext;
 import org.apache.commons.configuration.MapConfiguration;
+import org.apache.commons.lang3.StringUtils;
 import org.apache.commons.lang3.tuple.Pair;
 import org.apache.pinot.common.config.TlsConfig;
 import org.apache.pinot.common.utils.TlsUtils;
@@ -37,6 +39,7 @@ public class ConnectionUtils {
 
   public static final String INFO_HEADERS = "headers";
   public static final String PINOT_JAVA_TLS_PREFIX = "pinot.java_client.tls";
+  public static final int APP_ID_MAX_CHARS = 256;
 
   private ConnectionUtils() {
   }
@@ -56,14 +59,18 @@ public static SSLContext getSSLContextFromProperties(Properties properties) {
   }
 
 
-  public static String getUserAgentVersionFromClassPath(String userAgentKey) {
+  public static String getUserAgentVersionFromClassPath(String userAgentKey, @Nullable String appId) {
     Properties userAgentProperties = new Properties();
     try {
       userAgentProperties.load(ConnectionUtils.class.getClassLoader()
           .getResourceAsStream("version.properties"));
     } catch (IOException e) {
       LOGGER.warn("Unable to set user agent version");
     }
-    return userAgentProperties.getProperty(userAgentKey, "pinot-java");
+    String userAgentFromProperties = userAgentProperties.getProperty(userAgentKey, "unknown");
+    if (StringUtils.isNotEmpty(appId)) {
+      return appId.substring(0, Math.min(APP_ID_MAX_CHARS, appId.length())) + "-" + userAgentFromProperties;
+    }
+    return userAgentFromProperties;
   }
 }
diff --git a/...dbc-client/src/main/java/org/apache/pinot/client/controller/PinotControllerTransport.java b/...dbc-client/src/main/java/org/apache/pinot/client/controller/PinotControllerTransport.java
@@ -27,6 +27,7 @@
 import java.util.concurrent.Future;
 import javax.annotation.Nullable;
 import javax.net.ssl.SSLContext;
+import org.apache.commons.lang3.StringUtils;
 import org.apache.pinot.client.ConnectionTimeouts;
 import org.apache.pinot.client.PinotClientException;
 import org.apache.pinot.client.TlsProtocols;
@@ -50,9 +51,8 @@ public class PinotControllerTransport {
   private final String _scheme;
   private final AsyncHttpClient _httpClient;
 
-
-  public PinotControllerTransport(Map<String, String> headers, String scheme,
-    @Nullable SSLContext sslContext, ConnectionTimeouts connectionTimeouts, TlsProtocols tlsProtocols) {
+  public PinotControllerTransport(Map<String, String> headers, String scheme, @Nullable SSLContext sslContext,
+      ConnectionTimeouts connectionTimeouts, TlsProtocols tlsProtocols, @Nullable String appId) {
     _headers = headers;
     _scheme = scheme;
 
@@ -62,23 +62,29 @@ public PinotControllerTransport(Map<String, String> headers, String scheme,
     }
 
     builder.setReadTimeout(connectionTimeouts.getReadTimeoutMs())
-            .setConnectTimeout(connectionTimeouts.getConnectTimeoutMs())
-            .setHandshakeTimeout(connectionTimeouts.getHandshakeTimeoutMs())
-            .setUserAgent(getUserAgentVersionFromClassPath())
-            .setEnabledProtocols(tlsProtocols.getEnabledProtocols().toArray(new String[0]));
+        .setConnectTimeout(connectionTimeouts.getConnectTimeoutMs())
+        .setHandshakeTimeout(connectionTimeouts.getHandshakeTimeoutMs())
+        .setUserAgent(getUserAgentVersionFromClassPath(appId))
+        .setEnabledProtocols(tlsProtocols.getEnabledProtocols().toArray(new String[0]));
 
     _httpClient = Dsl.asyncHttpClient(builder.build());
   }
 
-  private String getUserAgentVersionFromClassPath() {
+  private String getUserAgentVersionFromClassPath(@Nullable String appId) {
     Properties userAgentProperties = new Properties();
     try {
-      userAgentProperties.load(PinotControllerTransport.class.getClassLoader()
-              .getResourceAsStream("version.properties"));
+      userAgentProperties.load(
+          PinotControllerTransport.class.getClassLoader().getResourceAsStream("version.properties"));
     } catch (IOException e) {
       LOGGER.warn("Unable to set user agent version");
     }
-    return userAgentProperties.getProperty("ua", "unknown");
+    String userAgentFromProperties = userAgentProperties.getProperty("ua", "unknown");
+    if (StringUtils.isNotEmpty(appId)) {
+      return
+          appId.substring(0, Math.min(org.apache.pinot.client.utils.ConnectionUtils.APP_ID_MAX_CHARS, appId.length()))
+              + "-" + userAgentFromProperties;
+    }
+    return userAgentFromProperties;
   }
 
   public TableResponse getAllTables(String controllerAddress) {

diff --git a/...ent/src/main/java/org/apache/pinot/client/controller/PinotControllerTransportFactory.java b/...ent/src/main/java/org/apache/pinot/client/controller/PinotControllerTransportFactory.java
@@ -41,12 +41,13 @@ public class PinotControllerTransportFactory {
   private int _readTimeoutMs = Integer.parseInt(DEFAULT_CONTROLLER_READ_TIMEOUT_MS);
   private int _connectTimeoutMs = Integer.parseInt(DEFAULT_CONTROLLER_CONNECT_TIMEOUT_MS);
   private int _handshakeTimeoutMs = Integer.parseInt(DEFAULT_CONTROLLER_HANDSHAKE_TIMEOUT_MS);
+  private String _appId = null;
 
   public PinotControllerTransport buildTransport() {
-    ConnectionTimeouts connectionTimeouts = ConnectionTimeouts.create(_readTimeoutMs, _connectTimeoutMs,
-            _handshakeTimeoutMs);
+    ConnectionTimeouts connectionTimeouts =
+        ConnectionTimeouts.create(_readTimeoutMs, _connectTimeoutMs, _handshakeTimeoutMs);
     TlsProtocols tlsProtocols = TlsProtocols.defaultProtocols(_tlsV10Enabled);
-    return new PinotControllerTransport(_headers, _scheme, _sslContext, connectionTimeouts, tlsProtocols);
+    return new PinotControllerTransport(_headers, _scheme, _sslContext, connectionTimeouts, tlsProtocols, _appId);
   }
 
   public Map<String, String> getHeaders() {
@@ -74,16 +75,17 @@ public void setSslContext(SSLContext sslContext) {
   }
 
   public PinotControllerTransportFactory withConnectionProperties(Properties properties) {
-    _readTimeoutMs = Integer.parseInt(properties.getProperty("controllerReadTimeoutMs",
-            DEFAULT_CONTROLLER_READ_TIMEOUT_MS));
-    _connectTimeoutMs = Integer.parseInt(properties.getProperty("controllerConnectTimeoutMs",
-            DEFAULT_CONTROLLER_CONNECT_TIMEOUT_MS));
-    _handshakeTimeoutMs = Integer.parseInt(properties.getProperty("controllerHandshakeTimeoutMs",
-            DEFAULT_CONTROLLER_HANDSHAKE_TIMEOUT_MS));
-    _tlsV10Enabled = Boolean.parseBoolean(properties.getProperty("controllerTlsV10Enabled",
-            DEFAULT_CONTROLLER_TLS_V10_ENABLED))
-            || Boolean.parseBoolean(System.getProperties().getProperty("controller.tlsV10Enabled",
-            DEFAULT_CONTROLLER_TLS_V10_ENABLED));
+    _readTimeoutMs =
+        Integer.parseInt(properties.getProperty("controllerReadTimeoutMs", DEFAULT_CONTROLLER_READ_TIMEOUT_MS));
+    _connectTimeoutMs =
+        Integer.parseInt(properties.getProperty("controllerConnectTimeoutMs", DEFAULT_CONTROLLER_CONNECT_TIMEOUT_MS));
+    _handshakeTimeoutMs = Integer.parseInt(
+        properties.getProperty("controllerHandshakeTimeoutMs", DEFAULT_CONTROLLER_HANDSHAKE_TIMEOUT_MS));
+    _appId = properties.getProperty("appId");
+    _tlsV10Enabled =
+        Boolean.parseBoolean(properties.getProperty("controllerTlsV10Enabled", DEFAULT_CONTROLLER_TLS_V10_ENABLED))
+            || Boolean.parseBoolean(
+            System.getProperties().getProperty("controller.tlsV10Enabled", DEFAULT_CONTROLLER_TLS_V10_ENABLED));
     return this;
   }
 }
diff --git a/...inot-jdbc-client/src/test/java/org/apache/pinot/client/DummyPinotControllerTransport.java b/...inot-jdbc-client/src/test/java/org/apache/pinot/client/DummyPinotControllerTransport.java
@@ -29,10 +29,10 @@
 
 public class DummyPinotControllerTransport extends PinotControllerTransport {
 
-  public DummyPinotControllerTransport(Map<String, String> headers, String scheme, @Nullable SSLContext sslContext) {
-    super(headers, scheme, sslContext,
-            ConnectionTimeouts.create(1000, 1000, 1000),
-            TlsProtocols.defaultProtocols(true));
+  public DummyPinotControllerTransport(Map<String, String> headers, String scheme, @Nullable SSLContext sslContext,
+      @Nullable String appId) {
+    super(headers, scheme, sslContext, ConnectionTimeouts.create(1000, 1000, 1000), TlsProtocols.defaultProtocols(true),
+        appId);
   }
 
   @Override
@@ -47,6 +47,10 @@ public ControllerTenantBrokerResponse getBrokersFromController(String controller
   }
 
   public static DummyPinotControllerTransport create() {
-    return new DummyPinotControllerTransport(null, null, null);
+    return create("dummy");
+  }
+
+  public static DummyPinotControllerTransport create(String appId) {
+    return new DummyPinotControllerTransport(null, null, null, appId);
   }
 }
diff --git a/...-clients/pinot-jdbc-client/src/test/java/org/apache/pinot/client/PinotConnectionTest.java b/...-clients/pinot-jdbc-client/src/test/java/org/apache/pinot/client/PinotConnectionTest.java
@@ -35,4 +35,32 @@ public void createStatementTest()
     Statement statement = pinotConnection.createStatement();
     Assert.assertNotNull(statement);
   }
+
+  @Test
+  public void setUserAgentTest()
+      throws Exception {
+    StringBuilder appId = new StringBuilder("appId-");
+    for (int i = 0; i < 256; i++) {
+       appId.append(i);
+    }
+    DummyPinotControllerTransport userAgentPinotControllerTransport = DummyPinotControllerTransport
+        .create(appId.toString());
+
+    PinotConnection pinotConnection =
+        new PinotConnection("dummy", _dummyPinotClientTransport, "dummy", userAgentPinotControllerTransport);
+    Statement statement = pinotConnection.createStatement();
+    Assert.assertNotNull(statement);
+  }
+
+  @Test
+  public void unsetUserAgentTest()
+      throws Exception {
+    DummyPinotControllerTransport userAgentPinotControllerTransport = DummyPinotControllerTransport
+        .create(null);
+
+    PinotConnection pinotConnection =
+        new PinotConnection("dummy", _dummyPinotClientTransport, "dummy", userAgentPinotControllerTransport);
+    Statement statement = pinotConnection.createStatement();
+    Assert.assertNotNull(statement);
+  }
 }