Skip to content

Commit a9f06ba

Browse files
joviegasjoviegas
authored
Merge pull request #3504 from aws/merge-dsql-customization
Merge dsql customization
2 parents fcaec42 + 1fcca71 commit a9f06ba

File tree

9 files changed

+811
-0
lines changed

9 files changed

+811
-0
lines changed

‎.changes/next-release/feature-AxdbFrontend-d6ce1e6.json‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,6 @@
1+
{
2+
"type": "feature",
3+
"category": "DSQL",
4+
"contributor": "APandher",
5+
"description": "Add IAM Token Generation Utility for DSQL"
6+
}

‎feature.metadata‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1 @@
1+
{"trebuchetFeatureArn":"arn:aws:trebuchet:::feature:v2:2a236723-7c96-48cd-8055-918b35a8308a","c2jModelsRevision":25,"messageId":1,"serviceId":"DSQL","serviceModule":"dsql"}

‎services/dsql/pom.xml‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -41,6 +41,11 @@
4141
</plugins>
4242
</build>
4343
<dependencies>
44+
<dependency>
45+
<groupId>nl.jqno.equalsverifier</groupId>
46+
<artifactId>equalsverifier</artifactId>
47+
<scope>test</scope>
48+
</dependency>
4449
<dependency>
4550
<groupId>software.amazon.awssdk</groupId>
4651
<artifactId>protocol-core</artifactId>

‎services/dsql/src/main/java/software/amazon/awssdk/services/dsql/DefaultDsqlUtilities.java‎

