You are viewing a plain text version of this content. The canonical link for it is here.
Posted to github@arrow.apache.org by "alexguo-db (via GitHub)" <gi...@apache.org> on 2025/07/18 23:13:58 UTC
[PR] feat(csharp/test/Drivers/Databricks): Support token refresh to extend connection lifetime [arrow-adbc]
alexguo-db opened a new pull request, #3177:
URL: https://github.com/apache/arrow-adbc/pull/3177
## Motivation
In scenarios like PowerBI dataset refresh, if a query runs longer than the OAuth token's expiration time (typically 1 hour for AAD tokens), the connection fails. PowerBI only refreshes access tokens if they have less than 20 minutes of expiration time and never updates tokens after a connection is opened.
This PR implements token refresh functionality in the Databricks ADBC driver using the Databricks token exchange API. When an OAuth token is about to expire within a configurable time limit, the driver automatically exchanges it for a new token with a longer expiration time.
## Key Components
1. **JWT Token Decoder**: Parses JWT tokens to extract expiration time
2. **Token Exchange Client**: Handles API calls to the Databricks token exchange endpoint
3. **Token Exchange Handler**: HTTP handler that intercepts requests and refreshes tokens when needed
## Changes
- Added new connection string parameter `adbc.databricks.token_renew_limit` to control when token renewal happens
- Implemented JWT token decoding to extract token expiration time
- Created token exchange client to handle API calls to Databricks token exchange endpoint
- Added HTTP handler to intercept requests and refresh tokens when needed
- Updated connection handling to create and configure the token exchange components
## Testing
- Unit tests for JWT token decoding, token exchange client, and token exchange handler
- End-to-end tests that verify token refresh functionality with real tokens
```
dotnet test --filter "FullyQualifiedName~JwtTokenDecoderTests"
dotnet test --filter "FullyQualifiedName~TokenExchangeClientTests"
dotnet test --filter "FullyQualifiedName~TokenExchangeDelegatingHandlerTests"
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] feat(csharp/test/Drivers/Databricks): Support token refresh to extend connection lifetime [arrow-adbc]
Posted by "CurtHagenlocher (via GitHub)" <gi...@apache.org>.
CurtHagenlocher commented on code in PR #3177:
URL: https://github.com/apache/arrow-adbc/pull/3177#discussion_r2220971029
##########
csharp/src/Drivers/Databricks/Auth/JwtTokenDecoder.cs:
##########
@@ -0,0 +1,96 @@
+/*
+* Licensed to the Apache Software Foundation (ASF) under one or more
+* contributor license agreements. See the NOTICE file distributed with
+* this work for additional information regarding copyright ownership.
+* The ASF licenses this file to You under the Apache License, Version 2.0
+* (the "License"); you may not use this file except in compliance with
+* the License. You may obtain a copy of the License at
+*
+* http://www.apache.org/licenses/LICENSE-2.0
+*
+* Unless required by applicable law or agreed to in writing, software
+* distributed under the License is distributed on an "AS IS" BASIS,
+* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+* See the License for the specific language governing permissions and
+* limitations under the License.
+*/
+
+using System;
+using System.Text;
+using System.Text.Json;
+
+namespace Apache.Arrow.Adbc.Drivers.Databricks.Auth
+{
+ /// <summary>
+ /// Utility class for decoding JWT tokens and extracting claims.
+ /// </summary>
+ internal static class JwtTokenDecoder
Review Comment:
I would have hand-rolled this as well.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] feat(csharp/test/Drivers/Databricks): Support token refresh to extend connection lifetime [arrow-adbc]
Posted by "alexguo-db (via GitHub)" <gi...@apache.org>.
alexguo-db commented on code in PR #3177:
URL: https://github.com/apache/arrow-adbc/pull/3177#discussion_r2220271743
##########
csharp/src/Drivers/Databricks/DatabricksConnection.cs:
##########
@@ -334,44 +336,78 @@ protected override HttpMessageHandler CreateHttpHandler()
baseHandler = new RetryHttpHandler(baseHandler, TemporarilyUnavailableRetryTimeout);
}
- // Add OAuth handler if OAuth authentication is being used
+ // Add OAuth client credentials handler if OAuth M2M authentication is being used
if (Properties.TryGetValue(SparkParameters.AuthType, out string? authType) &&
SparkAuthTypeParser.TryParse(authType, out SparkAuthType authTypeValue) &&
authTypeValue == SparkAuthType.OAuth &&
Properties.TryGetValue(DatabricksParameters.OAuthGrantType, out string? grantTypeStr) &&
DatabricksOAuthGrantTypeParser.TryParse(grantTypeStr, out DatabricksOAuthGrantType grantType) &&
grantType == DatabricksOAuthGrantType.ClientCredentials)
{
- // Note: We assume that properties have already been validated
- if (Properties.TryGetValue(SparkParameters.HostName, out string? host) && !string.IsNullOrEmpty(host))
- {
- // Use hostname directly if provided
- }
- else if (Properties.TryGetValue(AdbcOptions.Uri, out string? uri) && !string.IsNullOrEmpty(uri))
- {
- // Extract hostname from URI if URI is provided
- if (Uri.TryCreate(uri, UriKind.Absolute, out Uri? parsedUri))
- {
- host = parsedUri.Host;
- }
- }
+ string host = GetHost();
Properties.TryGetValue(DatabricksParameters.OAuthClientId, out string? clientId);
Properties.TryGetValue(DatabricksParameters.OAuthClientSecret, out string? clientSecret);
Properties.TryGetValue(DatabricksParameters.OAuthScope, out string? scope);
- HttpClient OauthHttpClient = new HttpClient(HiveServer2TlsImpl.NewHttpClientHandler(TlsOptions, _proxyConfigurator));
+ if (_httpClient == null)
+ {
+ _httpClient = new HttpClient(HiveServer2TlsImpl.NewHttpClientHandler(TlsOptions, _proxyConfigurator));
+ }
var tokenProvider = new OAuthClientCredentialsProvider(
- OauthHttpClient,
+ _httpClient,
clientId!,
clientSecret!,
host!,
scope: scope ?? "sql",
timeoutMinutes: 1
);
- return new OAuthDelegatingHandler(baseHandler, tokenProvider);
+ baseHandler = new OAuthDelegatingHandler(baseHandler, tokenProvider);
+ }
+
+ // Add token exchange handler if token renewal is enabled and the auth type is OAuth access token
Review Comment:
Yes, I will make that change
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] feat(csharp/test/Drivers/Databricks): Support token refresh to extend connection lifetime [arrow-adbc]
Posted by "alexguo-db (via GitHub)" <gi...@apache.org>.
alexguo-db commented on code in PR #3177:
URL: https://github.com/apache/arrow-adbc/pull/3177#discussion_r2220275645
##########
csharp/src/Drivers/Databricks/DatabricksConnection.cs:
##########
@@ -334,44 +336,78 @@ protected override HttpMessageHandler CreateHttpHandler()
baseHandler = new RetryHttpHandler(baseHandler, TemporarilyUnavailableRetryTimeout);
}
- // Add OAuth handler if OAuth authentication is being used
+ // Add OAuth client credentials handler if OAuth M2M authentication is being used
if (Properties.TryGetValue(SparkParameters.AuthType, out string? authType) &&
SparkAuthTypeParser.TryParse(authType, out SparkAuthType authTypeValue) &&
authTypeValue == SparkAuthType.OAuth &&
Properties.TryGetValue(DatabricksParameters.OAuthGrantType, out string? grantTypeStr) &&
DatabricksOAuthGrantTypeParser.TryParse(grantTypeStr, out DatabricksOAuthGrantType grantType) &&
grantType == DatabricksOAuthGrantType.ClientCredentials)
{
- // Note: We assume that properties have already been validated
- if (Properties.TryGetValue(SparkParameters.HostName, out string? host) && !string.IsNullOrEmpty(host))
- {
- // Use hostname directly if provided
- }
- else if (Properties.TryGetValue(AdbcOptions.Uri, out string? uri) && !string.IsNullOrEmpty(uri))
- {
- // Extract hostname from URI if URI is provided
- if (Uri.TryCreate(uri, UriKind.Absolute, out Uri? parsedUri))
- {
- host = parsedUri.Host;
- }
- }
+ string host = GetHost();
Properties.TryGetValue(DatabricksParameters.OAuthClientId, out string? clientId);
Properties.TryGetValue(DatabricksParameters.OAuthClientSecret, out string? clientSecret);
Properties.TryGetValue(DatabricksParameters.OAuthScope, out string? scope);
- HttpClient OauthHttpClient = new HttpClient(HiveServer2TlsImpl.NewHttpClientHandler(TlsOptions, _proxyConfigurator));
+ if (_httpClient == null)
+ {
+ _httpClient = new HttpClient(HiveServer2TlsImpl.NewHttpClientHandler(TlsOptions, _proxyConfigurator));
Review Comment:
That makes sense to me, I added the tracing and retry functionality to both auth HTTP clients
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] feat(csharp/test/Drivers/Databricks): Support token refresh to extend connection lifetime [arrow-adbc]
Posted by "alexguo-db (via GitHub)" <gi...@apache.org>.
alexguo-db commented on code in PR #3177:
URL: https://github.com/apache/arrow-adbc/pull/3177#discussion_r2220505611
##########
csharp/src/Drivers/Databricks/Auth/JwtTokenDecoder.cs:
##########
@@ -0,0 +1,96 @@
+/*
+* Licensed to the Apache Software Foundation (ASF) under one or more
+* contributor license agreements. See the NOTICE file distributed with
+* this work for additional information regarding copyright ownership.
+* The ASF licenses this file to You under the Apache License, Version 2.0
+* (the "License"); you may not use this file except in compliance with
+* the License. You may obtain a copy of the License at
+*
+* http://www.apache.org/licenses/LICENSE-2.0
+*
+* Unless required by applicable law or agreed to in writing, software
+* distributed under the License is distributed on an "AS IS" BASIS,
+* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+* See the License for the specific language governing permissions and
+* limitations under the License.
+*/
+
+using System;
+using System.Text;
+using System.Text.Json;
+
+namespace Apache.Arrow.Adbc.Drivers.Databricks.Auth
+{
+ /// <summary>
+ /// Utility class for decoding JWT tokens and extracting claims.
+ /// </summary>
+ internal static class JwtTokenDecoder
Review Comment:
`System.IdentityModel.Tokens.Jwt` is not built-in, we need to add it as an external dependency. I remember last time when we added `K4os.Compression.LZ4` dependency it took a long time to get security approval for builds, so not sure if we should add another dependency
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] feat(csharp/test/Drivers/Databricks): Support token refresh to extend connection lifetime [arrow-adbc]
Posted by "jadewang-db (via GitHub)" <gi...@apache.org>.
jadewang-db commented on code in PR #3177:
URL: https://github.com/apache/arrow-adbc/pull/3177#discussion_r2220531077
##########
csharp/src/Drivers/Databricks/DatabricksConnection.cs:
##########
@@ -321,57 +323,78 @@ private void ValidateProperties()
protected override HttpMessageHandler CreateHttpHandler()
{
HttpMessageHandler baseHandler = base.CreateHttpHandler();
+ HttpMessageHandler baseAuthHandler = HiveServer2TlsImpl.NewHttpClientHandler(TlsOptions, _proxyConfigurator);
// Add tracing handler to propagate W3C trace context if enabled
if (_tracePropagationEnabled)
{
baseHandler = new TracingDelegatingHandler(baseHandler, this, _traceParentHeaderName, _traceStateEnabled);
+ baseAuthHandler = new TracingDelegatingHandler(baseAuthHandler, this, _traceParentHeaderName, _traceStateEnabled);
}
if (TemporarilyUnavailableRetry)
{
// Add retry handler for 503 responses
baseHandler = new RetryHttpHandler(baseHandler, TemporarilyUnavailableRetryTimeout);
+ baseAuthHandler = new RetryHttpHandler(baseAuthHandler, TemporarilyUnavailableRetryTimeout);
Review Comment:
also we can have a follow up PR to make current retry handler richer, I don't see they are conflict
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] feat(csharp/test/Drivers/Databricks): Support token refresh to extend connection lifetime [arrow-adbc]
Posted by "alexguo-db (via GitHub)" <gi...@apache.org>.
alexguo-db commented on code in PR #3177:
URL: https://github.com/apache/arrow-adbc/pull/3177#discussion_r2220508724
##########
csharp/src/Drivers/Databricks/DatabricksConnection.cs:
##########
@@ -321,57 +323,78 @@ private void ValidateProperties()
protected override HttpMessageHandler CreateHttpHandler()
{
HttpMessageHandler baseHandler = base.CreateHttpHandler();
+ HttpMessageHandler baseAuthHandler = HiveServer2TlsImpl.NewHttpClientHandler(TlsOptions, _proxyConfigurator);
// Add tracing handler to propagate W3C trace context if enabled
if (_tracePropagationEnabled)
{
baseHandler = new TracingDelegatingHandler(baseHandler, this, _traceParentHeaderName, _traceStateEnabled);
+ baseAuthHandler = new TracingDelegatingHandler(baseAuthHandler, this, _traceParentHeaderName, _traceStateEnabled);
}
if (TemporarilyUnavailableRetry)
{
// Add retry handler for 503 responses
baseHandler = new RetryHttpHandler(baseHandler, TemporarilyUnavailableRetryTimeout);
+ baseAuthHandler = new RetryHttpHandler(baseAuthHandler, TemporarilyUnavailableRetryTimeout);
Review Comment:
Hmm yes we should retry on other codes, 502, 502, 408. I will create a different retry delegating handler for Auth operations, since it seems different enough from the retry logic of thrift operations.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] feat(csharp/test/Drivers/Databricks): Support token refresh to extend connection lifetime [arrow-adbc]
Posted by "alexguo-db (via GitHub)" <gi...@apache.org>.
alexguo-db commented on PR #3177:
URL: https://github.com/apache/arrow-adbc/pull/3177#issuecomment-3098734908
Please give a chance for @jadewang-db to take a look before merging, thanks
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] feat(csharp/test/Drivers/Databricks): Support token refresh to extend connection lifetime [arrow-adbc]
Posted by "alexguo-db (via GitHub)" <gi...@apache.org>.
alexguo-db commented on code in PR #3177:
URL: https://github.com/apache/arrow-adbc/pull/3177#discussion_r2220543800
##########
csharp/src/Drivers/Databricks/DatabricksConnection.cs:
##########
@@ -321,57 +323,78 @@ private void ValidateProperties()
protected override HttpMessageHandler CreateHttpHandler()
{
HttpMessageHandler baseHandler = base.CreateHttpHandler();
+ HttpMessageHandler baseAuthHandler = HiveServer2TlsImpl.NewHttpClientHandler(TlsOptions, _proxyConfigurator);
// Add tracing handler to propagate W3C trace context if enabled
if (_tracePropagationEnabled)
{
baseHandler = new TracingDelegatingHandler(baseHandler, this, _traceParentHeaderName, _traceStateEnabled);
+ baseAuthHandler = new TracingDelegatingHandler(baseAuthHandler, this, _traceParentHeaderName, _traceStateEnabled);
}
if (TemporarilyUnavailableRetry)
{
// Add retry handler for 503 responses
baseHandler = new RetryHttpHandler(baseHandler, TemporarilyUnavailableRetryTimeout);
+ baseAuthHandler = new RetryHttpHandler(baseAuthHandler, TemporarilyUnavailableRetryTimeout);
Review Comment:
Sure, these are the differences:
Auth retry HTTP handler:
- needs to retry on 502, 503, 504, 408
- needs to use exponential backoff for retries
Thrift retry HTTP handler:
- only retries on 503
- relies on the returned Retry-After header to determine how long to wait
I think we can unify in a future PR so that for both auth and Thrift requests:
- retries on 502, 503, 504, 408
- if Retry-After header is set, use that to determine how long to wait
- otherwise, use exponential backoff
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] feat(csharp/test/Drivers/Databricks): Support token refresh to extend connection lifetime [arrow-adbc]
Posted by "alexguo-db (via GitHub)" <gi...@apache.org>.
alexguo-db commented on code in PR #3177:
URL: https://github.com/apache/arrow-adbc/pull/3177#discussion_r2220508724
##########
csharp/src/Drivers/Databricks/DatabricksConnection.cs:
##########
@@ -321,57 +323,78 @@ private void ValidateProperties()
protected override HttpMessageHandler CreateHttpHandler()
{
HttpMessageHandler baseHandler = base.CreateHttpHandler();
+ HttpMessageHandler baseAuthHandler = HiveServer2TlsImpl.NewHttpClientHandler(TlsOptions, _proxyConfigurator);
// Add tracing handler to propagate W3C trace context if enabled
if (_tracePropagationEnabled)
{
baseHandler = new TracingDelegatingHandler(baseHandler, this, _traceParentHeaderName, _traceStateEnabled);
+ baseAuthHandler = new TracingDelegatingHandler(baseAuthHandler, this, _traceParentHeaderName, _traceStateEnabled);
}
if (TemporarilyUnavailableRetry)
{
// Add retry handler for 503 responses
baseHandler = new RetryHttpHandler(baseHandler, TemporarilyUnavailableRetryTimeout);
+ baseAuthHandler = new RetryHttpHandler(baseAuthHandler, TemporarilyUnavailableRetryTimeout);
Review Comment:
Hmm yes we should retry on other codes, 502, 504, 408. I will create a different retry delegating handler for Auth operations, since it seems different enough from the retry logic of thrift operations.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] feat(csharp/test/Drivers/Databricks): Support token refresh to extend connection lifetime [arrow-adbc]
Posted by "jadewang-db (via GitHub)" <gi...@apache.org>.
jadewang-db commented on code in PR #3177:
URL: https://github.com/apache/arrow-adbc/pull/3177#discussion_r2220479072
##########
csharp/src/Drivers/Databricks/Auth/JwtTokenDecoder.cs:
##########
@@ -0,0 +1,96 @@
+/*
+* Licensed to the Apache Software Foundation (ASF) under one or more
+* contributor license agreements. See the NOTICE file distributed with
+* this work for additional information regarding copyright ownership.
+* The ASF licenses this file to You under the Apache License, Version 2.0
+* (the "License"); you may not use this file except in compliance with
+* the License. You may obtain a copy of the License at
+*
+* http://www.apache.org/licenses/LICENSE-2.0
+*
+* Unless required by applicable law or agreed to in writing, software
+* distributed under the License is distributed on an "AS IS" BASIS,
+* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+* See the License for the specific language governing permissions and
+* limitations under the License.
+*/
+
+using System;
+using System.Text;
+using System.Text.Json;
+
+namespace Apache.Arrow.Adbc.Drivers.Databricks.Auth
+{
+ /// <summary>
+ /// Utility class for decoding JWT tokens and extracting claims.
+ /// </summary>
+ internal static class JwtTokenDecoder
Review Comment:
there is no builtin jwt token decoder in dotnet, what about System.IdentityModel.Tokens.Jwt
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] feat(csharp/test/Drivers/Databricks): Support token refresh to extend connection lifetime [arrow-adbc]
Posted by "jadewang-db (via GitHub)" <gi...@apache.org>.
jadewang-db commented on code in PR #3177:
URL: https://github.com/apache/arrow-adbc/pull/3177#discussion_r2220493096
##########
csharp/src/Drivers/Databricks/DatabricksConnection.cs:
##########
@@ -321,57 +323,78 @@ private void ValidateProperties()
protected override HttpMessageHandler CreateHttpHandler()
{
HttpMessageHandler baseHandler = base.CreateHttpHandler();
+ HttpMessageHandler baseAuthHandler = HiveServer2TlsImpl.NewHttpClientHandler(TlsOptions, _proxyConfigurator);
// Add tracing handler to propagate W3C trace context if enabled
if (_tracePropagationEnabled)
{
baseHandler = new TracingDelegatingHandler(baseHandler, this, _traceParentHeaderName, _traceStateEnabled);
+ baseAuthHandler = new TracingDelegatingHandler(baseAuthHandler, this, _traceParentHeaderName, _traceStateEnabled);
}
if (TemporarilyUnavailableRetry)
{
// Add retry handler for 503 responses
baseHandler = new RetryHttpHandler(baseHandler, TemporarilyUnavailableRetryTimeout);
+ baseAuthHandler = new RetryHttpHandler(baseAuthHandler, TemporarilyUnavailableRetryTimeout);
Review Comment:
currently retryhttphandler only handle 503, is it enough? should we retry on timed out?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] feat(csharp/test/Drivers/Databricks): Support token refresh to extend connection lifetime [arrow-adbc]
Posted by "CurtHagenlocher (via GitHub)" <gi...@apache.org>.
CurtHagenlocher merged PR #3177:
URL: https://github.com/apache/arrow-adbc/pull/3177
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] feat(csharp/test/Drivers/Databricks): Support token refresh to extend connection lifetime [arrow-adbc]
Posted by "CurtHagenlocher (via GitHub)" <gi...@apache.org>.
CurtHagenlocher commented on code in PR #3177:
URL: https://github.com/apache/arrow-adbc/pull/3177#discussion_r2219877656
##########
csharp/src/Drivers/Databricks/DatabricksConnection.cs:
##########
@@ -334,44 +336,78 @@ protected override HttpMessageHandler CreateHttpHandler()
baseHandler = new RetryHttpHandler(baseHandler, TemporarilyUnavailableRetryTimeout);
}
- // Add OAuth handler if OAuth authentication is being used
+ // Add OAuth client credentials handler if OAuth M2M authentication is being used
if (Properties.TryGetValue(SparkParameters.AuthType, out string? authType) &&
SparkAuthTypeParser.TryParse(authType, out SparkAuthType authTypeValue) &&
authTypeValue == SparkAuthType.OAuth &&
Properties.TryGetValue(DatabricksParameters.OAuthGrantType, out string? grantTypeStr) &&
DatabricksOAuthGrantTypeParser.TryParse(grantTypeStr, out DatabricksOAuthGrantType grantType) &&
grantType == DatabricksOAuthGrantType.ClientCredentials)
{
- // Note: We assume that properties have already been validated
- if (Properties.TryGetValue(SparkParameters.HostName, out string? host) && !string.IsNullOrEmpty(host))
- {
- // Use hostname directly if provided
- }
- else if (Properties.TryGetValue(AdbcOptions.Uri, out string? uri) && !string.IsNullOrEmpty(uri))
- {
- // Extract hostname from URI if URI is provided
- if (Uri.TryCreate(uri, UriKind.Absolute, out Uri? parsedUri))
- {
- host = parsedUri.Host;
- }
- }
+ string host = GetHost();
Properties.TryGetValue(DatabricksParameters.OAuthClientId, out string? clientId);
Properties.TryGetValue(DatabricksParameters.OAuthClientSecret, out string? clientSecret);
Properties.TryGetValue(DatabricksParameters.OAuthScope, out string? scope);
- HttpClient OauthHttpClient = new HttpClient(HiveServer2TlsImpl.NewHttpClientHandler(TlsOptions, _proxyConfigurator));
+ if (_httpClient == null)
+ {
+ _httpClient = new HttpClient(HiveServer2TlsImpl.NewHttpClientHandler(TlsOptions, _proxyConfigurator));
+ }
var tokenProvider = new OAuthClientCredentialsProvider(
- OauthHttpClient,
+ _httpClient,
clientId!,
clientSecret!,
host!,
scope: scope ?? "sql",
timeoutMinutes: 1
);
- return new OAuthDelegatingHandler(baseHandler, tokenProvider);
+ baseHandler = new OAuthDelegatingHandler(baseHandler, tokenProvider);
+ }
+
+ // Add token exchange handler if token renewal is enabled and the auth type is OAuth access token
Review Comment:
This is mutually exclusive with the `if` at line 340, yes? The former is for client_credentials while this is for a user token? I'm wondering whether there might be some way of making this more obvious in the code, because `_httpClient` is being initialized in both but differently and it looked a bit tricky to reason about. Perhaps this second `if` could instead be an `else if`?
##########
csharp/src/Drivers/Databricks/DatabricksConnection.cs:
##########
@@ -334,44 +336,78 @@ protected override HttpMessageHandler CreateHttpHandler()
baseHandler = new RetryHttpHandler(baseHandler, TemporarilyUnavailableRetryTimeout);
}
- // Add OAuth handler if OAuth authentication is being used
+ // Add OAuth client credentials handler if OAuth M2M authentication is being used
if (Properties.TryGetValue(SparkParameters.AuthType, out string? authType) &&
SparkAuthTypeParser.TryParse(authType, out SparkAuthType authTypeValue) &&
authTypeValue == SparkAuthType.OAuth &&
Properties.TryGetValue(DatabricksParameters.OAuthGrantType, out string? grantTypeStr) &&
DatabricksOAuthGrantTypeParser.TryParse(grantTypeStr, out DatabricksOAuthGrantType grantType) &&
grantType == DatabricksOAuthGrantType.ClientCredentials)
{
- // Note: We assume that properties have already been validated
- if (Properties.TryGetValue(SparkParameters.HostName, out string? host) && !string.IsNullOrEmpty(host))
- {
- // Use hostname directly if provided
- }
- else if (Properties.TryGetValue(AdbcOptions.Uri, out string? uri) && !string.IsNullOrEmpty(uri))
- {
- // Extract hostname from URI if URI is provided
- if (Uri.TryCreate(uri, UriKind.Absolute, out Uri? parsedUri))
- {
- host = parsedUri.Host;
- }
- }
+ string host = GetHost();
Properties.TryGetValue(DatabricksParameters.OAuthClientId, out string? clientId);
Properties.TryGetValue(DatabricksParameters.OAuthClientSecret, out string? clientSecret);
Properties.TryGetValue(DatabricksParameters.OAuthScope, out string? scope);
- HttpClient OauthHttpClient = new HttpClient(HiveServer2TlsImpl.NewHttpClientHandler(TlsOptions, _proxyConfigurator));
+ if (_httpClient == null)
Review Comment:
Can this method be called correctly more than once, and if so, then under what circumstances? If it can't be, then there's no need for the null check here. It could be replaced e.g. by a `Debug.Assert(_httpClient == null)`. This would make the method a little easier to reason about.
##########
csharp/src/Drivers/Databricks/DatabricksConnection.cs:
##########
@@ -334,44 +336,78 @@ protected override HttpMessageHandler CreateHttpHandler()
baseHandler = new RetryHttpHandler(baseHandler, TemporarilyUnavailableRetryTimeout);
}
- // Add OAuth handler if OAuth authentication is being used
+ // Add OAuth client credentials handler if OAuth M2M authentication is being used
if (Properties.TryGetValue(SparkParameters.AuthType, out string? authType) &&
SparkAuthTypeParser.TryParse(authType, out SparkAuthType authTypeValue) &&
authTypeValue == SparkAuthType.OAuth &&
Properties.TryGetValue(DatabricksParameters.OAuthGrantType, out string? grantTypeStr) &&
DatabricksOAuthGrantTypeParser.TryParse(grantTypeStr, out DatabricksOAuthGrantType grantType) &&
grantType == DatabricksOAuthGrantType.ClientCredentials)
{
- // Note: We assume that properties have already been validated
- if (Properties.TryGetValue(SparkParameters.HostName, out string? host) && !string.IsNullOrEmpty(host))
- {
- // Use hostname directly if provided
- }
- else if (Properties.TryGetValue(AdbcOptions.Uri, out string? uri) && !string.IsNullOrEmpty(uri))
- {
- // Extract hostname from URI if URI is provided
- if (Uri.TryCreate(uri, UriKind.Absolute, out Uri? parsedUri))
- {
- host = parsedUri.Host;
- }
- }
+ string host = GetHost();
Properties.TryGetValue(DatabricksParameters.OAuthClientId, out string? clientId);
Properties.TryGetValue(DatabricksParameters.OAuthClientSecret, out string? clientSecret);
Properties.TryGetValue(DatabricksParameters.OAuthScope, out string? scope);
- HttpClient OauthHttpClient = new HttpClient(HiveServer2TlsImpl.NewHttpClientHandler(TlsOptions, _proxyConfigurator));
+ if (_httpClient == null)
+ {
+ _httpClient = new HttpClient(HiveServer2TlsImpl.NewHttpClientHandler(TlsOptions, _proxyConfigurator));
Review Comment:
I know this was already creating a handler before, but should the credential-related handlers be adding the tracing-related functionality and/or the retry functionality? Right now, tracing is only added to the primary (data) handler and not to either of the credential handlers and retry is being added only to the token exchange path. The code is complex enough that it's not clear whether this is deliberate or accidental.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] feat(csharp/test/Drivers/Databricks): Support token refresh to extend connection lifetime [arrow-adbc]
Posted by "alexguo-db (via GitHub)" <gi...@apache.org>.
alexguo-db commented on code in PR #3177:
URL: https://github.com/apache/arrow-adbc/pull/3177#discussion_r2220274834
##########
csharp/src/Drivers/Databricks/DatabricksConnection.cs:
##########
@@ -334,44 +336,78 @@ protected override HttpMessageHandler CreateHttpHandler()
baseHandler = new RetryHttpHandler(baseHandler, TemporarilyUnavailableRetryTimeout);
}
- // Add OAuth handler if OAuth authentication is being used
+ // Add OAuth client credentials handler if OAuth M2M authentication is being used
if (Properties.TryGetValue(SparkParameters.AuthType, out string? authType) &&
SparkAuthTypeParser.TryParse(authType, out SparkAuthType authTypeValue) &&
authTypeValue == SparkAuthType.OAuth &&
Properties.TryGetValue(DatabricksParameters.OAuthGrantType, out string? grantTypeStr) &&
DatabricksOAuthGrantTypeParser.TryParse(grantTypeStr, out DatabricksOAuthGrantType grantType) &&
grantType == DatabricksOAuthGrantType.ClientCredentials)
{
- // Note: We assume that properties have already been validated
- if (Properties.TryGetValue(SparkParameters.HostName, out string? host) && !string.IsNullOrEmpty(host))
- {
- // Use hostname directly if provided
- }
- else if (Properties.TryGetValue(AdbcOptions.Uri, out string? uri) && !string.IsNullOrEmpty(uri))
- {
- // Extract hostname from URI if URI is provided
- if (Uri.TryCreate(uri, UriKind.Absolute, out Uri? parsedUri))
- {
- host = parsedUri.Host;
- }
- }
+ string host = GetHost();
Properties.TryGetValue(DatabricksParameters.OAuthClientId, out string? clientId);
Properties.TryGetValue(DatabricksParameters.OAuthClientSecret, out string? clientSecret);
Properties.TryGetValue(DatabricksParameters.OAuthScope, out string? scope);
- HttpClient OauthHttpClient = new HttpClient(HiveServer2TlsImpl.NewHttpClientHandler(TlsOptions, _proxyConfigurator));
+ if (_httpClient == null)
Review Comment:
No, it should only be called once per connection.
https://github.com/apache/arrow-adbc/blob/main/csharp/src/Drivers/Databricks/DatabricksDatabase.cs#L43
Specifically, it's called during OpenAsync, which is called right after initializing the DatabricksConnection
I added the change you suggested
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] feat(csharp/test/Drivers/Databricks): Support token refresh to extend connection lifetime [arrow-adbc]
Posted by "jadewang-db (via GitHub)" <gi...@apache.org>.
jadewang-db commented on code in PR #3177:
URL: https://github.com/apache/arrow-adbc/pull/3177#discussion_r2220530453
##########
csharp/src/Drivers/Databricks/Auth/JwtTokenDecoder.cs:
##########
@@ -0,0 +1,96 @@
+/*
+* Licensed to the Apache Software Foundation (ASF) under one or more
+* contributor license agreements. See the NOTICE file distributed with
+* this work for additional information regarding copyright ownership.
+* The ASF licenses this file to You under the Apache License, Version 2.0
+* (the "License"); you may not use this file except in compliance with
+* the License. You may obtain a copy of the License at
+*
+* http://www.apache.org/licenses/LICENSE-2.0
+*
+* Unless required by applicable law or agreed to in writing, software
+* distributed under the License is distributed on an "AS IS" BASIS,
+* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+* See the License for the specific language governing permissions and
+* limitations under the License.
+*/
+
+using System;
+using System.Text;
+using System.Text.Json;
+
+namespace Apache.Arrow.Adbc.Drivers.Databricks.Auth
+{
+ /// <summary>
+ /// Utility class for decoding JWT tokens and extracting claims.
+ /// </summary>
+ internal static class JwtTokenDecoder
Review Comment:
ok, defer to @CurtHagenlocher on this.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org