Skip to content

Commit 405bc54

Browse files
toddmeng-dbtoddmeng-db
committed
token refresh + test files
1 parent 136af43 commit 405bc54

2 files changed

Lines changed: 194 additions & 58 deletions

File tree

‎csharp/src/Drivers/Databricks/Auth/OAuthClientCredentialsService.cs‎

Lines changed: 134 additions & 58 deletions
Original file line numberDiff line numberDiff line change
@@ -28,128 +28,204 @@ namespace Apache.Arrow.Adbc.Drivers.Databricks.Auth
2828
/// <summary>
2929
/// Service for obtaining OAuth access tokens using the client credentials grant type.
3030
/// </summary>
31-
internal class OAuthClientCredentialsService
31+
internal class OAuthClientCredentialsService : IDisposable
3232
{
3333
private readonly Lazy<HttpClient> _httpClient;
3434
private readonly string _clientId;
3535
private readonly string _clientSecret;
3636
private readonly Uri _baseUri;
37-
private readonly string? _tenantId;
38-
private readonly string _scope;
3937
private readonly string _tokenEndpoint;
4038
private readonly int _timeoutMinutes;
39+
private readonly SemaphoreSlim _tokenLock = new SemaphoreSlim(1, 1);
40+
private TokenInfo? _cachedToken;
41+
42+
private class TokenInfo
43+
{
44+
public string? AccessToken { get; set; }
45+
public DateTime ExpiresAt { get; set; }
46+
47+
public bool IsExpired => DateTime.UtcNow >= ExpiresAt;
48+
49+
// Add buffer time to refresh token before actual expiration
50+
public bool NeedsRefresh => DateTime.UtcNow >= ExpiresAt.AddMinutes(-5);
51+
}
4152

4253
/// <summary>
4354
/// Initializes a new instance of the <see cref="OAuthClientCredentialsService"/> class.
4455
/// </summary>
4556
/// <param name="clientId">The OAuth client ID.</param>
4657
/// <param name="clientSecret">The OAuth client secret.</param>
4758
/// <param name="baseUri">The base URI of the Databricks workspace.</param>
48-
/// <param name="tenantId">The Azure AD tenant ID. Required for Azure Databricks.</param>
49-
/// <param name="scope">The OAuth scope to request. Default is "all-apis".</param>
50-
/// <param name="timeoutMinutes">The timeout in minutes for HTTP requests. Default is 5 minutes.</param>
5159
public OAuthClientCredentialsService(
5260
string clientId,
5361
string clientSecret,
5462
Uri baseUri,
55-
string? tenantId = null,
56-
string scope = "all-apis",
57-
int timeoutMinutes = 5)
63+
int timeoutMinutes = 1,
64+
HttpClient? httpClient = null)
5865
{
5966
_clientId = clientId ?? throw new ArgumentNullException(nameof(clientId));
6067
_clientSecret = clientSecret ?? throw new ArgumentNullException(nameof(clientSecret));
6168
_baseUri = baseUri ?? throw new ArgumentNullException(nameof(baseUri));
62-
_tenantId = tenantId;
63-
_scope = scope ?? "all-apis";
6469
_timeoutMinutes = timeoutMinutes;
6570
_tokenEndpoint = DetermineTokenEndpoint();
6671

67-
_httpClient = new Lazy<HttpClient>(() =>
68-
{
69-
var client = new HttpClient();
70-
client.Timeout = TimeSpan.FromMinutes(_timeoutMinutes);
71-
return client;
72-
});
72+
_httpClient = httpClient != null
73+
? new Lazy<HttpClient>(() => httpClient)
74+
: new Lazy<HttpClient>(() =>
75+
{
76+
var client = new HttpClient();
77+
client.Timeout = TimeSpan.FromMinutes(_timeoutMinutes);
78+
return client;
79+
});
7380
}
7481

7582
private HttpClient HttpClient => _httpClient.Value;
7683

