Skip to content

Commit 3dfb885

Browse files
alex268alex268
committed
Added option channelInitializer
1 parent 83165d2 commit 3dfb885

File tree

5 files changed

+201
-42
lines changed

5 files changed

+201
-42
lines changed

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

Lines changed: 82 additions & 29 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@
33
import java.lang.reflect.InvocationTargetException;
44
import java.sql.SQLException;
55
import java.util.Properties;
6+
import java.util.function.Consumer;
67
import java.util.function.Supplier;
78
import java.util.logging.Level;
89
import java.util.logging.Logger;
@@ -54,6 +55,10 @@ public class YdbConnectionProperties {
5455
static final YdbProperty<String> METADATA_URL = YdbProperty.content("metadataURL",
5556
"Custom URL for the metadata service authentication");
5657

58+
static final YdbProperty<Object> CHANNEL_INITIALIZER = YdbProperty.object("channelInitializer",
59+
"Custom GRPC channel initilizer, use object instance or class full name impementing"
60+
+ " Consumer<ManagedChannelBuilder>");
61+
5762
static final YdbProperty<Object> TOKEN_PROVIDER = YdbProperty.object("tokenProvider",
5863
"Custom token provider, use object instance or class full name impementing Supplier<String>");
5964

@@ -75,6 +80,7 @@ public class YdbConnectionProperties {
7580
private final YdbValue<String> iamEndpoint;
7681
private final YdbValue<String> metadataUrl;
7782
private final YdbValue<Object> tokenProvider;
83+
private final YdbValue<Object> channelInitializer;
7884
private final YdbValue<String> grpcCompression;
7985

8086
public YdbConnectionProperties(String username, String password, Properties props) throws SQLException {
@@ -92,6 +98,7 @@ public YdbConnectionProperties(String username, String password, Properties prop
9298
this.iamEndpoint = IAM_ENDPOINT.readValue(props);
9399
this.metadataUrl = METADATA_URL.readValue(props);
94100
this.tokenProvider = TOKEN_PROVIDER.readValue(props);
101+
this.channelInitializer = CHANNEL_INITIALIZER.readValue(props);
95102
this.grpcCompression = GRPC_COMPRESSION.readValue(props);
96103
}
97104

@@ -203,36 +210,12 @@ public GrpcTransportBuilder applyToGrpcTransport(GrpcTransportBuilder builder) t
203210
}
204211

205212
Object provider = tokenProvider.getValue();
206-
if (provider instanceof Supplier) {
207-
Supplier<?> prov = (Supplier<?>) provider;
208-
builder = builder.withAuthProvider((rpc) -> () -> prov.get().toString());
209-
} else if (provider instanceof AuthProvider) {
210-
AuthProvider prov = (AuthProvider) provider;
211-
builder = builder.withAuthProvider(prov);
212-
} else if (provider instanceof String) {
213-
String className = (String) provider;
214-
if (!FQCN.matcher(className).matches()) {
215-
throw new SQLException("tokenProvider must be full class name or instance of Supplier<String>");
216-
}
213+
builder = applyTokenProvider(builder, provider);
214+
}
217215

218-
try {
219-
Class<?> clazz = Class.forName(className);
220-
if (!Supplier.class.isAssignableFrom(clazz)) {
221-
throw new SQLException("tokenProvider " + className + " is not implement Supplier<String>");
222-
}
223-
Supplier<?> prov = clazz.asSubclass(Supplier.class)
224-
.getConstructor(new Class<?>[0])
225-
.newInstance(new Object[0]);
226-
builder = builder.withAuthProvider((rpc) -> () -> prov.get().toString());
227-
} catch (ClassNotFoundException ex) {
228-
throw new SQLException("tokenProvider " + className + " not found", ex);
229-
} catch (NoSuchMethodException | SecurityException | InstantiationException | IllegalAccessException
230-
| IllegalArgumentException | InvocationTargetException ex) {
231-
throw new SQLException("Cannot construct tokenProvider " + className, ex);
232-
}
233-
} else if (provider != null) {
234-
throw new SQLException("Cannot parse tokenProvider " + provider.getClass().getName());
235-
}
216+
if (channelInitializer.hasValue()) {
217+
Object initializer = channelInitializer.getValue();
218+
builder = applyChannelInitializer(builder, initializer);
236219
}
237220

238221
if (grpcCompression.hasValue()) {
@@ -248,4 +231,74 @@ public GrpcTransportBuilder applyToGrpcTransport(GrpcTransportBuilder builder) t
248231

249232
return builder;
250233
}
234+
235+
private GrpcTransportBuilder applyTokenProvider(GrpcTransportBuilder builder, Object provider) throws SQLException {
236+
if (provider instanceof Supplier) {
237+
Supplier<?> prov = (Supplier<?>) provider;
238+
builder = builder.withAuthProvider((rpc) -> () -> prov.get().toString());
239+
} else if (provider instanceof AuthProvider) {
240+
AuthProvider prov = (AuthProvider) provider;
241+
builder = builder.withAuthProvider(prov);
242+
} else if (provider instanceof String) {
243+
String className = (String) provider;
244+
if (!FQCN.matcher(className).matches()) {
245+
throw new SQLException("tokenProvider must be full class name or instance of Supplier<String>");
246+
}
247+
248+
try {
249+
Class<?> clazz = Class.forName(className);
250+
if (!Supplier.class.isAssignableFrom(clazz)) {
251+
throw new SQLException("tokenProvider " + className + " is not implement Supplier<String>");
252+
}
253+
Supplier<?> prov = clazz.asSubclass(Supplier.class)
254+
.getConstructor(new Class<?>[0])
255+
.newInstance(new Object[0]);
256+
builder = builder.withAuthProvider((rpc) -> () -> prov.get().toString());
257+
} catch (ClassNotFoundException ex) {
258+
throw new SQLException("tokenProvider " + className + " not found", ex);
259+
} catch (NoSuchMethodException | SecurityException | InstantiationException | IllegalAccessException
260+
| IllegalArgumentException | InvocationTargetException ex) {
261+
throw new SQLException("Cannot construct tokenProvider " + className, ex);
262+
}
263+
} else if (provider != null) {
264+
throw new SQLException("Cannot parse tokenProvider " + provider.getClass().getName());
265+
}
266+
return builder;
267+
}
268+
269+
private GrpcTransportBuilder applyChannelInitializer(GrpcTransportBuilder builder, Object initializer)
270+
throws SQLException {
271+
if (initializer instanceof Consumer) {
272+
@SuppressWarnings("unchecked")
273+
Consumer<Object> prov = (Consumer<Object>) initializer;
274+
builder = builder.addChannelInitializer(prov);
275+
} else if (initializer instanceof String) {
276+
String className = (String) initializer;
277+
if (!FQCN.matcher(className).matches()) {
278+
throw new SQLException("channelInitializer must be full class name or instance of "
279+
+ "Consumer<ManagedChannelBuilder>");
280+
}
281+
282+
try {
283+
Class<?> clazz = Class.forName(className);
284+
if (!Consumer.class.isAssignableFrom(clazz)) {
285+
throw new SQLException("channelInitializer " + className + " is not implement "
286+
+ "Consumer<ManagedChannelBuilder>");
287+
}
288+
@SuppressWarnings("unchecked")
289+
Consumer<Object> prov = clazz.asSubclass(Consumer.class)
290+
.getConstructor(new Class<?>[0])
291+
.newInstance(new Object[0]);
292+
builder = builder.addChannelInitializer(prov);
293+
} catch (ClassNotFoundException ex) {
294+
throw new SQLException("channelInitializer " + className + " not found", ex);
295+
} catch (NoSuchMethodException | SecurityException | InstantiationException | IllegalAccessException
296+
| IllegalArgumentException | InvocationTargetException ex) {
297+
throw new SQLException("Cannot construct channelInitializer " + className, ex);
298+
}
299+
} else if (initializer != null) {
300+
throw new SQLException("Cannot parse channelInitializer " + initializer.getClass().getName());
301+
}
302+
return builder;
303+
}
251304
}

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

Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,22 @@
1+
package tech.ydb.jdbc.settings;
2+
3+
import java.util.function.Consumer;
4+
5+
import io.grpc.ManagedChannelBuilder;
6+
7+
/**
8+
*
9+
* @author Aleksandr Gorshenin
10+
*/
11+
public class BadChannelInitilizer implements Consumer<ManagedChannelBuilder<?>> {
12+
private final String authority;
13+
14+
public BadChannelInitilizer(String authority) {
15+
this.authority = authority;
16+
}
17+
18+
@Override
19+
public void accept(ManagedChannelBuilder<?> t) {
20+
t.overrideAuthority(authority);
21+
}
22+
}

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

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,16 @@
1+
package tech.ydb.jdbc.settings;
2+
3+
import java.util.function.Consumer;
4+
5+
import io.grpc.ManagedChannelBuilder;
6+
7+
/**
8+
*
9+
* @author Aleksandr Gorshenin
10+
*/
11+
public class CustomChannelInitilizer implements Consumer<ManagedChannelBuilder<?>> {
12+
@Override
13+
public void accept(ManagedChannelBuilder<?> t) {
14+
t.usePlaintext();
15+
}
16+
}

