Support Snowflake OAuth authentication for Athena Snowflake connector

athena-federation-sdk/src/main/java/com/amazonaws/athena/connector/credentials/CredentialsProvider.java

@@ -19,15 +19,47 @@
  */
 package com.amazonaws.athena.connector.credentials;
 
+import java.util.Map;
+
 /**
  * JDBC username and password provider.
  */
 public interface CredentialsProvider
 {
+    String USER = "user";
+    String PASSWORD = "password";
+
     /**
      * Retrieves credential for database.
      *
      * @return JDBC credential. See {@link DefaultCredentials}.
      */
     DefaultCredentials getCredential();
+
+    /**
+     * Retrieves credential properties as a map for database connection.
+     *
+     * Default Behavior:
+     * The default implementation returns a map containing only the basic "user" and "password" 
+     * properties extracted from the {@link DefaultCredentials} object returned by {@link #getCredential()}.
+     * This maintains backward compatibility with existing JDBC connection patterns.
+     *
+     * Extended Behavior:
+     * Implementations can override this method to provide additional connection properties beyond 
+     * just username and password. This enables support for advanced authentication mechanisms.
+     *
+     * Usage:
+     * The returned map is directly applied to JDBC connection properties, allowing for seamless 
+     * integration with various database drivers and authentication schemes without requiring 
+     * custom connection factory implementations.
+     *
+     * @return Map containing credential properties for database connection. The default implementation
+     *         returns a map with "user" and "password" keys. Overriding implementations may return
+     *         additional properties as needed for their specific authentication requirements.
+     */
+    default Map<String, String> getCredentialMap()
+    {
+        DefaultCredentials credential = getCredential();
+        return Map.of(USER, credential.getUser(), PASSWORD, credential.getPassword());
+    }
 }

athena-federation-sdk/src/main/java/com/amazonaws/athena/connector/lambda/security/CachableSecretsManager.java

@@ -61,6 +61,16 @@ public CachableSecretsManager(SecretsManagerClient secretsManager)
         this.secretsManager = secretsManager;
     }
 
+    /**
+     * Gets the underlying SecretsManagerClient instance.
+     *
+     * @return The SecretsManagerClient instance.
+     */
+    public SecretsManagerClient getSecretsManager()
+    {
+        return secretsManager;
+    }
+
     /**
      * Resolves any secrets found in the supplied string, for example: MyString${WithSecret} would have ${WithSecret}
      * repalced by the corresponding value of the secret in AWS Secrets Manager with that name. If no such secret is found

athena-jdbc/src/main/java/com/amazonaws/athena/connectors/jdbc/connection/GenericJdbcConnectionFactory.java

@@ -83,8 +83,7 @@ public Connection getConnection(final CredentialsProvider credentialsProvider)
             Matcher secretMatcher = SECRET_NAME_PATTERN.matcher(databaseConnectionConfig.getJdbcConnectionString());
             derivedJdbcString = secretMatcher.replaceAll(Matcher.quoteReplacement(""));
 
-            jdbcProperties.put("user", credentialsProvider.getCredential().getUser());
-            jdbcProperties.put("password", credentialsProvider.getCredential().getPassword());
+            jdbcProperties.putAll(credentialsProvider.getCredentialMap());
         }
         else {
             derivedJdbcString = databaseConnectionConfig.getJdbcConnectionString();

athena-jdbc/src/main/java/com/amazonaws/athena/connectors/jdbc/manager/JdbcMetadataHandler.java

@@ -168,6 +168,11 @@ protected JdbcConnectionFactory getJdbcConnectionFactory()
         return jdbcConnectionFactory;
     }
 
+    protected DatabaseConnectionConfig getDatabaseConnectionConfig()
+    {
+        return databaseConnectionConfig;
+    }
+
     protected CredentialsProvider getCredentialProvider()
     {
         final String secretName = databaseConnectionConfig.getSecret();

athena-jdbc/src/main/java/com/amazonaws/athena/connectors/jdbc/manager/JdbcRecordHandler.java

@@ -132,6 +132,11 @@ protected JdbcConnectionFactory getJdbcConnectionFactory()
         return jdbcConnectionFactory;
     }
 
+    protected DatabaseConnectionConfig getDatabaseConnectionConfig()
+    {
+        return databaseConnectionConfig;
+    }
+
     protected CredentialsProvider getCredentialProvider()
     {
         final String secretName = this.databaseConnectionConfig.getSecret();

athena-snowflake/athena-snowflake-connection.yaml

@@ -102,6 +102,7 @@ Resources:
         Statement:
           - Action:
               - secretsmanager:GetSecretValue
+              - secretsmanager:PutSecretValue
             Effect: Allow
             Resource: !Sub 'arn:${AWS::Partition}:secretsmanager:${AWS::Region}:${AWS::AccountId}:secret:${SecretName}*'
           - Action:

athena-snowflake/athena-snowflake.yaml

@@ -85,6 +85,7 @@ Resources:
         - Statement:
             - Action:
                 - secretsmanager:GetSecretValue
+                - secretsmanager:PutSecretValue
               Effect: Allow
               Resource: !Sub 'arn:${AWS::Partition}:secretsmanager:${AWS::Region}:${AWS::AccountId}:secret:${SecretNamePrefix}*'
           Version: '2012-10-17'

athena-snowflake/pom.xml

@@ -32,6 +32,12 @@
             <artifactId>snowflake-jdbc</artifactId>
             <version>3.24.2</version>
         </dependency>
+        <!-- https://mvnrepository.com/artifact/org.json/json -->
+        <dependency>
+            <groupId>org.json</groupId>
+            <artifactId>json</artifactId>
+            <version>20250107</version>
+        </dependency>
         <!-- https://mvnrepository.com/artifact/software.amazon.awssdk/rds -->
         <dependency>
             <groupId>software.amazon.awssdk</groupId>

athena-snowflake/src/main/java/com/amazonaws/athena/connectors/snowflake/SnowflakeConstants.java

@@ -35,6 +35,11 @@ public final class SnowflakeConstants
      */
     public static final int SINGLE_SPLIT_LIMIT_COUNT = 10000;
     public static final String SNOWFLAKE_QUOTE_CHARACTER = "\"";