7784
private string DetermineTokenEndpoint()
7885
{
79-
string host = _baseUri.Host.ToLowerInvariant();
80-
if (host.Contains("azuredatabricks.net"))
86+
// For workspace URLs, the token endpoint is always /oidc/v1/token
87+
// TODO: Might be different for Azure AAD SPs
88+
return $"{_baseUri.Scheme}://{_baseUri.Host}/oidc/v1/token";
89+
}
90+
91+
private string? GetValidCachedToken()
92+
{
93+
return _cachedToken != null && !_cachedToken.NeedsRefresh && _cachedToken.AccessToken != null
94+
? _cachedToken.AccessToken
95+
: null;
96+
}
97+
98+
99+
private async Task<string> RefreshTokenInternalAsync(CancellationToken cancellationToken)
100+
{
101+
var request = CreateTokenRequest();
102+
103+
HttpResponseMessage response;
104+
try
81105
{
82-
if (string.IsNullOrEmpty(_tenantId))
83-
{
84-
throw new ArgumentException("Azure Databricks requires a tenantId to determine the token endpoint.");
85-
}
106+
response = await HttpClient.SendAsync(request, cancellationToken);
107+
response.EnsureSuccessStatusCode();
108+
}
109+
catch (HttpRequestException ex)
110+
{
111+
throw new DatabricksException($"Failed to acquire OAuth access token: {ex.Message}", ex);
112+
}
86113

87-
return $"https://login.microsoftonline.com/{_tenantId}/oauth2/v2.0/token";
114+
string content = await response.Content.ReadAsStringAsync();
115+
116+
try
117+
{
118+
_cachedToken = ParseTokenResponse(content);
119+
return _cachedToken.AccessToken!;
88120
}
89-
else
121+
catch (JsonException ex)
90122
{
91-
// Applies to AWS and GCP (if using Databricks-hosted IdP)
92-
return "https://accounts.cloud.databricks.com/oidc/token";
123+
throw new DatabricksException($"Failed to parse OAuth response: {ex.Message}", ex);
93124
}
94125
}
95126

96-
/// <summary>
97-
/// Gets an OAuth access token using the client credentials grant type.
98-
/// </summary>
99-
/// <param name="cancellationToken">A cancellation token to cancel the operation.</param>
100-
/// <returns>The access token.</returns>
101-
/// <exception cref="DatabricksException">Thrown when the token request fails or the response is invalid.</exception>
102-
public async Task<string> GetAccessTokenAsync(CancellationToken cancellationToken = default)
127+
private HttpRequestMessage CreateTokenRequest()
103128
{
104129
var requestContent = new FormUrlEncodedContent(new[]
105130
{
106131
new KeyValuePair<string, string>("grant_type", "client_credentials"),
107-
new KeyValuePair<string, string>("client_id", _clientId),
108-
new KeyValuePair<string, string>("client_secret", _clientSecret),
109-
new KeyValuePair<string, string>("scope", _scope)
132+
new KeyValuePair<string, string>("scope", "sql")
110133
});
111134

112135
var request = new HttpRequestMessage(HttpMethod.Post, _tokenEndpoint)
113136
{
114137
Content = requestContent
115138
};
116139

140+
// Use Basic Auth with client ID and secret
141+
var authHeader = Convert.ToBase64String(
142+
System.Text.Encoding.ASCII.GetBytes($"{_clientId}:{_clientSecret}"));
143+
request.Headers.Authorization = new AuthenticationHeaderValue("Basic", authHeader);
117144
request.Headers.Accept.Add(new MediaTypeWithQualityHeaderValue("application/json"));
118145

119-
HttpResponseMessage response;
120-
try
146+
return request;
147+
}
148+
149+
private TokenInfo ParseTokenResponse(string content)
150+
{
151+
using var jsonDoc = JsonDocument.Parse(content);
152+
153+
if (!jsonDoc.RootElement.TryGetProperty("access_token", out var accessTokenElement))
121154
{
122-
response = await HttpClient.SendAsync(request, cancellationToken);
123-
response.EnsureSuccessStatusCode();
155+
throw new DatabricksException("OAuth response did not contain an access_token");
124156
}
125-
catch (HttpRequestException ex)
157+
158+
string? accessToken = accessTokenElement.GetString();
159+
if (string.IsNullOrEmpty(accessToken))
126160
{
127-
throw new DatabricksException($"Failed to acquire OAuth access token: {ex.Message}", ex);
161+
throw new DatabricksException("OAuth access_token was null or empty");
128162
}
129163

130-
string content = await response.Content.ReadAsStringAsync();
131-
132-
try
164+
// Get expiration time from response
165+
if (!jsonDoc.RootElement.TryGetProperty("expires_in", out var expiresInElement))
133166
{
134-
using var jsonDoc = JsonDocument.Parse(content);
135-
136-
if (!jsonDoc.RootElement.TryGetProperty("access_token", out var accessTokenElement))
137-
{
138-
throw new DatabricksException("OAuth response did not contain an access_token");
139-
}
167+
throw new DatabricksException("OAuth response did not contain expires_in");
168+
}
169+
170+
int expiresIn = expiresInElement.GetInt32();
171+
if (expiresIn <= 0)
172+
{
173+
throw new DatabricksException("OAuth expires_in value must be positive");
174+
}
175+
176+
return new TokenInfo
177+
{
178+
AccessToken = accessToken!,
179+
ExpiresAt = DateTime.UtcNow.AddSeconds(expiresIn)
180+
};
181+
}
182+
183+
/// <summary>
184+
/// Gets an OAuth access token using the client credentials grant type.
185+
/// </summary>
186+
/// <param name="cancellationToken">A cancellation token to cancel the operation.</param>
187+
/// <returns>The access token.</returns>
188+
/// <exception cref="DatabricksException">Thrown when the token request fails or the response is invalid.</exception>
189+
public async Task<string> GetAccessTokenAsync(CancellationToken cancellationToken = default)
190+
{
191+
// First try to get cached token without acquiring lock
192+
if (GetValidCachedToken() is string cachedToken)
193+
{
194+
return cachedToken;
195+
}
196+
197+
// If token needs refresh, acquire lock with timeout
198+
var lockTimeout = TimeSpan.FromSeconds(30); // Reasonable timeout for lock acquisition
199+
if (!await _tokenLock.WaitAsync(lockTimeout, cancellationToken).ConfigureAwait(false))
200+
{
201+
throw new TimeoutException("Timeout waiting for token refresh lock");
202+
}
140203

141-
string? accessToken = accessTokenElement.GetString();
142-
if (string.IsNullOrEmpty(accessToken))
204+
try
205+
{
206+
// Double-check pattern in case another thread refreshed while we were waiting
207+
if (GetValidCachedToken() is string refreshedToken)
143208
{
144-
throw new DatabricksException("OAuth access_token was null or empty");
209+
return refreshedToken;
145210
}
146211

147-
return accessToken!;
212+
return await RefreshTokenInternalAsync(cancellationToken).ConfigureAwait(false);
148213
}
149-
catch (JsonException ex)
214+
finally
150215
{
151-
throw new DatabricksException($"Failed to parse OAuth response: {ex.Message}", ex);
216+
_tokenLock.Release();
152217
}
153218
}
219+
220+
221+
public void Dispose()
222+
{
223+
_tokenLock.Dispose();
224+
if (_httpClient.IsValueCreated)
225+
{
226+
HttpClient.Dispose();
227+
}
228+
}
229+
154230
}
155231
}

