Skip to content

Commit 64d1fff

Browse files
Jithendar12Jithendar12
committed
Address Review Comments
1 parent caa8e0b commit 64d1fff

File tree

2 files changed

+128
-26
lines changed

2 files changed

+128
-26
lines changed

‎athena-synapse/src/main/java/com/amazonaws/athena/connectors/synapse/SynapseCredentialsProvider.java‎

Lines changed: 124 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -23,14 +23,18 @@
2323
import com.amazonaws.athena.connector.credentials.DefaultCredentials;
2424
import com.amazonaws.athena.connector.lambda.exceptions.AthenaConnectorException;
2525
import com.amazonaws.athena.connector.lambda.security.CachableSecretsManager;
26+
import com.fasterxml.jackson.core.JsonProcessingException;
2627
import com.fasterxml.jackson.databind.JsonNode;
2728
import com.fasterxml.jackson.databind.ObjectMapper;
2829
import com.fasterxml.jackson.databind.node.ObjectNode;
2930
import software.amazon.awssdk.services.glue.model.ErrorDetails;
3031
import software.amazon.awssdk.services.glue.model.FederationSourceErrorCode;
3132
import software.amazon.awssdk.services.secretsmanager.SecretsManagerClient;
3233
import software.amazon.awssdk.services.secretsmanager.model.PutSecretValueRequest;
34+
import software.amazon.awssdk.services.secretsmanager.model.ResourceNotFoundException;
35+
import software.amazon.awssdk.services.secretsmanager.model.SecretsManagerException;
3336

37+
import java.io.IOException;
3438
import java.net.URI;
3539
import java.net.http.HttpClient;
3640
import java.net.http.HttpRequest;
@@ -42,14 +46,17 @@
4246

4347
public class SynapseCredentialsProvider implements CredentialsProvider
4448
{
49+
// Constants for OAuth token response fields
4550
private static final String ACCESS_TOKEN = "access_token";
4651
private static final String FETCHED_AT = "fetched_at";
4752
private static final String EXPIRES_IN = "expires_in";
4853

54+
// Constants for OAuth configuration fields
4955
private static final String CLIENT_ID = "client_id";
5056
private static final String CLIENT_SECRET = "client_secret";
5157
private static final String TENANT_ID = "tenant_id";
5258

59+
// Constants for basic authentication fields
5360
private static final String USER = "user";
5461
private static final String PASSWORD = "password";
5562
private static final String USERNAME = "username";
@@ -96,9 +103,25 @@ public Map<String, String> getCredentialMap()
96103

97104
return credentialMap;
98105
}
99-
catch (Exception e) {
106+
catch (JsonProcessingException e) {
100107
throw new AthenaConnectorException(
101-
"Failed to retrieve Synapse credentials",
108+
"Failed to parse Synapse credentials JSON",
109+
ErrorDetails.builder()
110+
.errorCode(FederationSourceErrorCode.INVALID_RESPONSE_EXCEPTION.toString())
111+
.build()
112+
);
113+
}
114+
catch (ResourceNotFoundException e) {
115+
throw new AthenaConnectorException(
116+
"Secret not found: " + secretName,
117+
ErrorDetails.builder()
118+
.errorCode(FederationSourceErrorCode.ENTITY_NOT_FOUND_EXCEPTION.toString())
119+
.build()
120+
);
121+
}
122+
catch (SecretsManagerException e) {
123+
throw new AthenaConnectorException(
124+
"Failed to retrieve Synapse credentials from Secrets Manager",
102125
ErrorDetails.builder()
103126
.errorCode(FederationSourceErrorCode.INTERNAL_SERVICE_EXCEPTION.toString())
104127
.build()
@@ -117,14 +140,47 @@ public String getOAuthAccessToken()
117140
}
118141
return null;
119142
}
120-
catch (Exception e) {
143+
catch (JsonProcessingException e) {
144+
throw new AthenaConnectorException(
145+
"Failed to parse OAuth configuration JSON",
146+
ErrorDetails.builder()
147+
.errorCode(FederationSourceErrorCode.INVALID_RESPONSE_EXCEPTION.toString())
148+
.build()
149+
);
150+
}
151+
catch (ResourceNotFoundException e) {
121152
throw new AthenaConnectorException(
122-
"Failed to get OAuth access token",
153+
"Secret not found: " + secretName,
154+
ErrorDetails.builder()
155+
.errorCode(FederationSourceErrorCode.ENTITY_NOT_FOUND_EXCEPTION.toString())
156+
.build()
157+
);
158+
}
159+
catch (SecretsManagerException e) {
160+
throw new AthenaConnectorException(
161+
"Failed to retrieve OAuth configuration from Secrets Manager",
123162
ErrorDetails.builder()
124163
.errorCode(FederationSourceErrorCode.INTERNAL_SERVICE_EXCEPTION.toString())
125164
.build()
126165
);
127166
}
167+
catch (IOException e) {
168+
throw new AthenaConnectorException(
169+
"Failed to communicate with OAuth endpoint",
170+
ErrorDetails.builder()
171+
.errorCode(FederationSourceErrorCode.OPERATION_TIMEOUT_EXCEPTION.toString())
172+
.build()
173+
);
174+
}
175+
catch (InterruptedException e) {
176+
Thread.currentThread().interrupt();
177+
throw new AthenaConnectorException(
178+
"OAuth token fetch interrupted",
179+
ErrorDetails.builder()
180+
.errorCode(FederationSourceErrorCode.OPERATION_TIMEOUT_EXCEPTION.toString())
181+
.build()
182+
);
183+
}
128184
}
129185

130186
private boolean isOAuthConfigured(Map<String, String> oauthConfig)
@@ -137,7 +193,7 @@ private boolean isOAuthConfigured(Map<String, String> oauthConfig)
137193
!oauthConfig.get(TENANT_ID).isEmpty();
138194
}
139195

