Add applicationNamePrefix property to JDBC driver

prestodb · Jul 2, 2018 · bbd2598 · bbd2598
1 parent e6ca22b
commit bbd2598
Show file tree

Hide file tree

Showing 5 changed files with 89 additions and 3 deletions.
diff --git a/presto-docs/src/main/sphinx/installation/jdbc.rst b/presto-docs/src/main/sphinx/installation/jdbc.rst
@@ -76,6 +76,10 @@ Name                              Description
 ``password``                      Password to use for LDAP authentication.
 ``socksProxy``                    SOCKS proxy host and port. Example: ``localhost:1080``
 ``httpProxy``                     HTTP proxy host and port. Example: ``localhost:8888``
+``applicationNamePrefix``         Prefix to append to any specified ``ApplicationName`` client info
+                                  property, which is used to set the source name for the Presto query.
+                                  If neither this property nor ``ApplicationName`` are set, the source
+                                  for the query will be ``presto-jdbc``.
 ``accessToken``                   Access token for token based authentication.
 ``SSL``                           Use HTTPS for connections
 ``SSLKeyStorePath``               The location of the Java KeyStore file that contains the certificate

diff --git a/presto-jdbc/src/main/java/com/facebook/presto/jdbc/ConnectionProperties.java b/presto-jdbc/src/main/java/com/facebook/presto/jdbc/ConnectionProperties.java
@@ -35,6 +35,7 @@ final class ConnectionProperties
     public static final ConnectionProperty<String> PASSWORD = new Password();
     public static final ConnectionProperty<HostAndPort> SOCKS_PROXY = new SocksProxy();
     public static final ConnectionProperty<HostAndPort> HTTP_PROXY = new HttpProxy();
+    public static final ConnectionProperty<String> APPLICATION_NAME_PREFIX = new ApplicationNamePrefix();
     public static final ConnectionProperty<Boolean> SSL = new Ssl();
     public static final ConnectionProperty<String> SSL_KEY_STORE_PATH = new SslKeyStorePath();
     public static final ConnectionProperty<String> SSL_KEY_STORE_PASSWORD = new SslKeyStorePassword();
@@ -53,6 +54,7 @@ final class ConnectionProperties
             .add(PASSWORD)
             .add(SOCKS_PROXY)
             .add(HTTP_PROXY)
+            .add(APPLICATION_NAME_PREFIX)
             .add(SSL)
             .add(SSL_KEY_STORE_PATH)
             .add(SSL_KEY_STORE_PASSWORD)
@@ -139,6 +141,15 @@ public HttpProxy()
         }
     }
 
