feat(go/adbc/driver/flightsql): Add OAuth Support to Flight Client #2651
Conversation
xborder
changed the title
|
I would suggest to update |
| }, nil | ||
| } | ||
|
|
||
| func (f *tokenExchange) GetToken(ctx context.Context) (*oauth2.Token, error) { |
There was a problem hiding this comment.
Could you please add RWLock here to make this method thread safe?
| }, nil | ||
| } | ||
|
|
||
| func (c *clientCredentials) GetToken(ctx context.Context) (*oauth2.Token, error) { |
There was a problem hiding this comment.
The same question is here - should it be thread safe? If so, could you please add rwlock here?
There was a problem hiding this comment.
maybe we can use https://pkg.go.dev/golang.org/x/[email protected]/clientcredentials#Config instead and just call https://pkg.go.dev/golang.org/x/[email protected]/clientcredentials#Config.TokenSource
If we do that, then most of this code can be entirely removed in favor of just using https://pkg.go.dev/google.golang.org/grpc/credentials/oauth#TokenSource and https://pkg.go.dev/google.golang.org/grpc#WithPerRPCCredentials which would manage the oauth flow for us entirely as long as we can construct the token source which performs refreshes as necessary (as per the clientcredentials config above)
Essentially, SetOptions can create the TokenSource and then getFlightClient just adds the token source as a dialoption via WithPerRPCCredentials. That would be much more convenient and less work/code for us to maintain
@zeroshade I was not aware of that, thank you for bringing this up. While I was testing this proposal I realized that TokenSource requires transport security. I think this conflicts with development and how to run tests making it mandatory to use tls. What alternatives do I have here? The options I see:
Do you have any suggestion? |
|
We have a TLS suite set up already for tests here: you can reuse that |
|
Also TokenSource on it's own shouldn't require transport security, it just has a method that returns a bool letting the caller know whether or not it requires it |
Replaced extra structs for OAuth in favor of grpc's TokenSource and grpc.DialOption WithPerRPCCredentials
* Changed OAuth tests to use TLS since it is a requirement from grpc's TokenSource and WithPerRPCCredentials * Split DoSetupSuit to setupFlightServer and setupDatabase so they can be used independently
* Removed adbc.flight.sql.token in favour of adbc.flight.sql.authorization_header when the client wants to pass a token or adbc.flight.sql.oauth.exchange.subject_token when client wants to do token exchange * Simplified options to set oauth flow. Now client can set client_credentials or token_exchange instead of integers
docs/source/driver/flight_sql.rst
Outdated
| @@ -159,6 +159,12 @@ few optional authentication schemes: | |||
| header will then be sent back as the ``authorization`` header on all | |||
| future requests. | |||
|
|
|||
| - (Go only) OAuth 2.0 authentication flows. | |||
There was a problem hiding this comment.
isn't it not Go only? Anything that uses the flightsql driver should be able to use the options that are being added. (We should add constants to the python adbc_driver_flightsql package)
There was a problem hiding this comment.
Does it make sense to create a separate PR for this?
There was a problem hiding this comment.
It makes sense to make a separate PR to add the option constants, but I would still say that the "Go only" should be removed as nothing would prevent any other binding from using these options.
|
Just to add to this (and this could be a follow-on PR) - what happens if https://github.com/apache/arrow-adbc/blob/a187ead78afebe85c75b466a06ad6e01ae4ac8c6/go/adbc/driver/flightsql/flightsql_statement.go#L194C21-L194C25 is part of a long running operation and the token expires? How does the token get refreshed and then continue the poll operation? By the way, this is a concern for all ADBC drivers in general, as we don't seem to have a standard way of failing "mid stream" and then "resume" the long running operation without starting it over. |
Or if runs long. |
And for context, #2655 is making an attempt to demonstrate this capability for BigQuery in the .NET driver when using Entra authentication. The Databricks driver also needs to do this. Essentially, any OAuth protected data source could support this type of retry behavior. |
|
@davidhcoe I think you can refer to this comment. |
Description
This pull request introduces OAuth support to the Flight client in the GO driver. The changes include the addition of OAuth access token support, implementation of token exchange and client credentials OAuth flows.
Related Issues
Changes Made
tokenas a database optiontokengets exchanged and the result is added to theAuthorizationheader as aBearertokenclient_idandclient_secretare used to obtain a access token that is added to theAuthorizationheader as aBearertokenHere's the markdown code for the OAuth 2.0 configuration options table:
markdown# OAuth 2.0 Configuration Options
adbc.flight.sql.oauth.flowclient_credentials,token_exchangeadbc.flight.sql.oauth.client_idadbc.flight.sql.oauth.client_secretadbc.flight.sql.oauth.token_uriadbc.flight.sql.oauth.scope"read.all offline_access")adbc.flight.sql.oauth.exchange.subject_tokenadbc.flight.sql.oauth.exchange.subject_token_typeadbc.flight.sql.oauth.exchange.actor_tokenadbc.flight.sql.oauth.exchange.actor_token_typeadbc.flight.sql.oauth.exchange.audadbc.flight.sql.oauth.exchange.resourceadbc.flight.sql.oauth.exchange.scopeadbc.flight.sql.oauth.exchange.requested_token_typeSupported token types:
urn:ietf:params:oauth:token-type:access_tokenurn:ietf:params:oauth:token-type:refresh_tokenurn:ietf:params:oauth:token-type:id_tokenurn:ietf:params:oauth:token-type:saml1urn:ietf:params:oauth:token-type:saml2urn:ietf:params:oauth:token-type:jwt