Lines changed: 166 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,166 @@
1+
/*
2+
* Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License").
5+
* You may not use this file except in compliance with the License.
6+
* A copy of the License is located at
7+
*
8+
* http://aws.amazon.com/apache2.0
9+
*
10+
* or in the "license" file accompanying this file. This file is distributed
11+
* on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
12+
* express or implied. See the License for the specific language governing
13+
* permissions and limitations under the License.
14+
*/
15+
16+
package software.amazon.awssdk.services.dsql;
17+
18+
import java.time.Clock;
19+
import java.time.Instant;
20+
import software.amazon.awssdk.annotations.Immutable;
21+
import software.amazon.awssdk.annotations.SdkInternalApi;
22+
import software.amazon.awssdk.annotations.SdkTestInternalApi;
23+
import software.amazon.awssdk.auth.credentials.AwsCredentials;
24+
import software.amazon.awssdk.auth.credentials.CredentialUtils;
25+
import software.amazon.awssdk.auth.signer.Aws4Signer;
26+
import software.amazon.awssdk.auth.signer.params.Aws4PresignerParams;
27+
import software.amazon.awssdk.awscore.client.config.AwsClientOption;
28+
import software.amazon.awssdk.core.client.config.SdkClientConfiguration;
29+
import software.amazon.awssdk.http.SdkHttpFullRequest;
30+
import software.amazon.awssdk.http.SdkHttpMethod;
31+
import software.amazon.awssdk.identity.spi.AwsCredentialsIdentity;
32+
import software.amazon.awssdk.identity.spi.IdentityProvider;
33+
import software.amazon.awssdk.regions.Region;
34+
import software.amazon.awssdk.services.dsql.model.GenerateAuthTokenRequest;
35+
import software.amazon.awssdk.utils.CompletableFutureUtils;
36+
import software.amazon.awssdk.utils.Logger;
37+
import software.amazon.awssdk.utils.StringUtils;
38+
39+
@Immutable
40+
@SdkInternalApi
41+
public final class DefaultDsqlUtilities implements DsqlUtilities {
42+
private static final Logger log = Logger.loggerFor(DsqlUtilities.class);
43+
private final Aws4Signer signer = Aws4Signer.create();
44+
private final Region region;
45+
private final IdentityProvider<? extends AwsCredentialsIdentity> credentialsProvider;
46+
private final Clock clock;
47+
48+
public DefaultDsqlUtilities(DefaultBuilder builder) {
49+
this(builder, Clock.systemUTC());
50+
}
51+
52+
/**
53+
* For testing purposes only
54+
*/
55+
@SdkTestInternalApi
56+
public DefaultDsqlUtilities(DefaultBuilder builder, Clock clock) {
57+
this.credentialsProvider = builder.credentialsProvider;
58+
this.region = builder.region;
59+
this.clock = clock;
60+
}
61+
62+
/**
63+
* Used by DSQL low-level client's utilities() method
64+
*/
65+
@SdkInternalApi
66+
static DsqlUtilities create(SdkClientConfiguration clientConfiguration) {
67+
return new DefaultBuilder().clientConfiguration(clientConfiguration).build();
68+
}
69+
70+
@Override
71+
public String generateDbConnectAuthToken(GenerateAuthTokenRequest request) {
72+
return generateAuthToken(request, false);
73+
}
74+
75+
@Override
76+
public String generateDbConnectAdminAuthToken(GenerateAuthTokenRequest request) {
77+
return generateAuthToken(request, true);
78+
}
79+
80+
private String generateAuthToken(GenerateAuthTokenRequest request, boolean isAdmin) {
81+
String action = isAdmin ? "DbConnectAdmin" : "DbConnect";
82+
83+
SdkHttpFullRequest httpRequest = SdkHttpFullRequest.builder()
84+
.method(SdkHttpMethod.GET)
85+
.protocol("https")
86+
.host(request.hostname())
87+
.encodedPath("/")
88+
.putRawQueryParameter("Action", action)
89+
.build();
90+
91+
Instant expirationTime = Instant.now(clock).plus(request.expiresIn());
92+
93+
Aws4PresignerParams presignRequest = Aws4PresignerParams.builder()
94+
.signingClockOverride(clock)
95+
.expirationTime(expirationTime)
96+
.awsCredentials(resolveCredentials(request))
97+
.signingName("dsql")
98+
.signingRegion(resolveRegion(request))
99+
.build();
100+
101+
SdkHttpFullRequest fullRequest = signer.presign(httpRequest, presignRequest);
102+
String signedUrl = fullRequest.getUri().toString();
103+
104+
log.debug(() -> "Generated DSQL authentication token with expiration of " + expirationTime);
105+
return StringUtils.replacePrefixIgnoreCase(signedUrl, "https://", "");
106+
}
107+
108+
private Region resolveRegion(GenerateAuthTokenRequest request) {
109+
if (request.region() != null) {
110+
return request.region();
111+
}
112+
113+
if (this.region != null) {
114+
return this.region;
115+
}
116+
117+
throw new IllegalArgumentException("Region must be provided in GenerateAuthTokenRequest or DsqlUtilities");
118+
}
119+
120+
// TODO: update this to use AwsCredentialsIdentity when we migrate Signers to accept the new type.
121+
private AwsCredentials resolveCredentials(GenerateAuthTokenRequest request) {
122+
if (request.credentialsIdentityProvider() != null) {
123+
return CredentialUtils.toCredentials(
124+
CompletableFutureUtils.joinLikeSync(request.credentialsIdentityProvider().resolveIdentity()));
125+
}
126+
127+
if (this.credentialsProvider != null) {
128+
return CredentialUtils.toCredentials(CompletableFutureUtils.joinLikeSync(this.credentialsProvider.resolveIdentity()));
129+
}
130+
131+
throw new IllegalArgumentException("CredentialsProvider must be provided in GenerateAuthTokenRequest or DsqlUtilities");
132+
}
133+
134+
@SdkInternalApi
135+
public static final class DefaultBuilder implements DsqlUtilities.Builder {
136+
private Region region;
137+
private IdentityProvider<? extends AwsCredentialsIdentity> credentialsProvider;
138+
139+
Builder clientConfiguration(SdkClientConfiguration clientConfiguration) {
140+
this.credentialsProvider = clientConfiguration.option(AwsClientOption.CREDENTIALS_IDENTITY_PROVIDER);
141+
this.region = clientConfiguration.option(AwsClientOption.AWS_REGION);
142+
143+
return this;
144+
}
145+
146+
@Override
147+
public Builder region(Region region) {
148+
this.region = region;
149+
return this;
150+
}
151+
152+
@Override
153+
public Builder credentialsProvider(IdentityProvider<? extends AwsCredentialsIdentity> credentialsProvider) {
154+
this.credentialsProvider = credentialsProvider;
155+
return this;
156+
}
157+
158+
/**
159+
* Construct a {@link DsqlUtilities} object.
160+
*/
161+
@Override
162+
public DsqlUtilities build() {
163+
return new DefaultDsqlUtilities(this);
164+
}
165+
}
166+
}

‎services/dsql/src/main/java/software/amazon/awssdk/services/dsql/DsqlUtilities.java‎

