Increase lease timeout for sso operations and provide arrow-adbc configuration hook for future adjustments (#13)

* Increase timeout to 30s to support headless sso and ide sso

* Add hooks for configuring lease timeouts without disturbing singleton design.

* Change timeout to millis.

No need to worry about cpu alignment. single entrypoint to value through setter and getter normalizes. Decisions on how to handle data made explicit for future readers.

Author: VersusFacit

authexternalbrowser.go

@@ -340,7 +340,8 @@ func doAuthenticateByExternalBrowser(
 func waitForSamlResponse(ctx context.Context, lease *Lease, l net.Listener, application string) (string, error) {
 	encodedChan := make(chan string, 1)
 	errChan := make(chan error, 1)
-	ticker := time.NewTicker(leaseTTL / 2)
+	ttl := leaseTTL()
+	ticker := time.NewTicker(ttl / 2)
 
 	go func() {
 		conn, err := l.Accept()
@@ -399,7 +400,7 @@ func waitForSamlResponse(ctx context.Context, lease *Lease, l net.Listener, appl
 	for {
 		select {
 		case <-ticker.C:
-			lease.Renew(leaseTTL)
+			lease.Renew(ttl)
 		case s := <-encodedChan:
 			ticker.Stop()
 			return s, nil

lease.go

@@ -12,6 +12,7 @@ import (
 	"os"
 	"path/filepath"
 	"strconv"
+	"sync/atomic"
 	"time"
 )
 
@@ -63,19 +64,38 @@ func (lease *Lease) Release() error {
 //
 // [1] https://en.wikipedia.org/wiki/Lease_(computer_science)
 type LeaseHandler struct {
-	path    string        // absolute path to the lease file
-	dir     string        // directory where the lease file is stored
-	timeout time.Duration // how long to keep trying to acquire or renew a lease
+	path          string // absolute path to the lease file
+	dir           string // directory where the lease file is stored
+	timeoutMillis int32  // atomic; how long to keep trying to acquire or renew a lease
 }
 
 func NewLeaseHandler(path string, timeout time.Duration) (*LeaseHandler, error) {
 	abspath, err := filepath.Abs(path)
 	if err != nil {
 		return nil, fmt.Errorf("lease: %w", err)
 	}
-	timeout = max(timeout, MinLeaseOperationTimeout)
 	dir := filepath.Dir(abspath)
-	return &LeaseHandler{path: abspath, dir: dir, timeout: timeout}, nil
+	h := &LeaseHandler{path: abspath, dir: dir}
+	h.SetTimeout(timeout)
+	return h, nil
+}
+
+// Normalize timeout to atomic-friendly i32 and store atomically
+// Clamp to prevent wraparound. Truncate to ms.
+func (l *LeaseHandler) SetTimeout(d time.Duration) {
+	if d < MinLeaseOperationTimeout {
+		d = MinLeaseOperationTimeout
+	}
+	ms := d.Milliseconds()
+	if ms > math.MaxInt32 {
+		ms = math.MaxInt32
+	}
+	atomic.StoreInt32(&l.timeoutMillis, int32(ms))
+}
+
+func (l *LeaseHandler) getTimeout() time.Duration {
+	ms := atomic.LoadInt32(&l.timeoutMillis)
+	return time.Duration(ms) * time.Millisecond
 }
 
 func randomLeaseId() (string, error) {
@@ -251,7 +271,7 @@ func (l *LeaseHandler) Acquire(ttl time.Duration) (*Lease, error) {
 
 	base := time.Duration(0)
 	m := gracePeriod
-	deadline := time.Now().Add(l.timeout)
+	deadline := time.Now().Add(l.getTimeout())
 	for time.Now().Before(deadline) {
 		data, err := l.read(base, m, min(deadline.Sub(time.Now())/2, maxPollInterval))
 		base, m = nextWait(base, m)
@@ -287,15 +307,15 @@ func (l *LeaseHandler) Acquire(ttl time.Duration) (*Lease, error) {
 		return &Lease{id: newLeaseId, expiry: expiry, handler: l}, nil
 	}
 
-	return nil, fmt.Errorf("timed out trying to acquire lease after %s: %s", l.timeout, l.path)
+	return nil, fmt.Errorf("timed out trying to acquire lease after %s: %s", l.getTimeout(), l.path)
 }
 
 func (l *LeaseHandler) renew(leaseId *string, ttl time.Duration, currentExpiry time.Time) (time.Time, error) {
 	// fmt.Fprintf(os.Stdout, "[%v] Renew('%s')\n", routine.Goid(), *leaseId)
 
 	base := time.Duration(0)
 	m := time.Duration(0)
-	deadline := time.Now().Add(l.timeout)
+	deadline := time.Now().Add(l.getTimeout())
 	for time.Now().Before(deadline) {
 		wait(base, m, min(deadline.Sub(time.Now())/2, maxPollInterval))
 		base, m = nextWait(base, m)

secure_storage_manager.go

@@ -11,18 +11,64 @@ import (
 	"strconv"
 	"strings"
 	"sync"
+	"sync/atomic"
 	"time"
 )
 
 const (
-	leaseTTL              = 10 * time.Second
-	leaseOperationTimeout = 60 * time.Second
+	// accessed through getters that provide an override hook, do not use directly
+	_defaultLeaseTTL              = 30 * time.Second
+	_defaultLeaseOperationTimeout = 90 * time.Second
 
 	credCacheDirEnv   = "SF_TEMPORARY_CREDENTIAL_CACHE_DIR"
 	credLeaseFileName = "credential_cache.lease"
 	credCacheFileName = "credential_cache_v1.json"
 )
 
+// --- EvalOnce hook to configure lease semantics ------------------------
+
+var (
+	_cfgOnce         sync.Once
+	_overrideTTL     atomic.Value // stores time.Duration
+	_overrideTimeout atomic.Value // ditto
+)
+
+// Call once per process.
+//
+// Ignore 0 values and keep defaults
+// Propogate changes to singleton credentialsStorage's LeaseHandler
+func ConfigureLeaseOnce(ttl, timeout time.Duration) {
+	_cfgOnce.Do(func() {
+		if ttl > 0 {
+			_overrideTTL.Store(ttl)
+		}
+		if timeout > 0 {
+			_overrideTimeout.Store(timeout)
+			// warn: will fail if of other type but this is not dbt's use
+			// of the api
+			if fb, ok := credentialsStorage.(*fileBasedSecureStorageManager); ok && fb.leaseHandler != nil {
+				fb.leaseHandler.SetTimeout(timeout)
+			}
+		}
+	})
+}
+
+func leaseTTL() time.Duration {
+	if v := _overrideTTL.Load(); v != nil {
+		return v.(time.Duration)
+	}
+	return _defaultLeaseTTL
+}
+
+func leaseOperationTimeout() time.Duration {
+	if v := _overrideTimeout.Load(); v != nil {
+		return v.(time.Duration)
+	}
+	return _defaultLeaseOperationTimeout
+}
+
+// --- CacheDir resolution ------------------------
+
 type cacheDirConf struct {
 	envVar       string
 	pathSegments []string
@@ -98,7 +144,7 @@ func newFileBasedSecureStorageManager() (*fileBasedSecureStorageManager, error)
 	if err != nil {
 		return nil, err
 	}
-	leaseHandler, err := NewLeaseHandler(filepath.Join(credDirPath, credLeaseFileName), leaseOperationTimeout)
+	leaseHandler, err := NewLeaseHandler(filepath.Join(credDirPath, credLeaseFileName), leaseOperationTimeout())
 	if err != nil {
 		return nil, err
 	}
@@ -173,11 +219,11 @@ func (ssm *fileBasedSecureStorageManager) getTokens(data map[string]any) map[str
 }
 
 func (ssm *fileBasedSecureStorageManager) acquireLease() (*Lease, error) {
-	return ssm.leaseHandler.Acquire(leaseTTL)
+	return ssm.leaseHandler.Acquire(leaseTTL())
 }
 
 func (ssm *fileBasedSecureStorageManager) withCacheFile(lease *Lease, action func(*os.File) error) error {
-	err := lease.Renew(leaseTTL / 2)
+	err := lease.Renew(leaseTTL() / 2)
 	if err != nil {
 		logger.Warnf("Unable to lease cache. %v", err)
 		return err