140-
private String fetchAccessToken(Map<String, String> oauthConfig) throws Exception
196+
private String fetchAccessToken(Map<String, String> oauthConfig) throws IOException, InterruptedException
141197
{
142198
String accessToken = oauthConfig.get(ACCESS_TOKEN);
143199

@@ -156,7 +212,7 @@ private String fetchAccessToken(Map<String, String> oauthConfig) throws Exceptio
156212
return fetchAndStoreNewToken(oauthConfig);
157213
}
158214

159-
private String fetchAndStoreNewToken(Map<String, String> oauthConfig) throws Exception
215+
private String fetchAndStoreNewToken(Map<String, String> oauthConfig) throws IOException, InterruptedException
160216
{
161217
String clientId = oauthConfig.get(CLIENT_ID);
162218
String clientSecret = oauthConfig.get(CLIENT_SECRET);
@@ -179,33 +235,78 @@ private String fetchAndStoreNewToken(Map<String, String> oauthConfig) throws Exc
179235
HttpResponse<String> response = client.send(request, HttpResponse.BodyHandlers.ofString());
180236

181237
if (response.statusCode() != 200) {
238+
String errorMessage = String.format("Failed to fetch access token. Status code: %d", response.statusCode());
239+
if (response.statusCode() == 401 || response.statusCode() == 403) {
240+
throw new AthenaConnectorException(
241+
errorMessage,
242+
ErrorDetails.builder()
243+
.errorCode(FederationSourceErrorCode.INVALID_CREDENTIALS_EXCEPTION.toString())
244+
.build()
245+
);
246+
}
247+
if (response.statusCode() == 429) {
248+
throw new AthenaConnectorException(
249+
errorMessage,
250+
ErrorDetails.builder()
251+
.errorCode(FederationSourceErrorCode.THROTTLING_EXCEPTION.toString())
252+
.build()
253+
);
254+
}
182255
throw new AthenaConnectorException(
183-
"Failed to fetch access token",
256+
errorMessage,
184257
ErrorDetails.builder()
185258
.errorCode(FederationSourceErrorCode.INTERNAL_SERVICE_EXCEPTION.toString())
186259
.build()
187260
);
188261
}
189262

190-
JsonNode tokenResponse = objectMapper.readTree(response.body());
191-
String accessToken = tokenResponse.get(ACCESS_TOKEN).asText();
192-
long expiresIn = tokenResponse.get(EXPIRES_IN).asLong();
193-
long fetchedAt = Instant.now().getEpochSecond();
263+
try {
264+
JsonNode tokenResponse = objectMapper.readTree(response.body());
265+
String accessToken = tokenResponse.get(ACCESS_TOKEN).asText();
266+
long expiresIn = tokenResponse.get(EXPIRES_IN).asLong();
267+
long fetchedAt = Instant.now().getEpochSecond();
194268

195-
oauthConfig.put(ACCESS_TOKEN, accessToken);
196-
oauthConfig.put(EXPIRES_IN, String.valueOf(expiresIn));
197-
oauthConfig.put(FETCHED_AT, String.valueOf(fetchedAt));
269+
oauthConfig.put(ACCESS_TOKEN, accessToken);
270+
oauthConfig.put(EXPIRES_IN, String.valueOf(expiresIn));
271+
oauthConfig.put(FETCHED_AT, String.valueOf(fetchedAt));
198272

199-
ObjectNode updatedSecretJson = objectMapper.createObjectNode();
200-
for (Map.Entry<String, String> entry : oauthConfig.entrySet()) {
201-
updatedSecretJson.put(entry.getKey(), entry.getValue());
202-
}
273+
ObjectNode updatedSecretJson = objectMapper.createObjectNode();
274+
for (Map.Entry<String, String> entry : oauthConfig.entrySet()) {
275+
updatedSecretJson.put(entry.getKey(), entry.getValue());
276+
}
203277

204-
secretsManager.getSecretsManager().putSecretValue(PutSecretValueRequest.builder()
205-
.secretId(secretName)
206-
.secretString(updatedSecretJson.toString())
207-
.build());
278+
try {
279+
secretsManager.getSecretsManager().putSecretValue(PutSecretValueRequest.builder()
280+
.secretId(secretName)
281+
.secretString(updatedSecretJson.toString())
282+
.build());
283+
}
284+
catch (ResourceNotFoundException e) {
285+
throw new AthenaConnectorException(
286+
"Secret not found: " + secretName,
287+
ErrorDetails.builder()
288+
.errorCode(FederationSourceErrorCode.ENTITY_NOT_FOUND_EXCEPTION.toString())
289+
.build()
290+
);
291+
}
292+
catch (SecretsManagerException e) {
293+
throw new AthenaConnectorException(
294+
"Failed to update access token in Secrets Manager",
295+
ErrorDetails.builder()
296+
.errorCode(FederationSourceErrorCode.INTERNAL_SERVICE_EXCEPTION.toString())
297+
.build()
298+
);
299+
}
208300

209-
return accessToken;
301+
return accessToken;
302+
}
303+
catch (JsonProcessingException e) {
304+
throw new AthenaConnectorException(
305+
"Failed to parse OAuth token response",
306+
ErrorDetails.builder()
307+
.errorCode(FederationSourceErrorCode.INVALID_RESPONSE_EXCEPTION.toString())
308+
.build()
309+
);
310+
}
210311
}
211312
}

