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 index 29ad24d6fef8a..8048e57835116 100644 --- 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 parseKeyValuePair(String keyValue) } } + protected static final class ListValidateConvertor + implements Converter + { + public static final Converter 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> { public static final ClassListConverter CLASS_LIST_CONVERTER = new ClassListConverter(); + private ClassListConverter() {} @Override @@ -221,6 +246,7 @@ protected static final class HttpProtocolConverter implements Converter> { 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 index 07717eac9c165..87d7d61f898f0 100644 --- 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 TIMEZONE_ID = new TimeZoneId(); public static final ConnectionProperty> EXTRA_CREDENTIALS = new ExtraCredentials(); public static final ConnectionProperty> CUSTOM_HEADERS = new CustomHeaders(); + public static final ConnectionProperty CLIENT_TAGS = new ClientTags(); public static final ConnectionProperty> SESSION_PROPERTIES = new SessionProperties(); public static final ConnectionProperty> HTTP_PROTOCOLS = new HttpProtocols(); public static final ConnectionProperty> 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 + { + public ClientTags() + { + super("clientTags", NOT_REQUIRED, ALLOWED, LIST_VALIDATE_CONVERTOR); + } + } + private static class DisableCompression extends AbstractConnectionProperty { 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 index d89ef658c5de8..e344e207aa879 100644 --- 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 index 5f34ccf50b90f..cf04159dfdac6 100644 --- 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 getCustomHeaders() return CUSTOM_HEADERS.getValue(properties).orElse(ImmutableMap.of()); } + public Optional getClientTags() + throws SQLException + { + return CLIENT_TAGS.getValue(properties); + } + public Map 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 index 4e611c4bb8c9e..cd03b0554d288 100644 --- 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 index 04ae27f3bb90e..5a18464b21a54 100644 --- 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)