Add support for passing clientTags in JDBC connection URI

prestodb · May 5, 2022 · 0c35954 · 0c35954
1 parent 67f09ca
commit 0c35954
Show file tree

Hide file tree

Showing 6 changed files with 77 additions and 1 deletion.
diff --git a/presto-jdbc/src/main/java/com/facebook/presto/jdbc/AbstractConnectionProperty.java b/presto-jdbc/src/main/java/com/facebook/presto/jdbc/AbstractConnectionProperty.java
@@ -25,6 +25,8 @@
 import java.util.Optional;
 import java.util.Properties;
 import java.util.function.Predicate;
+import java.util.regex.Pattern;
+import java.util.stream.Collectors;
 
 import static com.google.common.base.Preconditions.checkArgument;
 import static com.google.common.collect.ImmutableList.toImmutableList;
@@ -192,10 +194,33 @@ public List<String> parseKeyValuePair(String keyValue)
         }
     }
 
+    protected static final class ListValidateConvertor
+            implements Converter<String>
+    {
+        public static final Converter<String> LIST_VALIDATE_CONVERTOR = new ListValidateConvertor();
+
+        private ListValidateConvertor() {}
+
+        @Override
+        public String convert(String value)
+        {
+            return Splitter.on(',').trimResults().splitToList(value).stream().map(this::validatePattern).collect(Collectors.joining(","));
+        }
+
+        private String validatePattern(String value)
+        {
+            Pattern alphaNumericFilter = Pattern.compile("^[a-zA-Z0-9]+$");
+            boolean isAlphaNumeric = alphaNumericFilter.matcher(value).matches();
+            checkArgument(isAlphaNumeric, "Input client tag should contain only alphanumeric characters: %s", value);
+            return value;
+        }
+    }
+
     protected static final class ClassListConverter
             implements Converter<List<QueryInterceptor>>
     {
         public static final ClassListConverter CLASS_LIST_CONVERTER = new ClassListConverter();
+
         private ClassListConverter() {}
 
         @Override
@@ -221,6 +246,7 @@ protected static final class HttpProtocolConverter
             implements Converter<List<Protocol>>
     {
         public static final HttpProtocolConverter HTTP_PROTOCOL_CONVERTER = new HttpProtocolConverter();
+
         private HttpProtocolConverter() {}
 
         @Override

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
@@ -28,6 +28,7 @@
 
 import static com.facebook.presto.jdbc.AbstractConnectionProperty.ClassListConverter.CLASS_LIST_CONVERTER;
 import static com.facebook.presto.jdbc.AbstractConnectionProperty.HttpProtocolConverter.HTTP_PROTOCOL_CONVERTER;
+import static com.facebook.presto.jdbc.AbstractConnectionProperty.ListValidateConvertor.LIST_VALIDATE_CONVERTOR;
 import static com.facebook.presto.jdbc.AbstractConnectionProperty.StringMapConverter.STRING_MAP_CONVERTER;
 import static com.facebook.presto.jdbc.AbstractConnectionProperty.checkedPredicate;
 import static java.util.Collections.unmodifiableMap;
@@ -57,6 +58,7 @@ final class ConnectionProperties
     public static final ConnectionProperty<String> TIMEZONE_ID = new TimeZoneId();
     public static final ConnectionProperty<Map<String, String>> EXTRA_CREDENTIALS = new ExtraCredentials();
     public static final ConnectionProperty<Map<String, String>> CUSTOM_HEADERS = new CustomHeaders();
+    public static final ConnectionProperty<String> CLIENT_TAGS = new ClientTags();
     public static final ConnectionProperty<Map<String, String>> SESSION_PROPERTIES = new SessionProperties();
     public static final ConnectionProperty<List<Protocol>> HTTP_PROTOCOLS = new HttpProtocols();
     public static final ConnectionProperty<List<QueryInterceptor>> QUERY_INTERCEPTORS = new QueryInterceptors();
@@ -83,6 +85,7 @@ final class ConnectionProperties
             .add(TIMEZONE_ID)
             .add(EXTRA_CREDENTIALS)
             .add(CUSTOM_HEADERS)
+            .add(CLIENT_TAGS)
             .add(SESSION_PROPERTIES)
             .add(HTTP_PROTOCOLS)
             .add(QUERY_INTERCEPTORS)
@@ -169,6 +172,15 @@ public ApplicationNamePrefix()
         }
     }
 