‎athena-synapse/src/test/java/com/amazonaws/athena/connectors/synapse/SynapseCredentialsProviderTest.java‎

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -32,6 +32,7 @@
3232
import software.amazon.awssdk.services.secretsmanager.SecretsManagerClient;
3333
import software.amazon.awssdk.services.secretsmanager.model.GetSecretValueRequest;
3434
import software.amazon.awssdk.services.secretsmanager.model.GetSecretValueResponse;
35+
import software.amazon.awssdk.services.secretsmanager.model.ResourceNotFoundException;
3536

3637
import java.net.http.HttpClient;
3738
import java.net.http.HttpRequest;
@@ -194,7 +195,7 @@ public void testGetOAuthAccessTokenWithInvalidJson() throws Exception
194195
fail("Expected AthenaConnectorException");
195196
}
196197
catch (AthenaConnectorException e) {
197-
assertEquals(FederationSourceErrorCode.INTERNAL_SERVICE_EXCEPTION.toString(),
198+
assertEquals(FederationSourceErrorCode.INVALID_RESPONSE_EXCEPTION.toString(),
198199
e.getErrorDetails().errorCode());
199200
}
200201
}
@@ -204,12 +205,12 @@ public void testGetCredentialMapWithInvalidSecret()
204205
{
205206
try {
206207
when(mockSecretsClient.getSecretValue(any(GetSecretValueRequest.class)))
207-
.thenThrow(new RuntimeException("Secret not found"));
208+
.thenThrow(ResourceNotFoundException.builder().message("Secret not found").build());
208209
credentialsProvider.getCredentialMap();
209210
fail("Expected AthenaConnectorException");
210211
}
211212
catch (AthenaConnectorException e) {
212-
assertEquals(FederationSourceErrorCode.INTERNAL_SERVICE_EXCEPTION.toString(),
213+
assertEquals(FederationSourceErrorCode.ENTITY_NOT_FOUND_EXCEPTION.toString(),
213214
e.getErrorDetails().errorCode());
214215
}
215216
}

0 commit comments

Comments
 (0)