‎jdbc/src/test/java/tech/ydb/jdbc/settings/NonSupplierTokenProvider.java‎

Lines changed: 0 additions & 11 deletions
This file was deleted.

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

Lines changed: 81 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -2,8 +2,10 @@
22

33
import java.sql.SQLException;
44
import java.util.Properties;
5+
import java.util.function.Consumer;
56
import java.util.function.Supplier;
67

8+
import io.grpc.ManagedChannelBuilder;
79
import org.junit.jupiter.api.Assertions;
810
import org.junit.jupiter.api.Test;
911

@@ -100,10 +102,10 @@ public void tokenProviderInvalidClassTest() throws SQLException {
100102
@Test
101103
public void tokenProviderNonSupplierClassTest() throws SQLException {
102104
Properties props = new Properties();
103-
props.put("tokenProvider", "tech.ydb.jdbc.settings.NonSupplierTokenProvider");
105+
props.put("tokenProvider", "tech.ydb.jdbc.settings.CustomChannelInitilizer");
104106
YdbConnectionProperties cp = new YdbConnectionProperties(null, null, props);
105107
ExceptionAssert.sqlException(
106-
"tokenProvider tech.ydb.jdbc.settings.NonSupplierTokenProvider is not implement Supplier<String>",
108+
"tokenProvider tech.ydb.jdbc.settings.CustomChannelInitilizer is not implement Supplier<String>",
107109
() -> cp.applyToGrpcTransport(GrpcTransport.forConnectionString(YDB_URL))
108110
);
109111
}
@@ -119,4 +121,81 @@ public void tokenProviderWrongObjectTest() throws SQLException {
119121
() -> cp.applyToGrpcTransport(GrpcTransport.forConnectionString(YDB_URL))
120122
);
121123
}
124+
125+
@Test
126+
public void channelInitializerConsumerTest() throws SQLException {
127+
Properties props = new Properties();
128+
props.put("channelInitializer", (Consumer<ManagedChannelBuilder<?>>) (ManagedChannelBuilder<?> t) -> {
129+
t.maxTraceEvents(0);
130+
});
131+
YdbConnectionProperties cp = new YdbConnectionProperties(null, null, props);
132+
GrpcTransportBuilder builder = cp.applyToGrpcTransport(GrpcTransport.forConnectionString(YDB_URL));
133+
Assertions.assertFalse(builder.getChannelInitializers().isEmpty());
134+
}
135+
136+
@Test
137+
public void channelInitializerClassTest() throws SQLException {
138+
Properties props = new Properties();
139+
props.put("channelInitializer", "tech.ydb.jdbc.settings.CustomChannelInitilizer");
140+
YdbConnectionProperties cp = new YdbConnectionProperties(null, null, props);
141+
GrpcTransportBuilder builder = cp.applyToGrpcTransport(GrpcTransport.forConnectionString(YDB_URL));
142+
Assertions.assertFalse(builder.getChannelInitializers().isEmpty());
143+
}
144+
145+
@Test
146+
public void channelInitializerWrongClassNameTest() throws SQLException {
147+
Properties props = new Properties();
148+
props.put("channelInitializer", "1tech.ydb.jdbc.settings.StaticTokenProvider");
149+
YdbConnectionProperties cp = new YdbConnectionProperties(null, null, props);
150+
ExceptionAssert.sqlException(
151+
"channelInitializer must be full class name or instance of Consumer<ManagedChannelBuilder>",
152+
() -> cp.applyToGrpcTransport(GrpcTransport.forConnectionString(YDB_URL))
153+
);
154+
}
155+
156+
@Test
157+
public void channelInitializerUnknownClassTest() throws SQLException {
158+
Properties props = new Properties();
159+
props.put("channelInitializer", "tech.ydb.jdbc.settings.TestTokenProvider");
160+
YdbConnectionProperties cp = new YdbConnectionProperties(null, null, props);
161+
ExceptionAssert.sqlException(
162+
"channelInitializer tech.ydb.jdbc.settings.TestTokenProvider not found",
163+
() -> cp.applyToGrpcTransport(GrpcTransport.forConnectionString(YDB_URL))
164+
);
165+
}
166+
167+
@Test
168+
public void channelInitializerInvalidClassTest() throws SQLException {
169+
Properties props = new Properties();
170+
props.put("channelInitializer", "tech.ydb.jdbc.settings.BadChannelInitilizer");
171+
YdbConnectionProperties cp = new YdbConnectionProperties(null, null, props);
172+
ExceptionAssert.sqlException(
173+
"Cannot construct channelInitializer tech.ydb.jdbc.settings.BadChannelInitilizer",
174+
() -> cp.applyToGrpcTransport(GrpcTransport.forConnectionString(YDB_URL))
175+
);
176+
}
177+
178+
@Test
179+
public void channelInitializerNonConsumerClassTest() throws SQLException {
180+
Properties props = new Properties();
181+
props.put("channelInitializer", "tech.ydb.jdbc.settings.StaticTokenProvider");
182+
YdbConnectionProperties cp = new YdbConnectionProperties(null, null, props);
183+
ExceptionAssert.sqlException(
184+
"channelInitializer tech.ydb.jdbc.settings.StaticTokenProvider is not implement "
185+
+ "Consumer<ManagedChannelBuilder>",
186+
() -> cp.applyToGrpcTransport(GrpcTransport.forConnectionString(YDB_URL))
187+
);
188+
}
189+
190+
@Test
191+
public void channelInitializerWrongObjectTest() throws SQLException {
192+
Properties props = new Properties();
193+
props.put("channelInitializer", new Object());
194+
YdbConnectionProperties cp = new YdbConnectionProperties(null, null, props);
195+
196+
ExceptionAssert.sqlException(
197+
"Cannot parse channelInitializer java.lang.Object",
198+
() -> cp.applyToGrpcTransport(GrpcTransport.forConnectionString(YDB_URL))
199+
);
200+
}
122201
}

0 commit comments

Comments
 (0)