From 287218ec55bbebb3902a3983ec5a02b2356b7406 Mon Sep 17 00:00:00 2001 From: Colin Pistell Date: Fri, 3 Oct 2025 11:40:14 -0600 Subject: [PATCH 1/2] feat(source/bigquery): add service account impersonation support for bigquery --- docs/en/resources/sources/bigquery.md | 17 +-- internal/sources/bigquery/bigquery.go | 119 +++++++++++++++------ internal/sources/bigquery/bigquery_test.go | 20 ++++ 3 files changed, 116 insertions(+), 40 deletions(-) diff --git a/docs/en/resources/sources/bigquery.md b/docs/en/resources/sources/bigquery.md index 51c9e865bdb8..a45caf671386 100644 --- a/docs/en/resources/sources/bigquery.md +++ b/docs/en/resources/sources/bigquery.md @@ -122,6 +122,7 @@ sources: # allowedDatasets: # Optional: Restricts tool access to a specific list of datasets. # - "my_dataset_1" # - "other_project.my_dataset_2" + # impersonateServiceAccount: "service-account@project-id.iam.gserviceaccount.com" # Optional: Service account to impersonate ``` Initialize a BigQuery source that uses the client's access token: @@ -136,14 +137,16 @@ sources: # allowedDatasets: # Optional: Restricts tool access to a specific list of datasets. # - "my_dataset_1" # - "other_project.my_dataset_2" + # impersonateServiceAccount: "service-account@project-id.iam.gserviceaccount.com" # Optional: Service account to impersonate ``` ## Reference -| **field** | **type** | **required** | **description** | -|-----------------|:--------:|:------------:|---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------| -| kind | string | true | Must be "bigquery". | -| project | string | true | Id of the Google Cloud project to use for billing and as the default project for BigQuery resources. | -| 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) | -| 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. | -| useClientOAuth | bool | false | If true, forwards the client's OAuth access token from the "Authorization" header to downstream queries. | +| **field** | **type** | **required** | **description** | +|---------------------------|:--------:|:------------:|---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------| +| kind | string | true | Must be "bigquery". | +| project | string | true | Id of the Google Cloud project to use for billing and as the default project for BigQuery resources. | +| 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) | +| 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. | +| useClientOAuth | bool | false | If true, forwards the client's OAuth access token from the "Authorization" header to downstream queries. | +| 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) | diff --git a/internal/sources/bigquery/bigquery.go b/internal/sources/bigquery/bigquery.go index 7beffdb1bc03..9831df18778f 100644 --- a/internal/sources/bigquery/bigquery.go +++ b/internal/sources/bigquery/bigquery.go @@ -31,6 +31,7 @@ import ( "golang.org/x/oauth2/google" bigqueryrestapi "google.golang.org/api/bigquery/v2" "google.golang.org/api/googleapi" + "google.golang.org/api/impersonate" "google.golang.org/api/option" ) @@ -59,12 +60,13 @@ func newConfig(ctx context.Context, name string, decoder *yaml.Decoder) (sources type Config struct { // BigQuery configs - Name string `yaml:"name" validate:"required"` - Kind string `yaml:"kind" validate:"required"` - Project string `yaml:"project" validate:"required"` - Location string `yaml:"location"` - AllowedDatasets []string `yaml:"allowedDatasets"` - UseClientOAuth bool `yaml:"useClientOAuth"` + Name string `yaml:"name" validate:"required"` + Kind string `yaml:"kind" validate:"required"` + Project string `yaml:"project" validate:"required"` + Location string `yaml:"location"` + AllowedDatasets []string `yaml:"allowedDatasets"` + UseClientOAuth bool `yaml:"useClientOAuth"` + ImpersonateServiceAccount string `yaml:"impersonateServiceAccount"` } func (r Config) SourceConfigKind() string { @@ -86,7 +88,7 @@ func (r Config) Initialize(ctx context.Context, tracer trace.Tracer) (sources.So } } else { // Initializes a BigQuery Google SQL source - client, restService, tokenSource, err = initBigQueryConnection(ctx, tracer, r.Name, r.Project, r.Location) + client, restService, tokenSource, err = initBigQueryConnection(ctx, tracer, r.Name, r.Project, r.Location, r.ImpersonateServiceAccount) if err != nil { return nil, fmt.Errorf("error creating client from ADC: %w", err) } @@ -124,17 +126,18 @@ func (r Config) Initialize(ctx context.Context, tracer trace.Tracer) (sources.So } s := &Source{ - Name: r.Name, - Kind: SourceKind, - Project: r.Project, - Location: r.Location, - Client: client, - RestService: restService, - TokenSource: tokenSource, - MaxQueryResultRows: 50, - ClientCreator: clientCreator, - AllowedDatasets: allowedDatasets, - UseClientOAuth: r.UseClientOAuth, + Name: r.Name, + Kind: SourceKind, + Project: r.Project, + Location: r.Location, + Client: client, + RestService: restService, + TokenSource: tokenSource, + MaxQueryResultRows: 50, + ClientCreator: clientCreator, + AllowedDatasets: allowedDatasets, + UseClientOAuth: r.UseClientOAuth, + ImpersonateServiceAccount: r.ImpersonateServiceAccount, } s.makeDataplexCatalogClient = s.lazyInitDataplexClient(ctx, tracer) return s, nil @@ -156,6 +159,7 @@ type Source struct { ClientCreator BigqueryClientCreator AllowedDatasets map[string]struct{} UseClientOAuth bool + ImpersonateServiceAccount string makeDataplexCatalogClient func() (*dataplexapi.CatalogClient, DataplexClientCreator, error) } @@ -235,7 +239,7 @@ func (s *Source) lazyInitDataplexClient(ctx context.Context, tracer trace.Tracer return func() (*dataplexapi.CatalogClient, DataplexClientCreator, error) { once.Do(func() { - c, cc, e := initDataplexConnection(ctx, tracer, s.Name, s.Project, s.UseClientOAuth) + c, cc, e := initDataplexConnection(ctx, tracer, s.Name, s.Project, s.UseClientOAuth, s.ImpersonateServiceAccount) if e != nil { err = fmt.Errorf("failed to initialize dataplex client: %w", e) return @@ -253,34 +257,60 @@ func initBigQueryConnection( name string, project string, location string, + impersonateServiceAccount string, ) (*bigqueryapi.Client, *bigqueryrestapi.Service, oauth2.TokenSource, error) { ctx, span := sources.InitConnectionSpan(ctx, tracer, SourceKind, name) defer span.End() - cred, err := google.FindDefaultCredentials(ctx, bigqueryapi.Scope) - if err != nil { - return nil, nil, nil, fmt.Errorf("failed to find default Google Cloud credentials with scope %q: %w", bigqueryapi.Scope, err) - } - userAgent, err := util.UserAgentFromContext(ctx) if err != nil { return nil, nil, nil, err } + var tokenSource oauth2.TokenSource + var opts []option.ClientOption + + if impersonateServiceAccount != "" { + // Create impersonated credentials token source + ts, err := impersonate.CredentialsTokenSource(ctx, impersonate.CredentialsConfig{ + TargetPrincipal: impersonateServiceAccount, + Scopes: []string{bigqueryapi.Scope}, + }) + if err != nil { + return nil, nil, nil, fmt.Errorf("failed to create impersonated credentials for %q: %w", impersonateServiceAccount, err) + } + tokenSource = ts + opts = []option.ClientOption{ + option.WithUserAgent(userAgent), + option.WithTokenSource(ts), + } + } else { + // Use default credentials + cred, err := google.FindDefaultCredentials(ctx, bigqueryapi.Scope) + if err != nil { + return nil, nil, nil, fmt.Errorf("failed to find default Google Cloud credentials with scope %q: %w", bigqueryapi.Scope, err) + } + tokenSource = cred.TokenSource + opts = []option.ClientOption{ + option.WithUserAgent(userAgent), + option.WithCredentials(cred), + } + } + // Initialize the high-level BigQuery client - client, err := bigqueryapi.NewClient(ctx, project, option.WithUserAgent(userAgent), option.WithCredentials(cred)) + client, err := bigqueryapi.NewClient(ctx, project, opts...) if err != nil { return nil, nil, nil, fmt.Errorf("failed to create BigQuery client for project %q: %w", project, err) } client.Location = location // Initialize the low-level BigQuery REST service using the same credentials - restService, err := bigqueryrestapi.NewService(ctx, option.WithUserAgent(userAgent), option.WithCredentials(cred)) + restService, err := bigqueryrestapi.NewService(ctx, opts...) if err != nil { return nil, nil, nil, fmt.Errorf("failed to create BigQuery v2 service: %w", err) } - return client, restService, cred.TokenSource, nil + return client, restService, tokenSource, nil } // initBigQueryConnectionWithOAuthToken initialize a BigQuery client with an @@ -348,6 +378,7 @@ func initDataplexConnection( name string, project string, useClientOAuth bool, + impersonateServiceAccount string, ) (*dataplexapi.CatalogClient, DataplexClientCreator, error) { var client *dataplexapi.CatalogClient var clientCreator DataplexClientCreator @@ -356,11 +387,6 @@ func initDataplexConnection( ctx, span := sources.InitConnectionSpan(ctx, tracer, SourceKind, name) defer span.End() - cred, err := google.FindDefaultCredentials(ctx) - if err != nil { - return nil, nil, fmt.Errorf("failed to find default Google Cloud credentials: %w", err) - } - userAgent, err := util.UserAgentFromContext(ctx) if err != nil { return nil, nil, err @@ -369,7 +395,34 @@ func initDataplexConnection( if useClientOAuth { clientCreator = newDataplexClientCreator(ctx, project, userAgent) } else { - client, err = dataplexapi.NewCatalogClient(ctx, option.WithUserAgent(userAgent), option.WithCredentials(cred)) + var opts []option.ClientOption + + if impersonateServiceAccount != "" { + // Create impersonated credentials token source + ts, err := impersonate.CredentialsTokenSource(ctx, impersonate.CredentialsConfig{ + TargetPrincipal: impersonateServiceAccount, + Scopes: []string{"https://www.googleapis.com/auth/cloud-platform"}, + }) + if err != nil { + return nil, nil, fmt.Errorf("failed to create impersonated credentials for %q: %w", impersonateServiceAccount, err) + } + opts = []option.ClientOption{ + option.WithUserAgent(userAgent), + option.WithTokenSource(ts), + } + } else { + // Use default credentials + cred, err := google.FindDefaultCredentials(ctx) + if err != nil { + return nil, nil, fmt.Errorf("failed to find default Google Cloud credentials: %w", err) + } + opts = []option.ClientOption{ + option.WithUserAgent(userAgent), + option.WithCredentials(cred), + } + } + + client, err = dataplexapi.NewCatalogClient(ctx, opts...) if err != nil { return nil, nil, fmt.Errorf("failed to create Dataplex client for project %q: %w", project, err) } diff --git a/internal/sources/bigquery/bigquery_test.go b/internal/sources/bigquery/bigquery_test.go index c283b623ed24..197a3b050bc3 100644 --- a/internal/sources/bigquery/bigquery_test.go +++ b/internal/sources/bigquery/bigquery_test.go @@ -90,6 +90,26 @@ func TestParseFromYamlBigQuery(t *testing.T) { }, }, }, + { + desc: "with service account impersonation example", + in: ` + sources: + my-instance: + kind: bigquery + project: my-project + location: us + impersonateServiceAccount: service-account@my-project.iam.gserviceaccount.com + `, + want: server.SourceConfigs{ + "my-instance": bigquery.Config{ + Name: "my-instance", + Kind: bigquery.SourceKind, + Project: "my-project", + Location: "us", + ImpersonateServiceAccount: "service-account@my-project.iam.gserviceaccount.com", + }, + }, + }, } for _, tc := range tcs { t.Run(tc.desc, func(t *testing.T) { From baddecc336f75a502b2b831fba1bdf94cdab6fa1 Mon Sep 17 00:00:00 2001 From: Colin Pistell Date: Tue, 28 Oct 2025 17:23:33 -0600 Subject: [PATCH 2/2] addressing PR comments --- internal/sources/bigquery/bigquery.go | 54 ++++++++++++++++----------- 1 file changed, 33 insertions(+), 21 deletions(-) diff --git a/internal/sources/bigquery/bigquery.go b/internal/sources/bigquery/bigquery.go index 491bf4422db3..d9729e12d0bb 100644 --- a/internal/sources/bigquery/bigquery.go +++ b/internal/sources/bigquery/bigquery.go @@ -96,6 +96,10 @@ func (r Config) Initialize(ctx context.Context, tracer trace.Tracer) (sources.So return nil, fmt.Errorf("writeMode 'protected' cannot be used with useClientOAuth 'true'") } + if r.UseClientOAuth && r.ImpersonateServiceAccount != "" { + return nil, fmt.Errorf("useClientOAuth cannot be used with impersonateServiceAccount") + } + var client *bigqueryapi.Client var restService *bigqueryrestapi.Service var tokenSource oauth2.TokenSource @@ -149,17 +153,18 @@ func (r Config) Initialize(ctx context.Context, tracer trace.Tracer) (sources.So } s := &Source{ - Name: r.Name, - Kind: SourceKind, - Project: r.Project, - Location: r.Location, - Client: client, - RestService: restService, - TokenSource: tokenSource, - MaxQueryResultRows: 50, - WriteMode: r.WriteMode, - AllowedDatasets: allowedDatasets, - UseClientOAuth: r.UseClientOAuth, + Name: r.Name, + Kind: SourceKind, + Project: r.Project, + Location: r.Location, + Client: client, + RestService: restService, + TokenSource: tokenSource, + MaxQueryResultRows: 50, + WriteMode: r.WriteMode, + AllowedDatasets: allowedDatasets, + UseClientOAuth: r.UseClientOAuth, + ClientCreator: clientCreator, ImpersonateServiceAccount: r.ImpersonateServiceAccount, } s.SessionProvider = s.newBigQuerySessionProvider() @@ -329,6 +334,17 @@ func (s *Source) BigQueryTokenSource() oauth2.TokenSource { } func (s *Source) BigQueryTokenSourceWithScope(ctx context.Context, scope string) (oauth2.TokenSource, error) { + if s.ImpersonateServiceAccount != "" { + // Create impersonated credentials token source with the requested scope + ts, err := impersonate.CredentialsTokenSource(ctx, impersonate.CredentialsConfig{ + TargetPrincipal: s.ImpersonateServiceAccount, + Scopes: []string{scope}, + }) + if err != nil { + return nil, fmt.Errorf("failed to create impersonated credentials for %q with scope %q: %w", s.ImpersonateServiceAccount, scope, err) + } + return ts, nil + } return google.DefaultTokenSource(ctx, scope) } @@ -398,11 +414,6 @@ func initBigQueryConnection( ctx, span := sources.InitConnectionSpan(ctx, tracer, SourceKind, name) defer span.End() - cred, err := google.FindDefaultCredentials(ctx, "https://www.googleapis.com/auth/cloud-platform") - if err != nil { - return nil, nil, nil, fmt.Errorf("failed to find default Google Cloud credentials with scope %q: %w", bigqueryapi.Scope, err) - } - userAgent, err := util.UserAgentFromContext(ctx) if err != nil { return nil, nil, nil, err @@ -412,18 +423,19 @@ func initBigQueryConnection( var opts []option.ClientOption if impersonateServiceAccount != "" { - // Create impersonated credentials token source - ts, err := impersonate.CredentialsTokenSource(ctx, impersonate.CredentialsConfig{ + // Create impersonated credentials token source with cloud-platform scope + // This broader scope is needed for tools like conversational analytics + cloudPlatformTokenSource, err := impersonate.CredentialsTokenSource(ctx, impersonate.CredentialsConfig{ TargetPrincipal: impersonateServiceAccount, - Scopes: []string{bigqueryapi.Scope}, + Scopes: []string{"https://www.googleapis.com/auth/cloud-platform"}, }) if err != nil { return nil, nil, nil, fmt.Errorf("failed to create impersonated credentials for %q: %w", impersonateServiceAccount, err) } - tokenSource = ts + tokenSource = cloudPlatformTokenSource opts = []option.ClientOption{ option.WithUserAgent(userAgent), - option.WithTokenSource(ts), + option.WithTokenSource(cloudPlatformTokenSource), } } else { // Use default credentials