Add PAT support

snowflakedb · Dec 10, 2024 · c391329 · c391329
1 parent e8dc943
commit c391329
Show file tree

Hide file tree

Showing 3 changed files with 21 additions and 14 deletions.
diff --git a/src/main/java/net/snowflake/client/core/SessionUtil.java b/src/main/java/net/snowflake/client/core/SessionUtil.java
@@ -231,6 +231,8 @@ private static AuthenticatorType getAuthenticator(SFLoginInput loginInput) {
       } else if (loginInput.getAuthenticator().equalsIgnoreCase(AuthenticatorType.OAUTH.name())) {
         // OAuth access code Authentication
         return AuthenticatorType.OAUTH;
+      } else if (loginInput.getAuthenticator().equalsIgnoreCase(AuthenticatorType.PROGRAMMATIC_ACCESS_TOKEN.name())) {
+        return AuthenticatorType.PROGRAMMATIC_ACCESS_TOKEN;
       } else if (loginInput
           .getAuthenticator()
           .equalsIgnoreCase(AuthenticatorType.SNOWFLAKE_JWT.name())) {
@@ -290,17 +292,17 @@ static SFLoginOutput openSession(
     }
 
     final AuthenticatorType authenticator = getAuthenticator(loginInput);
-    if (!authenticator.equals(AuthenticatorType.OAUTH)) {
-      // OAuth does not require a username
-      AssertUtil.assertTrue(
-          loginInput.getUserName() != null, "missing user name for opening session");
-    } else {
-      // OAUTH needs either token or password
-      AssertUtil.assertTrue(
-          loginInput.getToken() != null || loginInput.getPassword() != null,
-          "missing token or password for opening session");
-    }
-    if (authenticator.equals(AuthenticatorType.EXTERNALBROWSER)) {
+      if (authenticator.equals(AuthenticatorType.OAUTH) || authenticator.equals(AuthenticatorType.PROGRAMMATIC_ACCESS_TOKEN)) {
+        // OAUTH and PAT needs either token or password
+        AssertUtil.assertTrue(
+            loginInput.getToken() != null || loginInput.getPassword() != null,
+            "missing token or password for opening session");
+      } else {
+        // OAuth does not require a username
+        AssertUtil.assertTrue(
+            loginInput.getUserName() != null, "missing user name for opening session");
+      }
+      if (authenticator.equals(AuthenticatorType.EXTERNALBROWSER)) {
       if ((Constants.getOS() == Constants.OS.MAC || Constants.getOS() == Constants.OS.WINDOWS)
           && loginInput.isEnableClientStoreTemporaryCredential()) {
         // force to set the flag for Mac/Windows users
@@ -506,7 +508,7 @@ private static SFLoginOutput newSession(
         }
       } else if (authenticatorType == AuthenticatorType.OKTA) {
         data.put(ClientAuthnParameter.RAW_SAML_RESPONSE.name(), tokenOrSamlResponse);
-      } else if (authenticatorType == AuthenticatorType.OAUTH) {
+      } else if (authenticatorType == AuthenticatorType.OAUTH || authenticatorType == AuthenticatorType.PROGRAMMATIC_ACCESS_TOKEN) {
         data.put(ClientAuthnParameter.AUTHENTICATOR.name(), authenticatorType.name());
 
         // Fix for HikariCP refresh token issue:SNOW-533673.

diff --git a/src/main/java/net/snowflake/client/core/auth/AuthenticatorType.java b/src/main/java/net/snowflake/client/core/auth/AuthenticatorType.java
@@ -51,5 +51,10 @@ public enum AuthenticatorType {
   /*
    * Client credentials flow with clientId and clientSecret as input
    */
-  OAUTH_CLIENT_CREDENTIALS
+  OAUTH_CLIENT_CREDENTIALS,
+
+  /*
+   * Authenticator to support PAT created in Snowflake
+   */
+  PROGRAMMATIC_ACCESS_TOKEN
 }
diff --git a/src/test/java/net/snowflake/client/AbstractDriverIT.java b/src/test/java/net/snowflake/client/AbstractDriverIT.java
@@ -325,7 +325,7 @@ public static Connection getConnection(
     properties.put("internal", Boolean.TRUE.toString()); // TODO: do we need this?
     properties.put("insecureMode", false); // use OCSP for all tests.
 
-    properties.put("authenticator", AuthenticatorType.OAUTH_CLIENT_CREDENTIALS.name());
+    properties.put("authenticator", AuthenticatorType.PROGRAMMATIC_ACCESS_TOKEN.name());
 
     if (injectSocketTimeout > 0) {
       properties.put("injectSocketTimeout", String.valueOf(injectSocketTimeout));