Lines changed: 143 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,143 @@
1+
/*
2+
* Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License").
5+
* You may not use this file except in compliance with the License.
6+
* A copy of the License is located at
7+
*
8+
* http://aws.amazon.com/apache2.0
9+
*
10+
* or in the "license" file accompanying this file. This file is distributed
11+
* on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
12+
* express or implied. See the License for the specific language governing
13+
* permissions and limitations under the License.
14+
*/
15+
16+
package software.amazon.awssdk.services.dsql;
17+
18+
import java.util.function.Consumer;
19+
import software.amazon.awssdk.annotations.SdkPublicApi;
20+
import software.amazon.awssdk.auth.credentials.AwsCredentialsProvider;
21+
import software.amazon.awssdk.identity.spi.AwsCredentialsIdentity;
22+
import software.amazon.awssdk.identity.spi.IdentityProvider;
23+
import software.amazon.awssdk.regions.Region;
24+
import software.amazon.awssdk.services.dsql.model.GenerateAuthTokenRequest;
25+
26+
/**
27+
* Utilities for working with DSQL. An instance of this class can be created by:
28+
* <p>
29+
* 1) Using the low-level client {@link DsqlClient#utilities()} (or {@link DsqlAsyncClient#utilities()}} method.
30+
* This is recommended as SDK will use the same configuration from the {@link DsqlClient} object to create the
31+
* {@link DsqlUtilities} object.
32+
*
33+
* @snippet :
34+
* {@code
35+
* DsqlClient dsqlClient = DsqlClient.create();
36+
* DsqlUtilities utilities = DsqlClient.utilities();
37+
* }
38+
*
39+
* <p>
40+
* 2) Directly using the {@link #builder()} method.
41+
*
42+
* @snippet :
43+
* {@code
44+
* DsqlUtilities utilities = DsqlUtilities.builder()
45+
* .credentialsProvider(DefaultCredentialsProvider.create())
46+
* .region(Region.US_WEST_2)
47+
* .build()
48+
* }
49+
*
50+
* Note: This class does not make network calls.
51+
*/
52+
@SdkPublicApi
53+
public interface DsqlUtilities {
54+
/**
55+
* Create a builder that can be used to configure and create a {@link DsqlUtilities}.
56+
*/
57+
static Builder builder() {
58+
return new DefaultDsqlUtilities.DefaultBuilder();
59+
}
60+
61+
/**
62+
* Generates an authentication token for IAM authentication to an DSQL database.
63+
*
64+
* @param request The request used to generate the authentication token
65+
* @return String to use as the DSQL authentication token
66+
* @throws IllegalArgumentException if the required parameters are not valid
67+
*/
68+
default String generateDbConnectAuthToken(Consumer<GenerateAuthTokenRequest.Builder> request) {
69+
return generateDbConnectAuthToken(GenerateAuthTokenRequest.builder().applyMutation(request).build());
70+
}
71+
72+
/**
73+
* Generates an authentication token for IAM authentication to an DSQL database.
74+
*
75+
* @param request The request used to generate the authentication token
76+
* @return String to use as the DSQL authentication token
77+
* @throws IllegalArgumentException if the required parameters are not valid
78+
*/
79+
default String generateDbConnectAuthToken(GenerateAuthTokenRequest request) {
80+
throw new UnsupportedOperationException();
81+
}
82+
83+
/**
84+
* Generates an admin authentication token for IAM authentication to an DSQL database.
85+
*
86+
* @param request The request used to generate the admin authentication token
87+
* @return String to use as the DSQL authentication token
88+
* @throws IllegalArgumentException if the required parameters are not valid
89+
*/
90+
default String generateDbConnectAdminAuthToken(Consumer<GenerateAuthTokenRequest.Builder> request) {
91+
return generateDbConnectAdminAuthToken(GenerateAuthTokenRequest.builder().applyMutation(request).build());
92+
}
93+
94+
/**
95+
* Generates an admin authentication token for IAM authentication to an DSQL database.
96+
*
97+
* @param request The request used to generate the admin authentication token
98+
* @return String to use as the DSQL authentication token
99+
* @throws IllegalArgumentException if the required parameters are not valid
100+
*/
101+
default String generateDbConnectAdminAuthToken(GenerateAuthTokenRequest request) {
102+
throw new UnsupportedOperationException();
103+
}
104+
105+
106+
/**
107+
* Builder for creating an instance of {@link DsqlUtilities}. It can be configured using
108+
* {@link DsqlUtilities#builder()}.
109+
* Once configured, the {@link DsqlUtilities} can created using {@link #build()}.
110+
*/
111+
@SdkPublicApi
112+
interface Builder {
113+
/**
114+
* The default region to use when working with the methods in {@link DsqlUtilities} class.
115+
*
116+
* @return This object for method chaining
117+
*/
118+
Builder region(Region region);
119+
120+
/**
121+
* The default credentials provider to use when working with the methods in {@link DsqlUtilities} class.
122+
*
123+
* @return This object for method chaining
124+
*/
125+
default Builder credentialsProvider(AwsCredentialsProvider credentialsProvider) {
126+
return credentialsProvider((IdentityProvider<? extends AwsCredentialsIdentity>) credentialsProvider);
127+
}
128+
129+
/**
130+
* The default credentials provider to use when working with the methods in {@link DsqlUtilities} class.
131+
*
132+
* @return This object for method chaining
133+
*/
134+
default Builder credentialsProvider(IdentityProvider<? extends AwsCredentialsIdentity> credentialsProvider) {
135+
throw new UnsupportedOperationException();
136+
}
137+
138+
/**
139+
* Create a {@link DsqlUtilities}
140+
*/
141+
DsqlUtilities build();
142+
}
143+
}

0 commit comments

Comments
 (0)