feat(go/adbc/driver/flightsql): Add OAuth Support to Flight Client by xborder · Pull Request #2651 · apache/arrow-adbc

xborder · 2025-03-25T17:14:29Z

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

Closes #2650

Changes Made

Added token as a database option
Added support for Token Exchange. If configured, token gets exchanged and the result is added to the Authorization header as a Bearer token
Added support for Client Credentials. If configured, client_id and client_secret are used to obtain a access token that is added to the Authorization header as a Bearer token
Added new driver options to allow third-party applications to configure oauth flows:
Added tests

Here's the markdown code for the OAuth 2.0 configuration options table:
markdown# OAuth 2.0 Configuration Options

Option	Description
`adbc.flight.sql.oauth.flow`	Specifies the OAuth 2.0 flow type to use. Possible values: `client_credentials`, `token_exchange`
`adbc.flight.sql.oauth.client_id`	Unique identifier issued to the client application by the authorization server
`adbc.flight.sql.oauth.client_secret`	Secret associated to the client_id. Used to authenticate the client application to the authorization server
`adbc.flight.sql.oauth.token_uri`	The endpoint URL where the client application requests tokens from the authorization server
`adbc.flight.sql.oauth.scope`	Space-separated list of permissions that the client is requesting access to (e.g `"read.all offline_access"`)
`adbc.flight.sql.oauth.exchange.subject_token`	The security token that the client application wants to exchange
`adbc.flight.sql.oauth.exchange.subject_token_type`	Identifier for the type of the subject token. Check list below for supported token types.
`adbc.flight.sql.oauth.exchange.actor_token`	A security token that represents the identity of the acting party
`adbc.flight.sql.oauth.exchange.actor_token_type`	Identifier for the type of the actor token. Check list below for supported token types.
`adbc.flight.sql.oauth.exchange.aud`	The intended audience for the requested security token
`adbc.flight.sql.oauth.exchange.resource`	The resource server where the client intends to use the requested security token
`adbc.flight.sql.oauth.exchange.scope`	Specific permissions requested for the new token
`adbc.flight.sql.oauth.exchange.requested_token_type`	The type of token the client wants to receive in exchange. Check list below for supported token types.

Supported token types:

urn:ietf:params:oauth:token-type:access_token
urn:ietf:params:oauth:token-type:refresh_token
urn:ietf:params:oauth:token-type:id_token
urn:ietf:params:oauth:token-type:saml1
urn:ietf:params:oauth:token-type:saml2
urn:ietf:params:oauth:token-type:jwt

jbonofre · 2025-03-27T13:03:46Z

I would suggest to update docs as part of this PR (with basically what you have in PR description).

howareyouman

Thank you, Helder, for this PR!
The only concerns I have are around thread safety.

howareyouman · 2025-03-27T13:44:20Z

