[peco-engineer-bot]

Summary

Automated fix for #475 — CloudFetch and M2M token exchange are unbounded and uninterruptible (v1.13.0).

Added two new ConnOptions in connector.go. WithCloudFetchHTTPClient sets the CloudFetch download client, so callers can set http.Client.Timeout without replacing the Thrift transport; it wins over WithTransport whichever comes first. WithCloudFetchMinTimeToExpiry sets MinTimeToExpiry and its comment says an expired link ends the read for good. In auth/oauth/m2m, GetTokenSource now uses an http.Client with a timeout (tokenExchangeTimeout, default 30s), and Authenticate no longer holds the mutex during a cached token refresh. The unit tests TestCloudFetchConnOptions and TestM2MTokenExchangeIsBounded pass, and make test and make lint (0 issues) are green. The two e2e tests now compile, but they SKIP in this shell because the tool scrubs the DATABRICKS_PECOTESTING_* credentials, so I could not see them pass against the live warehouse. A maintainer should run TestE2ECloudFetchMinTimeToExpiryIsHonored and TestE2ECloudFetchHTTPClientOption with creds before merging. The M2M part is offline-only: a live warehouse can't be made to hang its token endpoint. The expected behavior comes from the issue, which asks for the exchange to be bounded like the 10s endpoint-discovery cap.

Root cause & plan

Root cause: The issue has five parts. I checked each against the code:

(1) Context: arrowRows.go:126/128 does pass context.Background() to NewCloudBatchIterator/NewLocalBatchIterator instead of the scanner's ctx. But the ctx the scanner gets is rows.ctx, which NewRows deliberately detaches from the caller with context.WithoutCancel (ES-1934053 / #295 / #371). The existing e2e test TestE2EArrowBatchesSurviveQueryContextCancellation protects that. So "the caller's QueryContext cancellation should stop downloads" goes against a documented design decision. Passing ctx through only keeps the context values (conn, correlation and query IDs) and ties downloads to resultsCancel. rows.Close already cancels downloads through RowScanner.Close, so there is nothing that fails today to test.

(2) The download client is http.DefaultClient, which has no timeout. The only way to set cfg.HTTPClient is WithTransport, which also replaces the Thrift transport. There is no ConnOption for the CloudFetch http.Client, so a caller cannot set http.Client.Timeout.

(3) MinTimeToExpiry also has no ConnOption. An expired link returns the terminal error errors.New(ErrLinkExpired) and nothing re-issues links.

(4) m2m.authClient.Authenticate holds c.mx across tokenSource.Token(). GetTokenSource builds the token source with context.Background(), so clientcredentials uses http.DefaultClient with no timeout. A hung token endpoint therefore blocks Authenticate, and every request on the connector, indefinitely. Endpoint discovery is capped at 10s (oauth.hostConfigTimeout); the token exchange is not.

(5) Doc mismatch: already fixed. connector.go now says "WithCloudFetch ... Default is true."
Files: connector.go, auth/oauth/m2m/m2m.go, internal/rows/arrowbased/arrowRows.go, driver_e2e_test.go, connector_test.go, auth/oauth/m2m/m2m_test.go
Planned coverage:

