Skip to content

fix(auth): thread connector transport and User-Agent into OAuth M2M token requests - #498

Open
sevbanbayrak wants to merge 1 commit into
databricks:mainfrom
sevbanbayrak:fix/m2m-token-user-agent
Open

sevbanbayrak wants to merge 1 commit into
databricks:mainfrom
sevbanbayrak:fix/m2m-token-user-agent

Conversation

@sevbanbayrak

Copy link
Copy Markdown

Fixes #413

Problem

With WithClientCredentials the OAuth M2M token exchange (and the OIDC discovery before it) is issued through http.DefaultClient, so system.access.audit records mintOAuthToken and workspaceInHouseOAuthClientAuthentication as Go-http-client/2.0 instead of the driver User-Agent and the WithUserAgentEntry value. Query traffic is attributed correctly; only the token path is anonymous, which matters for ISV partner attribution.

Root cause

The M2M authenticator is constructed when its option is applied (WithClientCredentials, or the DSN path), before WithUserAgentEntry / WithTransport may have been applied, and it created its token source on a bare context.Background(), so no HTTP client could reach oauth2.

Fix

  • internal/client.PooledClient runs once the config is final. It now hands any authenticator that exposes SetHTTPClient(*http.Client) a client on the same base transport (WithTransport or the pooled one) wrapped with BuildUserAgent(cfg). It deliberately wraps the base transport, not the auth Transport, so a token fetch cannot recurse into Authenticate.
  • m2m.authClient implements SetHTTPClient (structural, same pattern as M2MCredentials for the kernel backend) and puts that client on the oauth2.HTTPClient context used for GetConfig and TokenSource.
  • oauth.GetEndpoint reuses the transport of an oauth2.HTTPClient found on ctx for the /.well-known/databricks-config lookup and OIDC discovery, closing the in-tree NOTE about WithTransport not being threaded into endpoint resolution.

Authenticators without SetHTTPClient (PAT, U2M, token providers) are unchanged. No public API change; option order does not matter.

Tests

  • m2m: Authenticate on an Azure host (offline endpoint) sends the token request through the injected client and reuses the cached token source.
  • oauth: GetEndpoint performs discovery through the ctx client's transport and resolves the served token_endpoint.
  • client: PooledClient injects a client whose transport carries BuildUserAgent(cfg) over cfg.Transport, without mutating the caller's request; authenticators without the hook are untouched.

go test ./... and golangci-lint run are clean locally.

Verification against a workspace

Ran SELECT 1 through a connector configured with WithClientCredentials + WithUserAgentEntry("Insider_UATest413/…") against an AWS workspace, then queried system.access.audit:

action_name before (driver 1.13.0 traffic, same window) after (this branch)
workspaceInHouseOAuthClientAuthentication Go-http-client/2.0 godatabrickssqlconnector/1.16.0 (Insider_UATest413/…)
mintOAuthToken Go-http-client/2.0 godatabrickssqlconnector/1.16.0 (Insider_UATest413/…)
oidcTokenAuthorization already correct unchanged

…oken requests

With WithClientCredentials the token exchange (and the OIDC discovery before
it) was issued through http.DefaultClient, so system.access.audit recorded
mintOAuthToken and workspaceInHouseOAuthClientAuthentication as
"Go-http-client/2.0" instead of the driver User-Agent and WithUserAgentEntry.

The M2M authenticator is built when its option is applied, before the
User-Agent entry or transport may be set, so the driver now injects an
*http.Client once the config is final: PooledClient passes a client on the
same base transport, with BuildUserAgent, to any authenticator exposing
SetHTTPClient. The M2M authenticator stores it on the oauth2 context used for
GetConfig and TokenSource, and GetEndpoint reuses that client's transport for
discovery. Authenticators without SetHTTPClient are unaffected.

Fixes databricks#413

Signed-off-by: Sevban Bayrak <svbnbyrk@gmail.com>

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

OAuth M2M: token requests ignore the connector's User-Agent (WithUserAgentEntry not threaded into the token exchange)

1 participant