+	}, nil
+}
+
+func (f *tokenExchange) GetToken(ctx context.Context) (*oauth2.Token, error) {


Could you please add RWLock here to make this method thread safe?

howareyouman · 2025-03-27T13:49:01Z

+	}, nil
+}
+
+func (c *clientCredentials) GetToken(ctx context.Context) (*oauth2.Token, error) {


The same question is here - should it be thread safe? If so, could you please add rwlock here?

zeroshade

maybe we can use https://pkg.go.dev/golang.org/x/oauth2@v0.28.0/clientcredentials#Config instead and just call https://pkg.go.dev/golang.org/x/oauth2@v0.28.0/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

…ken is set

xborder · 2025-04-07T16:19:35Z

maybe we can use https://pkg.go.dev/golang.org/x/oauth2@v0.28.0/clientcredentials#Config instead and just call https://pkg.go.dev/golang.org/x/oauth2@v0.28.0/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:

grpc-go seems to have some test certs for its examples. Would it make sense to have this for these tests?
Implement the PerRPCCredentials to set RequireTransportSecurity to return false on certain conditions. Besides that, I think the implementation would not be that much different. Just more aligned with grpc DialOptions.

Do you have any suggestion?

zeroshade · 2025-04-07T16:30:41Z

We have a TLS suite set up already for tests here:

arrow-adbc/go/adbc/driver/flightsql/flightsql_adbc_test.go

Line 857 in 854d31e

type TLSTests struct {

you can reuse that

zeroshade · 2025-04-07T16:32:07Z

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

…ntials interface

* 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

zeroshade · 2025-04-11T17:49:09Z

  header will then be sent back as the ``authorization`` header on all
  future requests.

+- (Go only) OAuth 2.0 authentication flows.


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)

Does it make sense to create a separate PR for this?

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.

davidhcoe · 2025-04-12T13:44:22Z

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.

davidhcoe · 2025-04-12T13:53:16Z

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

arrow-adbc/go/adbc/driver/flightsql/flightsql_statement.go

Line 77 in a187ead

return cnxn.execute(ctx, s.sqlQuery, opts...)

runs long.

davidhcoe · 2025-04-12T13:57:39Z

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

arrow-adbc/go/adbc/driver/flightsql/flightsql_statement.go

Line 77 in a187ead

return cnxn.execute(ctx, s.sqlQuery, opts...)

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.

xborder · 2025-04-14T12:28:19Z

@davidhcoe I think you can refer to this comment.
TokenSource returns the obtain token or triggers a refresh if it expired.
I tested locally and could confirm that it was triggering refresh calls to the IdP in between flight calls. Even if the request runs long, it should still be valid until the request is finished/aborted. For the next request it should refresh.
It will require a better setup to test those edge cases but for most common scenarios this PR should address it.

zeroshade

just a couple nits, but otherwise this looks good to me! Thanks for this!

…pache#2651) ## 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 - Closes #[2650](apache#2650) ## Changes Made 1. Added `token` as a database option 1. Added support for [Token Exchange](https://datatracker.ietf.org/doc/html/rfc8693). If configured, `token` gets exchanged and the result is added to the `Authorization` header as a `Bearer` token 1. Added support for [Client Credentials](https://datatracker.ietf.org/doc/html/rfc6749#section-4.4). If configured, `client_id` and `client_secret` are used to obtain a access token that is added to the `Authorization` header as a `Bearer` token 1. Added new driver options to allow third-party applications to configure oauth flows: 1. Added tests Here's the markdown code for the OAuth 2.0 configuration options table: markdown# OAuth 2.0 Configuration Options | Option | Description | |--------|-------------| | `adbc.flight.sql.oauth.flow` | Specifies the OAuth 2.0 flow type to use. Possible values: `client_credentials`, `token_exchange` | | `adbc.flight.sql.oauth.client_id` | Unique identifier issued to the client application by the authorization server | | `adbc.flight.sql.oauth.client_secret` | Secret associated to the client_id. Used to authenticate the client application to the authorization server | | `adbc.flight.sql.oauth.token_uri` | The endpoint URL where the client application requests tokens from the authorization server | | `adbc.flight.sql.oauth.scope` | Space-separated list of permissions that the client is requesting access to (e.g `"read.all offline_access"`) | | `adbc.flight.sql.oauth.exchange.subject_token` | The security token that the client application wants to exchange | | `adbc.flight.sql.oauth.exchange.subject_token_type` | Identifier for the type of the subject token. Check list below for supported token types. | | `adbc.flight.sql.oauth.exchange.actor_token` | A security token that represents the identity of the acting party | | `adbc.flight.sql.oauth.exchange.actor_token_type` | Identifier for the type of the actor token. Check list below for supported token types. | | `adbc.flight.sql.oauth.exchange.aud` | The intended audience for the requested security token | | `adbc.flight.sql.oauth.exchange.resource` | The resource server where the client intends to use the requested security token | | `adbc.flight.sql.oauth.exchange.scope` | Specific permissions requested for the new token | | `adbc.flight.sql.oauth.exchange.requested_token_type` | The type of token the client wants to receive in exchange. Check list below for supported token types. | **Supported token types:** * `urn:ietf:params:oauth:token-type:access_token` * `urn:ietf:params:oauth:token-type:refresh_token` * `urn:ietf:params:oauth:token-type:id_token` * `urn:ietf:params:oauth:token-type:saml1` * `urn:ietf:params:oauth:token-type:saml2` * `urn:ietf:params:oauth:token-type:jwt`

…3849) This pull request adds OAuth options introduced in #2651 to Python. ## Changes Made - Added enums in python with available OAuth options - Updated documentation - Added python recipes with working mock oauth server ## Related Issues Closes #2714.

…pache#3849) This pull request adds OAuth options introduced in apache#2651 to Python. ## Changes Made - Added enums in python with available OAuth options - Updated documentation - Added python recipes with working mock oauth server ## Related Issues Closes apache#2714.

xborder added 5 commits March 20, 2025 21:51

initial attempt. working with token exchange

39da2c4

add oauth support for flight client

c10af80

add oauth tests

7758432

add missing configurations to token exchange flow

bb8a1f1

add missing configurations to client credentials

10cd341

xborder changed the title ~~feat(go/adbc/flightsql): Add OAuth Support to Flight Client~~ feat(go/adbc/driver/flightsql): Add OAuth Support to Flight Client Mar 25, 2025

xborder added 3 commits March 25, 2025 17:21

revert changes to example_usage_test

775e779

fix(go/adbc/driver/flightsql): tests linting

3f83c43

fix(go/adbc/driver/flightsql): reset test suit db instead of server

e74fb32

xborder marked this pull request as ready for review March 25, 2025 23:42

xborder requested a review from zeroshade as a code owner March 25, 2025 23:42

