Skip to content

Intraprocess credential caching #21

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 4 commits into
base: v1.17.0+dbt
Choose a base branch
from
+167 −37
Open

Intraprocess credential caching #21

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
7ce08bf
Intraprocess credential caching
jasonlin45 Nov 25, 2025
5d801d7
go fmt
jasonlin45 Nov 25, 2025
be41471
Centralize cache behavior + defensive lease handling
jasonlin45 Nov 25, 2025
11195ef
Optional lease
jasonlin45 Nov 25, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
175 changes: 97 additions & 78 deletions auth.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -430,15 +430,15 @@ func authenticate(
if !respd.Success {
logger.WithContext(ctx).Errorln("Authentication FAILED")
sc.rest.TokenAccessor.SetTokens("", "", -1)
if sessionParameters[clientRequestMfaToken] == true {
credentialsStorage.deleteCredential(lease, newMfaTokenSpec(sc.cfg.Host, sc.cfg.User))
}
if sessionParameters[clientStoreTemporaryCredential] == true && sc.cfg.Authenticator == AuthTypeExternalBrowser {
credentialsStorage.deleteCredential(lease, newIDTokenSpec(sc.cfg.Host, sc.cfg.User))
}
if sessionParameters[clientStoreTemporaryCredential] == true && sc.cfg.Authenticator.isOauthNativeFlow() {
credentialsStorage.deleteCredential(lease, newOAuthAccessTokenSpec(sc.cfg.OauthTokenRequestURL, sc.cfg.User))
}
if sessionParameters[clientRequestMfaToken] == true {
_ = deleteCredentialWithLease(newMfaTokenSpec(sc.cfg.Host, sc.cfg.User))
Copy link
Collaborator

@felipecrv felipecrv Nov 25, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is more like "delete credential without lease" because you're not requiring a lease, right?

Copy link
Collaborator

@felipecrv felipecrv Nov 25, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

How does this work since this function has a lease acquired already? I don't remember if I made lease acquisition recursive (even if I did, it's bad to rely on recursive locks/leases).

Copy link
Author

@jasonlin45 jasonlin45 Nov 25, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Deleted and deferring to existing deleteCredential on the storage manager

The lease provided is nil a lot of the time. There is indeed a risk of deadlock. A lease was acquired within the function itself.

}
if sessionParameters[clientStoreTemporaryCredential] == true && sc.cfg.Authenticator == AuthTypeExternalBrowser {
_ = deleteCredentialWithLease(newIDTokenSpec(sc.cfg.Host, sc.cfg.User))
}
if sessionParameters[clientStoreTemporaryCredential] == true && sc.cfg.Authenticator.isOauthNativeFlow() {
_ = deleteCredentialWithLease(newOAuthAccessTokenSpec(sc.cfg.OauthTokenRequestURL, sc.cfg.User))
}
Copy link
Collaborator

@felipecrv felipecrv Nov 25, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think for all of these here you could just pass the lease. deletion is a very destructive operation.

Copy link
Author

@jasonlin45 jasonlin45 Nov 25, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Agree - got rid of this and went back to the existing one

code, err := strconv.Atoi(respd.Code)
if err != nil {
return nil, err
Expand All @@ -452,19 +452,19 @@ func authenticate(
logger.WithContext(ctx).Info("Authentication SUCCESS")
sc.rest.TokenAccessor.SetTokens(respd.Data.Token, respd.Data.MasterToken, respd.Data.SessionID)

if sessionParameters[clientRequestMfaToken] == true && sc.cfg.Authenticator == AuthTypeUsernamePasswordMFA {
token := respd.Data.MfaToken
credentialsStorage.setCredential(lease, newMfaTokenSpec(sc.cfg.Host, sc.cfg.User), token)
}

if sessionParameters[clientStoreTemporaryCredential] == true && sc.cfg.Authenticator == AuthTypeExternalBrowser {
token := respd.Data.IDToken
// XXX: for some reason, token is empty here some times and we
// don't want to clear the cache, so let's skip it if it's empty
if token != "" {
credentialsStorage.setCredential(lease, newIDTokenSpec(sc.cfg.Host, sc.cfg.User), token)
}
}
if sessionParameters[clientRequestMfaToken] == true && sc.cfg.Authenticator == AuthTypeUsernamePasswordMFA {
token := respd.Data.MfaToken
_ = setCredentialWithLease(newMfaTokenSpec(sc.cfg.Host, sc.cfg.User), token)
}

if sessionParameters[clientStoreTemporaryCredential] == true && sc.cfg.Authenticator == AuthTypeExternalBrowser {
token := respd.Data.IDToken
// XXX: for some reason, token is empty here some times and we
// don't want to clear the cache, so let's skip it if it's empty
if token != "" {
_ = setCredentialWithLease(newIDTokenSpec(sc.cfg.Host, sc.cfg.User), token)
}
}
Copy link
Collaborator

@felipecrv felipecrv Nov 25, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Same concern with setters. You already have a lease here. Pass the lease.

return &respd.Data, nil
}

Expand Down Expand Up @@ -633,13 +633,13 @@ func authenticateByAuthorizationCode(sc *snowflakeConn, lease *Lease) (string, e
lockKey := newOAuthAuthorizationCodeLockKey(oauthClient.tokenURL(), sc.cfg.User)
valueAwaiter := valueAwaitHolder.get(lockKey)
defer valueAwaiter.resumeOne()
token, err := awaitValue(valueAwaiter, func() (string, error) {
return credentialsStorage.getCredential(lease, newOAuthAccessTokenSpec(oauthClient.tokenURL(), sc.cfg.User))
}, func(s string, err error) bool {
return s != ""
}, func() string {
return ""
})
token, err := awaitValue(valueAwaiter, func() (string, error) {
Copy link
Collaborator

@felipecrv felipecrv Nov 25, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

You need to go fmt to make these spaces become tabs. Make the diff cleaner.

Copy link
Author

@jasonlin45 jasonlin45 Nov 25, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is my bad, messed up my editor extensions.
ran a manual go fmt and this is a lot cleaner now.

return getCredentialFast(newOAuthAccessTokenSpec(oauthClient.tokenURL(), sc.cfg.User))
Copy link
Collaborator

@felipecrv felipecrv Nov 25, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Can you add getCredentialFast to the same interface that contains getCredential? You would still take the lease, but not try to renew it right before making the query.

Copy link
Collaborator

@felipecrv felipecrv Nov 25, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

UPDATE: getCredentialRelaxed is a better name

Copy link
Author

@jasonlin45 jasonlin45 Nov 25, 2025 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Changed to getCredentialRelaxed and moved to the storage manager itself

}, func(s string, err error) bool {
return s != ""
}, func() string {
return ""
})
if err != nil || token != "" {
return token, err
}
Expand Down Expand Up @@ -742,11 +742,14 @@ func authenticateWithConfig(sc *snowflakeConn) error {
mfaTokenLockKey := newMfaTokenLockKey(sc.cfg.Host, sc.cfg.User)
idTokenLockKey := newIDTokenLockKey(sc.cfg.Host, sc.cfg.User)

lease, err := credentialsStorage.acquireLease()
if err != nil {
return err
}
defer lease.Release()
// Avoid acquiring the global lease up front; use fast-path lookups that
// consult the in-memory cache first and only acquire the lease on misses.
var lease *Lease
defer func() {
if lease != nil {
lease.Release()
}
}()

key := extBrowserBackoffKey(sc.cfg.Host, sc.cfg.User)

Expand All @@ -762,7 +765,7 @@ func authenticateWithConfig(sc *snowflakeConn) error {
sc.cfg.IDToken, _ = awaitValue(
valueAwaiter,
func() (string, error) {
tok, _ := credentialsStorage.getCredential(lease, newIDTokenSpec(sc.cfg.Host, sc.cfg.User))
tok, _ := getCredentialFast(newIDTokenSpec(sc.cfg.Host, sc.cfg.User))
return tok, nil
},
func(s string, err error) bool {
Expand All @@ -773,10 +776,10 @@ func authenticateWithConfig(sc *snowflakeConn) error {
},
)

} else if sc.cfg.ClientStoreTemporaryCredential == ConfigBoolTrue {
tok, _ := credentialsStorage.getCredential(lease, newIDTokenSpec(sc.cfg.Host, sc.cfg.User))
sc.cfg.IDToken = tok
}
} else if sc.cfg.ClientStoreTemporaryCredential == ConfigBoolTrue {
tok, _ := getCredentialFast(newIDTokenSpec(sc.cfg.Host, sc.cfg.User))
Copy link
Collaborator

@felipecrv felipecrv Nov 25, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Call it getCredentialsStale or getCredentialRelaxed. You're implementing something that is analogous to relaxed atomic reads.

https://www.google.com/search?hl=en&q=relaxed%20memory%20model

Copy link
Collaborator

@felipecrv felipecrv Nov 25, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is a "Relaxed" read. All the writes sets and delete must be "Release" (i.e. they must have already "Acquire"d a lease).

If the "Relaxed" read is bad, we "Acquire" a lease before querying it again.

sc.cfg.IDToken = tok
}
// Disable console login by default
if sc.cfg.DisableConsoleLogin == configBoolNotSet {
sc.cfg.DisableConsoleLogin = ConfigBoolTrue
Expand All @@ -793,7 +796,7 @@ func authenticateWithConfig(sc *snowflakeConn) error {
sc.cfg.MfaToken, _ = awaitValue(
valueAwaiter,
func() (string, error) {
tok, err := credentialsStorage.getCredential(lease, newMfaTokenSpec(sc.cfg.Host, sc.cfg.User))
tok, err := getCredentialFast(newMfaTokenSpec(sc.cfg.Host, sc.cfg.User))
if err != nil {
logger.WithContext(sc.ctx).Warnf("failed to get MFA token from credential storage: %v", err)
}
Expand All @@ -804,13 +807,13 @@ func authenticateWithConfig(sc *snowflakeConn) error {
return ""
},
)
} else if sc.cfg.ClientRequestMfaToken == ConfigBoolTrue {
tok, err := credentialsStorage.getCredential(lease, newMfaTokenSpec(sc.cfg.Host, sc.cfg.User))
if err != nil {
logger.WithContext(sc.ctx).Warnf("failed to get MFA token from credential storage: %v", err)
}
sc.cfg.MfaToken = tok
}
} else if sc.cfg.ClientRequestMfaToken == ConfigBoolTrue {
tok, err := getCredentialFast(newMfaTokenSpec(sc.cfg.Host, sc.cfg.User))
if err != nil {
logger.WithContext(sc.ctx).Warnf("failed to get MFA token from credential storage: %v", err)
}
sc.cfg.MfaToken = tok
}
}
}

Expand All @@ -829,28 +832,36 @@ func authenticateWithConfig(sc *snowflakeConn) error {
lastFail.Delete(key)
}

samlResponse, proofKey, err = authenticateByExternalBrowser(
sc.ctx,
lease,
sc.rest,
sc.cfg.Authenticator.String(),
sc.cfg.Application,
sc.cfg.Account,
sc.cfg.User,
sc.cfg.ExternalBrowserTimeout,
sc.cfg.DisableConsoleLogin)
// External browser path requires a live lease for periodic renewals
if lease == nil {
lease, err = credentialsStorage.acquireLease()
Copy link
Collaborator

@felipecrv felipecrv Nov 25, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

If you got all the way here without a lease, what you should probably start from the beginning, but with a lease. This can be easily implemented with recursion (max depth will be 2).

You change func authenticateWithConfig(sc *snowflakeConn) error { to take an optionalLease parameter and you change the code here to be:

if optionalLease == nil {
   lease = acquire the lease
   return authenticateWithConfig(sc, lease)
}

Copy link
Author

@jasonlin45 jasonlin45 Nov 25, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Recursion can't be done here because of interactions with the valueAwaiter the Snowflake folks added. Recursing out of the function releases other goroutine to go ahead, and then I had a deadlock.

if err != nil {
sc.cleanup()
return err
}
}
samlResponse, proofKey, err = authenticateByExternalBrowser(
sc.ctx,
lease,
sc.rest,
sc.cfg.Authenticator.String(),
sc.cfg.Application,
sc.cfg.Account,
sc.cfg.User,
sc.cfg.ExternalBrowserTimeout,
sc.cfg.DisableConsoleLogin)
if err != nil {
sc.cleanup()
return err
}
}
}
authData, err = authenticate(
sc.ctx,
lease,
sc,
samlResponse,
proofKey)
authData, err = authenticate(
sc.ctx,
lease,
Copy link
Collaborator

@felipecrv felipecrv Nov 25, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

you're passing a nil lease to a function that didn't expect nil

Copy link
Collaborator

@felipecrv felipecrv Nov 25, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

you might have to rename the parameter to optionalLease

Copy link
Author

@jasonlin45 jasonlin45 Nov 25, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Good point. This can be nil in some cases. Renamed accordingly!

sc,
samlResponse,
proofKey)
if err != nil {
if errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) {
sc.cleanup()
Expand All @@ -861,9 +872,9 @@ func authenticateWithConfig(sc *snowflakeConn) error {

switch {
// Case 1: cached ID token failed -> clear + try one interactive refresh
case sc.cfg.Authenticator == AuthTypeExternalBrowser && sc.cfg.IDToken != "":
credentialsStorage.deleteCredential(lease, newIDTokenSpec(sc.cfg.Host, sc.cfg.User))
sc.cfg.IDToken = ""
case sc.cfg.Authenticator == AuthTypeExternalBrowser && sc.cfg.IDToken != "":
_ = deleteCredentialWithLease(newIDTokenSpec(sc.cfg.Host, sc.cfg.User))
sc.cfg.IDToken = ""

// NOTE on tabstorms:
// We intentionally do NOT pre-gate the cached-ID-token fallback with a LoadOrStore
Expand All @@ -874,11 +885,19 @@ func authenticateWithConfig(sc *snowflakeConn) error {
// considered low probability and low impact

// reenter interactive flow
samlResponse, proofKey, err = authenticateByExternalBrowser(
sc.ctx, lease, sc.rest, sc.cfg.Authenticator.String(),
sc.cfg.Application, sc.cfg.Account, sc.cfg.User,
sc.cfg.ExternalBrowserTimeout, sc.cfg.DisableConsoleLogin,
)
if lease == nil {
lease, err = credentialsStorage.acquireLease()
if err != nil {
// let SAML-phase failure remain retryable; do not set backoff
sc.cleanup()
return err
}
}
samlResponse, proofKey, err = authenticateByExternalBrowser(
sc.ctx, lease, sc.rest, sc.cfg.Authenticator.String(),
sc.cfg.Application, sc.cfg.Account, sc.cfg.User,
sc.cfg.ExternalBrowserTimeout, sc.cfg.DisableConsoleLogin,
)
if err != nil {
// let SAML-phase failure remain retryable; do not set backoff
sc.cleanup()
Expand All @@ -892,12 +911,12 @@ func authenticateWithConfig(sc *snowflakeConn) error {
fallthrough

// Case 2: still failing, but could be an OAuth refreshable error
case errors.As(err, &se) && slices.Contains(refreshOAuthTokenErrorCodes, strconv.Itoa(se.Number)):
credentialsStorage.deleteCredential(lease, newOAuthAccessTokenSpec(sc.cfg.OauthTokenRequestURL, sc.cfg.User))
case errors.As(err, &se) && slices.Contains(refreshOAuthTokenErrorCodes, strconv.Itoa(se.Number)):
_ = deleteCredentialWithLease(newOAuthAccessTokenSpec(sc.cfg.OauthTokenRequestURL, sc.cfg.User))

if sc.cfg.Authenticator == AuthTypeOAuthAuthorizationCode {
doRefreshTokenWithLease(sc, lease)
}
if sc.cfg.Authenticator == AuthTypeOAuthAuthorizationCode {
doRefreshTokenWithLease(sc, lease)
}

// if refreshing succeeds for authorization code, we will take a token from cache
// if it fails, we will just run the full flow
Expand Down Expand Up @@ -941,7 +960,7 @@ func doRefreshTokenWithLock(sc *snowflakeConn, lease *Lease) {
if _, err = getValueWithLock(chooseLockerForAuth(sc.cfg), lockKey, func() (string, error) {
if err = oauthClient.refreshToken(lease); err != nil {
logger.Warnf("cannot refresh token. %v", err)
credentialsStorage.deleteCredential(lease, newOAuthRefreshTokenSpec(sc.cfg.OauthTokenRequestURL, sc.cfg.User))
_ = deleteCredentialWithLease(newOAuthRefreshTokenSpec(sc.cfg.OauthTokenRequestURL, sc.cfg.User))
return "", err
}
return "", nil
Expand All @@ -957,7 +976,7 @@ func doRefreshTokenWithLease(sc *snowflakeConn, lease *Lease) {
} else {
if rfErr := oauthClient.refreshToken(lease); rfErr != nil {
logger.Warnf("cannot refresh token. %v", rfErr)
credentialsStorage.deleteCredential(lease, newOAuthRefreshTokenSpec(sc.cfg.OauthTokenRequestURL, sc.cfg.User))
_ = deleteCredentialWithLease(newOAuthRefreshTokenSpec(sc.cfg.OauthTokenRequestURL, sc.cfg.User))
}
}

Expand Down
Loading
Loading