+    private static class ApplicationNamePrefix
+            extends AbstractConnectionProperty<String>
+    {
+        public ApplicationNamePrefix()
+        {
+            super("applicationNamePrefix", NOT_REQUIRED, ALLOWED, STRING_CONVERTER);
+        }
+    }
+
     private static class Ssl
             extends AbstractConnectionProperty<Boolean>
     {

diff --git a/presto-jdbc/src/main/java/com/facebook/presto/jdbc/PrestoConnection.java b/presto-jdbc/src/main/java/com/facebook/presto/jdbc/PrestoConnection.java
@@ -54,7 +54,6 @@
 import java.util.concurrent.atomic.AtomicLong;
 import java.util.concurrent.atomic.AtomicReference;
 
-import static com.google.common.base.MoreObjects.firstNonNull;
 import static com.google.common.base.Preconditions.checkArgument;
 import static com.google.common.base.Strings.nullToEmpty;
 import static com.google.common.collect.Maps.fromProperties;
@@ -83,6 +82,7 @@ public class PrestoConnection
     private final URI jdbcUri;
     private final URI httpUri;
     private final String user;
+    private final Optional<String> applicationNamePrefix;
     private final Map<String, String> clientInfo = new ConcurrentHashMap<>();
     private final Map<String, String> sessionProperties = new ConcurrentHashMap<>();
     private final Map<String, String> preparedStatements = new ConcurrentHashMap<>();
@@ -98,6 +98,7 @@ public class PrestoConnection
         this.schema.set(uri.getSchema());
         this.catalog.set(uri.getCatalog());
         this.user = uri.getUser();
+        this.applicationNamePrefix = uri.getApplicationNamePrefix();
 
         this.queryExecutor = requireNonNull(queryExecutor, "queryExecutor is null");
 
@@ -636,7 +637,18 @@ String getStartTransactionSql()
 
     StatementClient startQuery(String sql, Map<String, String> sessionPropertiesOverride)
     {
-        String source = firstNonNull(clientInfo.get("ApplicationName"), "presto-jdbc");
+        String source = "presto-jdbc";
+        String applicationName = clientInfo.get("ApplicationName");
+        if (applicationNamePrefix.isPresent()) {
+            source = applicationNamePrefix.get();
+            if (applicationName != null) {
+                source += applicationName;
+            }
+        }
+        else if (applicationName != null) {
+            source = applicationName;
+        }
+
         Optional<String> traceToken = Optional.ofNullable(clientInfo.get("TraceToken"));
         Iterable<String> clientTags = Splitter.on(',').trimResults().omitEmptyStrings()
                 .split(nullToEmpty(clientInfo.get("ClientTags")));

diff --git a/presto-jdbc/src/main/java/com/facebook/presto/jdbc/PrestoDriverUri.java b/presto-jdbc/src/main/java/com/facebook/presto/jdbc/PrestoDriverUri.java
@@ -39,6 +39,7 @@
 import static com.facebook.presto.client.OkHttpUtil.setupSsl;
 import static com.facebook.presto.client.OkHttpUtil.tokenAuth;
 import static com.facebook.presto.jdbc.ConnectionProperties.ACCESS_TOKEN;
+import static com.facebook.presto.jdbc.ConnectionProperties.APPLICATION_NAME_PREFIX;
 import static com.facebook.presto.jdbc.ConnectionProperties.HTTP_PROXY;
 import static com.facebook.presto.jdbc.ConnectionProperties.KERBEROS_CONFIG_PATH;
 import static com.facebook.presto.jdbc.ConnectionProperties.KERBEROS_CREDENTIAL_CACHE_PATH;
@@ -125,6 +126,12 @@ public String getUser()
         return USER.getRequiredValue(properties);
     }
 
+    public Optional<String> getApplicationNamePrefix()
+            throws SQLException
+    {
+        return APPLICATION_NAME_PREFIX.getValue(properties);
+    }
+
     public Properties getProperties()
     {
         return properties;

diff --git a/presto-jdbc/src/test/java/com/facebook/presto/jdbc/TestJdbcConnection.java b/presto-jdbc/src/test/java/com/facebook/presto/jdbc/TestJdbcConnection.java
@@ -25,6 +25,7 @@
 
 import java.sql.Connection;
 import java.sql.DriverManager;
+import java.sql.PreparedStatement;
 import java.sql.ResultSet;
 import java.sql.SQLException;
 import java.sql.Statement;
@@ -33,6 +34,8 @@
 import static com.facebook.presto.jdbc.TestPrestoDriver.closeQuietly;
 import static java.lang.String.format;
 import static org.assertj.core.api.Assertions.assertThat;
+import static org.testng.Assert.assertFalse;
+import static org.testng.Assert.assertTrue;
 
 public class TestJdbcConnection
 {
@@ -167,10 +170,39 @@ public void testSession()
         }
     }
 
+    @Test
+    public void testApplicationName()
+            throws SQLException
+    {
+        try (Connection connection = createConnection()) {
+            assertConnectionSource(connection, "presto-jdbc");
+        }
+
+        try (Connection connection = createConnection()) {
+            connection.setClientInfo("ApplicationName", "testing");
+            assertConnectionSource(connection, "testing");
+        }
+
+        try (Connection connection = createConnection("applicationNamePrefix=fruit:")) {
+            assertConnectionSource(connection, "fruit:");
+        }
+
+        try (Connection connection = createConnection("applicationNamePrefix=fruit:")) {
+            connection.setClientInfo("ApplicationName", "testing");
+            assertConnectionSource(connection, "fruit:testing");
+        }
+    }
+
     private Connection createConnection()
             throws SQLException
     {
-        String url = format("jdbc:presto://%s/hive/default", server.getAddress());
+        return createConnection("");
+    }
+
+    private Connection createConnection(String extra)
+            throws SQLException
+    {
+        String url = format("jdbc:presto://%s/hive/default?%s", server.getAddress(), extra);
         return DriverManager.getConnection(url, "test", null);
     }
 
@@ -202,4 +234,24 @@ private static Set<String> listSession(Connection connection)
         }
         return set.build();
     }
+
+    private static void assertConnectionSource(Connection connection, String expectedSource)
+            throws SQLException
+    {
+        String queryId;
+        try (Statement statement = connection.createStatement();
+                ResultSet rs = statement.executeQuery("SELECT 123")) {
+            queryId = rs.unwrap(PrestoResultSet.class).getQueryId();
+        }
+
+        try (PreparedStatement statement = connection.prepareStatement(
+                "SELECT source FROM system.runtime.queries WHERE query_id = ?")) {
+            statement.setString(1, queryId);
+            try (ResultSet rs = statement.executeQuery()) {
+                assertTrue(rs.next());
+                assertThat(rs.getString("source")).isEqualTo(expectedSource);
+                assertFalse(rs.next());
+            }
+        }
+    }
 }