Skip to content

Conversation

@alexguo-db
Copy link
Contributor

@alexguo-db alexguo-db commented Jul 23, 2025

Motivation

Databricks will eventually require that all non-inhouse OAuth tokens be exchanged for Databricks OAuth tokens before accessing resources. This change implements mandatory token exchange before sending Thrift requests. This check and exchange is performed in the background for now to reduce latency, but it will eventually need to be blocking if non-inhouse OAuth tokens will fail to access Databricks resources in the future.

Key Components

  1. JWT Token Decoder - Decodes JWT tokens to inspect the issuer claim and determine if token exchange is necessary
  2. MandatoryTokenExchangeDelegatingHandler - HTTP handler that intercepts requests and performs token exchange when required
  3. TokenExchangeClient - Handles the token exchange logic with the same /oidc/v1/token endpoint as token refresh, with slightly different parameters

Changes

  • Added new connection string parameter: IdentityFederationClientId for service principal workload identity federation scenarios
  • Implemented token exchange logic that checks JWT issuer against workspace host
  • Introduced fallback behavior to maintain backward compatibility if token exchange fails

Testing

dotnet test --filter "FullyQualifiedName~MandatoryTokenExchangeDelegatingHandlerTests"

[xUnit.net 00:00:00.00] xUnit.net VSTest Adapter v3.1.1+bf6400fd51 (64-bit .NET 8.0.7)
[xUnit.net 00:00:00.06]   Discovering: Apache.Arrow.Adbc.Tests.Drivers.Databricks
[xUnit.net 00:00:00.15]   Discovered:  Apache.Arrow.Adbc.Tests.Drivers.Databricks
[xUnit.net 00:00:00.16]   Starting:    Apache.Arrow.Adbc.Tests.Drivers.Databricks
[xUnit.net 00:00:01.77]   Finished:    Apache.Arrow.Adbc.Tests.Drivers.Databricks
  Apache.Arrow.Adbc.Tests.Drivers.Databricks test net8.0 succeeded (2.6s)

Test summary: total: 11, failed: 0, succeeded: 11, skipped: 0, duration: 2.6s

dotnet test --filter "FullyQualifiedName~TokenExchangeClientTests"

[xUnit.net 00:00:00.00] xUnit.net VSTest Adapter v3.1.1+bf6400fd51 (64-bit .NET 8.0.7)
[xUnit.net 00:00:00.06]   Discovering: Apache.Arrow.Adbc.Tests.Drivers.Databricks
[xUnit.net 00:00:00.14]   Discovered:  Apache.Arrow.Adbc.Tests.Drivers.Databricks
[xUnit.net 00:00:00.15]   Starting:    Apache.Arrow.Adbc.Tests.Drivers.Databricks
[xUnit.net 00:00:00.23]   Finished:    Apache.Arrow.Adbc.Tests.Drivers.Databricks
  Apache.Arrow.Adbc.Tests.Drivers.Databricks test net8.0 succeeded (0.8s)

Test summary: total: 19, failed: 0, succeeded: 19, skipped: 0, duration: 0.8s

dotnet test --filter "FullyQualifiedName~JwtTokenDecoderTests"

[xUnit.net 00:00:00.00] xUnit.net VSTest Adapter v3.1.1+bf6400fd51 (64-bit .NET 8.0.7)
[xUnit.net 00:00:00.06]   Discovering: Apache.Arrow.Adbc.Tests.Drivers.Databricks
[xUnit.net 00:00:00.14]   Discovered:  Apache.Arrow.Adbc.Tests.Drivers.Databricks
[xUnit.net 00:00:00.15]   Starting:    Apache.Arrow.Adbc.Tests.Drivers.Databricks
[xUnit.net 00:00:00.19]   Finished:    Apache.Arrow.Adbc.Tests.Drivers.Databricks
  Apache.Arrow.Adbc.Tests.Drivers.Databricks test net8.0 succeeded (0.8s)

Test summary: total: 10, failed: 0, succeeded: 10, skipped: 0, duration: 0.8s

Also tested E2E manually with AAD tokens for Azure Databricks workspaces, AAD tokens for AWS Databricks workspaces, and service principal workload identity federation tokens

@alexguo-db alexguo-db force-pushed the alex-guo_data/add-mandatory-token-exchange branch from fff5a7c to 8079738 Compare July 24, 2025 00:09
@alexguo-db alexguo-db force-pushed the alex-guo_data/add-mandatory-token-exchange branch from 8079738 to 162c114 Compare July 24, 2025 00:32
@alexguo-db alexguo-db force-pushed the alex-guo_data/add-mandatory-token-exchange branch from 97803c3 to fc75b07 Compare July 25, 2025 17:15
@alexguo-db alexguo-db force-pushed the alex-guo_data/add-mandatory-token-exchange branch from fc75b07 to e80892e Compare July 25, 2025 17:38
@alexguo-db alexguo-db marked this pull request as ready for review July 25, 2025 17:45
@github-actions github-actions bot added this to the ADBC Libraries 20 milestone Jul 25, 2025
/// The client ID of the service principal when using workload identity federation.
/// Default value is empty if not specified.
/// </summary>
public const string IdentityFederationClientId = "adbc.databricks.identity_federation_client_id";
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Can this also be added to the documentation at csharp\src\Drivers\Databricks\readme.md?

I haven't always been paying attention to that, so there may be a bit of a gap now which could be addressed all at once with a separate PR.

@CurtHagenlocher CurtHagenlocher merged commit 7ff3364 into apache:main Jul 31, 2025
8 checks passed
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.

2 participants