Skip to content

Commit

Permalink
SNOW-931952: Provide IncludeRetryReason config param
Browse files Browse the repository at this point in the history
  • Loading branch information
sfc-gh-pfus committed Oct 6, 2023
1 parent 207dec5 commit 6315c32
Show file tree
Hide file tree
Showing 13 changed files with 191 additions and 46 deletions.
4 changes: 2 additions & 2 deletions authokta.go
Original file line number Diff line number Diff line change
Expand Up @@ -216,7 +216,7 @@ func postAuthSAML(
fullURL := sr.getFullURL(authenticatorRequestPath, params)

logger.Infof("fullURL: %v", fullURL)
resp, err := sr.FuncPost(ctx, sr, fullURL, headers, body, timeout, true, defaultTimeProvider)
resp, err := sr.FuncPost(ctx, sr, fullURL, headers, body, timeout, true, defaultTimeProvider, nil)
if err != nil {
return nil, err
}
Expand Down Expand Up @@ -274,7 +274,7 @@ func postAuthOKTA(
if err != nil {
return nil, err
}
resp, err := sr.FuncPost(ctx, sr, targetURL, headers, body, timeout, false, defaultTimeProvider)
resp, err := sr.FuncPost(ctx, sr, targetURL, headers, body, timeout, false, defaultTimeProvider, nil)
if err != nil {
return nil, err
}
Expand Down
4 changes: 2 additions & 2 deletions chunk_downloader.go
Original file line number Diff line number Diff line change
Expand Up @@ -264,7 +264,7 @@ func getChunk(
if err != nil {
return nil, err
}
return newRetryHTTP(ctx, sc.rest.Client, http.NewRequest, u, headers, timeout, sc.currentTimeProvider).execute()
return newRetryHTTP(ctx, sc.rest.Client, http.NewRequest, u, headers, timeout, sc.currentTimeProvider, sc.cfg).execute()
}

func (scd *snowflakeChunkDownloader) startArrowBatches() error {
Expand Down Expand Up @@ -636,7 +636,7 @@ func (f *httpStreamChunkFetcher) fetch(URL string, rows chan<- []*string) error
if err != nil {
return err
}
res, err := newRetryHTTP(context.Background(), f.client, http.NewRequest, fullURL, f.headers, 0, defaultTimeProvider).execute()
res, err := newRetryHTTP(context.Background(), f.client, http.NewRequest, fullURL, f.headers, 0, defaultTimeProvider, nil).execute()
if err != nil {
return err
}
Expand Down
2 changes: 1 addition & 1 deletion client.go
Original file line number Diff line number Diff line change
Expand Up @@ -35,5 +35,5 @@ func (cli *httpClient) Post(
timeout time.Duration,
raise4xx bool,
currentTimeProvider currentTimeProvider) (*http.Response, error) {
return cli.sr.FuncPost(ctx, cli.sr, url, headers, body, timeout, raise4xx, currentTimeProvider)
return cli.sr.FuncPost(ctx, cli.sr, url, headers, body, timeout, raise4xx, currentTimeProvider, nil)
}
20 changes: 20 additions & 0 deletions dsn.go
Original file line number Diff line number Diff line change
Expand Up @@ -99,6 +99,8 @@ type Config struct {
ClientStoreTemporaryCredential ConfigBool // When true the ID token is cached in the credential manager. True by default in Windows/OSX. False for Linux.

DisableQueryContextCache bool // Should HTAP query context cache be disabled

IncludeRetryReason ConfigBool // Should retried request contain retry reason
}

// Validate enables testing if config is correct.
Expand Down Expand Up @@ -235,6 +237,9 @@ func DSN(cfg *Config) (dsn string, err error) {
if cfg.DisableQueryContextCache {
params.Add("disableQueryContextCache", "true")
}
if cfg.IncludeRetryReason == ConfigBoolFalse {
params.Add("includeRetryReason", "false")
}

params.Add("ocspFailOpen", strconv.FormatBool(cfg.OCSPFailOpen != OCSPFailOpenFalse))

Expand Down Expand Up @@ -473,6 +478,10 @@ func fillMissingConfigParameters(cfg *Config) error {
cfg.ValidateDefaultParameters = ConfigBoolTrue
}

if cfg.IncludeRetryReason == configBoolNotSet {
cfg.IncludeRetryReason = ConfigBoolTrue
}

if strings.HasSuffix(cfg.Host, defaultDomain) && len(cfg.Host) == len(defaultDomain) {
return &SnowflakeError{
Number: ErrCodeFailedToParseHost,
Expand Down Expand Up @@ -714,6 +723,17 @@ func parseDSNParams(cfg *Config, params string) (err error) {
return
}
cfg.DisableQueryContextCache = b
case "includeRetryReason":
var vv bool
vv, err = strconv.ParseBool(value)
if err != nil {
return
}
if vv {
cfg.IncludeRetryReason = ConfigBoolTrue
} else {
cfg.IncludeRetryReason = ConfigBoolFalse
}
default:
if cfg.Params == nil {
cfg.Params = make(map[string]*string)
Expand Down
Loading

0 comments on commit 6315c32

Please sign in to comment.