Added option channelInitializer

jdbc/src/main/java/tech/ydb/jdbc/settings/YdbConnectionProperties.java

@@ -3,6 +3,7 @@
 import java.lang.reflect.InvocationTargetException;
 import java.sql.SQLException;
 import java.util.Properties;
+import java.util.function.Consumer;
 import java.util.function.Supplier;
 import java.util.logging.Level;
 import java.util.logging.Logger;
@@ -54,6 +55,10 @@ public class YdbConnectionProperties {
     static final YdbProperty<String> METADATA_URL = YdbProperty.content("metadataURL",
             "Custom URL for the metadata service authentication");
 
+    static final YdbProperty<Object> CHANNEL_INITIALIZER = YdbProperty.object("channelInitializer",
+            "Custom GRPC channel initilizer, use object instance or class full name impementing"
+                    + " Consumer<ManagedChannelBuilder>");
+
     static final YdbProperty<Object> TOKEN_PROVIDER = YdbProperty.object("tokenProvider",
             "Custom token provider, use object instance or class full name impementing Supplier<String>");
 
@@ -75,6 +80,7 @@ public class YdbConnectionProperties {
     private final YdbValue<String> iamEndpoint;
     private final YdbValue<String> metadataUrl;
     private final YdbValue<Object> tokenProvider;
+    private final YdbValue<Object> channelInitializer;
     private final YdbValue<String> grpcCompression;
 
     public YdbConnectionProperties(String username, String password, Properties props) throws SQLException {
@@ -92,6 +98,7 @@ public YdbConnectionProperties(String username, String password, Properties prop
         this.iamEndpoint = IAM_ENDPOINT.readValue(props);
         this.metadataUrl = METADATA_URL.readValue(props);
         this.tokenProvider = TOKEN_PROVIDER.readValue(props);
+        this.channelInitializer = CHANNEL_INITIALIZER.readValue(props);
         this.grpcCompression = GRPC_COMPRESSION.readValue(props);
     }
 
@@ -203,36 +210,12 @@ public GrpcTransportBuilder applyToGrpcTransport(GrpcTransportBuilder builder) t
             }
 
             Object provider = tokenProvider.getValue();
-            if (provider instanceof Supplier) {
-                Supplier<?> prov = (Supplier<?>) provider;
-                builder = builder.withAuthProvider((rpc) -> () -> prov.get().toString());
-            } else if (provider instanceof AuthProvider) {
-                AuthProvider prov = (AuthProvider) provider;
-                builder = builder.withAuthProvider(prov);
-            } else if (provider instanceof String) {
-                String className = (String) provider;
-                if (!FQCN.matcher(className).matches()) {
-                    throw new SQLException("tokenProvider must be full class name or instance of Supplier<String>");
-                }
+            builder = applyTokenProvider(builder, provider);
+        }
 
-                try {
-                    Class<?> clazz = Class.forName(className);
-                    if (!Supplier.class.isAssignableFrom(clazz)) {
-                        throw new SQLException("tokenProvider " + className + " is not implement Supplier<String>");
-                    }
-                    Supplier<?> prov = clazz.asSubclass(Supplier.class)
-                            .getConstructor(new Class<?>[0])
-                            .newInstance(new Object[0]);
-                    builder = builder.withAuthProvider((rpc) -> () -> prov.get().toString());
-                } catch (ClassNotFoundException ex) {
-                    throw new SQLException("tokenProvider " + className + " not found", ex);
-                } catch (NoSuchMethodException | SecurityException | InstantiationException | IllegalAccessException
-                        | IllegalArgumentException | InvocationTargetException ex) {
-                    throw new SQLException("Cannot construct tokenProvider " + className, ex);
-                }
-            } else if (provider != null) {
-                throw new SQLException("Cannot parse tokenProvider " + provider.getClass().getName());
-            }
+        if (channelInitializer.hasValue()) {
+            Object initializer = channelInitializer.getValue();
+            builder = applyChannelInitializer(builder, initializer);
         }
 
         if (grpcCompression.hasValue()) {
@@ -248,4 +231,89 @@ public GrpcTransportBuilder applyToGrpcTransport(GrpcTransportBuilder builder) t
 
         return builder;
     }
+
+    private GrpcTransportBuilder applyTokenProvider(GrpcTransportBuilder builder, Object provider) throws SQLException {
+        if (provider instanceof Supplier) {
+            Supplier<?> prov = (Supplier<?>) provider;
+            builder = builder.withAuthProvider((rpc) -> () -> prov.get().toString());
+        } else if (provider instanceof AuthProvider) {
+            AuthProvider prov = (AuthProvider) provider;
+            builder = builder.withAuthProvider(prov);
+        } else if (provider instanceof String) {
+            String className = (String) provider;
+            if (!FQCN.matcher(className).matches()) {
+                throw new SQLException("tokenProvider must be full class name or instance of Supplier<String>");
+            }
+
+            try {
+                Class<?> clazz = Class.forName(className);
+                if (!Supplier.class.isAssignableFrom(clazz)) {
+                    throw new SQLException("tokenProvider " + className + " is not implement Supplier<String>");
+                }
+                Supplier<?> prov = clazz.asSubclass(Supplier.class)
+                        .getConstructor(new Class<?>[0])
+                        .newInstance(new Object[0]);
+                builder = builder.withAuthProvider((rpc) -> () -> prov.get().toString());
+            } catch (ClassNotFoundException ex) {
+                throw new SQLException("tokenProvider " + className + " not found", ex);
+            } catch (NoSuchMethodException | SecurityException | InstantiationException | IllegalAccessException
+                    | IllegalArgumentException | InvocationTargetException ex) {
+                throw new SQLException("Cannot construct tokenProvider " + className, ex);
+            }
+        } else if (provider != null) {
+            throw new SQLException("Cannot parse tokenProvider " + provider.getClass().getName());
+        }
+        return builder;
+    }
+
+    private GrpcTransportBuilder applyChannelInitializer(GrpcTransportBuilder builder, Object initializer)
+            throws SQLException {
+        if (initializer instanceof Consumer) {
+            @SuppressWarnings("unchecked")
+            Consumer<Object> prov = (Consumer<Object>) initializer;
+            builder = builder.addChannelInitializer(prov);
+        } else if (initializer instanceof String) {
+            String className = (String) initializer;
+
+            if (FQCN.matcher(className.trim()).matches()) {
+                builder.addChannelInitializer(newInitializerInstance(className.trim()));
+            } else {
+                String[] classNames = className.split(",");
+                if (classNames.length < 2) {
+                    throw new SQLException("channelInitializer must be full class name or instance of "
+                            + "Consumer<ManagedChannelBuilder>");
+                }
+
+                for (String name: classNames) {
+                    if (!FQCN.matcher(name.trim()).matches()) {
+                        throw new SQLException("channelInitializer must be full class name or instance of "
+                                + "Consumer<ManagedChannelBuilder>");
+                    }
+                    builder.addChannelInitializer(newInitializerInstance(name.trim()));
+                }
+            }
+        } else if (initializer != null) {
+            throw new SQLException("Cannot parse channelInitializer " + initializer.getClass().getName());
+        }
+        return builder;
+    }
+
+    @SuppressWarnings("unchecked")
+    private Consumer<Object> newInitializerInstance(String className) throws SQLException {
+        try {
+            Class<?> clazz = Class.forName(className);
+            if (!Consumer.class.isAssignableFrom(clazz)) {
+                throw new SQLException("channelInitializer " + className + " is not implement "
+                        + "Consumer<ManagedChannelBuilder>");
+            }
+            return clazz.asSubclass(Consumer.class)
+                    .getConstructor(new Class<?>[0])
+                    .newInstance(new Object[0]);
+        } catch (ClassNotFoundException ex) {
+            throw new SQLException("channelInitializer " + className + " not found", ex);
+        } catch (NoSuchMethodException | SecurityException | InstantiationException | IllegalAccessException
+                | IllegalArgumentException | InvocationTargetException ex) {
+            throw new SQLException("Cannot construct channelInitializer " + className, ex);
+        }
+    }
 }

