SNOW-1825495 OAuth flows implementation (#2135)

Co-authored-by: Michał Hofman <[email protected]>
Co-authored-by: Piotr Bulawa <[email protected]>
Co-authored-by: Maxim Mishchenko <[email protected]>
Co-authored-by: Mikołaj Kubik <[email protected]>
Co-authored-by: Yijun Xie <[email protected]>
Co-authored-by: Zexin Yao <[email protected]>
Co-authored-by: Jakub Szczerbiński <[email protected]>
Co-authored-by: Patryk Cyrek <[email protected]>

Author: 9 people

.github/workflows/parameters/private/parameters_aws_auth_tests.json.gpg

@@ -13,6 +13,7 @@ Source code is also available at: https://github.com/snowflakedb/snowflake-conne
   - Dropped support for Python 3.8.
   - Basic decimal floating-point type support.
   - Added handling of PAT provided in `password` field.
+  - Added experimental support for OAuth authorization code and client credentials flows.
   - Improved error message for client-side query cancellations due to timeouts.
   - Added support of GCS regional endpoints.
   - Added `gcs_use_virtual_endpoints` connection property that forces the usage of the virtual GCS usage. Thanks to this it should be possible to set up private DNS entry for the GCS endpoint. See more: https://cloud.google.com/storage/docs/request-endpoints#xml-api

.github/workflows/parameters/private/rsa_keys/rsa_key.p8.gpg

@@ -35,29 +35,46 @@ timestamps {
         string(name: 'parent_job', value: env.JOB_NAME),
         string(name: 'parent_build_number', value: env.BUILD_NUMBER)
       ]
-      stage('Test') {
-          try {
-          def commit_hash = "main" // default which we want to override
-          def bptp_tag = "bptp-stable"
-          def response = authenticatedGithubCall("https://api.github.com/repos/snowflakedb/snowflake/git/ref/tags/${bptp_tag}")
-          commit_hash = response.object.sha
-          // Append the bptp-stable commit sha to params
-          params += [string(name: 'svn_revision', value: commit_hash)]
-          } catch(Exception e) {
-          println("Exception computing commit hash from: ${response}")
+      parallel(
+      'Test': {
+        stage('Test') {
+            try {
+            def commit_hash = "main" // default which we want to override
+            def bptp_tag = "bptp-stable"
+            def response = authenticatedGithubCall("https://api.github.com/repos/snowflakedb/snowflake/git/ref/tags/${bptp_tag}")
+            commit_hash = response.object.sha
+            // Append the bptp-stable commit sha to params
+            params += [string(name: 'svn_revision', value: commit_hash)]
+            } catch(Exception e) {
+            println("Exception computing commit hash from: ${response}")
+            }
+          parallel (
+            'Test Python 39': { build job: 'RT-PyConnector39-PC',parameters: params},
+            'Test Python 310': { build job: 'RT-PyConnector310-PC',parameters: params},
+            'Test Python 311': { build job: 'RT-PyConnector311-PC',parameters: params},
+            'Test Python 312': { build job: 'RT-PyConnector312-PC',parameters: params},
+            'Test Python 313': { build job: 'RT-PyConnector313-PC',parameters: params},
+            'Test Python 39 OldDriver': { build job: 'RT-PyConnector39-OldDriver-PC',parameters: params},
+            'Test Python 39 FIPS': { build job: 'RT-FIPS-PyConnector39',parameters: params},
+            )
+          }
+        },
+      'Test Authentication': {
+        stage('Test Authentication') {
+          withCredentials([
+            string(credentialsId: 'a791118f-a1ea-46cd-b876-56da1b9bc71c', variable: 'NEXUS_PASSWORD'),
+            string(credentialsId: 'sfctest0-parameters-secret', variable: 'PARAMETERS_SECRET')
+          ]) {
+            sh '''\
+            |#!/bin/bash -e
+            |$WORKSPACE/ci/test_authentication.sh
+            '''.stripMargin()
           }
-        parallel (
-          'Test Python 39': { build job: 'RT-PyConnector39-PC',parameters: params},
-          'Test Python 310': { build job: 'RT-PyConnector310-PC',parameters: params},
-          'Test Python 311': { build job: 'RT-PyConnector311-PC',parameters: params},
-          'Test Python 312': { build job: 'RT-PyConnector312-PC',parameters: params},
-          'Test Python 313': { build job: 'RT-PyConnector313-PC',parameters: params},
-          'Test Python 39 OldDriver': { build job: 'RT-PyConnector39-OldDriver-PC',parameters: params},
-          'Test Python 39 FIPS': { build job: 'RT-FIPS-PyConnector39',parameters: params},
-          )
         }
       }
-    }
+    )
+  }
+}
 
 
 pipeline {

.github/workflows/parameters/private/rsa_keys/rsa_key_invalid.p8.gpg

@@ -0,0 +1,24 @@
+#!/bin/bash -e
+
+set -o pipefail
+
+
+export WORKSPACE=${WORKSPACE:-/mnt/workspace}
+export SOURCE_ROOT=${SOURCE_ROOT:-/mnt/host}
+
+MVNW_EXE=$SOURCE_ROOT/mvnw
+AUTH_PARAMETER_FILE=./.github/workflows/parameters/private/parameters_aws_auth_tests.json
+eval $(jq -r '.authtestparams | to_entries | map("export \(.key)=\(.value|tostring)")|.[]' $AUTH_PARAMETER_FILE)
+
+export SNOWFLAKE_AUTH_TEST_PRIVATE_KEY_PATH=./.github/workflows/parameters/private/rsa_keys/rsa_key.p8
+export SNOWFLAKE_AUTH_TEST_INVALID_PRIVATE_KEY_PATH=./.github/workflows/parameters/private/rsa_keys/rsa_key_invalid.p8
+
+export SF_OCSP_TEST_MODE=true
+export SF_ENABLE_EXPERIMENTAL_AUTHENTICATION=true
+export RUN_AUTH_TESTS=true
+export AUTHENTICATION_TESTS_ENV="docker"
+export PYTHONPATH=$SOURCE_ROOT
+
+python3 -m pip install --break-system-packages -e .
+
+python3 -m pytest test/auth/*

DESCRIPTION.md

@@ -0,0 +1,27 @@
+#!/bin/bash -e
+
+set -o pipefail
+
+
+export THIS_DIR="$( cd "$( dirname "${BASH_SOURCE[0]}" )" && pwd )"
+export WORKSPACE=${WORKSPACE:-/tmp}
+
+CI_DIR="$( cd "$( dirname "${BASH_SOURCE[0]}" )" && pwd )"
+if [[ -n "$JENKINS_HOME" ]]; then
+  ROOT_DIR="$(cd "${CI_DIR}/.." && pwd)"
+  export WORKSPACE=${WORKSPACE:-/tmp}
+  echo "Use /sbin/ip"
+  IP_ADDR=$(/sbin/ip -4 addr show scope global dev eth0 | grep inet | awk '{print $2}' | cut -d / -f 1)
+
+fi
+
+gpg --quiet --batch --yes --decrypt --passphrase="$PARAMETERS_SECRET" --output $THIS_DIR/../.github/workflows/parameters/private/parameters_aws_auth_tests.json "$THIS_DIR/../.github/workflows/parameters/private/parameters_aws_auth_tests.json.gpg"
+gpg --quiet --batch --yes --decrypt --passphrase="$PARAMETERS_SECRET" --output $THIS_DIR/../.github/workflows/parameters/private/rsa_keys/rsa_key.p8 "$THIS_DIR/../.github/workflows/parameters/private/rsa_keys/rsa_key.p8.gpg"
+gpg --quiet --batch --yes --decrypt --passphrase="$PARAMETERS_SECRET" --output $THIS_DIR/../.github/workflows/parameters/private/rsa_keys/rsa_key_invalid.p8 "$THIS_DIR/../.github/workflows/parameters/private/rsa_keys/rsa_key_invalid.p8.gpg"
+
+docker run \
+  -v $(cd $THIS_DIR/.. && pwd):/mnt/host \
+  -v $WORKSPACE:/mnt/workspace \
+  --rm \
+  nexus.int.snowflakecomputing.com:8086/docker/snowdrivers-test-external-browser-python:1 \
+  "/mnt/host/ci/container/test_authentication.sh"

Jenkinsfile

@@ -7,6 +7,8 @@
 from .keypair import AuthByKeyPair
 from .no_auth import AuthNoAuth
 from .oauth import AuthByOAuth
+from .oauth_code import AuthByOauthCode
+from .oauth_credentials import AuthByOauthCredentials
 from .okta import AuthByOkta
 from .pat import AuthByPAT
 from .usrpwdmfa import AuthByUsrPwdMfa
@@ -18,6 +20,8 @@
         AuthByDefault,
         AuthByKeyPair,
         AuthByOAuth,
+        AuthByOauthCode,
+        AuthByOauthCredentials,
         AuthByOkta,
         AuthByUsrPwdMfa,
         AuthByWebBrowser,
@@ -34,6 +38,8 @@
     "AuthByKeyPair",
     "AuthByPAT",
     "AuthByOAuth",
+    "AuthByOauthCode",
+    "AuthByOauthCredentials",
     "AuthByOkta",
     "AuthByUsrPwdMfa",
     "AuthByWebBrowser",

ci/container/test_authentication.sh

@@ -47,6 +47,7 @@
     ACCEPT_TYPE_APPLICATION_SNOWFLAKE,
     CONTENT_TYPE_APPLICATION_JSON,
     ID_TOKEN_INVALID_LOGIN_REQUEST_GS_CODE,
+    OAUTH_ACCESS_TOKEN_EXPIRED_GS_CODE,
     PYTHON_CONNECTOR_USER_AGENT,
     ReauthenticationRequest,
 )
@@ -86,7 +87,7 @@ class Auth:
 
     def __init__(self, rest) -> None:
         self._rest = rest
-        self.token_cache = TokenCache.make()
+        self._token_cache: TokenCache | None = None
 
     @staticmethod
     def base_auth_data(
@@ -350,7 +351,7 @@ def post_request_wrapper(self, url, headers, body) -> None:
                 # clear stored id_token if failed to connect because of id_token
                 # raise an exception for reauth without id_token
                 self._rest.id_token = None
-                self.delete_temporary_credential(
+                self._delete_temporary_credential(
                     self._rest._host, user, TokenType.ID_TOKEN
                 )
                 raise ReauthenticationRequest(
@@ -360,6 +361,14 @@ def post_request_wrapper(self, url, headers, body) -> None:
                         sqlstate=SQLSTATE_CONNECTION_WAS_NOT_ESTABLISHED,
                     )
                 )
+            elif errno == OAUTH_ACCESS_TOKEN_EXPIRED_GS_CODE:
+                raise ReauthenticationRequest(
+                    ProgrammingError(
+                        msg=ret["message"],
+                        errno=int(errno),
+                        sqlstate=SQLSTATE_CONNECTION_WAS_NOT_ESTABLISHED,
+                    )
+                )
 
             from . import AuthByKeyPair
 
@@ -374,7 +383,7 @@ def post_request_wrapper(self, url, headers, body) -> None:
             from . import AuthByUsrPwdMfa
 
             if isinstance(auth_instance, AuthByUsrPwdMfa):
-                self.delete_temporary_credential(
+                self._delete_temporary_credential(
                     self._rest._host, user, TokenType.MFA_TOKEN
                 )
             Error.errorhandler_wrapper(
@@ -466,7 +475,7 @@ def _read_temporary_credential(
         user: str,
         cred_type: TokenType,
     ) -> str | None:
-        return self.token_cache.retrieve(TokenKey(host, user, cred_type))
+        return self.get_token_cache().retrieve(TokenKey(host, user, cred_type))
 
     def read_temporary_credentials(
         self,
@@ -500,7 +509,7 @@ def _write_temporary_credential(
                 "no credential is given when try to store temporary credential"
             )
             return
-        self.token_cache.store(TokenKey(host, user, cred_type), cred)
+        self.get_token_cache().store(TokenKey(host, user, cred_type), cred)
 
     def write_temporary_credentials(
         self,
@@ -524,10 +533,15 @@ def write_temporary_credentials(
                 host, user, TokenType.MFA_TOKEN, response["data"].get("mfaToken")
             )
 
-    def delete_temporary_credential(
+    def _delete_temporary_credential(
         self, host: str, user: str, cred_type: TokenType
     ) -> None:
-        self.token_cache.remove(TokenKey(host, user, cred_type))
+        self.get_token_cache().remove(TokenKey(host, user, cred_type))
+
+    def get_token_cache(self) -> TokenCache:
+        if self._token_cache is None:
+            self._token_cache = TokenCache.make()
+        return self._token_cache
 
 
 def get_token_from_private_key(