[gabepesco]

Description

#5859 fixed #5858 by giving U2M OAuth a single shared connection. That stops the browser login racing (#5646), but every query then goes through one connection, and the connector serializes requests on it. So concurrent_tasks has no effect for U2M users.

This PR keeps one connection per thread and shares the login instead:

• One login for all connections. SharedU2MCredentialsProvider is passed to every connection as credentials_provider. It creates the connector's own OAuth provider once, under a lock, so the browser flow runs a single time per workspace in each process.
• One token refresh at a time. Header reads are serialized. Databricks can issue single-use refresh tokens, so a separate token copy per connection would break on the first refresh.
• Recovery after a failed login. If a token read fails (refresh token expired or revoked), the login is discarded and the next request logs in again. Threads waiting on a login that fails get its error instead of each opening another browser window. A network failure during a refresh keeps the login, so the next request retries the refresh without a browser flow.
• Fallback for older connectors. Connectors without the internals this relies on (< 4.1) keep fix: Set DatabricksConnectionConfig to only use a shared_connection with U2M OAuth #5859's single shared connection, so shared_connection is now True only in that case.

PAT and M2M auth are unchanged. No state or schema changes, so no migration is needed.

Known gap: auth_type is still passed to the connector because its kernel backend rejects a custom credentials_provider. On warehouses where connector 4.6+ falls back to that backend, each connection would run its own login. I haven't been able to verify this path.

Related: #5646, #5746, #5858, #5859

Test Plan

• New unit tests in tests/utils/test_databricks.py. They cover a single browser flow across threads, serialized token reads, retry after a failed login, waiting threads not retrying a failed login, and re-login after a failed token read, and keeping the login after a network failure during refresh. Two tests run the real connector with only OAuthManager.get_tokens mocked; one builds per-connection auth exactly as connect() does, with both auth_type and the shared provider. Each test fails when the fix it covers is reverted.
• Updated tests/core/test_connection_config.py: U2M now uses ThreadLocalConnectionPool with the shared provider, and the old-connector fallback uses the shared connection.
• Live check against a Databricks workspace with U2M (databricks-oauth), running 4 threads with 15 small queries each. There was one browser login and 4 distinct connections, and it finished in 6.5s vs 13.9s with the shared connection (2.1x).
• Not verified live: token refresh after the ~1h access-token expiry (covered by mocked tests only), and the kernel backend path noted above.

Checklist