jdbc/src/test/java/tech/ydb/jdbc/settings/BadChannelInitilizer.java

@@ -0,0 +1,22 @@
+package tech.ydb.jdbc.settings;
+
+import java.util.function.Consumer;
+
+import io.grpc.ManagedChannelBuilder;
+
+/**
+ *
+ * @author Aleksandr Gorshenin
+ */
+public class BadChannelInitilizer implements Consumer<ManagedChannelBuilder<?>> {
+    private final String authority;
+
+    public BadChannelInitilizer(String authority) {
+        this.authority = authority;
+    }
+
+    @Override
+    public void accept(ManagedChannelBuilder<?> t) {
+        t.overrideAuthority(authority);
+    }
+}

jdbc/src/test/java/tech/ydb/jdbc/settings/CustomChannelInitilizer.java

@@ -0,0 +1,16 @@
+package tech.ydb.jdbc.settings;
+
+import java.util.function.Consumer;
+
+import io.grpc.ManagedChannelBuilder;
+
+/**
+ *
+ * @author Aleksandr Gorshenin
+ */
+public class CustomChannelInitilizer implements Consumer<ManagedChannelBuilder<?>> {
+    @Override
+    public void accept(ManagedChannelBuilder<?> t) {
+        t.usePlaintext();
+    }
+}