+    public static final String AUTH_CODE = "auth_code";
+    public static final String CLIENT_ID = "client_id";
+    public static final String TOKEN_URL = "token_url";
+    public static final String REDIRECT_URI = "redirect_uri";
+    public static final String CLIENT_SECRET = "client_secret";
 
     private SnowflakeConstants() {}
 }

athena-snowflake/src/main/java/com/amazonaws/athena/connectors/snowflake/SnowflakeCredentialsProvider.java

@@ -0,0 +1,237 @@
+/*-
+ * #%L
+ * athena-snowflake
+ * %%
+ * Copyright (C) 2019 - 2025 Amazon Web Services
+ * %%
+ * 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.
+ * #L%
+ */
+package com.amazonaws.athena.connectors.snowflake;
+
+import com.amazonaws.athena.connector.credentials.CredentialsProvider;
+import com.amazonaws.athena.connector.credentials.DefaultCredentials;
+import com.amazonaws.athena.connector.lambda.security.CachableSecretsManager;
+import com.fasterxml.jackson.databind.ObjectMapper;
+import com.google.common.annotations.VisibleForTesting;
+import org.json.JSONObject;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+import software.amazon.awssdk.services.secretsmanager.SecretsManagerClient;
+import software.amazon.awssdk.utils.Validate;
+
+import java.io.BufferedReader;
+import java.io.IOException;
+import java.io.InputStream;
+import java.io.InputStreamReader;
+import java.io.OutputStream;
+import java.net.HttpURLConnection;
+import java.net.URL;
+import java.net.URLEncoder;
+import java.nio.charset.StandardCharsets;
+import java.util.Base64;
+import java.util.HashMap;
+import java.util.Map;
+
+/**
+ * Snowflake OAuth credentials provider that manages OAuth token lifecycle.
+ * This provider handles token refresh, expiration, and provides credential properties
+ * for Snowflake OAuth connections.
+ */
+public class SnowflakeCredentialsProvider implements CredentialsProvider
+{
+    private static final Logger LOGGER = LoggerFactory.getLogger(SnowflakeCredentialsProvider.class);
+
+    public static final String ACCESS_TOKEN = "access_token";
+    public static final String FETCHED_AT = "fetched_at";
+    public static final String REFRESH_TOKEN = "refresh_token";
+    public static final String EXPIRES_IN = "expires_in";
+    public static final String USERNAME = "username";
+    public static final String PASSWORD = "password";
+    public static final String USER = "user";
+
+    private final String oauthSecretName;
+    private final CachableSecretsManager secretsManager;
+    private final ObjectMapper objectMapper;
+
+    public SnowflakeCredentialsProvider(String oauthSecretName)
+    {
+        this(oauthSecretName, SecretsManagerClient.create());
+    }
+
+    @VisibleForTesting
+    public SnowflakeCredentialsProvider(String oauthSecretName, SecretsManagerClient secretsClient)
+    {
+        this.oauthSecretName = Validate.notNull(oauthSecretName, "oauthSecretName must not be null");
+        this.secretsManager = new CachableSecretsManager(secretsClient);
+        this.objectMapper = new ObjectMapper();
+    }
+
+    @Override
+    public DefaultCredentials getCredential()
+    {
+        Map<String, String> credentialMap = getCredentialMap();
+        return new DefaultCredentials(
+            credentialMap.get(USER),
+            credentialMap.get(PASSWORD)
+        );
+    }
+
+    @Override
+    public Map<String, String> getCredentialMap()
+    {
+        try {
+            String secretString = secretsManager.getSecret(oauthSecretName);
+            Map<String, String> oauthConfig = objectMapper.readValue(secretString, Map.class);
+
+            if (oauthConfig.containsKey(SnowflakeConstants.AUTH_CODE) && !oauthConfig.get(SnowflakeConstants.AUTH_CODE).isEmpty()) {
+                // OAuth flow
+                String accessToken = fetchAccessTokenFromSecret(oauthConfig);
+
+                Map<String, String> credentialMap = new HashMap<>();
+                credentialMap.put(USER, oauthConfig.get(USERNAME));
+                credentialMap.put(PASSWORD, accessToken);
+                credentialMap.put("authenticator", "oauth");
+
+                return credentialMap;
+            }
+            else {
+                // Fallback to standard credentials
+                return Map.of(
+                        USER, oauthConfig.get(USERNAME),
+                        PASSWORD, oauthConfig.get(PASSWORD)
+                );
+            }
+        }
+        catch (Exception ex) {
+            throw new RuntimeException("Error retrieving Snowflake credentials: " + ex.getMessage(), ex);
+        }
+    }
+
+    private String loadTokenFromSecretsManager(Map<String, String> oauthConfig)
+    {
+        if (oauthConfig.containsKey(ACCESS_TOKEN)) {
+            return oauthConfig.get(ACCESS_TOKEN);
+        }
+        return null;
+    }
+
+    private void saveTokenToSecretsManager(JSONObject tokenJson, Map<String, String> oauthConfig)
+    {
+        // Update token related fields
+        tokenJson.put(FETCHED_AT, System.currentTimeMillis() / 1000);
+        oauthConfig.put(ACCESS_TOKEN, tokenJson.getString(ACCESS_TOKEN));
+        oauthConfig.put(REFRESH_TOKEN, tokenJson.getString(REFRESH_TOKEN));
+        oauthConfig.put(EXPIRES_IN, String.valueOf(tokenJson.getInt(EXPIRES_IN)));
+        oauthConfig.put(FETCHED_AT, String.valueOf(tokenJson.getLong(FETCHED_AT)));
+
+        // Save updated secret
+        secretsManager.getSecretsManager().putSecretValue(builder -> builder
+                .secretId(this.oauthSecretName)
+                .secretString(String.valueOf(new JSONObject(oauthConfig)))
+                .build());
+    }
+
+    private String fetchAccessTokenFromSecret(Map<String, String> oauthConfig) throws Exception
+    {
+        String accessToken;
+        String clientId = Validate.notNull(oauthConfig.get(SnowflakeConstants.CLIENT_ID), "Missing required property: client_id");
+        String tokenEndpoint = Validate.notNull(oauthConfig.get(SnowflakeConstants.TOKEN_URL), "Missing required property: token_url");
+        String redirectUri = Validate.notNull(oauthConfig.get(SnowflakeConstants.REDIRECT_URI), "Missing required property: redirect_uri");
+        String clientSecret = Validate.notNull(oauthConfig.get(SnowflakeConstants.CLIENT_SECRET), "Missing required property: client_secret");
+        String authCode = Validate.notNull(oauthConfig.get(SnowflakeConstants.AUTH_CODE), "Missing required property: auth_code");
+
+        accessToken = loadTokenFromSecretsManager(oauthConfig);
+
+        if (accessToken == null) {
+            LOGGER.debug("First time auth. Using authorization_code...");
+            JSONObject tokenJson = getTokenFromAuthCode(authCode, redirectUri, tokenEndpoint, clientId, clientSecret);
+            saveTokenToSecretsManager(tokenJson, oauthConfig);
+            accessToken = tokenJson.getString(ACCESS_TOKEN);
+        }
+        else {
+            long expiresIn = Long.parseLong(oauthConfig.get(EXPIRES_IN));
+            long fetchedAt = Long.parseLong(oauthConfig.getOrDefault(FETCHED_AT, String.valueOf(0L)));
+            long now = System.currentTimeMillis() / 1000;
+
+            if ((now - fetchedAt) < expiresIn - 60) {
+                LOGGER.debug("Access token still valid.");
+            }
+            else {
+                LOGGER.debug("Access token expired. Using refresh_token...");
+                JSONObject refreshed = refreshAccessToken(oauthConfig.get(REFRESH_TOKEN), tokenEndpoint, clientId, clientSecret);
+                refreshed.put(REFRESH_TOKEN, oauthConfig.get(REFRESH_TOKEN));
+                saveTokenToSecretsManager(refreshed, oauthConfig);
+                accessToken = refreshed.getString(ACCESS_TOKEN);
+            }
+        }
+        return accessToken;
+    }
+
+    private JSONObject getTokenFromAuthCode(String authCode, String redirectUri, String tokenEndpoint, String clientId, String clientSecret) throws Exception
+    {
+        String body = "grant_type=authorization_code"
+                + "&code=" + authCode
+                + "&redirect_uri=" + redirectUri;
+
+        return requestToken(body, tokenEndpoint, clientId, clientSecret);
+    }
+
+    private JSONObject refreshAccessToken(String refreshToken, String tokenEndpoint, String clientId, String clientSecret) throws Exception
+    {
+        String body = "grant_type=refresh_token"
+                + "&refresh_token=" + URLEncoder.encode(refreshToken, StandardCharsets.UTF_8);
+
+        return requestToken(body, tokenEndpoint, clientId, clientSecret);
+    }
+
+    private JSONObject requestToken(String requestBody, String tokenEndpoint, String clientId, String clientSecret) throws Exception
+    {
+        HttpURLConnection conn = getHttpURLConnection(tokenEndpoint, clientId, clientSecret);
+
+        try (OutputStream os = conn.getOutputStream()) {
+            os.write(requestBody.getBytes(StandardCharsets.UTF_8));
+        }
+
+        int responseCode = conn.getResponseCode();
+        InputStream is = (responseCode >= 200 && responseCode < 300) ?
+                conn.getInputStream() : conn.getErrorStream();
+
+        String response = new BufferedReader(new InputStreamReader(is))
+                .lines()
+                .reduce("", (acc, line) -> acc + line);
+
+        if (responseCode != 200) {
+            throw new RuntimeException("Failed: " + responseCode + " - " + response);
+        }
+
+        JSONObject tokenJson = new JSONObject(response);
+        tokenJson.put(FETCHED_AT, System.currentTimeMillis() / 1000);
+        return tokenJson;
+    }
+
+    static HttpURLConnection getHttpURLConnection(String tokenEndpoint, String clientId, String clientSecret) throws IOException
+    {
+        URL url = new URL(tokenEndpoint);
+        HttpURLConnection conn = (HttpURLConnection) url.openConnection();
+
+        String authHeader = Base64.getEncoder()
+                .encodeToString((clientId + ":" + clientSecret).getBytes(StandardCharsets.UTF_8));
+
+        conn.setRequestMethod("POST");
+        conn.setRequestProperty("Authorization", "Basic " + authHeader);
+        conn.setRequestProperty("Content-Type", "application/x-www-form-urlencoded");
+        conn.setDoOutput(true);
+        return conn;
+    }
+}