diff --git a/jdbc/src/main/java/tech/ydb/jdbc/settings/YdbConnectionProperties.java b/jdbc/src/main/java/tech/ydb/jdbc/settings/YdbConnectionProperties.java index 938634b..6deab80 100644 --- a/jdbc/src/main/java/tech/ydb/jdbc/settings/YdbConnectionProperties.java +++ b/jdbc/src/main/java/tech/ydb/jdbc/settings/YdbConnectionProperties.java @@ -8,6 +8,7 @@ import java.util.logging.Logger; import java.util.regex.Pattern; +import tech.ydb.auth.AuthProvider; import tech.ydb.auth.TokenAuthProvider; import tech.ydb.auth.iam.CloudAuthHelper; import tech.ydb.core.auth.StaticCredentials; @@ -205,6 +206,9 @@ public GrpcTransportBuilder applyToGrpcTransport(GrpcTransportBuilder builder) t 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()) { @@ -226,7 +230,7 @@ public GrpcTransportBuilder applyToGrpcTransport(GrpcTransportBuilder builder) t | IllegalArgumentException | InvocationTargetException ex) { throw new SQLException("Cannot construct tokenProvider " + className, ex); } - } else { + } else if (provider != null) { throw new SQLException("Cannot parse tokenProvider " + provider.getClass().getName()); } } diff --git a/jdbc/src/test/java/tech/ydb/jdbc/settings/YdbConnectionPropertiesTest.java b/jdbc/src/test/java/tech/ydb/jdbc/settings/YdbConnectionPropertiesTest.java index a3f7633..c436894 100644 --- a/jdbc/src/test/java/tech/ydb/jdbc/settings/YdbConnectionPropertiesTest.java +++ b/jdbc/src/test/java/tech/ydb/jdbc/settings/YdbConnectionPropertiesTest.java @@ -8,6 +8,7 @@ import org.junit.jupiter.api.Test; import tech.ydb.auth.AuthIdentity; +import tech.ydb.auth.AuthProvider; import tech.ydb.core.grpc.GrpcTransport; import tech.ydb.core.grpc.GrpcTransportBuilder; import tech.ydb.jdbc.impl.helper.ExceptionAssert; @@ -20,7 +21,18 @@ public class YdbConnectionPropertiesTest { private static final String YDB_URL = "grpc://localhost/local"; @Test - public void tokenProviderObject() throws SQLException { + public void tokenProviderTest() throws SQLException { + Properties props = new Properties(); + props.put("tokenProvider", (AuthProvider) () -> () -> "AUTH_PROVIDER"); + YdbConnectionProperties cp = new YdbConnectionProperties(null, null, props); + GrpcTransportBuilder builder = cp.applyToGrpcTransport(GrpcTransport.forConnectionString(YDB_URL)); + try (AuthIdentity identity = builder.getAuthProvider().createAuthIdentity(null)) { + Assertions.assertEquals("AUTH_PROVIDER", identity.getToken()); + } + } + + @Test + public void tokenProviderSupplierTest() throws SQLException { Properties props = new Properties(); props.put("tokenProvider", (Supplier) () -> "SUPPLIER"); YdbConnectionProperties cp = new YdbConnectionProperties(null, null, props);