jdbc/src/test/java/tech/ydb/jdbc/settings/YdbConnectionPropertiesTest.java

@@ -2,8 +2,10 @@
 
 import java.sql.SQLException;
 import java.util.Properties;
+import java.util.function.Consumer;
 import java.util.function.Supplier;
 
+import io.grpc.ManagedChannelBuilder;
 import org.junit.jupiter.api.Assertions;
 import org.junit.jupiter.api.Test;
 
@@ -100,10 +102,10 @@ public void tokenProviderInvalidClassTest() throws SQLException {
     @Test
     public void tokenProviderNonSupplierClassTest() throws SQLException {
         Properties props = new Properties();
-        props.put("tokenProvider", "tech.ydb.jdbc.settings.NonSupplierTokenProvider");
+        props.put("tokenProvider", "tech.ydb.jdbc.settings.CustomChannelInitilizer");
         YdbConnectionProperties cp = new YdbConnectionProperties(null, null, props);
         ExceptionAssert.sqlException(
-                "tokenProvider tech.ydb.jdbc.settings.NonSupplierTokenProvider is not implement Supplier<String>",
+                "tokenProvider tech.ydb.jdbc.settings.CustomChannelInitilizer is not implement Supplier<String>",
                 () -> cp.applyToGrpcTransport(GrpcTransport.forConnectionString(YDB_URL))
         );
     }
@@ -119,4 +121,103 @@ public void tokenProviderWrongObjectTest() throws SQLException {
                 () -> cp.applyToGrpcTransport(GrpcTransport.forConnectionString(YDB_URL))
         );
     }
