Skip to content

Commit 5a56d89

Browse files
alex268alex268
authored
Added option 'tokenProvider' to use custom auth provider (#107)
2 parents 4a6d7d3 + f909f29 commit 5a56d89

File tree

8 files changed

+222
-11
lines changed

8 files changed

+222
-11
lines changed

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

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -202,6 +202,7 @@ public DriverPropertyInfo[] toPropertyInfo() throws SQLException {
202202
YdbConnectionProperties.USE_METADATA.toInfo(properties),
203203
YdbConnectionProperties.IAM_ENDPOINT.toInfo(properties),
204204
YdbConnectionProperties.METADATA_URL.toInfo(properties),
205+
YdbConnectionProperties.TOKEN_PROVIDER.toInfo(properties),
205206
YdbConnectionProperties.GRPC_COMPRESSION.toInfo(properties),
206207

207208
YdbClientProperties.KEEP_QUERY_TEXT.toInfo(properties),

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

Lines changed: 55 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,12 @@
11
package tech.ydb.jdbc.settings;
22

3+
import java.lang.reflect.InvocationTargetException;
34
import java.sql.SQLException;
45
import java.util.Properties;
6+
import java.util.function.Supplier;
57
import java.util.logging.Level;
68
import java.util.logging.Logger;
9+
import java.util.regex.Pattern;
710

811
import tech.ydb.auth.TokenAuthProvider;
912
import tech.ydb.auth.iam.CloudAuthHelper;
@@ -16,6 +19,8 @@
1619

1720
public class YdbConnectionProperties {
1821
private static final Logger LOGGER = Logger.getLogger(YdbDriver.class.getName());
22+
private static final String ID_PATTERN = "\\p{javaJavaIdentifierStart}\\p{javaJavaIdentifierPart}*";
23+
private static final Pattern FQCN = Pattern.compile(ID_PATTERN + "(\\." + ID_PATTERN + ")*");
1924

2025
static final YdbProperty<String> TOKEN = YdbProperty.content(YdbConfig.TOKEN_KEY, "Authentication token");
2126

@@ -48,6 +53,9 @@ public class YdbConnectionProperties {
4853
static final YdbProperty<String> METADATA_URL = YdbProperty.content("metadataURL",
4954
"Custom URL for the metadata service authentication");
5055

56+
static final YdbProperty<Object> TOKEN_PROVIDER = YdbProperty.object("tokenProvider",
57+
"Custom token provider, use object instance or class full name impementing Supplier<String>");
58+
5159
static final YdbProperty<String> GRPC_COMPRESSION = YdbProperty.string(
5260
"grpcCompression", "Use specified GRPC compressor (supported only none and gzip)"
5361
);
@@ -65,13 +73,12 @@ public class YdbConnectionProperties {
6573
private final YdbValue<Boolean> useMetadata;
6674
private final YdbValue<String> iamEndpoint;
6775
private final YdbValue<String> metadataUrl;
76+
private final YdbValue<Object> tokenProvider;
6877
private final YdbValue<String> grpcCompression;
6978

70-
public YdbConnectionProperties(YdbConfig config) throws SQLException {
71-
this.username = config.getUsername();
72-
this.password = config.getPassword();
73-
74-
Properties props = config.getProperties();
79+
public YdbConnectionProperties(String username, String password, Properties props) throws SQLException {
80+
this.username = username;
81+
this.password = password;
7582

7683
this.localDatacenter = LOCAL_DATACENTER.readValue(props);
7784
this.useSecureConnection = USE_SECURE_CONNECTION.readValue(props);
@@ -83,9 +90,14 @@ public YdbConnectionProperties(YdbConfig config) throws SQLException {
8390
this.useMetadata = USE_METADATA.readValue(props);
8491
this.iamEndpoint = IAM_ENDPOINT.readValue(props);
8592
this.metadataUrl = METADATA_URL.readValue(props);
93+
this.tokenProvider = TOKEN_PROVIDER.readValue(props);
8694
this.grpcCompression = GRPC_COMPRESSION.readValue(props);
8795
}
8896

97+
public YdbConnectionProperties(YdbConfig config) throws SQLException {
98+
this(config.getUsername(), config.getPassword(), config.getProperties());
99+
}
100+
89101
String getLocalDataCenter() {
90102
return localDatacenter.getValue();
91103
}
@@ -98,7 +110,7 @@ byte[] getSecureConnectionCert() {
98110
return secureConnectionCertificate.getValue();
99111
}
100112

101-
public GrpcTransportBuilder applyToGrpcTransport(GrpcTransportBuilder builder) {
113+
public GrpcTransportBuilder applyToGrpcTransport(GrpcTransportBuilder builder) throws SQLException {
102114
if (localDatacenter.hasValue()) {
103115
builder = builder.withBalancingSettings(BalancingSettings.fromLocation(localDatacenter.getValue()));
104116
}
@@ -180,6 +192,43 @@ public GrpcTransportBuilder applyToGrpcTransport(GrpcTransportBuilder builder) {
180192
} else {
181193
builder = builder.withAuthProvider(CloudAuthHelper.getServiceAccountJsonAuthProvider(json));
182194
}
195+
usedProvider = "service account file credentitals";
196+
}
197+
198+
if (tokenProvider.hasValue()) {
199+
if (usedProvider != null) {
200+
LOGGER.log(Level.WARNING, "Dublicate authentication config! Token Provider credentials replaces {0}",
201+
usedProvider);
202+
}
203+
204+
Object provider = tokenProvider.getValue();
205+
if (provider instanceof Supplier) {
206+
Supplier<?> prov = (Supplier<?>) provider;
207+
builder = builder.withAuthProvider((rpc) -> () -> prov.get().toString());
208+
} else if (provider instanceof String) {
209+
String className = (String) provider;
210+
if (!FQCN.matcher(className).matches()) {
211+
throw new SQLException("tokenProvider must be full class name or instance of Supplier<String>");
212+
}
213+
214+
try {
215+
Class<?> clazz = Class.forName(className);
216+
if (!Supplier.class.isAssignableFrom(clazz)) {
217+
throw new SQLException("tokenProvider " + className + " is not implement Supplier<String>");
218+
}
219+
Supplier<?> prov = clazz.asSubclass(Supplier.class)
220+
.getConstructor(new Class<?>[0])
221+
.newInstance(new Object[0]);
222+
builder = builder.withAuthProvider((rpc) -> () -> prov.get().toString());
223+
} catch (ClassNotFoundException ex) {
224+
throw new SQLException("tokenProvider " + className + " not found", ex);
225+
} catch (NoSuchMethodException | SecurityException | InstantiationException | IllegalAccessException
226+
| IllegalArgumentException | InvocationTargetException ex) {
227+
throw new SQLException("Cannot construct tokenProvider " + className, ex);
228+
}
229+
} else {
230+
throw new SQLException("Cannot parse tokenProvider " + provider.getClass().getName());
231+
}
183232
}
184233

185234
if (grpcCompression.hasValue()) {

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

Lines changed: 9 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -43,6 +43,11 @@ public YdbValue<T> readValue(Properties props) throws SQLException {
4343
}
4444
}
4545

46+
if (clazz == Object.class) {
47+
T typed = clazz.cast(value);
48+
return new YdbValue<>(true, typed.toString(), typed);
49+
}
50+
4651
if (value instanceof String) {
4752
try {
4853
String stringValue = (String) value;
@@ -69,6 +74,10 @@ DriverPropertyInfo toInfo(Properties values) throws SQLException {
6974
return info;
7075
}
7176

77+
public static YdbProperty<Object> object(String name, String description) {
78+
return new YdbProperty<>(name, description, null, Object.class, v -> v);
79+
}
80+
7281
public static YdbProperty<String> string(String name, String description) {
7382
return string(name, description, null);
7483
}
@@ -143,11 +152,6 @@ public static YdbProperty<Duration> duration(String name, String description, St
143152
});
144153
}
145154

146-
@Deprecated
147-
public static YdbProperty<byte[]> bytes(String name, String description) {
148-
return new YdbProperty<>(name, description, null, byte[].class, YdbLookup::byteFileReference);
149-
}
150-
151155
public static YdbProperty<String> content(String name, String description) {
152156
return new YdbProperty<>(name, description, null, String.class, YdbLookup::stringFileReference);
153157
}

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

Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,20 @@
1+
package tech.ydb.jdbc.settings;
2+
3+
import java.util.function.Supplier;
4+
5+
/**
6+
*
7+
* @author Aleksandr Gorshenin
8+
*/
9+
public class BadTokenProvider implements Supplier<String> {
10+
private final String token;
11+
12+
public BadTokenProvider(String token) {
13+
this.token = token;
14+
}
15+
16+
@Override
17+
public String get() {
18+
return token;
19+
}
20+
}

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

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,11 @@
1+
package tech.ydb.jdbc.settings;
2+
3+
/**
4+
*
5+
* @author Aleksandr Gorshenin
6+
*/
7+
public class NonSupplierTokenProvider {
8+
public String getToken() {
9+
return "STATIC";
10+
}
11+
}

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

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,14 @@
1+
package tech.ydb.jdbc.settings;
2+
3+
import java.util.function.Supplier;
4+
5+
/**
6+
*
7+
* @author Aleksandr Gorshenin
8+
*/
9+
public class StaticTokenProvider implements Supplier<String>{
10+
@Override
11+
public String get() {
12+
return "STATIC";
13+
}
14+
}

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

Lines changed: 110 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,110 @@
1+
package tech.ydb.jdbc.settings;
2+
3+
import java.sql.SQLException;
4+
import java.util.Properties;
5+
import java.util.function.Supplier;
6+
7+
import org.junit.jupiter.api.Assertions;
8+
import org.junit.jupiter.api.Test;
9+
10+
import tech.ydb.auth.AuthIdentity;
11+
import tech.ydb.core.grpc.GrpcTransport;
12+
import tech.ydb.core.grpc.GrpcTransportBuilder;
13+
import tech.ydb.jdbc.impl.helper.ExceptionAssert;
14+
15+
/**
16+
*
17+
* @author Aleksandr Gorshenin
18+
*/
19+
public class YdbConnectionPropertiesTest {
20+
private static final String YDB_URL = "grpc://localhost/local";
21+
22+
@Test
23+
public void tokenProviderObject() throws SQLException {
24+
Properties props = new Properties();
25+
props.put("tokenProvider", (Supplier<String>) () -> "SUPPLIER");
26+
YdbConnectionProperties cp = new YdbConnectionProperties(null, null, props);
27+
GrpcTransportBuilder builder = cp.applyToGrpcTransport(GrpcTransport.forConnectionString(YDB_URL));
28+
try (AuthIdentity identity = builder.getAuthProvider().createAuthIdentity(null)) {
29+
Assertions.assertEquals("SUPPLIER", identity.getToken());
30+
}
31+
}
32+
33+
@Test
34+
public void tokenProviderToStringTest() throws SQLException {
35+
Properties props = new Properties();
36+
props.put("tokenProvider", (Supplier<Object>) () -> new StringBuilder("test"));
37+
YdbConnectionProperties cp = new YdbConnectionProperties(null, null, props);
38+
GrpcTransportBuilder builder = cp.applyToGrpcTransport(GrpcTransport.forConnectionString(YDB_URL));
39+
try (AuthIdentity identity = builder.getAuthProvider().createAuthIdentity(null)) {
40+
Assertions.assertEquals("test", identity.getToken());
41+
}
42+
}
43+
44+
@Test
45+
public void tokenProviderClassTest() throws SQLException {
46+
Properties props = new Properties();
47+
props.put("tokenProvider", "tech.ydb.jdbc.settings.StaticTokenProvider");
48+
YdbConnectionProperties cp = new YdbConnectionProperties(null, null, props);
49+
GrpcTransportBuilder builder = cp.applyToGrpcTransport(GrpcTransport.forConnectionString(YDB_URL));
50+
try (AuthIdentity identity = builder.getAuthProvider().createAuthIdentity(null)) {
51+
Assertions.assertEquals("STATIC", identity.getToken());
52+
}
53+
}
54+
55+
@Test
56+
public void tokenProviderWrongClassNameTest() throws SQLException {
57+
Properties props = new Properties();
58+
props.put("tokenProvider", "1tech.ydb.jdbc.settings.StaticTokenProvider");
59+
YdbConnectionProperties cp = new YdbConnectionProperties(null, null, props);
60+
ExceptionAssert.sqlException(
61+
"tokenProvider must be full class name or instance of Supplier<String>",
62+
() -> cp.applyToGrpcTransport(GrpcTransport.forConnectionString(YDB_URL))
63+
);
64+
}
65+
66+
@Test
67+
public void tokenProviderUnknownClassTest() throws SQLException {
68+
Properties props = new Properties();
69+
props.put("tokenProvider", "tech.ydb.jdbc.settings.TestTokenProvider");
70+
YdbConnectionProperties cp = new YdbConnectionProperties(null, null, props);
71+
ExceptionAssert.sqlException(
72+
"tokenProvider tech.ydb.jdbc.settings.TestTokenProvider not found",
73+
() -> cp.applyToGrpcTransport(GrpcTransport.forConnectionString(YDB_URL))
74+
);
75+
}
76+
77+
@Test
78+
public void tokenProviderInvalidClassTest() throws SQLException {
79+
Properties props = new Properties();
80+
props.put("tokenProvider", "tech.ydb.jdbc.settings.BadTokenProvider");
81+
YdbConnectionProperties cp = new YdbConnectionProperties(null, null, props);
82+
ExceptionAssert.sqlException(
83+
"Cannot construct tokenProvider tech.ydb.jdbc.settings.BadTokenProvider",
84+
() -> cp.applyToGrpcTransport(GrpcTransport.forConnectionString(YDB_URL))
85+
);
86+
}
87+
88+
@Test
89+
public void tokenProviderNonSupplierClassTest() throws SQLException {
90+
Properties props = new Properties();
91+
props.put("tokenProvider", "tech.ydb.jdbc.settings.NonSupplierTokenProvider");
92+
YdbConnectionProperties cp = new YdbConnectionProperties(null, null, props);
93+
ExceptionAssert.sqlException(
94+
"tokenProvider tech.ydb.jdbc.settings.NonSupplierTokenProvider is not implement Supplier<String>",
95+
() -> cp.applyToGrpcTransport(GrpcTransport.forConnectionString(YDB_URL))
96+
);
97+
}
98+
99+
@Test
100+
public void tokenProviderWrongObjectTest() throws SQLException {
101+
Properties props = new Properties();
102+
props.put("tokenProvider", new Object());
103+
YdbConnectionProperties cp = new YdbConnectionProperties(null, null, props);
104+
105+
ExceptionAssert.sqlException(
106+
"Cannot parse tokenProvider java.lang.Object",
107+
() -> cp.applyToGrpcTransport(GrpcTransport.forConnectionString(YDB_URL))
108+
);
109+
}
110+
}

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

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -316,6 +316,7 @@ static DriverPropertyInfo[] defaultPropertyInfo(@Nullable String localDatacenter
316316
new DriverPropertyInfo("useMetadata", ""),
317317
new DriverPropertyInfo("iamEndpoint", ""),
318318
new DriverPropertyInfo("metadataURL", ""),
319+
new DriverPropertyInfo("tokenProvider", ""),
319320
new DriverPropertyInfo("grpcCompression", ""),
320321
new DriverPropertyInfo("keepQueryText", ""),
321322
new DriverPropertyInfo("sessionKeepAliveTime", ""),
@@ -362,6 +363,7 @@ static DriverPropertyInfo[] customizedPropertyInfo() {
362363
new DriverPropertyInfo("useMetadata", "true"),
363364
new DriverPropertyInfo("iamEndpoint", "iam.endpoint.com"),
364365
new DriverPropertyInfo("metadataURL", "https://metadata.com"),
366+
new DriverPropertyInfo("tokenProvider", "tech.ydb.jdbc.settings.StaticTokenProvider"),
365367
new DriverPropertyInfo("grpcCompression", "gzip"),
366368
new DriverPropertyInfo("keepQueryText", "true"),
367369
new DriverPropertyInfo("sessionKeepAliveTime", "15m"),

0 commit comments

Comments
 (0)