‎csharp/test/Drivers/Databricks/Auth/OAuthClientCredentialsServiceTests.cs‎

Lines changed: 60 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,60 @@
1+
using System;
2+
using System.Threading;
3+
using System.Threading.Tasks;
4+
using Apache.Arrow.Adbc.Drivers.Databricks.Auth;
5+
using Apache.Arrow.Adbc.Tests;
6+
using Apache.Arrow.Adbc.Tests.Drivers.Apache.Spark;
7+
using Xunit;
8+
using Xunit.Abstractions;
9+
10+
namespace Apache.Arrow.Adbc.Tests.Drivers.Databricks.Auth
11+
{
12+
public class OAuthClientCredentialsServiceTests : TestBase<SparkTestConfiguration, SparkTestEnvironment>, IDisposable
13+
{
14+
private readonly OAuthClientCredentialsService _service;
15+
16+
public OAuthClientCredentialsServiceTests(ITestOutputHelper? outputHelper)
17+
: base(outputHelper, new SparkTestEnvironment.Factory())
18+
{
19+
_service = new OAuthClientCredentialsService(
20+
TestConfiguration.Username,
21+
TestConfiguration.Password,
22+
new Uri(TestConfiguration.Uri),
23+
timeoutMinutes: 1);
24+
}
25+
26+
[Fact]
27+
public async Task GetAccessToken_WithValidCredentials_ReturnsToken()
28+
{
29+
var token = await _service.GetAccessTokenAsync(CancellationToken.None);
30+
31+
Assert.NotNull(token);
32+
Assert.NotEmpty(token);
33+
}
34+
35+
[Fact]
36+
public async Task GetAccessToken_WithCancellation_ThrowsOperationCanceledException()
37+
{
38+
using var cts = new CancellationTokenSource();
39+
cts.Cancel();
40+
41+
await Assert.ThrowsAsync<OperationCanceledException>(() =>
42+
_service.GetAccessTokenAsync(cts.Token));
43+
}
44+
45+
[Fact]
46+
public async Task GetAccessToken_MultipleCalls_ReusesCachedToken()
47+
{
48+
var token1 = await _service.GetAccessTokenAsync(CancellationToken.None);
49+
var token2 = await _service.GetAccessTokenAsync(CancellationToken.None);
50+
51+
Assert.Equal(token1, token2);
52+
}
53+
54+
void IDisposable.Dispose()
55+
{
56+
_service.Dispose();
57+
base.Dispose();
58+
}
59+
}
60+
}

0 commit comments

Comments
 (0)