feat: add quotaProject support for BigQuery and Looker conversational analytics (#2610)

Fixes #2543

## Summary
- Add optional `quotaProject` to BigQuery source config
- Apply `option.WithQuotaProject(quotaProject)` when creating BigQuery
clients/services
- Add optional `quotaProject` to Looker source config
- For `looker-conversational-analytics`, set `X-Goog-User-Project`
header when configured

## Why
When `useClientOAuth: true`, quota/billing can be attributed to the
token-issuing project. This allows explicit quota attribution to a
deployment-controlled project.

## Validation
- go test ./internal/sources/bigquery
- go test ./internal/sources/looker
- go test ./internal/tools/looker/lookerconversationalanalytics

---------

Co-authored-by: Yuan Teoh <45984206+Yuan325@users.noreply.github.com>

Author: Deeven-Seru

internal/sources/bigquery/bigquery.go

@@ -88,6 +88,7 @@ type Config struct {
 	WriteMode                 string              `yaml:"writeMode"`
 	AllowedDatasets           StringOrStringSlice `yaml:"allowedDatasets"`
 	UseClientOAuth            string              `yaml:"useClientOAuth"`
+	QuotaProject              string              `yaml:"quotaProject"`
 	ImpersonateServiceAccount string              `yaml:"impersonateServiceAccount"`
 	Scopes                    StringOrStringSlice `yaml:"scopes"`
 	MaxQueryResultRows        int                 `yaml:"maxQueryResultRows"`
@@ -166,7 +167,7 @@ func (r Config) Initialize(ctx context.Context, tracer trace.Tracer) (sources.So
 
 	if strings.ToLower(r.UseClientOAuth) == "false" || r.UseClientOAuth == "" {
 		// Initializes a BigQuery Google SQL source
-		client, restService, tokenSource, err = initBigQueryConnection(ctx, tracer, r.Name, r.Project, r.Location, r.ImpersonateServiceAccount, r.Scopes)
+		client, restService, tokenSource, err = initBigQueryConnection(ctx, tracer, r.Name, r.Project, r.Location, r.QuotaProject, r.ImpersonateServiceAccount, r.Scopes)
 		if err != nil {
 			return nil, fmt.Errorf("error creating client from ADC: %w", err)
 		}
@@ -183,7 +184,7 @@ func (r Config) Initialize(ctx context.Context, tracer trace.Tracer) (sources.So
 			s.AuthTokenHeaderName = r.UseClientOAuth
 		}
 		// use client OAuth
-		baseClientCreator, err := newBigQueryClientCreator(ctx, tracer, r.Project, r.Location, r.Name)
+		baseClientCreator, err := newBigQueryClientCreator(ctx, tracer, r.Project, r.Location, r.QuotaProject, r.Name)
 		if err != nil {
 			return nil, fmt.Errorf("error constructing client creator: %w", err)
 		}
@@ -446,6 +447,10 @@ func (s *Source) BigQueryLocation() string {
 	return s.Location
 }
 
+func (s *Source) BigQueryQuotaProject() string {
+	return s.QuotaProject
+}
+
 func (s *Source) BigQueryTokenSource() oauth2.TokenSource {
 	return s.TokenSource
 }
@@ -690,6 +695,7 @@ func initBigQueryConnection(
 	name string,
 	project string,
 	location string,
+	quotaProject string,
 	impersonateServiceAccount string,
 	scopes []string,
 ) (*bigqueryapi.Client, *bigqueryrestapi.Service, oauth2.TokenSource, error) {
@@ -741,6 +747,10 @@ func initBigQueryConnection(
 		}
 	}
 
+	if quotaProject != "" {
+		opts = append(opts, option.WithQuotaProject(quotaProject))
+	}
+
 	// Initialize the high-level BigQuery client
 	client, err := bigqueryapi.NewClient(ctx, project, opts...)
 	if err != nil {
@@ -764,6 +774,7 @@ func initBigQueryConnectionWithOAuthToken(
 	tracer trace.Tracer,
 	project string,
 	location string,
+	quotaProject string,
 	name string,
 	userAgent string,
 	tokenString string,
@@ -777,16 +788,24 @@ func initBigQueryConnectionWithOAuthToken(
 	}
 	ts := oauth2.StaticTokenSource(token)
 
+	opts := []option.ClientOption{
+		option.WithUserAgent(userAgent),
+		option.WithTokenSource(ts),
+	}
+	if quotaProject != "" {
+		opts = append(opts, option.WithQuotaProject(quotaProject))
+	}
+
 	// Initialize the BigQuery client with tokenSource
-	client, err := bigqueryapi.NewClient(ctx, project, option.WithUserAgent(userAgent), option.WithTokenSource(ts))
+	client, err := bigqueryapi.NewClient(ctx, project, opts...)
 	if err != nil {
 		return nil, nil, fmt.Errorf("failed to create BigQuery client for project %q: %w", project, err)
 	}
 	client.Location = location
 
 	if wantRestService {
 		// Initialize the low-level BigQuery REST service using the same credentials
-		restService, err := bigqueryrestapi.NewService(ctx, option.WithUserAgent(userAgent), option.WithTokenSource(ts))
+		restService, err := bigqueryrestapi.NewService(ctx, opts...)
 		if err != nil {
 			return nil, nil, fmt.Errorf("failed to create BigQuery v2 service: %w", err)
 		}
@@ -804,6 +823,7 @@ func newBigQueryClientCreator(
 	tracer trace.Tracer,
 	project string,
 	location string,
+	quotaProject string,
 	name string,
 ) (func(string, bool) (*bigqueryapi.Client, *bigqueryrestapi.Service, error), error) {
 	userAgent, err := util.UserAgentFromContext(ctx)
@@ -812,7 +832,7 @@ func newBigQueryClientCreator(
 	}
 
 	return func(tokenString string, wantRestService bool) (*bigqueryapi.Client, *bigqueryrestapi.Service, error) {
-		return initBigQueryConnectionWithOAuthToken(ctx, tracer, project, location, name, userAgent, tokenString, wantRestService)
+		return initBigQueryConnectionWithOAuthToken(ctx, tracer, project, location, quotaProject, name, userAgent, tokenString, wantRestService)
 	}, nil
 }
 

internal/sources/bigquery/bigquery_test.go

@@ -154,6 +154,28 @@ func TestParseFromYamlBigQuery(t *testing.T) {
 				},
 			},
 		},
+		{
+			desc: "quota project with client auth example",
+			in: `
+			kind: source
+			name: my-instance
+			type: bigquery
+			project: my-project
+			location: us
+			useClientOAuth: true
+			quotaProject: billing-project
+			`,
+			want: map[string]sources.SourceConfig{
+				"my-instance": bigquery.Config{
+					Name:           "my-instance",
+					Type:           bigquery.SourceType,
+					Project:        "my-project",
+					Location:       "us",
+					UseClientOAuth: "true",
+					QuotaProject:   "billing-project",
+				},
+			},
+		},
 		{
 			desc: "with allowed datasets example",
 			in: `

internal/sources/looker/looker.go

@@ -76,6 +76,7 @@ type Config struct {
 	ShowHiddenFields   bool   `yaml:"show_hidden_fields"`
 	Project            string `yaml:"project"`
 	Location           string `yaml:"location"`
+	QuotaProject       string `yaml:"quotaProject"`
 	SessionLength      int64  `yaml:"sessionLength"`
 }
 
@@ -178,6 +179,10 @@ func (s *Source) GoogleCloudLocation() string {
 	return s.Location
 }
 
+func (s *Source) GoogleCloudQuotaProject() string {
+	return s.QuotaProject
+}
+
 func (s *Source) GoogleCloudTokenSource() oauth2.TokenSource {
 	return s.TokenSource
 }

internal/sources/looker/looker_test.go

@@ -65,6 +65,36 @@ func TestParseFromYamlLooker(t *testing.T) {
 				},
 			},
 		},
+		{
+			desc: "with quota project",
+			in: `
+			kind: source
+			name: my-looker-instance
+			type: looker
+			base_url: http://example.looker.com/
+			client_id: jasdl;k;tjl
+			client_secret: sdakl;jgflkasdfkfg
+			quotaProject: billing-project
+			`,
+			want: map[string]sources.SourceConfig{
+				"my-looker-instance": looker.Config{
+					Name:               "my-looker-instance",
+					Type:               looker.SourceType,
+					BaseURL:            "http://example.looker.com/",
+					ClientId:           "jasdl;k;tjl",
+					ClientSecret:       "sdakl;jgflkasdfkfg",
+					Timeout:            "600s",
+					SslVerification:    true,
+					UseClientOAuth:     "false",
+					ShowHiddenModels:   true,
+					ShowHiddenExplores: true,
+					ShowHiddenFields:   true,
+					Location:           "us",
+					QuotaProject:       "billing-project",
+					SessionLength:      1200,
+				},
+			},
+		},
 	}
 	for _, tc := range tcs {
 		t.Run(tc.desc, func(t *testing.T) {

internal/tools/bigquery/bigqueryconversationalanalytics/bigqueryconversationalanalytics.go

@@ -64,6 +64,7 @@ type compatibleSource interface {
 	BigQueryTokenSourceWithScope(ctx context.Context, scopes []string) (oauth2.TokenSource, error)
 	BigQueryProject() string
 	BigQueryLocation() string
+	BigQueryQuotaProject() string
 	GetMaxQueryResultRows() int
 	UseClientAuthorization() bool
 	GetAuthTokenHeaderName() string
@@ -216,6 +217,9 @@ func (t Tool) Invoke(ctx context.Context, primitiveMgr tools.SourceProvider, par
 		"Content-Type":      "application/json",
 		"X-Goog-API-Client": util.GDAClientID,
 	}
+	if quotaProject := source.BigQueryQuotaProject(); quotaProject != "" {
+		headers["X-Goog-User-Project"] = quotaProject
+	}
 
 	payload := CAPayload{
 		Messages: []Message{{UserMessage: UserMessage{Text: finalQueryText}}},

internal/tools/looker/lookerconversationalanalytics/lookerconversationalanalytics.go

@@ -59,6 +59,7 @@ type compatibleSource interface {
 	GoogleCloudTokenSourceWithScope(ctx context.Context, scope string) (oauth2.TokenSource, error)
 	GoogleCloudProject() string
 	GoogleCloudLocation() string
+	GoogleCloudQuotaProject() string
 	UseClientAuthorization() bool
 	GetAuthTokenHeaderName() string
 	LookerApiSettings() *rtl.ApiSettings
@@ -259,6 +260,9 @@ func (t Tool) Invoke(ctx context.Context, primitiveMgr tools.SourceProvider, par
 		"Content-Type":      "application/json",
 		"X-Goog-API-Client": util.GDAClientID,
 	}
+	if quotaProject := source.GoogleCloudQuotaProject(); quotaProject != "" {
+		headers["X-Goog-User-Project"] = quotaProject
+	}
 
 	payload := CAPayload{
 		Messages: []Message{{UserMessage: UserMessage{Text: userQuery}}},