• e2e in driver_e2e_test.go (Thrift path, pecoTestingCreds): run a query big enough to return CloudFetch result links, e.g. SELECT * FROM range(0, 5000000) with WithCloudFetch(true), using a new WithCloudFetchMinTimeToExpiry(24*time.Hour) ConnOption. That margin is longer than any real presigned URL's lifetime. Iterate the rows and assert iteration fails with an error containing dbsqlerr.ErrLinkExpired. Running the same query without the option must return the exact row count. Red today: the option doesn't exist, the setting can't be reached, and cfg.MinTimeToExpiry stays 0. The test exercises the real server-issued link.ExpiryTime. (Fix the module name #3 MinTimeToExpiry cannot be set and has no effect on live CloudFetch downloads)
• e2e in driver_e2e_test.go: open a connector with a new WithCloudFetchHTTPClient(&http.Client{Timeout: 2*time.Minute, Transport: countingRT}). countingRT wraps http.DefaultTransport and counts requests. Run the large CloudFetch query, read every row, and assert: the exact row count, countingRT saw at least one request (the presigned-URL GETs went through the supplied client), and every counted request host differs from the workspace host (the Thrift transport was not replaced). Add a connector_test.go unit test: the option sets cfg.HTTPClient (same pointer, Timeout preserved) and leaves cfg.Transport nil. Also add one where WithTransport and WithCloudFetchHTTPClient are both given in either order; the explicit client must win. (Update references to new repo location #2 no way to give CloudFetch downloads their own http.Client (and so a Timeout) without replacing the Thrift transport)
• Unit test in auth/oauth/m2m (offline-only carve-out: the live warehouse can't be made to hang its token endpoint; the expectation comes from the issue, matching the 10s cap discovery already has). Use an httptest server that hangs on the token URL and build an authClient whose config points TokenURL at it, bypassing discovery through a small seam (e.g. a package-level tokenExchangeTimeout var and/or setting tokenSource from a clientcredentials.Config). Call Authenticate in a goroutine with the timeout shortened to ~200ms and assert it returns a non-nil error within a few seconds instead of hanging. Fix: make GetTokenSource bind the source to a context carrying oauth2.HTTPClient = &http.Client{Timeout: tokenExchangeTimeout} (default about 10–30s), so each exchange is bounded. (Revert "Fix the module name" #4 M2M token exchange is unbounded: a hung token endpoint blocks Authenticate forever)

Files changed

• driver_e2e_test.go
• connector_test.go
• auth/oauth/m2m/m2m_test.go
• connector.go
• auth/oauth/m2m/m2m.go

Test plan

• TestCloudFetchConnOptions — fails (red) against the original code, passes (green) after the fix
• TestM2MTokenExchangeIsBounded — fails (red) against the original code, passes (green) after the fix

🤖 Generated by engineer-bot (bug-fix flow) — review before merge.

Not changed in this PR (out of scope)

These facets could not be fixed here. Reason given for each — please follow up as needed:

• Live e2e green run of TestE2ECloudFetchMinTimeToExpiryIsHonored and TestE2ECloudFetchHTTPClientOption — This shell scrubs the DATABRICKS_PECOTESTING_* credentials, so both tests compile but skip here.

  Before this fix they failed because the new ConnOptions didn't exist. The fix routes cfg.MinTimeToExpiry and cfg.HTTPClient into the existing batchloader code (isLinkExpired, cloudIPCStreamIterator httpClient), but that still needs to be checked by running these tests with warehouse creds set.

• Caller QueryContext cancellation does not stop CloudFetch downloads (issue part 1) — The scanner's ctx (rows.ctx) is deliberately detached from caller cancellation via context.WithoutCancel (Fix context loss in polling and connection close operations #295/[ES-1934053] Detach result streaming from QueryContext cancellation #371, protected by TestE2EArrowBatchesSurviveQueryContextCancellation). Replacing context.Background() in arrowRows.go:126/128 would not add QueryContext cancellation, and changing the design is a maintainer decision.

  rows.Close already cancels downloads through RowScanner.Close. WithCloudFetchHTTPClient (with a Timeout) is now the supported way to bound a stalled download.

• Re-fetching result links on expiry (issue part 3) — This would be a new feature (re-issuing links through FetchResults), well beyond a minimal bug fix.

  The WithCloudFetchMinTimeToExpiry doc comment now says that an expired link is terminal and the query must be re-run.

[peco-review-bot]

Verdict: 1 Medium · 1 Low · 2 Nit

The fix is mostly sound, with 1 medium concern: a negative WithCloudFetchMinTimeToExpiry value skips the default clamp and makes the expiry check accept links that have already expired. The other findings are low-severity: the new options are silently ignored on the kernel backend, the mutex change has no test, and the token-exchange transport is hard-coded.

Other findings

• ⚪ Nit — The lock-scope change (no longer holding c.mx across a cached ts.Token()) has no test. TestM2MTokenExchangeIsBounded sets c.tokenSource directly and calls Authenticate once, so it only checks the HTTP timeout. A regression back to holding the lock for the whole call would still pass. A test that runs two concurrent Authenticate calls against a slow token server and checks that both finish in about one timeout (or that the second one isn't serialized behind the first) would pin this behavior. Separately, the first-fetch path (tokenSource == nil) still holds c.mx through GetConfig plus Token(). That is bounded now (10s + 30s), but other callers can still be blocked for up to about 40s.
• ⚪ Nit — The token-exchange client is a fresh &http.Client{Timeout: ...}, which uses http.DefaultTransport. So the M2M token request ignores any WithTransport / TLS / proxy settings on the connector, same as before. Pinning the timeout in now makes it a fixed 30s that callers can't tune. That's fine as a safe default, but if a corporate proxy or a custom TLS setup needs the connector's transport, this path can't use it. Consider at least documenting the limitation on GetTokenSource.

[peco-review-bot]

Verdict: 1 Low

Looks good, with one low-severity edge case: passing nil to WithCloudFetchHTTPClient wipes the CloudFetch client that WithTransport set up. The M2M timeout and the change to stop holding the mutex during a cached token refresh are correct and covered by tests that fail without the fix. The e2e tests still need a maintainer to run them against a live warehouse with credentials, as the description says.

[peco-review-bot]

Verdict: 2 Low

Looks good overall. The two ConnOptions are correct, WithTransport only fills HTTPClient when it's nil so precedence works in either option order, and cfg.HTTPClient is only used by the CloudFetch batch loader. The 30s token-exchange client timeout fixes the hang. I left two low-severity notes: the mutex change doesn't do what its comment claims, and the test doesn't cover the first-call Authenticate branch. The e2e tests still need a maintainer run with credentials, as the PR description says.

[peco-review-bot]

Verdict: 1 Medium · 1 Low

The new ConnOptions behave as their docs say, and the precedence and nil handling are tested. The M2M token exchange now has a timeout. One medium issue: the PR description says the Authenticate mutex was narrowed, but the code still holds it across Token().

[peco-review-bot]

Verdict: 1 Low

Looks good overall, with one low-severity issue: a comment in m2m.go misdescribes how concurrent token refreshes behave. The ConnOption precedence logic matches its docs and tests. The token-exchange timeout bound is correct and has offline tests that would fail without the fix. As the PR notes, the two e2e tests still need a run with live credentials.

[peco-review-bot]

🔵 Low — The new comment says that because c.mx is no longer held, "a slow token exchange does not serialize every request on the connector". That isn't accurate. The ReuseTokenSource returned by clientcredentials.Config.TokenSource holds its own mutex for the whole refresh (s.mu.Lock(); … s.new.Token()). So every concurrent Authenticate on this connector still waits behind a slow or hung exchange. They just wait on oauth2's lock now instead of c.mx. The fix that actually works is the tokenExchangeTimeout bound in GetTokenSource. Dropping c.mx on the cached path changes nothing for concurrency. Please reword the comment (and the PR description's claim) so a later reader doesn't count on non-serialization that isn't there. Something like: "c.mx only guards lazy init; concurrent refreshes are still serialized by ReuseTokenSource, but each is bounded by tokenExchangeTimeout."

[peco-engineer-bot]

⛔ NEEDS HUMAN DECISION — the bots can't resolve this thread; a maintainer's input is required.

I agree, and the code comment already says this at HEAD, so I made no edit. In auth/oauth/m2m/m2m.go, the comment in Authenticate just above ts.Token() now reads: "c.mx guards only the lazy init above. Token refreshes are still serialized inside the oauth2.ReuseTokenSource (it holds its own mutex for the whole exchange), so concurrent callers wait on an in-flight refresh; that wait is bounded by tokenExchangeTimeout." The comment on getOrInitTokenSource also says c.mx guards only the lazy init, and the tokenExchangeTimeout doc names the timeout as the actual fix. The PR description still makes the "does not serialize every request" claim, and I can't edit it from here. A human needs to reword that line to match the code comment.

[peco-review-bot]

Verdict: 1 Low

Looks good, with 1 low-severity doc concern about how Client.Timeout interacts with CloudFetch retries. The precedence between the new ConnOption and WithTransport holds because WithTransport only sets HTTPClient when it is nil. The M2M lock narrowing is safe because oauth2.ReuseTokenSource still serializes refreshes. Nit: the PR lists internal/rows/arrowbased/arrowRows.go as a touched file, but the diff doesn't change it, so the description should be corrected.

[peco-review-bot]

Verdict: 1 Low

Looks good, with one low concern: the M2M exchange can take about twice the documented bound because of oauth2's AuthStyle auto-detect retry. I checked that the precedence between WithCloudFetchHTTPClient and WithTransport holds in both option orders (WithTransport only sets HTTPClient when it is nil). The lazy init under the mutex and the token refresh outside it look correct.

[peco-review-bot]

Verdict: 1 Medium · 1 Low

Looks good overall, with one medium doc/contract concern. The new WithCloudFetchHTTPClient doc tells users to rely on the query context for an overall deadline, but rows deliberately strip the caller's cancellation from download contexts. Nit: the PR description lists internal/rows/arrowbased/arrowRows.go under Files, but the diff doesn't touch it.

[peco-review-bot]

Verdict: 1 Low

Looks good, with one low-severity concern. The new ConnOptions do what their doc comments say, which I checked against batchloader.go (retry loop, per-attempt timeout, one-second expiry resolution, cfg.HTTPClient used only for CloudFetch). The M2M change is also correct: the lazy init is locked, refreshes are serialized by ReuseTokenSource, and the timeout is bounded per request. The e2e tests still need a maintainer run with credentials before merge, as the PR description says.

[peco-review-bot]

Verdict: 1 Low

Looks good: the new ConnOptions are wired correctly, and the M2M timeout and lock-scope changes are sound. They match how the retry loop in fetchBatchBytes and the mutex inside oauth2.ReuseTokenSource actually behave. One low finding: an e2e assertion that can never fail. One nit: the PR description lists internal/rows/arrowbased/arrowRows.go as a changed file, but the diff doesn't touch it.

[peco-review-bot]

Verdict: 1 Low

Looks good, with one low-severity contract note. Option precedence (WithTransport sets HTTPClient only when it is nil), the retry/expiry wording in the doc comments (matches batchloader.go:404-453), and the M2M lock narrowing (ReuseTokenSource still serializes refreshes) all check out. Nit: the PR description lists internal/rows/arrowbased/arrowRows.go as a changed file, but the diff doesn't touch it. The root-cause section explains why the context change was dropped, so the file list should be updated to match.

[peco-review-bot]

✅ No issues identified by the review bot.