CurtHagenlocher commented on code in PR #2655:
URL: https://github.com/apache/arrow-adbc/pull/2655#discussion_r2052539701
##########
csharp/src/Drivers/BigQuery/BigQueryParameters.cs:
##########
@@ -22,6 +22,8 @@ namespace Apache.Arrow.Adbc.Drivers.BigQuery
/// </summary>
public class BigQueryParameters
{
+ public const string AccessToken = "adbc.bigquery.access_token";
+ public const string AudienceUri = "adbc.bigquery.audience_uri";
Review Comment:
Are these properties just for EntraId or can they also be used for Google
OAuth? If the former, should they be named in a way to make that clear?
Otherwise I might think I can supply a Google OAuth access token this way.
##########
csharp/src/Drivers/BigQuery/BigQueryParameters.cs:
##########
@@ -36,6 +38,8 @@ public class BigQueryParameters
public const string Scopes = "adbc.bigquery.scopes";
public const string IncludeConstraintsWithGetObjects =
"adbc.bigquery.include_constraints_getobjects";
public const string ClientTimeout = "adbc.bigquery.client.timeout";
+ public const string MaximumRetryAttempts =
"adbc.bigquery.maximum_retries";
+ public const string RetryDelayMs = "adbc.bigquery.retry_delay_ms";
Review Comment:
Is the "retry" here specific to refreshing the credential, and if so, should
the property name reflect that?
##########
csharp/src/Drivers/BigQuery/BigQueryConnection.cs:
##########
@@ -36,17 +37,15 @@ namespace Apache.Arrow.Adbc.Drivers.BigQuery
/// <summary>
/// BigQuery-specific implementation of <see cref="AdbcConnection"/>
/// </summary>
- public class BigQueryConnection : AdbcConnection
+ public class BigQueryConnection : AdbcConnection, ITokenProtectedResource
Review Comment:
I know it's a "preexisting condition" but In general we've made the
implementation types all be internal. Is there a specific API required on this
class that's not exposed through `ITokenProtectedResource`?
##########
csharp/src/Drivers/BigQuery/BigQueryConnection.cs:
##########
@@ -36,17 +37,15 @@ namespace Apache.Arrow.Adbc.Drivers.BigQuery
/// <summary>
/// BigQuery-specific implementation of <see cref="AdbcConnection"/>
/// </summary>
- public class BigQueryConnection : AdbcConnection
+ public class BigQueryConnection : AdbcConnection, ITokenProtectedResource
{
- readonly IReadOnlyDictionary<string, string> properties;
- BigQueryClient? client;
- GoogleCredential? credential;
+ IReadOnlyDictionary<string, string> properties;
Review Comment:
I feel like there's a lot of unnecessary complexity around keeping this as
an `IReadOnlyDictionary`. If it were instead a `Dictionary<string, string>`
then it could retain its `readonly`-ness and initialized only once in the
constructor, and we wouldn't need to make a copy of it every time `SetOption`
was called.
##########
csharp/src/Drivers/BigQuery/ITokenProtectedResource.cs:
##########
@@ -0,0 +1,40 @@
+/*
+ * 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.Threading.Tasks;
+
+namespace Apache.Arrow.Adbc.Drivers.BigQuery
+{
+ /// <summary>
+ /// Common interface for a token protected resource.
+ /// </summary>
+ public interface ITokenProtectedResource
+ {
+ /// <summary>
+ /// The function to call when updating the token.
+ /// </summary>
+ public Func<Task>? UpdateToken { get; set; }
+
+ /// <summary>
+ /// Determines the token needs to be updated.
+ /// </summary>
+ /// <param name="ex">The exception that occurs.</param>
+ /// <returns>True/False indicating a refresh is needed.</returns>
+ public bool TokenRequiresUpdate(Exception ex);
Review Comment:
How is client code expected to use this?
##########
csharp/src/Drivers/BigQuery/BigQueryParameters.cs:
##########
@@ -50,10 +54,18 @@ public class BigQueryParameters
public class BigQueryConstants
{
public const string UserAuthenticationType = "user";
+ public const string EntraIdAuthenticationType = "aad";
public const string ServiceAccountAuthenticationType = "service";
public const string TokenEndpoint =
"https://accounts.google.com/o/oauth2/token";;
public const string TreatLargeDecimalAsString = "true";
+ // Entra ID / Azure AD constants
+ public const string GrantType =
"urn:ietf:params:oauth:grant-type:token-exchange";
+ public const string SubjectTokenType =
"urn:ietf:params:oauth:token-type:id_token";
+ public const string RequestedTokenType =
"urn:ietf:params:oauth:token-type:access_token";
+ public const string EntraIdScope =
"https://www.googleapis.com/auth/cloud-platform";;
+ public const string StsTokenEndpoint =
"https://sts.googleapis.com/v1/token";;
Review Comment:
These are being exposed as public members to user code. How does user code
use these values?
##########
csharp/src/Drivers/BigQuery/ITokenProtectedResource.cs:
##########
@@ -0,0 +1,40 @@
+/*
+ * 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.Threading.Tasks;
+
+namespace Apache.Arrow.Adbc.Drivers.BigQuery
+{
+ /// <summary>
+ /// Common interface for a token protected resource.
+ /// </summary>
+ public interface ITokenProtectedResource
+ {
+ /// <summary>
+ /// The function to call when updating the token.
+ /// </summary>
+ public Func<Task>? UpdateToken { get; set; }
+
+ /// <summary>
+ /// Determines the token needs to be updated.
+ /// </summary>
+ /// <param name="ex">The exception that occurs.</param>
+ /// <returns>True/False indicating a refresh is needed.</returns>
+ public bool TokenRequiresUpdate(Exception ex);
Review Comment:
It feels like this class is serving as both the user-facing API and an
internal-facing API. My vague expectation would be that there be one interface
which allows the `UpdateToken` handler to be set and a different interface
which generic credential-handling code uses to implement credential logic.
##########
csharp/src/Drivers/BigQuery/BigQueryUtils.cs:
##########
@@ -0,0 +1,45 @@
+/*
+ * 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 Google;
+
+namespace Apache.Arrow.Adbc.Drivers.BigQuery
+{
+ internal class BigQueryUtils
+ {
+ public static bool TokenRequiresUpdate(Exception ex)
+ {
+ bool result = false;
+
+ switch (ex)
+ {
+ case GoogleApiException apiException:
+ if (apiException.HttpStatusCode ==
System.Net.HttpStatusCode.Unauthorized)
+ {
+ result = true;
Review Comment:
Single exit is overrated :P.
##########
csharp/src/Drivers/BigQuery/BigQueryParameters.cs:
##########
@@ -22,6 +22,8 @@ namespace Apache.Arrow.Adbc.Drivers.BigQuery
/// </summary>
public class BigQueryParameters
{
+ public const string AccessToken = "adbc.bigquery.access_token";
+ public const string AudienceUri = "adbc.bigquery.audience_uri";
Review Comment:
(I think it's okay to say that this is aspirationally for both -- but then
we should also file a work item to complete the implementation.)
##########
csharp/src/Drivers/BigQuery/BigQueryStatement.cs:
##########
@@ -36,88 +37,155 @@ namespace Apache.Arrow.Adbc.Drivers.BigQuery
/// <summary>
/// BigQuery-specific implementation of <see cref="AdbcStatement"/>
/// </summary>
- public class BigQueryStatement : AdbcStatement
+ class BigQueryStatement : AdbcStatement, ITokenProtectedResource,
IDisposable
{
- readonly BigQueryClient client;
- readonly GoogleCredential credential;
+ readonly BigQueryConnection bigQueryConnection;
- public BigQueryStatement(BigQueryClient client, GoogleCredential
credential)
+ public BigQueryStatement(BigQueryConnection bigQueryConnection)
{
- this.client = client;
- this.credential = credential;
+ if (bigQueryConnection == null) { throw new
AdbcException($"{nameof(bigQueryConnection)} cannot be null",
AdbcStatusCode.InvalidArgument); }
+
+ // pass on the handler since this isn't accessible publicly
+ UpdateToken = bigQueryConnection.UpdateToken;
+
+ this.bigQueryConnection = bigQueryConnection;
}
+ public Func<Task>? UpdateToken { get; set; }
+
public IReadOnlyDictionary<string, string>? Options { get; set; }
Review Comment:
Same comment as for Connection: if the options are logically mutable then
let's store them in a mutable data structure instead of having to rebuild the
dictionary every time an option is set.
##########
csharp/src/Drivers/BigQuery/BigQueryConnection.cs:
##########
@@ -309,6 +368,23 @@ public override IArrowArrayStream GetObjects(
return new
BigQueryInfoArrowStream(StandardSchemas.GetObjectsSchema, dataArrays);
}
+ /// <summary>
+ /// Renews the internal BigQueryClient with updated credentials.
+ /// </summary>
+ internal void UpdateClientToken()
+ {
+ // there isn't a way to set the credentials, just need to open a
new client
+ Client = Open();
+ }
+
+ /// <summary>
+ /// Determines if the token needs to be updated.
+ /// </summary>
+ public bool TokenRequiresUpdate(Exception ex) =>
BigQueryUtils.TokenRequiresUpdate(ex);
+
+ async Task<T> ExecuteWithRetriesAsync<T>(Func<Task<T>> action) =>
await RetryManager.ExecuteWithRetriesAsync<T>(this, action, MaxRetryAttempts,
RetryDelayMs);
+
+
Review Comment:
nit: extra blank line
##########
csharp/src/Drivers/BigQuery/RetryManager.cs:
##########
@@ -0,0 +1,90 @@
+
+/*
+* 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.Threading.Tasks;
+
+namespace Apache.Arrow.Adbc.Drivers.BigQuery
+{
+ /// <summary>
+ /// Class that will retry calling a method with an exponential backoff.
+ /// </summary>
+ public class RetryManager
Review Comment:
Is there a reason for this to be public?
--
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-unsubscr...@arrow.apache.org
For queries about this service, please contact Infrastructure at:
us...@infra.apache.org
