Skip to content

Commit c300ae5

Browse files
Genesis929dumians
Genesis929
authored and
dumians
committed
feat(bigquery): support custom oauth header name (googleapis#2564)
Adds support for custom OAuth header and updates documentation. ## Description > Should include a concise description of the changes (bug or feature), it's > impact, along with a summary of the solution ## PR Checklist > Thank you for opening a Pull Request! Before submitting your PR, there are a > few things you can do to make sure it goes smoothly: - [ ] Make sure you reviewed [CONTRIBUTING.md](https://github.com/googleapis/genai-toolbox/blob/main/CONTRIBUTING.md) - [ ] Make sure to open an issue as a [bug/issue](https://github.com/googleapis/genai-toolbox/issues/new/choose) before writing your code! That way we can discuss the change, evaluate designs, and agree on the general idea - [ ] Ensure the tests and linter pass - [ ] Code coverage does not decrease (if any source code was changed) - [ ] Appropriate docs were updated (if necessary) - [ ] Make sure to add `!` if this involve a breaking change 🛠️ Fixes #<issue_number_goes_here>
1 parent bfcb792 commit c300ae5

17 files changed

Lines changed: 224 additions & 46 deletions

File tree

‎docs/BIGQUERY_README.md‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -74,7 +74,7 @@ The BigQuery MCP server is configured using environment variables.
7474
```bash
7575
export BIGQUERY_PROJECT="<your-gcp-project-id>"
7676
export BIGQUERY_LOCATION="<your-dataset-location>" # Optional
77-
export BIGQUERY_USE_CLIENT_OAUTH="true" # Optional
77+
export BIGQUERY_USE_CLIENT_OAUTH="true" # Optional: true, false, or a custom header name
7878
export BIGQUERY_SCOPES="<comma-separated-scopes>" # Optional
7979
```
8080

‎docs/en/resources/sources/bigquery.md‎

Lines changed: 10 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -104,10 +104,15 @@ Common scopes include `https://www.googleapis.com/auth/bigquery` or
104104
### Authentication via User's OAuth Access Token
105105

106106
If the `useClientOAuth` parameter is set to `true`, Toolbox will instead use the
107-
OAuth access token for authentication. This token is parsed from the
108-
`Authorization` header passed in with the tool invocation request. This method
109-
allows Toolbox to make queries to [BigQuery][bigquery-docs] on behalf of the
110-
client or the end-user.
107+
OAuth access token for authentication. By default, this token is parsed from the
108+
`Authorization` header passed in with the tool invocation request.
109+
110+
If you need to use a non-standard header for the access token (e.g., to avoid
111+
conflicts with other services like Cloud Run), you can specify the header name
112+
in the `useClientOAuth` field (e.g., `useClientOAuth: X-BigQuery-Auth`).
113+
114+
This method allows Toolbox to make queries to [BigQuery][bigquery-docs] on behalf
115+
of the client or the end-user.
111116

112117
When using this on-behalf-of authentication, you must ensure that the
113118
identity used has been granted the correct IAM permissions.
@@ -166,7 +171,7 @@ useClientOAuth: true
166171
| location | string | false | Specifies the location (e.g., 'us', 'asia-northeast1') in which to run the query job. This location must match the location of any tables referenced in the query. Defaults to the table's location or 'US' if the location cannot be determined. [Learn More](https://cloud.google.com/bigquery/docs/locations) |
167172
| writeMode | string | false | Controls the write behavior for tools. `allowed` (default): All queries are permitted. `blocked`: Only `SELECT` statements are allowed for the `bigquery-execute-sql` tool. `protected`: Enables session-based execution where all tools associated with this source instance share the same [BigQuery session](https://cloud.google.com/bigquery/docs/sessions-intro). This allows for stateful operations using temporary tables (e.g., `CREATE TEMP TABLE`). For `bigquery-execute-sql`, `SELECT` statements can be used on all tables, but write operations are restricted to the session's temporary dataset. For tools like `bigquery-sql`, `bigquery-forecast`, and `bigquery-analyze-contribution`, the `writeMode` restrictions do not apply, but they will operate within the shared session. **Note:** The `protected` mode cannot be used with `useClientOAuth: true`. It is also not recommended for multi-user server environments, as all users would share the same session. A session is terminated automatically after 24 hours of inactivity or after 7 days, whichever comes first. A new session is created on the next request, and any temporary data from the previous session will be lost. |
168173
| allowedDatasets | []string | false | An optional list of dataset IDs that tools using this source are allowed to access. If provided, any tool operation attempting to access a dataset not in this list will be rejected. To enforce this, two types of operations are also disallowed: 1) Dataset-level operations (e.g., `CREATE SCHEMA`), and 2) operations where table access cannot be statically analyzed (e.g., `EXECUTE IMMEDIATE`, `CREATE PROCEDURE`). If a single dataset is provided, it will be treated as the default for prebuilt tools. |
169-
| useClientOAuth | bool | false | If true, forwards the client's OAuth access token from the "Authorization" header to downstream queries. **Note:** This cannot be used with `writeMode: protected`. |
174+
| useClientOAuth | string | false | If set to `'true'`, forwards the client's OAuth access token from the default `Authorization` header. If set to a custom header name (e.g., `X-My-Auth`), that header will be used instead. An empty string or `'false'` disables this feature. Defaults to `""` (disabled). |
170175
| scopes | []string | false | A list of OAuth 2.0 scopes to use for the credentials. If not provided, default scopes are used. |
171176
| impersonateServiceAccount | string | false | Service account email to impersonate when making BigQuery and Dataplex API calls. The authenticated principal must have the `roles/iam.serviceAccountTokenCreator` role on the target service account. [Learn More](https://cloud.google.com/iam/docs/service-account-impersonation) |
172177
| maxQueryResultRows | int | false | The maximum number of rows to return from a query. Defaults to 50. |

‎docs/en/resources/sources/looker.md‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -101,7 +101,7 @@ instead of hardcoding your secrets into the configuration file.
101101
| project | string | false | The project id to use in Google Cloud. |
102102
| location | string | false | The location to use in Google Cloud. (default: us) |
103103
| timeout | string | false | Maximum time to wait for query execution (e.g. "30s", "2m"). By default, 120s is applied. |
104-
| use_client_oauth | string | false | Use OAuth tokens instead of client_id and client_secret. (default: false) If a header name is provided, it will be used instead of "Authorization". |
104+
| use_client_oauth | string | false | If set to `'true'`, forwards the client's OAuth access token from the default `Authorization` header. If set to a custom header name (e.g., `X-Looker-Auth`), that header will be used instead. An empty string or `'false'` disables this feature. Defaults to `""` (disabled). |
105105
| show_hidden_models | string | false | Show or hide hidden models. (default: true) |
106106
| show_hidden_explores | string | false | Show or hide hidden explores. (default: true) |
107107
| show_hidden_fields | string | false | Show or hide hidden fields. (default: true) |

‎internal/sources/bigquery/bigquery.go‎

Lines changed: 36 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -86,7 +86,7 @@ type Config struct {
8686
Location string `yaml:"location"`
8787
WriteMode string `yaml:"writeMode"`
8888
AllowedDatasets StringOrStringSlice `yaml:"allowedDatasets"`
89-
UseClientOAuth bool `yaml:"useClientOAuth"`
89+
UseClientOAuth string `yaml:"useClientOAuth"`
9090
ImpersonateServiceAccount string `yaml:"impersonateServiceAccount"`
9191
Scopes StringOrStringSlice `yaml:"scopes"`
9292
MaxQueryResultRows int `yaml:"maxQueryResultRows"`
@@ -123,6 +123,7 @@ func (r Config) SourceConfigType() string {
123123
// Returns BigQuery source type
124124
return SourceType
125125
}
126+
126127
func (r Config) Initialize(ctx context.Context, tracer trace.Tracer) (sources.Source, error) {
127128
if r.WriteMode == "" {
128129
r.WriteMode = WriteModeAllowed
@@ -132,15 +133,15 @@ func (r Config) Initialize(ctx context.Context, tracer trace.Tracer) (sources.So
132133
r.MaxQueryResultRows = 50
133134
}
134135

135-
if r.WriteMode == WriteModeProtected && r.UseClientOAuth {
136+
if r.WriteMode == WriteModeProtected && strings.ToLower(r.UseClientOAuth) != "false" && r.UseClientOAuth != "" {
136137
// The protected mode only allows write operations to the session's temporary datasets.
137138
// when using client OAuth, a new session is created every
138139
// time a BigQuery tool is invoked. Therefore, no session data can
139140
// be preserved as needed by the protected mode.
140-
return nil, fmt.Errorf("writeMode 'protected' cannot be used with useClientOAuth 'true'")
141+
return nil, fmt.Errorf("writeMode 'protected' cannot be used with useClientOAuth enabled")
141142
}
142143

143-
if r.UseClientOAuth && r.ImpersonateServiceAccount != "" {
144+
if strings.ToLower(r.UseClientOAuth) != "false" && r.UseClientOAuth != "" && r.ImpersonateServiceAccount != "" {
144145
return nil, fmt.Errorf("useClientOAuth cannot be used with impersonateServiceAccount")
145146
}
146147

@@ -151,23 +152,16 @@ func (r Config) Initialize(ctx context.Context, tracer trace.Tracer) (sources.So
151152
var err error
152153

153154
s := &Source{
154-
Config: r,
155-
Client: client,
156-
RestService: restService,
157-
TokenSource: tokenSource,
158-
MaxQueryResultRows: r.MaxQueryResultRows,
159-
ClientCreator: clientCreator,
155+
Config: r,
156+
Client: client,
157+
RestService: restService,
158+
TokenSource: tokenSource,
159+
MaxQueryResultRows: r.MaxQueryResultRows,
160+
ClientCreator: clientCreator,
161+
AuthTokenHeaderName: "Authorization",
160162
}
161163

162-
if r.UseClientOAuth {
163-
// use client OAuth
164-
baseClientCreator, err := newBigQueryClientCreator(ctx, tracer, r.Project, r.Location, r.Name)
165-
if err != nil {
166-
return nil, fmt.Errorf("error constructing client creator: %w", err)
167-
}
168-
setupClientCaching(s, baseClientCreator)
169-
170-
} else {
164+
if strings.ToLower(r.UseClientOAuth) == "false" || r.UseClientOAuth == "" {
171165
// Initializes a BigQuery Google SQL source
172166
client, restService, tokenSource, err = initBigQueryConnection(ctx, tracer, r.Name, r.Project, r.Location, r.ImpersonateServiceAccount, r.Scopes)
173167
if err != nil {
@@ -176,6 +170,21 @@ func (r Config) Initialize(ctx context.Context, tracer trace.Tracer) (sources.So
176170
s.Client = client
177171
s.RestService = restService
178172
s.TokenSource = tokenSource
173+
174+
if r.WriteMode == WriteModeProtected {
175+
// session-based connections
176+
s.SessionProvider = s.newBigQuerySessionProvider()
177+
}
178+
} else {
179+
if strings.ToLower(r.UseClientOAuth) != "true" {
180+
s.AuthTokenHeaderName = r.UseClientOAuth
181+
}
182+
// use client OAuth
183+
baseClientCreator, err := newBigQueryClientCreator(ctx, tracer, r.Project, r.Location, r.Name)
184+
if err != nil {
185+
return nil, fmt.Errorf("error constructing client creator: %w", err)
186+
}
187+
setupClientCaching(s, baseClientCreator)
179188
}
180189

181190
allowedDatasets := make(map[string]struct{})
@@ -280,6 +289,7 @@ type Source struct {
280289
Client *bigqueryapi.Client
281290
RestService *bigqueryrestapi.Service
282291
TokenSource oauth2.TokenSource
292+
AuthTokenHeaderName string
283293
MaxQueryResultRows int
284294
ClientCreator BigqueryClientCreator
285295
AllowedDatasets map[string]struct{}
@@ -417,7 +427,11 @@ func (s *Source) newBigQuerySessionProvider() BigQuerySessionProvider {
417427
}
418428

419429
func (s *Source) UseClientAuthorization() bool {
420-
return s.UseClientOAuth
430+
return strings.ToLower(s.UseClientOAuth) != "false" && s.UseClientOAuth != ""
431+
}
432+
433+
func (s *Source) GetAuthTokenHeaderName() string {
434+
return s.AuthTokenHeaderName
421435
}
422436

423437
func (s *Source) BigQueryProject() string {
@@ -497,15 +511,15 @@ func (s *Source) lazyInitDataplexClient(ctx context.Context, tracer trace.Tracer
497511

498512
return func() (*dataplexapi.CatalogClient, DataplexClientCreator, error) {
499513
once.Do(func() {
500-
c, cc, e := initDataplexConnection(ctx, tracer, s.Name, s.Project, s.UseClientOAuth, s.ImpersonateServiceAccount, s.Scopes)
514+
c, cc, e := initDataplexConnection(ctx, tracer, s.Name, s.Project, s.UseClientAuthorization(), s.ImpersonateServiceAccount, s.Scopes)
501515
if e != nil {
502516
err = fmt.Errorf("failed to initialize dataplex client: %w", e)
503517
return
504518
}
505519
client = c
506520

507521
// If using OAuth, wrap the provided client creator (cc) with caching logic
508-
if s.UseClientOAuth && cc != nil {
522+
if s.UseClientAuthorization() && cc != nil {
509523
clientCreator = func(tokenString string) (*dataplexapi.CatalogClient, error) {
510524
// Check cache
511525
if val, found := s.dataplexCache.Get(tokenString); found {

‎internal/sources/bigquery/bigquery_test.go‎

Lines changed: 64 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -70,7 +70,7 @@ func TestParseFromYamlBigQuery(t *testing.T) {
7070
Project: "my-project",
7171
Location: "asia",
7272
WriteMode: "blocked",
73-
UseClientOAuth: false,
73+
UseClientOAuth: "",
7474
},
7575
},
7676
},
@@ -82,6 +82,46 @@ func TestParseFromYamlBigQuery(t *testing.T) {
8282
type: bigquery
8383
project: my-project
8484
location: us
85+
useClientOAuth: "true"
86+
`,
87+
want: map[string]sources.SourceConfig{
88+
"my-instance": bigquery.Config{
89+
Name: "my-instance",
90+
Type: bigquery.SourceType,
91+
Project: "my-project",
92+
Location: "us",
93+
UseClientOAuth: "true",
94+
},
95+
},
96+
},
97+
{
98+
desc: "with custom auth header name example",
99+
in: `
100+
kind: sources
101+
name: my-instance
102+
type: bigquery
103+
project: my-project
104+
location: us
105+
useClientOAuth: X-Custom-Auth
106+
`,
107+
want: map[string]sources.SourceConfig{
108+
"my-instance": bigquery.Config{
109+
Name: "my-instance",
110+
Type: bigquery.SourceType,
111+
Project: "my-project",
112+
Location: "us",
113+
UseClientOAuth: "X-Custom-Auth",
114+
},
115+
},
116+
},
117+
{
118+
desc: "use client auth with unquoted true",
119+
in: `
120+
kind: sources
121+
name: my-instance
122+
type: bigquery
123+
project: my-project
124+
location: us
85125
useClientOAuth: true
86126
`,
87127
want: map[string]sources.SourceConfig{
@@ -90,7 +130,27 @@ func TestParseFromYamlBigQuery(t *testing.T) {
90130
Type: bigquery.SourceType,
91131
Project: "my-project",
92132
Location: "us",
93-
UseClientOAuth: true,
133+
UseClientOAuth: "true",
134+
},
135+
},
136+
},
137+
{
138+
desc: "use client auth with unquoted false",
139+
in: `
140+
kind: sources
141+
name: my-instance
142+
type: bigquery
143+
project: my-project
144+
location: us
145+
useClientOAuth: false
146+
`,
147+
want: map[string]sources.SourceConfig{
148+
"my-instance": bigquery.Config{
149+
Name: "my-instance",
150+
Type: bigquery.SourceType,
151+
Project: "my-project",
152+
Location: "us",
153+
UseClientOAuth: "false",
94154
},
95155
},
96156
},
@@ -253,7 +313,7 @@ func TestInitialize_MaxQueryResultRows(t *testing.T) {
253313
Name: "test-default",
254314
Type: bigquery.SourceType,
255315
Project: "test-project",
256-
UseClientOAuth: true,
316+
UseClientOAuth: "true",
257317
},
258318
want: 50,
259319
},
@@ -263,7 +323,7 @@ func TestInitialize_MaxQueryResultRows(t *testing.T) {
263323
Name: "test-configured",
264324
Type: bigquery.SourceType,
265325
Project: "test-project",
266-
UseClientOAuth: true,
326+
UseClientOAuth: "true",
267327
MaxQueryResultRows: 100,
268328
},
269329
want: 100,

‎internal/tools/bigquery/bigqueryanalyzecontribution/bigqueryanalyzecontribution.go‎

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -52,6 +52,7 @@ func newConfig(ctx context.Context, name string, decoder *yaml.Decoder) (tools.T
5252
type compatibleSource interface {
5353
BigQueryClient() *bigqueryapi.Client
5454
UseClientAuthorization() bool
55+
GetAuthTokenHeaderName() string
5556
IsDatasetAllowed(projectID, datasetID string) bool
5657
BigQueryAllowedDatasets() []string
5758
BigQuerySession() bigqueryds.BigQuerySessionProvider
@@ -340,7 +341,11 @@ func (t Tool) RequiresClientAuthorization(resourceMgr tools.SourceProvider) (boo
340341
}
341342

342343
func (t Tool) GetAuthTokenHeaderName(resourceMgr tools.SourceProvider) (string, error) {
343-
return "Authorization", nil
344+
source, err := tools.GetCompatibleSource[compatibleSource](resourceMgr, t.Source, t.Name, t.Type)
345+
if err != nil {
346+
return "", err
347+
}
348+
return source.GetAuthTokenHeaderName(), nil
344349
}
345350

346351
func (t Tool) GetParameters() parameters.Parameters {

‎internal/tools/bigquery/bigqueryconversationalanalytics/bigqueryconversationalanalytics.go‎

Lines changed: 9 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -63,6 +63,7 @@ type compatibleSource interface {
6363
BigQueryLocation() string
6464
GetMaxQueryResultRows() int
6565
UseClientAuthorization() bool
66+
GetAuthTokenHeaderName() string
6667
IsDatasetAllowed(projectID, datasetID string) bool
6768
BigQueryAllowedDatasets() []string
6869
}
@@ -241,9 +242,9 @@ func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, para
241242
caURL := fmt.Sprintf(gdaURLFormat, projectID, location)
242243

243244
headers := map[string]string{
244-
"Authorization": fmt.Sprintf("Bearer %s", tokenStr),
245-
"Content-Type": "application/json",
246-
"X-Goog-API-Client": util.GDAClientID,
245+
source.GetAuthTokenHeaderName(): fmt.Sprintf("Bearer %s", tokenStr),
246+
"Content-Type": "application/json",
247+
"X-Goog-API-Client": util.GDAClientID,
247248
}
248249

249250
payload := CAPayload{
@@ -578,7 +579,11 @@ func appendMessage(messages []map[string]any, newMessage map[string]any) []map[s
578579
}
579580

580581
func (t Tool) GetAuthTokenHeaderName(resourceMgr tools.SourceProvider) (string, error) {
581-
return "Authorization", nil
582+
source, err := tools.GetCompatibleSource[compatibleSource](resourceMgr, t.Source, t.Name, t.Type)
583+
if err != nil {
584+
return "", err
585+
}
586+
return source.GetAuthTokenHeaderName(), nil
582587
}
583588

584589
func (t Tool) GetParameters() parameters.Parameters {

‎internal/tools/bigquery/bigqueryexecutesql/bigqueryexecutesql.go‎

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -54,6 +54,7 @@ type compatibleSource interface {
5454
BigQuerySession() bigqueryds.BigQuerySessionProvider
5555
BigQueryWriteMode() string
5656
UseClientAuthorization() bool
57+
GetAuthTokenHeaderName() string
5758
IsDatasetAllowed(projectID, datasetID string) bool
5859
BigQueryAllowedDatasets() []string
5960
RetrieveClientAndService(tools.AccessToken) (*bigqueryapi.Client, *bigqueryrestapi.Service, error)
@@ -311,7 +312,11 @@ func (t Tool) RequiresClientAuthorization(resourceMgr tools.SourceProvider) (boo
311312
}
312313

313314
func (t Tool) GetAuthTokenHeaderName(resourceMgr tools.SourceProvider) (string, error) {
314-
return "Authorization", nil
315+
source, err := tools.GetCompatibleSource[compatibleSource](resourceMgr, t.Source, t.Name, t.Type)
316+
if err != nil {
317+
return "", err
318+
}
319+
return source.GetAuthTokenHeaderName(), nil
315320
}
316321

317322
func (t Tool) GetParameters() parameters.Parameters {

‎internal/tools/bigquery/bigqueryforecast/bigqueryforecast.go‎

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -51,6 +51,7 @@ func newConfig(ctx context.Context, name string, decoder *yaml.Decoder) (tools.T
5151
type compatibleSource interface {
5252
BigQueryClient() *bigqueryapi.Client
5353
UseClientAuthorization() bool
54+
GetAuthTokenHeaderName() string
5455
IsDatasetAllowed(projectID, datasetID string) bool
5556
BigQueryAllowedDatasets() []string
5657
BigQuerySession() bigqueryds.BigQuerySessionProvider
@@ -301,7 +302,11 @@ func (t Tool) RequiresClientAuthorization(resourceMgr tools.SourceProvider) (boo
301302
}
302303

303304
func (t Tool) GetAuthTokenHeaderName(resourceMgr tools.SourceProvider) (string, error) {
304-
return "Authorization", nil
305+
source, err := tools.GetCompatibleSource[compatibleSource](resourceMgr, t.Source, t.Name, t.Type)
306+
if err != nil {
307+
return "", err
308+
}
309+
return source.GetAuthTokenHeaderName(), nil
305310
}
306311

307312
func (t Tool) GetParameters() parameters.Parameters {

0 commit comments

Comments
 (0)