SNOW-1825790 Token cache refactor - v2

snowflakedb · Feb 4, 2025 · 3a3de37 · 3a3de37
1 parent e926883
commit 3a3de37
Show file tree

Hide file tree

Showing 3 changed files with 156 additions and 152 deletions.
diff --git a/auth.go b/auth.go
@@ -26,8 +26,6 @@ const (
 )
 
 const (
-	idToken                        = "ID_TOKEN"
-	mfaToken                       = "MFATOKEN"
 	clientStoreTemporaryCredential = "CLIENT_STORE_TEMPORARY_CREDENTIAL"
 	clientRequestMfaToken          = "CLIENT_REQUEST_MFA_TOKEN"
 	idTokenAuthenticator           = "ID_TOKEN"
@@ -365,10 +363,10 @@ func authenticate(
 		logger.WithContext(ctx).Errorln("Authentication FAILED")
 		sc.rest.TokenAccessor.SetTokens("", "", -1)
 		if sessionParameters[clientRequestMfaToken] == true {
-			credentialsStorage.deleteCredential(sc, mfaToken)
+			credentialsStorage.deleteCredential(newMfaTokenSpec(sc.cfg.Host, sc.cfg.User))
 		}
 		if sessionParameters[clientStoreTemporaryCredential] == true {
-			credentialsStorage.deleteCredential(sc, idToken)
+			credentialsStorage.deleteCredential(newIDTokenSpec(sc.cfg.Host, sc.cfg.User))
 		}
 		code, err := strconv.Atoi(respd.Code)
 		if err != nil {
@@ -384,11 +382,11 @@ func authenticate(
 	sc.rest.TokenAccessor.SetTokens(respd.Data.Token, respd.Data.MasterToken, respd.Data.SessionID)
 	if sessionParameters[clientRequestMfaToken] == true {
 		token := respd.Data.MfaToken
-		credentialsStorage.setCredential(sc, mfaToken, token)
+		credentialsStorage.setCredential(newMfaTokenSpec(sc.cfg.Host, sc.cfg.User), token)
 	}
 	if sessionParameters[clientStoreTemporaryCredential] == true {
 		token := respd.Data.IDToken
-		credentialsStorage.setCredential(sc, idToken, token)
+		credentialsStorage.setCredential(newIDTokenSpec(sc.cfg.Host, sc.cfg.User), token)
 	}
 	return &respd.Data, nil
 }
@@ -523,7 +521,7 @@ func authenticateWithConfig(sc *snowflakeConn) error {
 			sc.cfg.ClientStoreTemporaryCredential = ConfigBoolTrue
 		}
 		if sc.cfg.ClientStoreTemporaryCredential == ConfigBoolTrue {
-			fillCachedIDToken(sc)
+			sc.cfg.IDToken = credentialsStorage.getCredential(newIDTokenSpec(sc.cfg.Host, sc.cfg.User))
 		}
 		// Disable console login by default
 		if sc.cfg.DisableConsoleLogin == configBoolNotSet {
@@ -536,7 +534,7 @@ func authenticateWithConfig(sc *snowflakeConn) error {
 			sc.cfg.ClientRequestMfaToken = ConfigBoolTrue
 		}
 		if sc.cfg.ClientRequestMfaToken == ConfigBoolTrue {
-			fillCachedMfaToken(sc)
+			sc.cfg.MfaToken = credentialsStorage.getCredential(newMfaTokenSpec(sc.cfg.Host, sc.cfg.User))
 		}
 	}
 
@@ -573,11 +571,3 @@ func authenticateWithConfig(sc *snowflakeConn) error {
 	sc.ctx = context.WithValue(sc.ctx, SFSessionIDKey, authData.SessionID)
 	return nil
 }
-
-func fillCachedIDToken(sc *snowflakeConn) {
-	credentialsStorage.getCredential(sc, idToken)
-}
-
-func fillCachedMfaToken(sc *snowflakeConn) {
-	credentialsStorage.getCredential(sc, mfaToken)
-}