• I have run make style and fixed any issues (ruff and mypy run directly; the pre-commit hooks need bash, which I don't have on Windows)
• I have added tests for my changes (if applicable)
• All existing tests pass (make fast-test) (related suites pass; I didn't run the full fast-test suite locally)
• My commits are signed off (git commit -s) per the DCO

🤖 Generated with Claude Code

Summary by CodeRabbit

• New Features
  • Databricks U2M OAuth can reuse sign-in credentials across connections when the connector supports shared login, reducing repeated browser sign-ins.
• Improvements
  • When shared login is unavailable, Databricks connections use the connector’s existing shared-connection behavior.
  • Temporary sign-in or credential-refresh failures can be retried, while network errors do not trigger a new browser sign-in.
  • Credentials are reused across equivalent workspace hostnames and authentication types.

[coderabbitai]

No actionable comments were generated in the recent review.

Walkthrough

Databricks U2M OAuth connections now use a shared, lazily initialized credentials provider when connector support is available. The provider serializes authentication and header generation, and is reused by normalized workspace and auth-type keys. Unsupported connectors retain shared-connection behavior.

Changes: Databricks U2M OAuth

Layer / File(s)	Summary
Shared provider lifecycle and validation sqlmesh/utils/databricks.py, tests/utils/test_databricks.py	A registry reuses providers by normalized workspace and auth type. Each provider creates its auth provider lazily, serializes header generation, and allows later attempts after failures. Tests cover concurrent callers, failure handling, provider reuse, and connector integration.
Connection selection and provider wiring sqlmesh/core/config/connection.py, tests/core/test_connection_config.py	Configuration uses a thread-local pool and passes a shared credentials provider when the connector supports shared U2M login. Tests cover provider reuse and shared-connection fallback when support is unavailable.

Estimated code review effort: 3 (Moderate) | ~25 minutes

Change: Bug fix · Severity of issue fixed: Medium

Sequence Diagram(s)

sequenceDiagram
  participant DatabricksConnectionConfig
  participant DatabricksSQLConnector
  participant SharedU2MCredentialsProvider
  participant DatabricksAuthProvider
  DatabricksConnectionConfig->>DatabricksSQLConnector: create connection with auth_type and provider
  DatabricksSQLConnector->>SharedU2MCredentialsProvider: request authorization headers
  SharedU2MCredentialsProvider->>DatabricksAuthProvider: create provider and read headers
  DatabricksAuthProvider-->>SharedU2MCredentialsProvider: authorization headers
  SharedU2MCredentialsProvider-->>DatabricksSQLConnector: authorization headers

Loading

Suggested reviewers: stuffbyyuki, mday-io

Merge Risk: Minimal · up to dc9ad

No actionable merge-blocking risk remains from the reviewed evidence.

Security architecture risk: Moderate · up to dc9ad

The shared provider adds important protections against competing logins and token refreshes. However, enabling concurrent connections depends on a capability check that does not establish whether the selected connector backend shares credentials. Some supported deployments may therefore lose the previous single-login protection. No credential theft or cross-user access was demonstrated.

Retained concerns

• Medium · security · inferred: Per-thread connections are enabled when a connector module exists, without checking whether the effective backend honors the shared credentials provider. The disclosed kernel path may therefore run independent browser logins, reintroducing the authentication-state races that the base avoided through one shared connection. The kernel outcome is unverified; the auth-factory test does not resolve backend dispatch.

Security review details

Security Blast Radius

• observed — The sharing scope is every U2M adapter in one process using the same normalized hostname and auth type, including different HTTP paths. Other hostnames and auth types receive distinct providers. Client-secret OAuth remains on its separate service-principal branch.

Security Findings and Attack Paths

• inferred — If backend dispatch bypasses the supplied provider, concurrent connections can initiate independent OAuth browser flows despite the local provider lock. This threatens authentication-state coordination and failure containment, not a demonstrated CSRF bypass, privilege escalation, or credential-exfiltration path. The connector implementation needed to settle this condition was unavailable.

Trust Boundaries and Controls

• observed — Identity sharing is intentional within the documented single-user process boundary. The provider delegates OAuth login to the connector rather than implementing a new authorization flow. A missing oauth_client_id component in the registry is not an established regression: U2M did not forward that field at the base either.

Resilience and Maintainability Implications

• observed — The lock bounds refresh concurrency to one header read per shared provider, and authentication errors propagate rather than producing unauthenticated headers. Recovery can reauthenticate through existing callbacks after invalidation. Tests assert serialization and recovery, but use simulated refresh failures rather than real single-use-token rotation or interrupted refresh responses.

Hardening Proposals

• proposed — Make concurrent-connection eligibility depend on the effective backend honoring shared credentials, with shared-connection fallback otherwise. Validate that boundary through actual connection/backend selection, not only the lower-level auth factory.

Pre-merge checks: 4 · 1

Failed checks: 1 warning

Check name	Status	Explanation	Resolution
Docstring Coverage		Docstring coverage is 46.15% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff with cyclomatic complexity of at least 2. Documenta…	Write docstrings for the functions missing them to satisfy the coverage threshold.

Passed checks: 4 passed

Check name	Status	Explanation
Title check		The title is 'fix(databricks): run U2M OAuth queries concurrently with a shared login'. It matches the main change: one connection per thread for U2M OAuth, with the login shared through SharedU2MCred…
Description check		The description follows the template. It has a Description section that explains the change: one connection per thread, a shared login, serialized token refresh, recovery after failed login, and a fal…
Linked Issues check		[#5858] is the only directly linked issue. It is closed and completed, so it gives historical context only. No active directly linked targets remain, so no linked-issue coding requirements apply. For …
Out of Scope Changes check		The whole-PR diff (merge base c4dbac4 to head dc9adec) changes four files: sqlmesh/core/config/connection.py, sqlmesh/utils/databricks.py (new), tests/core/test_connection_config.py, and tests/utils/t…

Full details: Title check

Explanation

The title is 'fix(databricks): run U2M OAuth queries concurrently with a shared login'. It matches the main change: one connection per thread for U2M OAuth, with the login shared through SharedU2MCredentialsProvider. The title is specific and follows the conventional-commit style. It is a little long (71 characters), but it is clear.

Full details: Description check

Explanation

The description follows the template. It has a Description section that explains the change: one connection per thread, a shared login, serialized token refresh, recovery after failed login, and a fallback for older connectors. It has a Test Plan section that lists unit tests, a updated connection config test, and a live check with timings. The Checklist is present. Two gaps remain: the 'All existing tests pass' item is unchecked, and the author states that the full make fast-test suite was not run. The author also documents known gaps: the kernel-backend path and live token refresh after expiry were not verified.

Full details: Linked Issues check

Explanation

[#5858] is the only directly linked issue. It is closed and completed, so it gives historical context only. No active directly linked targets remain, so no linked-issue coding requirements apply. For context, #5858 asked for U2M OAuth queries to run concurrently instead of sequentially. The diff addresses that context. DatabricksConnectionConfig.shared_connection returns False for U2M OAuth when supports_shared_u2m_login() is True (sqlmesh/core/config/connection.py). SharedU2MCredentialsProvider in sqlmesh/utils/databricks.py passes one shared login to each connection as credentials_provider. Older connectors keep a shared connection. The tests cover these paths. The kernel-backend limitation noted in the PR description is a known gap. It is not an active requirement of #5858.

Full details: Out of Scope Changes check

Explanation

The whole-PR diff (merge base c4dbac4 to head dc9adec) changes four files: sqlmesh/core/config/connection.py, sqlmesh/utils/databricks.py (new), tests/core/test_connection_config.py, and tests/utils/test_databricks.py. Every change maps to the PR purpose of per-thread U2M connections with one shared login. This covers the shared provider and its locking, serialized token reads, login retry and error propagation, the supports_shared_u2m_login() fallback, and the updated tests. The _uses_u2m_oauth helper narrows the U2M check to auth types containing "oauth" without a client secret. This check is part of the same U2M path, so it is in scope. No PAT, M2M, schema, or state changes appear in the diff.

Full details: Docstring Coverage

Explanation

Docstring coverage is 46.15% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff with cyclomatic complexity of at least 2. Documentation-only edits and deletions within surviving functions are included; unchanged diff context and fully deleted functions are excluded. Analyzed 13 functions across 4 files.

Resolution

Write docstrings for the functions missing them to satisfy the coverage threshold.

Finishing Touches

Generate unit tests (beta)

• Create a new PR

• Autofix · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts

Tip

🎯 Customize your risk assessments

Guide how CodeRabbit assesses merge risk and security architecture risk. Highlight sensitive changes and explain your security boundaries using Merge Risk Instructions and Security Risk Instructions in your review settings. Security architecture risk assessments require the Advanced plan or above.

✅ Automatic approvals, with your rules

Let CodeRabbit approve pull requests once required threads are resolved, pre-merge checks pass, and your risk limits are met. Enable Automatic approvals in review settings, with overrides for teams or file paths.

CodeRabbit Verify

• Run CodeRabbit Verify to test the changes in this pull request.

Recent review info

Run configuration

• Source mode: Auto
• Repository: Unavailable
• Branch: Unavailable
• File: No configuration file from this repository applied
• Commit: Unavailable
• Layers: defaults
• Review profile: CHILL
• Plan: Advanced
• Run ID: a825aee5-ef78-42fd-b1c5-c74e3dc12826

Commits

Reviewing files that changed from the base of the PR and between b8bdb47 and dc9adec.

Files selected for processing: 2

• sqlmesh/utils/databricks.py
• tests/utils/test_databricks.py

Included review availability: This review used your included allowance. Your plan provides up to 10 included reviews per hour; 9 remain after this review.

[coderabbitai]

• Fix comments: Fix CodeRabbit comments on this PR