+
+    @Test
+    public void channelInitializerConsumerTest() throws SQLException {
+        Properties props = new Properties();
+        props.put("channelInitializer", (Consumer<ManagedChannelBuilder<?>>) (ManagedChannelBuilder<?> t) -> {
+            t.maxTraceEvents(0);
+        });
+        YdbConnectionProperties cp = new YdbConnectionProperties(null, null, props);
+        GrpcTransportBuilder builder = cp.applyToGrpcTransport(GrpcTransport.forConnectionString(YDB_URL));
+        Assertions.assertFalse(builder.getChannelInitializers().isEmpty());
+    }
+
+    @Test
+    public void channelInitializerClassTest() throws SQLException {
+        Properties props = new Properties();
+        props.put("channelInitializer", "tech.ydb.jdbc.settings.CustomChannelInitilizer");
+        YdbConnectionProperties cp = new YdbConnectionProperties(null, null, props);
+        GrpcTransportBuilder builder = cp.applyToGrpcTransport(GrpcTransport.forConnectionString(YDB_URL));
+        Assertions.assertEquals(1, builder.getChannelInitializers().size());
+    }
+
+    @Test
+    public void channelInitializerClassListTest() throws SQLException {
+        Properties props = new Properties();
+        props.put("channelInitializer", "tech.ydb.jdbc.settings.CustomChannelInitilizer,   "
+                + "tech.ydb.jdbc.settings.CustomChannelInitilizer");
+        YdbConnectionProperties cp = new YdbConnectionProperties(null, null, props);
+        GrpcTransportBuilder builder = cp.applyToGrpcTransport(GrpcTransport.forConnectionString(YDB_URL));
+        Assertions.assertEquals(2, builder.getChannelInitializers().size());
+    }
+
+    @Test
+    public void channelInitializerWrongClassNameTest() throws SQLException {
+        Properties props = new Properties();
+        props.put("channelInitializer", "1tech.ydb.jdbc.settings.StaticTokenProvider");
+        YdbConnectionProperties cp = new YdbConnectionProperties(null, null, props);
+        ExceptionAssert.sqlException(
+                "channelInitializer must be full class name or instance of Consumer<ManagedChannelBuilder>",
+                () -> cp.applyToGrpcTransport(GrpcTransport.forConnectionString(YDB_URL))
+        );
+    }
+
+    @Test
+    public void channelInitializerWrongClassNameListTest() throws SQLException {
+        Properties props = new Properties();
+        props.put("channelInitializer", "tech.ydb.jdbc.settings.CustomChannelInitilizer,  "
+                + "1tech.ydb.jdbc.settings.StaticTokenProvider");
+        YdbConnectionProperties cp = new YdbConnectionProperties(null, null, props);
+        ExceptionAssert.sqlException(
+                "channelInitializer must be full class name or instance of Consumer<ManagedChannelBuilder>",
+                () -> cp.applyToGrpcTransport(GrpcTransport.forConnectionString(YDB_URL))
+        );
+    }
+
+    @Test
+    public void channelInitializerUnknownClassTest() throws SQLException {
+        Properties props = new Properties();
+        props.put("channelInitializer", "tech.ydb.jdbc.settings.TestTokenProvider");
+        YdbConnectionProperties cp = new YdbConnectionProperties(null, null, props);
+        ExceptionAssert.sqlException(
+                "channelInitializer tech.ydb.jdbc.settings.TestTokenProvider not found",
+                () -> cp.applyToGrpcTransport(GrpcTransport.forConnectionString(YDB_URL))
+        );
+    }
+
+    @Test
+    public void channelInitializerInvalidClassTest() throws SQLException {
+        Properties props = new Properties();
+        props.put("channelInitializer", "tech.ydb.jdbc.settings.BadChannelInitilizer");
+        YdbConnectionProperties cp = new YdbConnectionProperties(null, null, props);
+        ExceptionAssert.sqlException(
+                "Cannot construct channelInitializer tech.ydb.jdbc.settings.BadChannelInitilizer",
+                () -> cp.applyToGrpcTransport(GrpcTransport.forConnectionString(YDB_URL))
+        );
+    }
+
+    @Test
+    public void channelInitializerNonConsumerClassTest() throws SQLException {
+        Properties props = new Properties();
+        props.put("channelInitializer", "tech.ydb.jdbc.settings.StaticTokenProvider");
+        YdbConnectionProperties cp = new YdbConnectionProperties(null, null, props);
+        ExceptionAssert.sqlException(
+                "channelInitializer tech.ydb.jdbc.settings.StaticTokenProvider is not implement "
+                        + "Consumer<ManagedChannelBuilder>",
+                () -> cp.applyToGrpcTransport(GrpcTransport.forConnectionString(YDB_URL))
+        );
+    }
+
+    @Test
+    public void channelInitializerWrongObjectTest() throws SQLException {
+        Properties props = new Properties();
+        props.put("channelInitializer", new Object());
+        YdbConnectionProperties cp = new YdbConnectionProperties(null, null, props);
+
+        ExceptionAssert.sqlException(
+                "Cannot parse channelInitializer java.lang.Object",
+                () -> cp.applyToGrpcTransport(GrpcTransport.forConnectionString(YDB_URL))
+        );
+    }
 }