Skip to content

feat(source/bigquery): add service account impersonation support for bigquery #1641

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
drstrangelooker merged 7 commits into from
Oct 31, 2025
Merged

feat(source/bigquery): add service account impersonation support for bigquery #1641

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
7 commits
Select commit Hold shift + click to select a range
287218e
feat(source/bigquery): add service account impersonation support for …
jcpistell-gf Oct 3, 2025
0cda0dc
merge main into feature branch
jcpistell-gf Oct 28, 2025
baddecc
addressing PR comments
jcpistell-gf Oct 28, 2025
efdd332
Merge branch 'main' into google/feat/bigquery_sa_impersonation
jcpistell-gf Oct 29, 2025
9a1d4b9
Merge branch 'main' into google/feat/bigquery_sa_impersonation
jcpistell-gf Oct 29, 2025
cef81bc
Merge branch 'main' into google/feat/bigquery_sa_impersonation
jcpistell-gf Oct 30, 2025
f389a6b
Merge branch 'main' into google/feat/bigquery_sa_impersonation
jcpistell-gf Oct 31, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
17 changes: 10 additions & 7 deletions docs/en/resources/sources/bigquery.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand All @@ -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) |
119 changes: 86 additions & 33 deletions internal/sources/bigquery/bigquery.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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"
)

Expand Down Expand Up @@ -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 {
Expand All @@ -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)
Comment thread
jcpistell-gf marked this conversation as resolved.
if err != nil {
return nil, fmt.Errorf("error creating client from ADC: %w", err)
}
Expand Down Expand Up @@ -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
Expand All @@ -156,6 +159,7 @@ type Source struct {
ClientCreator BigqueryClientCreator
AllowedDatasets map[string]struct{}
UseClientOAuth bool
ImpersonateServiceAccount string
makeDataplexCatalogClient func() (*dataplexapi.CatalogClient, DataplexClientCreator, error)
}

Expand Down Expand Up @@ -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
Expand All @@ -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
Comment thread
jcpistell-gf marked this conversation as resolved.
}

// initBigQueryConnectionWithOAuthToken initialize a BigQuery client with an
Expand Down Expand Up @@ -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
Expand All @@ -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
Expand All @@ -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)
}
Expand Down
20 changes: 20 additions & 0 deletions internal/sources/bigquery/bigquery_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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) {
Expand Down