add plugin for DSQL auth tokens

danielfrankcom · danielfrankcom · commit 8c3d0a7b1fef · 2025-07-23T11:54:44.000-07:00
diff --git a/wrapper/build.gradle.kts b/wrapper/build.gradle.kts
@@ -37,6 +37,7 @@ dependencies {
     compileOnly("software.amazon.awssdk:rds:2.31.78")
     compileOnly("software.amazon.awssdk:auth:2.31.45") // Required for IAM (light implementation)
     compileOnly("software.amazon.awssdk:http-client-spi:2.31.60") // Required for IAM (light implementation)
+    compileOnly("software.amazon.awssdk:dsql:2.31.78")
     compileOnly("software.amazon.awssdk:sts:2.31.78")
     compileOnly("com.zaxxer:HikariCP:4.0.3") // Version 4.+ is compatible with Java 8
     compileOnly("com.mchange:c3p0:0.11.0")
@@ -73,6 +74,7 @@ dependencies {
     testImplementation("software.amazon.awssdk:rds:2.31.78")
     testImplementation("software.amazon.awssdk:auth:2.31.45") // Required for IAM (light implementation)
     testImplementation("software.amazon.awssdk:http-client-spi:2.31.60") // Required for IAM (light implementation)
+    testImplementation("software.amazon.awssdk:dsql:2.31.78")
     testImplementation("software.amazon.awssdk:ec2:2.31.78")
     testImplementation("software.amazon.awssdk:secretsmanager:2.31.12")
     testImplementation("software.amazon.awssdk:sts:2.31.78")
diff --git a/wrapper/src/main/java/software/amazon/jdbc/ConnectionPluginChainBuilder.java b/wrapper/src/main/java/software/amazon/jdbc/ConnectionPluginChainBuilder.java
@@ -45,6 +45,7 @@
 import software.amazon.jdbc.plugin.failover.FailoverConnectionPluginFactory;
 import software.amazon.jdbc.plugin.federatedauth.FederatedAuthPluginFactory;
 import software.amazon.jdbc.plugin.federatedauth.OktaAuthPluginFactory;
+import software.amazon.jdbc.plugin.iam.DsqlIamConnectionPluginFactory;
 import software.amazon.jdbc.plugin.iam.IamAuthConnectionPluginFactory;
 import software.amazon.jdbc.plugin.limitless.LimitlessConnectionPluginFactory;
 import software.amazon.jdbc.plugin.readwritesplitting.ReadWriteSplittingPluginFactory;
@@ -75,6 +76,7 @@ public class ConnectionPluginChainBuilder {
           put("failover", new FailoverConnectionPluginFactory());
           put("failover2", new software.amazon.jdbc.plugin.failover2.FailoverConnectionPluginFactory());
           put("iam", new IamAuthConnectionPluginFactory());
+          put("iamDsql", new DsqlIamConnectionPluginFactory());
           put("awsSecretsManager", new AwsSecretsManagerConnectionPluginFactory());
           put("federatedAuth", new FederatedAuthPluginFactory());
           put("okta", new OktaAuthPluginFactory());
@@ -114,6 +116,7 @@ public class ConnectionPluginChainBuilder {
           put(FastestResponseStrategyPluginFactory.class, 900);
           put(LimitlessConnectionPluginFactory.class, 950);
           put(IamAuthConnectionPluginFactory.class, 1000);
+          put(DsqlIamConnectionPluginFactory.class, 1001);
           put(AwsSecretsManagerConnectionPluginFactory.class, 1100);
           put(FederatedAuthPluginFactory.class, 1200);
           put(LogQueryConnectionPluginFactory.class, 1300);
diff --git a/wrapper/src/main/java/software/amazon/jdbc/plugin/federatedauth/FederatedAuthPlugin.java b/wrapper/src/main/java/software/amazon/jdbc/plugin/federatedauth/FederatedAuthPlugin.java
@@ -121,7 +121,7 @@ public Set<String> getSubscribedMethods() {
 
   public FederatedAuthPlugin(final PluginService pluginService,
       final CredentialsProviderFactory credentialsProviderFactory) {
-    this(pluginService, credentialsProviderFactory, new RdsUtils(), IamAuthUtils.getTokenUtility());
+    this(pluginService, credentialsProviderFactory, new RdsUtils(), IamAuthUtils.getRdsTokenUtility());
   }
 
   FederatedAuthPlugin(
diff --git a/wrapper/src/main/java/software/amazon/jdbc/plugin/federatedauth/OktaAuthPlugin.java b/wrapper/src/main/java/software/amazon/jdbc/plugin/federatedauth/OktaAuthPlugin.java
@@ -95,7 +95,7 @@ public class OktaAuthPlugin extends AbstractConnectionPlugin {
   private final TelemetryCounter fetchTokenCounter;
 
   public OktaAuthPlugin(PluginService pluginService, CredentialsProviderFactory credentialsProviderFactory) {
-    this(pluginService, credentialsProviderFactory, new RdsUtils(), IamAuthUtils.getTokenUtility());
+    this(pluginService, credentialsProviderFactory, new RdsUtils(), IamAuthUtils.getRdsTokenUtility());
   }
 
   OktaAuthPlugin(
diff --git a/wrapper/src/main/java/software/amazon/jdbc/plugin/iam/DsqlIamConnectionPluginFactory.java b/wrapper/src/main/java/software/amazon/jdbc/plugin/iam/DsqlIamConnectionPluginFactory.java
@@ -0,0 +1,34 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License").
+ * You may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package software.amazon.jdbc.plugin.iam;
+
+import java.util.Properties;
+import org.checkerframework.checker.nullness.qual.NonNull;
+import software.amazon.jdbc.ConnectionPlugin;
+import software.amazon.jdbc.ConnectionPluginFactory;
+import software.amazon.jdbc.PluginService;
+import software.amazon.jdbc.util.IamAuthUtils;
+
+/**
+ * Provides {@link ConnectionPlugin} instances which can be used to connect to Amazon Aurora DSQL.
+ */
+public class DsqlIamConnectionPluginFactory implements ConnectionPluginFactory {
+  @Override
+  public ConnectionPlugin getInstance(@NonNull final PluginService pluginService, @NonNull final Properties props) {
+    return new IamAuthConnectionPlugin(pluginService, IamAuthUtils.getDsqlTokenUtility());
+  }
+}
diff --git a/wrapper/src/main/java/software/amazon/jdbc/plugin/iam/DsqlTokenUtility.java b/wrapper/src/main/java/software/amazon/jdbc/plugin/iam/DsqlTokenUtility.java
@@ -0,0 +1,53 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License").
+ * You may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package software.amazon.jdbc.plugin.iam;
+
+import org.checkerframework.checker.nullness.qual.NonNull;
+import software.amazon.awssdk.auth.credentials.AwsCredentialsProvider;
+import software.amazon.awssdk.regions.Region;
+import software.amazon.awssdk.services.dsql.DsqlUtilities;
+
+/**
+ * Represents an {@link IamTokenUtility} which provides auth tokens for connecting to Amazon Aurora DSQL.
+ */
+public class DsqlTokenUtility implements IamTokenUtility {
+
+  public DsqlTokenUtility() { }
+
+  @Override
+  public String generateAuthenticationToken(
+      @NonNull final AwsCredentialsProvider credentialsProvider,
+      @NonNull final Region region,
+      @NonNull final String hostname,
+      final int port,
+      @NonNull final String username) {
+    final DsqlUtilities utilities = DsqlUtilities.builder()
+        .credentialsProvider(credentialsProvider)
+        .region(region)
+        .build();
+
+    if (username.equals("admin")) {
+      return utilities.generateDbConnectAdminAuthToken((builder) ->
+          builder.hostname(hostname).region(region)
+      );
+    } else {
+      return utilities.generateDbConnectAuthToken((builder) ->
+          builder.hostname(hostname).region(region)
+      );
+    }
+  }
+}
diff --git a/wrapper/src/main/java/software/amazon/jdbc/plugin/iam/IamAuthConnectionPlugin.java b/wrapper/src/main/java/software/amazon/jdbc/plugin/iam/IamAuthConnectionPlugin.java
@@ -86,11 +86,7 @@ public class IamAuthConnectionPlugin extends AbstractConnectionPlugin {
 
   private final IamTokenUtility iamTokenUtility;
 
-  public IamAuthConnectionPlugin(final @NonNull PluginService pluginService) {
-    this(pluginService, IamAuthUtils.getTokenUtility());
-  }
-
-  IamAuthConnectionPlugin(final @NonNull PluginService pluginService, IamTokenUtility utility) {
+  IamAuthConnectionPlugin(final @NonNull PluginService pluginService, final IamTokenUtility utility) {
     this.iamTokenUtility = utility;
     this.pluginService = pluginService;
     this.telemetryFactory = pluginService.getTelemetryFactory();
diff --git a/wrapper/src/main/java/software/amazon/jdbc/plugin/iam/IamAuthConnectionPluginFactory.java b/wrapper/src/main/java/software/amazon/jdbc/plugin/iam/IamAuthConnectionPluginFactory.java
@@ -20,10 +20,11 @@
 import software.amazon.jdbc.ConnectionPlugin;
 import software.amazon.jdbc.ConnectionPluginFactory;
 import software.amazon.jdbc.PluginService;
+import software.amazon.jdbc.util.IamAuthUtils;
 
 public class IamAuthConnectionPluginFactory implements ConnectionPluginFactory {
   @Override
   public ConnectionPlugin getInstance(final PluginService pluginService, final Properties props) {
-    return new IamAuthConnectionPlugin(pluginService);
+    return new IamAuthConnectionPlugin(pluginService, IamAuthUtils.getRdsTokenUtility());
   }
 }
diff --git a/wrapper/src/main/java/software/amazon/jdbc/util/IamAuthUtils.java b/wrapper/src/main/java/software/amazon/jdbc/util/IamAuthUtils.java
@@ -21,6 +21,7 @@
 import software.amazon.awssdk.regions.Region;
 import software.amazon.jdbc.HostSpec;
 import software.amazon.jdbc.PluginService;
+import software.amazon.jdbc.plugin.iam.DsqlTokenUtility;
 import software.amazon.jdbc.plugin.iam.IamTokenUtility;
 import software.amazon.jdbc.plugin.iam.LightRdsUtility;
 import software.amazon.jdbc.plugin.iam.RegularRdsUtility;
@@ -59,7 +60,7 @@ public static String getCacheKey(
     return String.format("%s:%s:%d:%s", region, hostname, port, user);
   }
 
-  public static IamTokenUtility getTokenUtility() {
+  public static IamTokenUtility getRdsTokenUtility() {
     try {
       // RegularRdsUtility requires AWS Java SDK RDS v2.x to be presented in classpath.
       Class.forName("software.amazon.awssdk.services.rds.RdsUtilities");
@@ -81,6 +82,15 @@ public static IamTokenUtility getTokenUtility() {
     }
   }
 
+  public static IamTokenUtility getDsqlTokenUtility() {
+    try {
+      Class.forName("software.amazon.awssdk.services.dsql.DsqlUtilities");
+      return new DsqlTokenUtility();
+    } catch (final ClassNotFoundException e) {
+      throw new RuntimeException(Messages.get("AuthenticationToken.javaDsqlSdkNotInClasspath"), e);
+    }
+  }
+
   public static String generateAuthenticationToken(
       final IamTokenUtility tokenUtils,
       final PluginService pluginService,
diff --git a/wrapper/src/main/java/software/amazon/jdbc/util/RdsUtils.java b/wrapper/src/main/java/software/amazon/jdbc/util/RdsUtils.java
@@ -144,6 +144,14 @@ public class RdsUtils {
               + "\\.(amazonaws\\.com\\.?|c2s\\.ic\\.gov\\.?|sc2s\\.sgov\\.gov\\.?))$",
           Pattern.CASE_INSENSITIVE);
 
+  private static final Pattern AURORA_DSQL_CLUSTER_PATTERN =
+      Pattern.compile(
+          "^(?<instance>[^.]+)\\."
+              + "(?<dns>dsql(?:-[^.]+)?)\\."
+              + "(?<domain>(?<region>[a-zA-Z0-9\\-]+)"
+              + "\\.on\\.aws\\.?)$",
+          Pattern.CASE_INSENSITIVE);
+
   private static final Pattern ELB_PATTERN =
       Pattern.compile(
           "^(?<instance>.+)\\.elb\\."
@@ -259,14 +267,25 @@ public String getRdsInstanceHostPattern(final String host) {
     return group == null ? "?" : "?." + group;
   }
 
+  public String getDsqlInstanceId(final String host) {
+    final String preparedHost = getPreparedHost(host);
+    if (StringUtils.isNullOrEmpty(preparedHost)) {
+      return null;
+    }
+
+    final Matcher matcher = cacheMatcher(preparedHost, AURORA_DSQL_CLUSTER_PATTERN);
+    return getRegexGroup(matcher, INSTANCE_GROUP);
+  }
+
   public String getRdsRegion(final String host) {
     final String preparedHost = getPreparedHost(host);
     if (StringUtils.isNullOrEmpty(preparedHost)) {
       return null;
     }
 
     final Matcher matcher = cacheMatcher(preparedHost,
-        AURORA_DNS_PATTERN, AURORA_CHINA_DNS_PATTERN, AURORA_OLD_CHINA_DNS_PATTERN, AURORA_GOV_DNS_PATTERN);
+        AURORA_DNS_PATTERN, AURORA_CHINA_DNS_PATTERN, AURORA_OLD_CHINA_DNS_PATTERN, AURORA_GOV_DNS_PATTERN,
+        AURORA_DSQL_CLUSTER_PATTERN);
     final String group = getRegexGroup(matcher, REGION_GROUP);
     if (group != null) {
       return group;
@@ -294,6 +313,11 @@ public boolean isLimitlessDbShardGroupDns(final String host) {
     return dnsGroup != null && dnsGroup.equalsIgnoreCase("shardgrp-");
   }
 
+  public boolean isDsqlCluster(final String host) {
+    final String instanceId = getDsqlInstanceId(host);
+    return instanceId != null;
+  }
+
   public String getRdsClusterHostUrl(final String host) {
     final String preparedHost = getPreparedHost(host);
     if (StringUtils.isNullOrEmpty(preparedHost)) {
diff --git a/wrapper/src/main/resources/aws_advanced_jdbc_wrapper_messages.properties b/wrapper/src/main/resources/aws_advanced_jdbc_wrapper_messages.properties
@@ -47,6 +47,7 @@ AwsSdk.unsupportedRegion=Unsupported AWS region ''{0}''. For supported regions p
 AwsSecretsManagerConnectionPlugin.endpointOverrideMisconfigured=The provided endpoint is invalid and could not be used to create a URI: `{0}`.
 AwsSecretsManagerConnectionPlugin.endpointOverrideInvalidConnection=A connection to the provided endpoint could not be established: `{0}`.
 AwsSecretsManagerConnectionPlugin.javaSdkNotInClasspath=Required dependency 'AWS Java SDK for AWS Secrets Manager' is not on the classpath.
+AuthenticationToken.javaDsqlSdkNotInClasspath=Required dependency 'AWS Java SDK for DSQL v2.x' is not on the classpath.
 AwsSecretsManagerConnectionPlugin.jacksonDatabindNotInClasspath=Required dependency 'Jackson Databind' is not on the classpath.
 AwsSecretsManagerConnectionPlugin.failedToFetchDbCredentials=Was not able to either fetch or read the database credentials from AWS Secrets Manager. Ensure the correct secretId and region properties have been provided.
 AwsSecretsManagerConnectionPlugin.missingRequiredConfigParameter=Configuration parameter ''{0}'' is required.
diff --git a/wrapper/src/test/java/software/amazon/jdbc/plugin/iam/DsqlIamConnectionPluginTest.java b/wrapper/src/test/java/software/amazon/jdbc/plugin/iam/DsqlIamConnectionPluginTest.java
diff --git a/wrapper/src/test/java/software/amazon/jdbc/plugin/iam/DsqlTokenUtilityTest.java b/wrapper/src/test/java/software/amazon/jdbc/plugin/iam/DsqlTokenUtilityTest.java
diff --git a/wrapper/src/test/java/software/amazon/jdbc/util/RdsUtilsTests.java b/wrapper/src/test/java/software/amazon/jdbc/util/RdsUtilsTests.java

Original file line number	Diff line number	Diff line change
`@@ -121,7 +121,7 @@ public Set<String> getSubscribedMethods() {`
`121`	`121`
`122`	`122`	`public FederatedAuthPlugin(final PluginService pluginService,`
`123`	`123`	`final CredentialsProviderFactory credentialsProviderFactory) {`
`124`		`- this(pluginService, credentialsProviderFactory, new RdsUtils(), IamAuthUtils.getTokenUtility());`
	`124`	`+ this(pluginService, credentialsProviderFactory, new RdsUtils(), IamAuthUtils.getRdsTokenUtility());`
`125`	`125`	`}`
`126`	`126`
`127`	`127`	`FederatedAuthPlugin(`
Original file line number	Diff line number	Diff line change
`@@ -95,7 +95,7 @@ public class OktaAuthPlugin extends AbstractConnectionPlugin {`
`95`	`95`	`private final TelemetryCounter fetchTokenCounter;`
`96`	`96`
`97`	`97`	`public OktaAuthPlugin(PluginService pluginService, CredentialsProviderFactory credentialsProviderFactory) {`
`98`		`- this(pluginService, credentialsProviderFactory, new RdsUtils(), IamAuthUtils.getTokenUtility());`
	`98`	`+ this(pluginService, credentialsProviderFactory, new RdsUtils(), IamAuthUtils.getRdsTokenUtility());`
`99`	`99`	`}`
`100`	`100`
`101`	`101`	`OktaAuthPlugin(`
Original file line number	Diff line number	Diff line change
`@@ -20,10 +20,11 @@`
`20`	`20`	`import software.amazon.jdbc.ConnectionPlugin;`
`21`	`21`	`import software.amazon.jdbc.ConnectionPluginFactory;`
`22`	`22`	`import software.amazon.jdbc.PluginService;`
	`23`	`+import software.amazon.jdbc.util.IamAuthUtils;`
`23`	`24`
`24`	`25`	`public class IamAuthConnectionPluginFactory implements ConnectionPluginFactory {`
`25`	`26`	`@Override`
`26`	`27`	`public ConnectionPlugin getInstance(final PluginService pluginService, final Properties props) {`
`27`		`- return new IamAuthConnectionPlugin(pluginService);`
	`28`	`+ return new IamAuthConnectionPlugin(pluginService, IamAuthUtils.getRdsTokenUtility());`
`28`	`29`	`}`
`29`	`30`	`}`