github-actions Bot added this to the ADBC Libraries 18 milestone Mar 25, 2025

xborder mentioned this pull request Mar 25, 2025

go/adbc/flightsql: OAuth 2.0 flows support in flight sql clients #2650

Closed

jbonofre self-requested a review March 27, 2025 12:59

howareyouman reviewed Mar 27, 2025

View reviewed changes

zeroshade requested changes Mar 27, 2025

View reviewed changes

xborder added 8 commits April 7, 2025 11:19

chore(go/adbc/driver/flightsql): refactor set auth header to a function

823e6b8

chore(go/adbc/driver/flightsql): refactor flightsql_oauth implementation

a0ad006

chore(go/adbc/driver/flightsql): small fixes

737e207

chore(go/adbc/driver/flightsql): moved token to flightsq configuration

9ff438c

chore(go/adbc/driver/flightsql): fix token key

f8ce583

chore(go/adbc/driver/flightsql): simplify setting token in SetOptions

78fc5fc

fix(go/adbc/driver/flightsql): throw error if triggering oauth but to…

b1f4a4d

…ken is set

test(go/adbc/driver/flightsql): test to fail oauth if token is set

1fc80e7

xborder closed this Apr 7, 2025

xborder reopened this Apr 7, 2025

xborder added 6 commits April 8, 2025 22:50

chore(go/adbc/driver/flightsql): replace oauth impl

2a84059

Replaced extra structs for OAuth in favor of grpc's TokenSource and grpc.DialOption WithPerRPCCredentials

test(go/adbc/driver/flightsql): adapt OAuth tests with tls

a2b2fe3

* 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

chore(go/adbc/driver/flightsql): replace TokenSource with PerRPCCrede…

bded2f0

…ntials interface

chore(go/adbc/driver/flightsql): remove atoi

76fc327

docs(source/driver): document oauth options

4d763b9

xborder requested a review from lidavidm as a code owner April 11, 2025 17:47

zeroshade reviewed Apr 11, 2025

View reviewed changes

xborder added 2 commits April 14, 2025 11:26

chore(go/adbc/driver/flightsql): code improvements

e7dc917

fix(go/adbc/driver/flightsql): fix expected message

64b8143

cultcargo reviewed Apr 14, 2025

View reviewed changes

Comment thread docs/source/driver/flight_sql.rst Outdated

Comment thread go/adbc/driver/flightsql/flightsql_database.go

howareyouman approved these changes Apr 15, 2025

View reviewed changes

docs(docs/source/driver): typo

9e98498

xborder requested a review from zeroshade April 16, 2025 13:13

zeroshade reviewed Apr 16, 2025

View reviewed changes

Comment thread go/adbc/driver/flightsql/flightsql_database.go Outdated

zeroshade approved these changes Apr 16, 2025

View reviewed changes

xborder added 2 commits April 17, 2025 10:20

docs(docs/source/driver): remove go only from docs

9a665d0

fix(go/adbc/driver/flightsql): small nits

c94f2ee

xborder mentioned this pull request Apr 17, 2025

python: add oauth options #2714

Closed

zeroshade merged commit 2ea2fcb into apache:main Apr 17, 2025

xborder mentioned this pull request May 16, 2025

go/adbc/driver/flightsql: TLS_ROOT_CERTS on OAuth flows #2830

Closed

xborder mentioned this pull request Jan 2, 2026

feat(python/adbc_driver_flightsql): add constants for OAuth options #3849

Merged

Uh oh!

Conversation

xborder commented Mar 25, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Related Issues

Changes Made

Uh oh!

jbonofre commented Mar 27, 2025

Uh oh!

howareyouman left a comment

Choose a reason for hiding this comment

Uh oh!

howareyouman Mar 27, 2025

Choose a reason for hiding this comment

Uh oh!

howareyouman Mar 27, 2025

Choose a reason for hiding this comment

Uh oh!

zeroshade left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

xborder commented Apr 7, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

zeroshade commented Apr 7, 2025

Uh oh!

zeroshade commented Apr 7, 2025

Uh oh!

zeroshade Apr 11, 2025

Choose a reason for hiding this comment

Uh oh!

xborder Apr 14, 2025

Choose a reason for hiding this comment

Uh oh!

zeroshade Apr 16, 2025

Choose a reason for hiding this comment

Uh oh!

xborder Apr 17, 2025

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

davidhcoe commented Apr 12, 2025

Uh oh!

davidhcoe commented Apr 12, 2025

Uh oh!

davidhcoe commented Apr 12, 2025

Uh oh!

xborder commented Apr 14, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

zeroshade left a comment

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

6 participants

xborder commented Mar 25, 2025 •

edited

Loading

xborder commented Apr 7, 2025 •

edited

Loading

xborder commented Apr 14, 2025 •

edited

Loading