+    private static class ClientTags
+            extends AbstractConnectionProperty<String>
+    {
+        public ClientTags()
+        {
+            super("clientTags", NOT_REQUIRED, ALLOWED, LIST_VALIDATE_CONVERTOR);
+        }
+    }
+
     private static class DisableCompression
             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
@@ -116,6 +116,7 @@ public class PrestoConnection
         this.sessionProperties = new ConcurrentHashMap<>(uri.getSessionProperties());
         this.connectionProperties = uri.getProperties();
         this.queryExecutor = requireNonNull(queryExecutor, "queryExecutor is null");
+        uri.getClientTags().ifPresent(tags -> clientInfo.put("ClientTags", tags));
 
         timeZoneId.set(uri.getTimeZoneId());
         locale.set(Locale.getDefault());

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
@@ -48,6 +48,7 @@
 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.CLIENT_TAGS;
 import static com.facebook.presto.jdbc.ConnectionProperties.CUSTOM_HEADERS;
 import static com.facebook.presto.jdbc.ConnectionProperties.DISABLE_COMPRESSION;
 import static com.facebook.presto.jdbc.ConnectionProperties.EXTRA_CREDENTIALS;
@@ -179,6 +180,12 @@ public Map<String, String> getCustomHeaders()
         return CUSTOM_HEADERS.getValue(properties).orElse(ImmutableMap.of());
     }
 
+    public Optional<String> getClientTags()
+            throws SQLException
+    {
+        return CLIENT_TAGS.getValue(properties);
+    }
+
     public Map<String, String> getSessionProperties()
             throws SQLException
     {

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
@@ -283,6 +283,15 @@ public void testCustomHeadersWithSpecialCharacters(String testHeaderValue)
         assertEquals(prestoConnection.getCustomHeaders(), customHeadersMap);
     }
 
+    @Test
+    public void testClientTags()
+            throws SQLException
+    {
+        try (Connection connection = createConnection("clientTags=c2,c3")) {
+            assertEquals(connection.getClientInfo("ClientTags"), "c2,c3");
+        }
+    }
+
     @Test
     public void testQueryInterceptors()
             throws SQLException

diff --git a/presto-jdbc/src/test/java/com/facebook/presto/jdbc/TestPrestoDriverUri.java b/presto-jdbc/src/test/java/com/facebook/presto/jdbc/TestPrestoDriverUri.java
@@ -22,6 +22,7 @@
 import java.sql.SQLException;
 import java.util.Properties;
 
+import static com.facebook.presto.jdbc.ConnectionProperties.CLIENT_TAGS;
 import static com.facebook.presto.jdbc.ConnectionProperties.CUSTOM_HEADERS;
 import static com.facebook.presto.jdbc.ConnectionProperties.DISABLE_COMPRESSION;
 import static com.facebook.presto.jdbc.ConnectionProperties.EXTRA_CREDENTIALS;
@@ -268,6 +269,26 @@ public void testUriWithCustomHeaders()
         assertEquals(properties.getProperty(CUSTOM_HEADERS.getKey()), customHeaders);
     }
 
+    @Test
+    public void testUriWithClientTags()
+            throws SQLException
+    {
+        String clientTags = "c1,c2";
+        PrestoDriverUri parameters = createDriverUri("presto://localhost:8080?clientTags=" + clientTags);
+        Properties properties = parameters.getProperties();
+        assertEquals(properties.getProperty(CLIENT_TAGS.getKey()), clientTags);
+    }
+
+    @Test(expectedExceptions = SQLException.class)
+    public void assertNonAlphanumericClientTags()
+            throws SQLException
+    {
+        String clientTags = "d1,@d2,d3";
+        PrestoDriverUri parameters = createDriverUri("presto://localhost:8080?clientTags=" + clientTags);
+        Properties properties = parameters.getProperties();
+        assertEquals(properties.getProperty(CLIENT_TAGS.getKey()), clientTags);
+    }
+
     @Test
     public void testUriWithSessionProperties()
             throws SQLException
@@ -299,7 +320,7 @@ public void testUriWithQueryInterceptors()
     }
 
     public static class TestForUriQueryInterceptor
-                implements QueryInterceptor
+            implements QueryInterceptor
     {}
 
     private static void assertUriPortScheme(PrestoDriverUri parameters, int port, String scheme)