This is an automated email from the ASF dual-hosted git repository.
claudevdm pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/beam.git
The following commit(s) were added to refs/heads/master by this push:
new 40df757ea39 Support attributing BigQuery quota and billing to a custom
GCP project (Python, Java, Go) (#39203)
40df757ea39 is described below
commit 40df757ea39be36ac71beebf849a5a7997ae0ff9
Author: Shahar Epstein <[email protected]>
AuthorDate: Mon Oct 19 06:09:18 2026 +0300
Support attributing BigQuery quota and billing to a custom GCP project
(Python, Java, Go) (#39203)
* [Python] Support a custom quota project for BigQuery IO
Add a quota_project_id parameter to ReadFromBigQuery and a
--quota_project_id pipeline option (GoogleCloudOptions) that attribute
BigQuery API quota and billing to a project other than the one the data
resides in.
The quota project is sent as the x-goog-user-project header on the
apitools BigQuery client, and through client_options on the
google-cloud-bigquery Client (streaming inserts) and the BigQuery
Storage Read client. The shared pipeline credentials are left untouched
instead of being copied with with_quota_project(), which for some
credential types drops the cached token. The header is only added to
requests to the BigQuery API hosts: the credentials refresh their token
through the same Http, and the token endpoint, metadata server and IAM
Credentials API must not receive it.
It is honored by ReadFromBigQuery (EXPORT and DIRECT_READ, including
temporary dataset cleanup), ReadAllFromBigQuery, WriteToBigQuery with
FILE_LOADS or STREAMING_INSERTS, and ReadBigQueryChangeHistory. The
value is resolved once at expand() time and passed to each DoFn so it
reaches separately deserialized SDF copies. STORAGE_WRITE_API is
Java-backed and not covered yet.
Generated-By: Claude Opus 5.5 <[email protected]>
* [Java] Add bigQueryQuotaProjectId pipeline option to BigQueryIO
Add BigQueryOptions.bigQueryQuotaProjectId to attribute BigQuery API
quota and billing to a project other than the one the data resides in.
For the REST client, the x-goog-user-project header is set by an
execute interceptor, so it is re-applied on every attempt, including
after HttpCredentialsAdapter re-initializes headers on a 401 refresh.
The Storage Read and Write clients set it through setQuotaProjectId.
Generated-By: Claude Opus 5.5 <[email protected]>
* [Go] Add WithQuotaProject option to bigqueryio
Add a WithQuotaProject query option that attributes BigQuery API quota
and billing to a project other than the one the data resides in, via
option.WithQuotaProject on the BigQuery client.
Read now accepts query options too. Only WithQuotaProject applies to it;
UseStandardSQL and WithQueryParameters panic, because Read generates a
legacy SQL statement without parameters.
Generated-By: Claude Opus 5.5 <[email protected]>
* Add CHANGES.md entry for BigQuery quota project support
Generated-By: Claude Opus 5.5 <[email protected]>
---
CHANGES.md | 2 +-
sdks/go/pkg/beam/io/bigqueryio/bigquery.go | 54 +++-
sdks/go/pkg/beam/io/bigqueryio/bigquery_test.go | 125 ++++++++
.../beam/sdk/io/gcp/bigquery/BigQueryOptions.java | 10 +
.../sdk/io/gcp/bigquery/BigQueryServicesImpl.java | 23 +-
.../io/gcp/bigquery/BigQueryServicesImplTest.java | 54 ++++
sdks/python/apache_beam/io/gcp/bigquery.py | 121 +++++--
.../apache_beam/io/gcp/bigquery_change_history.py | 42 ++-
.../io/gcp/bigquery_change_history_test.py | 75 +++++
.../apache_beam/io/gcp/bigquery_file_loads.py | 46 ++-
.../apache_beam/io/gcp/bigquery_file_loads_test.py | 47 +++
.../apache_beam/io/gcp/bigquery_read_internal.py | 11 +-
sdks/python/apache_beam/io/gcp/bigquery_test.py | 355 +++++++++++++++++++++
sdks/python/apache_beam/io/gcp/bigquery_tools.py | 162 +++++++++-
.../apache_beam/io/gcp/bigquery_tools_test.py | 232 +++++++++++++-
.../python/apache_beam/options/pipeline_options.py | 11 +
16 files changed, 1282 insertions(+), 88 deletions(-)
diff --git a/CHANGES.md b/CHANGES.md
index d9ca6f0b83d..c10b7262209 100644
--- a/CHANGES.md
+++ b/CHANGES.md
@@ -66,7 +66,7 @@
* Support for X source added (Java/Python)
([#X](https://github.com/apache/beam/issues/X)).
* (Java) FirestoreIO: added
`RpcQosOptions.Builder.withRampupThrottlingDisabled()`, mirroring the existing
DatastoreIO option, to opt out of the 500/50/5 write ramp-up
([#27170](https://github.com/apache/beam/issues/27170)).
-
+* Added support for attributing BigQuery API quota and billing to a specific
GCP project (quota project): `quota_project_id` parameter in `ReadFromBigQuery`
or `--quota_project_id` pipeline option, honored by the Python BigQuery IO
transforms (`WriteToBigQuery` with `FILE_LOADS` or `STREAMING_INSERTS`; not yet
`STORAGE_WRITE_API`, which is Java-backed), `--bigQueryQuotaProjectId` pipeline
option (Java), and `bigqueryio.WithQuotaProject` read/query option (Go)
([#37431](https://github.com [...]
## New Features / Improvements
diff --git a/sdks/go/pkg/beam/io/bigqueryio/bigquery.go
b/sdks/go/pkg/beam/io/bigqueryio/bigquery.go
index bbe951969b9..98818ccf1af 100644
--- a/sdks/go/pkg/beam/io/bigqueryio/bigquery.go
+++ b/sdks/go/pkg/beam/io/bigqueryio/bigquery.go
@@ -34,6 +34,7 @@ import (
bq "google.golang.org/api/bigquery/v2"
"google.golang.org/api/googleapi"
"google.golang.org/api/iterator"
+ "google.golang.org/api/option"
)
// writeSizeLimit is the maximum number of rows allowed by BQ in a write.
@@ -88,14 +89,22 @@ func NewQualifiedTableName(s string) (QualifiedTableName,
error) {
// Read reads all rows from the given table. The table must have a schema
// compatible with the given type, t, and Read returns a PCollection<t>. If the
// table has more rows than t, then Read is implicitly a projection.
-func Read(s beam.Scope, project, table string, t reflect.Type)
beam.PCollection {
+// Of the options, only WithQuotaProject applies to Read; UseStandardSQL and
+// WithQueryParameters panic, because the generated statement is legacy SQL and
+// takes no parameters.
+func Read(s beam.Scope, project, table string, t reflect.Type, options
...func(*QueryOptions) error) beam.PCollection {
mustParseTable(table)
s = s.Scope("bigquery.Read")
+ qo := applyQueryOptions(options)
+ if qo.UseStandardSQL || len(qo.parameters) > 0 {
+ panic("bigqueryio.Read: UseStandardSQL and WithQueryParameters
are only supported by Query")
+ }
+
stmt := constructSelectStatement(t, bigQueryTag, table)
- return query(s, project, stmt, t)
+ return queryWithOptions(s, project, stmt, t, qo)
}
func constructSelectStatement(t reflect.Type, tagKey string, table string)
string {
@@ -114,6 +123,9 @@ func constructSelectStatement(t reflect.Type, tagKey
string, table string) strin
type QueryOptions struct {
// UseStandardSQL enables BigQuery's Standard SQL dialect when
executing a query.
UseStandardSQL bool
+ // QuotaProject is the GCP project ID used for quota and billing
attribution
+ // of the BigQuery API calls, if different from the project the data
resides in.
+ QuotaProject string
// Parameters are the query parameters for parameterized queries.
// In the current implementation, user-defines types are not supported
in Value field.
// Use *bigquery.QueryParameterValue to build STRUCT/ARRAY parameters
@@ -129,6 +141,17 @@ func UseStandardSQL() func(qo *QueryOptions) error {
}
}
+// WithQuotaProject sets the GCP project ID used for quota and billing
+// attribution of the BigQuery API calls, if different from the project the
+// data resides in. The credentials used must have the
+// serviceusage.services.use permission on that project.
+func WithQuotaProject(project string) func(qo *QueryOptions) error {
+ return func(qo *QueryOptions) error {
+ qo.QuotaProject = project
+ return nil
+ }
+}
+
// WithQueryParameters sets the query parameters for parameterized queries.
// See QueryOptions.parameters for the list of supported Value types.
func WithQueryParameters(params ...bigquery.QueryParameter) func(qo
*QueryOptions) error {
@@ -138,6 +161,15 @@ func WithQueryParameters(params
...bigquery.QueryParameter) func(qo *QueryOption
}
}
+// clientOptions returns the BigQuery client options implied by qo.
+func clientOptions(qo QueryOptions) []option.ClientOption {
+ var opts []option.ClientOption
+ if qo.QuotaProject != "" {
+ opts = append(opts, option.WithQuotaProject(qo.QuotaProject))
+ }
+ return opts
+}
+
// Query executes a query. The output must have a schema compatible with the
given
// type, t. It returns a PCollection<t>.
func Query(s beam.Scope, project, q string, t reflect.Type, options
...func(*QueryOptions) error) beam.PCollection {
@@ -145,15 +177,22 @@ func Query(s beam.Scope, project, q string, t
reflect.Type, options ...func(*Que
return query(s, project, q, t, options...)
}
-func query(s beam.Scope, project, query string, t reflect.Type, options
...func(*QueryOptions) error) beam.PCollection {
- mustInferSchema(t)
-
+func applyQueryOptions(options []func(*QueryOptions) error) QueryOptions {
queryOptions := QueryOptions{}
for _, opt := range options {
if err := opt(&queryOptions); err != nil {
panic(err)
}
}
+ return queryOptions
+}
+
+func query(s beam.Scope, project, query string, t reflect.Type, options
...func(*QueryOptions) error) beam.PCollection {
+ return queryWithOptions(s, project, query, t,
applyQueryOptions(options))
+}
+
+func queryWithOptions(s beam.Scope, project, query string, t reflect.Type,
queryOptions QueryOptions) beam.PCollection {
+ mustInferSchema(t)
imp := beam.Impulse(s)
queryParameters, err := encodeQueryParameters(queryOptions.parameters)
@@ -181,8 +220,11 @@ type queryFn struct {
Options QueryOptions `json:"options"`
}
+// newClient creates the BigQuery client used by queryFn. Tests replace it.
+var newClient = bigquery.NewClient
+
func (f *queryFn) ProcessElement(ctx context.Context, _ []byte, emit
func(beam.X)) error {
- client, err := bigquery.NewClient(ctx, f.Project)
+ client, err := newClient(ctx, f.Project, clientOptions(f.Options)...)
if err != nil {
return err
}
diff --git a/sdks/go/pkg/beam/io/bigqueryio/bigquery_test.go
b/sdks/go/pkg/beam/io/bigqueryio/bigquery_test.go
index 9f83eeb847f..4cf0c46b9f6 100644
--- a/sdks/go/pkg/beam/io/bigqueryio/bigquery_test.go
+++ b/sdks/go/pkg/beam/io/bigqueryio/bigquery_test.go
@@ -16,11 +16,16 @@
package bigqueryio
import (
+ "context"
+ "errors"
+ "fmt"
"reflect"
+ "strings"
"testing"
"cloud.google.com/go/bigquery"
"github.com/apache/beam/sdks/v2/go/pkg/beam"
+ "google.golang.org/api/option"
)
func TestNewQualifiedTableName(t *testing.T) {
@@ -149,6 +154,88 @@ func Test_mustInferSchema(t *testing.T) {
}
}
+func TestWithQuotaProject(t *testing.T) {
+ qo := QueryOptions{}
+ if err := WithQuotaProject("quota-project")(&qo); err != nil {
+ t.Fatalf("WithQuotaProject() returned err: %v", err)
+ }
+ if got, want := qo.QuotaProject, "quota-project"; got != want {
+ t.Errorf("qo.QuotaProject = %q, want %q", got, want)
+ }
+}
+
+func TestClientOptions(t *testing.T) {
+ if got := clientOptions(QueryOptions{}); len(got) != 0 {
+ t.Errorf("clientOptions(no quota) = %v, want none", got)
+ }
+ got := clientOptions(QueryOptions{QuotaProject: "quota-project"})
+ if want :=
[]option.ClientOption{option.WithQuotaProject("quota-project")};
!reflect.DeepEqual(got, want) {
+ t.Errorf("clientOptions(with quota) = %v, want %v", got, want)
+ }
+}
+
+func TestQueryFn_passesQuotaProjectToClient(t *testing.T) {
+ var gotProject string
+ var gotOpts []option.ClientOption
+ errStop := errors.New("stop")
+ defer func(orig func(context.Context, string, ...option.ClientOption)
(*bigquery.Client, error)) {
+ newClient = orig
+ }(newClient)
+ newClient = func(_ context.Context, project string, opts
...option.ClientOption) (*bigquery.Client, error) {
+ gotProject, gotOpts = project, opts
+ return nil, errStop
+ }
+
+ fn := &queryFn{Project: "project", Options: QueryOptions{QuotaProject:
"quota-project"}}
+ if err := fn.ProcessElement(context.Background(), nil, nil);
!errors.Is(err, errStop) {
+ t.Fatalf("ProcessElement() err = %v, want %v", err, errStop)
+ }
+ if gotProject != "project" {
+ t.Errorf("client project = %q, want %q", gotProject, "project")
+ }
+ if want :=
[]option.ClientOption{option.WithQuotaProject("quota-project")};
!reflect.DeepEqual(gotOpts, want) {
+ t.Errorf("client options = %v, want %v", gotOpts, want)
+ }
+}
+
+// queryFnOf returns the queryFn in pipeline p.
+func queryFnOf(t *testing.T, p *beam.Pipeline) *queryFn {
+ t.Helper()
+ edges, _, err := p.Build()
+ if err != nil {
+ t.Fatalf("Build() failed: %v", err)
+ }
+ for _, e := range edges {
+ if e.DoFn != nil {
+ if fn, ok := e.DoFn.Recv.(*queryFn); ok {
+ return fn
+ }
+ }
+ }
+ t.Fatal("no queryFn in pipeline")
+ return nil
+}
+
+func TestReadAndQuery_passQuotaProjectToQueryFn(t *testing.T) {
+ type quotaSchema struct {
+ Name bigquery.NullString `bigquery:"name"`
+ }
+ beam.RegisterType(reflect.TypeOf((*quotaSchema)(nil)).Elem())
+ typ := reflect.TypeOf(quotaSchema{})
+
+ p, s := beam.NewPipelineWithRoot()
+ Read(s, "project", "project:dataset.table", typ,
WithQuotaProject("quota-project"))
+ if got := queryFnOf(t, p).Options.QuotaProject; got != "quota-project" {
+ t.Errorf("Read: queryFn.Options.QuotaProject = %q, want %q",
got, "quota-project")
+ }
+
+ p, s = beam.NewPipelineWithRoot()
+ Query(s, "project", "SELECT 1", typ, WithQuotaProject("quota-project"))
+ if got := queryFnOf(t, p).Options.QuotaProject; got != "quota-project" {
+ t.Errorf("Query: queryFn.Options.QuotaProject = %q, want %q",
got, "quota-project")
+ }
+}
+
func TestWithQueryParameters(t *testing.T) {
t.Run("WithQueryParameters sets parameters correctly", func(t
*testing.T) {
params := []bigquery.QueryParameter{
@@ -183,3 +270,41 @@ func TestWithQueryParameters(t *testing.T) {
}
})
}
+
+func TestRead_queryOnlyOptions(t *testing.T) {
+ type readSchema struct {
+ Name bigquery.NullString `bigquery:"name"`
+ }
+ beam.RegisterType(reflect.TypeOf((*readSchema)(nil)).Elem())
+ typ := reflect.TypeOf(readSchema{})
+
+ mustPanic := func(t *testing.T, opt func(*QueryOptions) error) {
+ t.Helper()
+ defer func() {
+ r := recover()
+ if r == nil {
+ t.Errorf("Read() with query-only option does
not panic")
+ } else if !strings.Contains(fmt.Sprint(r), "only
supported by Query") {
+ t.Errorf("Read() panicked with unexpected
value: %v", r)
+ }
+ }()
+ _, s := beam.NewPipelineWithRoot()
+ Read(s, "project", "project:dataset.table", typ, opt)
+ }
+
+ t.Run("UseStandardSQL panics", func(t *testing.T) {
+ mustPanic(t, UseStandardSQL())
+ })
+ t.Run("WithQueryParameters panics", func(t *testing.T) {
+ mustPanic(t, WithQueryParameters(bigquery.QueryParameter{Name:
"p", Value: 1}))
+ })
+ t.Run("WithQuotaProject does not panic", func(t *testing.T) {
+ defer func() {
+ if r := recover(); r != nil {
+ t.Errorf("Read() with WithQuotaProject
panicked: %v", r)
+ }
+ }()
+ _, s := beam.NewPipelineWithRoot()
+ Read(s, "project", "project:dataset.table", typ,
WithQuotaProject("quota-project"))
+ })
+}
diff --git
a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryOptions.java
b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryOptions.java
index 1826f4a1fd2..64382d64bb5 100644
---
a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryOptions.java
+++
b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryOptions.java
@@ -197,6 +197,16 @@ public interface BigQueryOptions
void setBigQueryProject(String value);
+ @Description(
+ "GCP project ID used for quota and billing attribution of BigQuery API
requests "
+ + "(sets the X-Goog-User-Project header), if different from the
project the data "
+ + "resides in. If unspecified, the project associated with the
credentials is used. "
+ + "The credentials used must have the serviceusage.services.use
permission on this "
+ + "project.")
+ String getBigQueryQuotaProjectId();
+
+ void setBigQueryQuotaProjectId(String value);
+
@Description("Maximum (best effort) size of a single append to the storage
API.")
@Default.Integer(2 * 1024 * 1024)
Integer getStorageApiAppendThresholdBytes();
diff --git
a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryServicesImpl.java
b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryServicesImpl.java
index 7a924c6aa02..a921be3700e 100644
---
a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryServicesImpl.java
+++
b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryServicesImpl.java
@@ -1800,7 +1800,8 @@ public class BigQueryServicesImpl implements
BigQueryServices {
}
/** Returns a BigQuery client builder using the specified {@link
BigQueryOptions}. */
- private static Bigquery.Builder newBigQueryClient(BigQueryOptions options) {
+ @VisibleForTesting
+ static Bigquery.Builder newBigQueryClient(BigQueryOptions options) {
// Do not log 404. It clutters the output and is possibly even required by
the
// caller.
RetryHttpRequestInitializer httpRequestInitializer =
@@ -1818,6 +1819,18 @@ public class BigQueryServicesImpl implements
BigQueryServices {
initBuilder.add(new
LatencyRecordingHttpRequestInitializer(API_METRIC_LABEL));
initBuilder.add(httpRequestInitializer);
+ // Set the quota project as a request header instead of deriving
credentials: a derived
+ // credential inherits the shared credential's cached request metadata,
which lacks
+ // x-goog-user-project until the next token refresh. Applied by an execute
interceptor so it
+ // is re-applied on every attempt; HttpCredentialsAdapter re-initializes
the request headers
+ // from the credential after a 401 refresh, which would otherwise override
it.
+ @Nullable String quotaProjectId = options.getBigQueryQuotaProjectId();
+ if (!Strings.isNullOrEmpty(quotaProjectId)) {
+ initBuilder.add(
+ request ->
+ request.setInterceptor(
+ r -> r.getHeaders().set("x-goog-user-project",
quotaProjectId)));
+ }
HttpRequestInitializer chainInitializer =
new ChainingHttpRequestInitializer(
Iterables.toArray(initBuilder.build(),
HttpRequestInitializer.class));
@@ -1848,6 +1861,10 @@ public class BigQueryServicesImpl implements
BigQueryServices {
if (!Strings.isNullOrEmpty(endpoint)) {
builder.setEndpoint(trimSchemaIfNecessary(endpoint));
}
+ @Nullable String quotaProjectId = options.getBigQueryQuotaProjectId();
+ if (!Strings.isNullOrEmpty(quotaProjectId)) {
+ builder.setQuotaProjectId(quotaProjectId);
+ }
return BigQueryWriteClient.create(
builder
.setCredentialsProvider(() ->
options.as(GcpOptions.class).getGcpCredential())
@@ -1973,6 +1990,10 @@ public class BigQueryServicesImpl implements
BigQueryServices {
if (!Strings.isNullOrEmpty(endpoint)) {
settingsBuilder.setEndpoint(trimSchemaIfNecessary(endpoint));
}
+ @Nullable String quotaProjectId = options.getBigQueryQuotaProjectId();
+ if (!Strings.isNullOrEmpty(quotaProjectId)) {
+ settingsBuilder.setQuotaProjectId(quotaProjectId);
+ }
UnaryCallSettings.Builder<CreateReadSessionRequest, ReadSession>
createReadSessionSettings =
settingsBuilder.getStubSettingsBuilder().createReadSessionSettings();
diff --git
a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryServicesImplTest.java
b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryServicesImplTest.java
index 90b2e1fb872..642d35b4390 100644
---
a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryServicesImplTest.java
+++
b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryServicesImplTest.java
@@ -41,6 +41,8 @@ import com.google.api.client.googleapis.json.GoogleJsonError;
import com.google.api.client.googleapis.json.GoogleJsonError.ErrorInfo;
import com.google.api.client.googleapis.json.GoogleJsonErrorContainer;
import com.google.api.client.googleapis.json.GoogleJsonResponseException;
+import com.google.api.client.http.GenericUrl;
+import com.google.api.client.http.HttpRequest;
import com.google.api.client.http.HttpResponseException;
import com.google.api.client.http.LowLevelHttpResponse;
import com.google.api.client.json.GenericJson;
@@ -2311,6 +2313,58 @@ public class BigQueryServicesImplTest {
assertEquals(123456, (long)
container.getCounter(metricName).getCumulative());
}
+ @Test
+ public void testQuotaProjectIdOverrides() throws IOException {
+ BigQueryOptions options =
PipelineOptionsFactory.create().as(BigQueryOptions.class);
+ options.setBigQueryQuotaProjectId("my-quota-project");
+
+ assertEquals(
+ "my-quota-project",
+ new BigQueryServicesImpl.StorageClientImpl(options)
+ .getClient()
+ .getSettings()
+ .getQuotaProjectId());
+ assertEquals(
+ "my-quota-project",
+ new BigQueryServicesImpl.WriteStreamServiceImpl(options)
+ .getClient()
+ .getSettings()
+ .getQuotaProjectId());
+ }
+
+ @Test
+ public void testQuotaProjectIdSetsRequestHeader() throws IOException {
+ BigQueryOptions options =
PipelineOptionsFactory.create().as(BigQueryOptions.class);
+ options.setGcpCredential(null);
+ GenericUrl url = new
GenericUrl("https://bigquery.googleapis.com/bigquery/v2/projects";);
+
+ HttpRequest request =
+ BigQueryServicesImpl.newBigQueryClient(options)
+ .build()
+ .getRequestFactory()
+ .buildGetRequest(url);
+ request.getInterceptor().intercept(request);
+
assertNull(request.getHeaders().getFirstHeaderStringValue("x-goog-user-project"));
+
+ options.setBigQueryQuotaProjectId("my-quota-project");
+ request =
+ BigQueryServicesImpl.newBigQueryClient(options)
+ .build()
+ .getRequestFactory()
+ .buildGetRequest(url);
+ // The interceptor runs before every attempt, including retries.
+ request.getInterceptor().intercept(request);
+ assertEquals(
+ "my-quota-project",
request.getHeaders().getFirstHeaderStringValue("x-goog-user-project"));
+
+ // Simulate HttpCredentialsAdapter re-initializing headers from a
credential that carries
+ // its own quota project after a 401 refresh: the configured value must
win on the retry.
+ request.getHeaders().set("x-goog-user-project", "credential-project");
+ request.getInterceptor().intercept(request);
+ assertEquals(
+ "my-quota-project",
request.getHeaders().getFirstHeaderStringValue("x-goog-user-project"));
+ }
+
@Test
public void testEndpointOverrides() throws IOException {
BigQueryOptions options =
PipelineOptionsFactory.create().as(BigQueryOptions.class);
diff --git a/sdks/python/apache_beam/io/gcp/bigquery.py
b/sdks/python/apache_beam/io/gcp/bigquery.py
index 2a52b31edf7..f3d64bfcfa6 100644
--- a/sdks/python/apache_beam/io/gcp/bigquery.py
+++ b/sdks/python/apache_beam/io/gcp/bigquery.py
@@ -720,7 +720,8 @@ class _CustomBigQuerySource(BoundedSource):
step_name=None,
unique_id=None,
temp_dataset=None,
- query_priority=BigQueryQueryPriority.BATCH):
+ query_priority=BigQueryQueryPriority.BATCH,
+ quota_project_id=None):
if table is not None and query is not None:
raise ValueError(
'Both a BigQuery table and a query were specified.'
@@ -754,6 +755,7 @@ class _CustomBigQuerySource(BoundedSource):
self.use_json_exports = use_json_exports
self.temp_dataset = temp_dataset
self.query_priority = query_priority
+ self.quota_project_id = quota_project_id
self._job_name = job_name or 'BQ_EXPORT_JOB'
self._step_name = step_name
self._source_uuid = unique_id
@@ -773,12 +775,15 @@ class _CustomBigQuerySource(BoundedSource):
'use_legacy_sql': self.use_legacy_sql,
'bigquery_job_labels': json.dumps(self.bigquery_job_labels),
'export_file_format': export_format,
+ 'quota_project_id': DisplayDataItem(
+ self._get_quota_project_id(),
+ label='BigQuery Quota Project').drop_if_none(),
'launchesBigQueryJobs': DisplayDataItem(
True, label="This Dataflow job launches bigquery jobs."),
}
def estimate_size(self):
- bq = bigquery_tools.BigQueryWrapper.from_pipeline_options(self.options)
+ bq = self._create_bq_wrapper()
if self.table_reference is not None:
table_ref = self.table_reference
if (isinstance(self.table_reference, vp.ValueProvider) and
@@ -844,6 +849,23 @@ class _CustomBigQuerySource(BoundedSource):
project = self.project
return project
+ def _get_quota_project_id(self):
+ """Returns the explicit quota project ID, else the pipeline option."""
+ quota_project_id = bigquery_tools._quota_project_id_from_options(
+ self.options)
+ return self.quota_project_id or quota_project_id
+
+ def _create_bq_wrapper(self, **kwargs):
+ """Creates a BigQueryWrapper for the API calls made by this source.
+
+ Every BigQuery client this source uses is built here, so that quota
+ attribution cannot be forgotten at an individual call site.
+ """
+ return bigquery_tools.BigQueryWrapper(
+ pipeline_options=self.options,
+ quota_project_id=self.quota_project_id,
+ **kwargs)
+
def _create_source(self, path, coder):
if not self.use_json_exports:
return create_avro_source(path, validate=self.validate)
@@ -858,10 +880,9 @@ class _CustomBigQuerySource(BoundedSource):
def split(self, desired_bundle_size, start_position=None,
stop_position=None):
if self.export_result is None:
- bq = bigquery_tools.BigQueryWrapper(
+ bq = self._create_bq_wrapper(
temp_dataset_id=(
- self.temp_dataset.datasetId if self.temp_dataset else None),
- client=bigquery_tools.BigQueryWrapper._bigquery_client(self.options))
+ self.temp_dataset.datasetId if self.temp_dataset else None))
if self.query is not None:
self._setup_temporary_dataset(bq)
@@ -1051,7 +1072,8 @@ class _CustomBigQueryStorageSource(BoundedSource):
temp_dataset: Optional[DatasetReference] = None,
temp_table: Optional[TableReference] = None,
use_native_datetime: Optional[bool] = False,
- timeout: Optional[float] = None):
+ timeout: Optional[float] = None,
+ quota_project_id: Optional[str] = None):
if table is not None and query is not None:
raise ValueError(
@@ -1090,6 +1112,7 @@ class _CustomBigQueryStorageSource(BoundedSource):
self._job_name = job_name or 'BQ_DIRECT_READ_JOB'
self._step_name = step_name
self._source_uuid = unique_id
+ self.quota_project_id = quota_project_id
def _get_project(self):
"""Returns the project that queries and exports will be billed to."""
@@ -1101,6 +1124,23 @@ class _CustomBigQueryStorageSource(BoundedSource):
return project
return self.project
+ def _get_quota_project_id(self):
+ """Returns the explicit quota project ID, else the pipeline option."""
+ quota_project_id = bigquery_tools._quota_project_id_from_options(
+ self.pipeline_options)
+ return self.quota_project_id or quota_project_id
+
+ def _create_bq_wrapper(self, **kwargs):
+ """Creates a BigQueryWrapper for the API calls made by this source.
+
+ Every BigQuery client this source uses is built here, so that quota
+ attribution cannot be forgotten at an individual call site.
+ """
+ return bigquery_tools.BigQueryWrapper(
+ pipeline_options=self.pipeline_options,
+ quota_project_id=self.quota_project_id,
+ **kwargs)
+
def _get_parent_project(self):
"""Returns the project that will be billed."""
if self.temp_table:
@@ -1188,14 +1228,16 @@ class _CustomBigQueryStorageSource(BoundedSource):
'use_native_datetime': self.use_native_datetime,
'selected_fields': str(self.selected_fields),
'row_restriction': str(self.row_restriction),
+ 'quota_project_id': DisplayDataItem(
+ self._get_quota_project_id(),
+ label='BigQuery Quota Project').drop_if_none(),
'launchesBigQueryJobs': DisplayDataItem(
True, label="This Dataflow job launches bigquery jobs."),
}
def estimate_size(self):
# Returns the pre-filtering size of the (temporary) table being read.
- bq = bigquery_tools.BigQueryWrapper.from_pipeline_options(
- self.pipeline_options)
+ bq = self._create_bq_wrapper()
if self.table_reference is not None:
table_ref = self.table_reference
if (isinstance(self.table_reference, vp.ValueProvider) and
@@ -1242,10 +1284,8 @@ class _CustomBigQueryStorageSource(BoundedSource):
def split(self, desired_bundle_size, start_position=None,
stop_position=None):
if self.split_result is None:
- bq = bigquery_tools.BigQueryWrapper(
- temp_table_ref=(self.temp_table if self.temp_table else None),
- client=bigquery_tools.BigQueryWrapper._bigquery_client(
- self.pipeline_options))
+ bq = self._create_bq_wrapper(
+ temp_table_ref=(self.temp_table if self.temp_table else None))
if self.query is not None:
self._setup_temporary_dataset(bq)
@@ -1278,7 +1318,8 @@ class _CustomBigQueryStorageSource(BoundedSource):
if self.row_restriction is not None:
requested_session.read_options.row_restriction = self.row_restriction
- storage_client = bq_storage.BigQueryReadClient()
+ storage_client = bigquery_tools.create_bq_storage_client(
+ self._get_quota_project_id())
stream_count = self._get_stream_count(bq, desired_bundle_size)
parent = 'projects/{}'.format(self.table_reference.projectId)
@@ -1302,8 +1343,10 @@ class _CustomBigQueryStorageSource(BoundedSource):
self.split_result = [
_CustomBigQueryStorageStreamSource(
- stream.name, self.use_native_datetime, self.timeout)
- for stream in read_session.streams
+ stream.name,
+ self.use_native_datetime,
+ self.timeout,
+ self._get_quota_project_id()) for stream in read_session.streams
]
for source in self.split_result:
@@ -1337,10 +1380,12 @@ class _CustomBigQueryStorageStreamSource(BoundedSource):
self,
read_stream_name: str,
use_native_datetime: Optional[bool] = True,
- timeout: Optional[float] = None):
+ timeout: Optional[float] = None,
+ quota_project_id: Optional[str] = None):
self.read_stream_name = read_stream_name
self.use_native_datetime = use_native_datetime
self.timeout = timeout
+ self.quota_project_id = quota_project_id
def display_data(self):
return {
@@ -1363,7 +1408,10 @@ class _CustomBigQueryStorageStreamSource(BoundedSource):
return SourceBundle(
weight=1.0,
source=_CustomBigQueryStorageStreamSource(
- self.read_stream_name, self.use_native_datetime),
+ self.read_stream_name,
+ self.use_native_datetime,
+ self.timeout,
+ self.quota_project_id),
start_position=None,
stop_position=None)
@@ -1399,7 +1447,8 @@ class _CustomBigQueryStorageStreamSource(BoundedSource):
def read_arrow(self):
- storage_client = bq_storage.BigQueryReadClient()
+ storage_client = bigquery_tools.create_bq_storage_client(
+ self.quota_project_id)
read_rows_kwargs = {'retry_delay_callback': self.retry_delay_callback}
if self.timeout is not None:
read_rows_kwargs['timeout'] = self.timeout
@@ -1418,7 +1467,8 @@ class _CustomBigQueryStorageStreamSource(BoundedSource):
yield py_row
def read_avro(self):
- storage_client = bq_storage.BigQueryReadClient()
+ storage_client = bigquery_tools.create_bq_storage_client(
+ self.quota_project_id)
read_rows_kwargs = {'retry_delay_callback': self.retry_delay_callback}
if self.timeout is not None:
read_rows_kwargs['timeout'] = self.timeout
@@ -1499,7 +1549,8 @@ class BigQueryWriteFn(DoFn):
with_batched_input=False,
ignore_unknown_columns=False,
max_retries=MAX_INSERT_RETRIES,
- max_insert_payload_size=MAX_INSERT_PAYLOAD_SIZE):
+ max_insert_payload_size=MAX_INSERT_PAYLOAD_SIZE,
+ quota_project_id=None):
"""Initialize a WriteToBigQuery transform.
Args:
@@ -1552,9 +1603,12 @@ class BigQueryWriteFn(DoFn):
backoffs (effectively retry forever).
max_insert_payload_size: The maximum byte size for a BigQuery legacy
streaming insert payload.
+ quota_project_id: Optional GCP project ID to use for quota and billing of
+ BigQuery API requests.
"""
self.schema = schema
self.test_client = test_client
+ self.quota_project_id = quota_project_id
self.create_disposition = create_disposition
self.write_disposition = write_disposition
if write_disposition in (BigQueryDisposition.WRITE_EMPTY,
@@ -1633,7 +1687,7 @@ class BigQueryWriteFn(DoFn):
if not self.bigquery_wrapper:
self.bigquery_wrapper = bigquery_tools.BigQueryWrapper(
- client=self.test_client)
+ client=self.test_client, quota_project_id=self.quota_project_id)
(
bigquery_tools.BigQueryWrapper.HISTOGRAM_METRIC_LOGGER.
@@ -1958,6 +2012,8 @@ class _StreamToBigQuery(PTransform):
yield (key, (value, insert_id))
def expand(self, input):
+ quota_project_id = bigquery_tools._quota_project_id_from_options(
+ input.pipeline.options)
bigquery_write_fn = BigQueryWriteFn(
schema=self.schema,
batch_size=self.batch_size,
@@ -1971,7 +2027,8 @@ class _StreamToBigQuery(PTransform):
ignore_unknown_columns=self.ignore_unknown_columns,
with_batched_input=self.with_auto_sharding,
max_retries=self._max_retries,
- max_insert_payload_size=self._max_insert_payload_size)
+ max_insert_payload_size=self._max_insert_payload_size,
+ quota_project_id=quota_project_id)
def _add_random_shard(element):
key = element[0]
@@ -3043,6 +3100,12 @@ class ReadFromBigQuery(PTransform):
PCollection with a schema and yielding Beam Rows via the option
`BEAM_ROW`. For more information on schemas, see
https://beam.apache.org/documentation/programming-guide/#what-is-a-schema)
+ quota_project_id (str): The GCP project ID to use for quota and billing
+ of BigQuery API requests issued by this transform, if different from
+ the project the data resides in. Falls back to the
+ ``--quota_project_id`` pipeline option if not set. The credentials
+ used must have the ``serviceusage.services.use`` permission on that
+ project.
query_output_schema: Required when output_type is 'BEAM_ROW' and a query
is specified. A BigQuery schema describing the query result columns,
since the schema cannot be auto-derived from an existing table when
@@ -3143,10 +3206,12 @@ class ReadFromBigQuery(PTransform):
'%s: table must be of type string'
'; got a callable instead' % self.__class__.__name__)
return output_pcollection | bigquery_schema_tools.convert_to_usertype(
- bigquery_tools.BigQueryWrapper().get_table(
- project_id=table_details.projectId,
- dataset_id=table_details.datasetId,
- table_id=table_details.tableId).schema,
+ bigquery_tools.BigQueryWrapper(
+ pipeline_options=output_pcollection.pipeline.options,
+ quota_project_id=self._kwargs.get('quota_project_id')).get_table(
+ project_id=table_details.projectId,
+ dataset_id=table_details.datasetId,
+ table_id=table_details.tableId).schema,
self._kwargs.get('selected_fields', None))
else:
raise ValueError(
@@ -3228,7 +3293,9 @@ class ReadFromBigQuery(PTransform):
bigquery_dataset_labels=self.bigquery_dataset_labels,
*self._args,
**self._kwargs))
- | _PassThroughThenCleanupTempDatasets(project_to_cleanup_pcoll))
+ | _PassThroughThenCleanupTempDatasets(
+ project_to_cleanup_pcoll,
+ quota_project_id=self._kwargs.get('quota_project_id')))
class ReadFromBigQueryRequest:
diff --git a/sdks/python/apache_beam/io/gcp/bigquery_change_history.py
b/sdks/python/apache_beam/io/gcp/bigquery_change_history.py
index 1e90c297477..9f836190d9c 100644
--- a/sdks/python/apache_beam/io/gcp/bigquery_change_history.py
+++ b/sdks/python/apache_beam/io/gcp/bigquery_change_history.py
@@ -416,7 +416,8 @@ class _PollChangeHistoryFn(beam.DoFn,
beam.transforms.core.RestrictionProvider):
start_time: Timestamp,
stop_time: Timestamp,
poll_interval: Duration,
- location: Optional[str] = None) -> None:
+ location: Optional[str] = None,
+ quota_project_id: Optional[str] = None) -> None:
self._table = table
self._project = project
self._change_function = change_function
@@ -425,9 +426,11 @@ class _PollChangeHistoryFn(beam.DoFn,
beam.transforms.core.RestrictionProvider):
self._stop_time = stop_time
self._poll_interval = poll_interval
self._location = location
+ self._quota_project_id = quota_project_id
def setup(self) -> None:
- self._bq_wrapper = bigquery_tools.BigQueryWrapper()
+ self._bq_wrapper = bigquery_tools.BigQueryWrapper(
+ quota_project_id=self._quota_project_id)
if self._location is None:
table_ref = bigquery_tools.parse_table_reference(
self._table, project=self._project)
@@ -570,7 +573,8 @@ class _ExecuteQueryFn(beam.DoFn):
change_type_column: str = 'change_type',
change_timestamp_column: str = 'change_timestamp',
columns: Optional[list[str]] = None,
- row_filter: Optional[str] = None) -> None:
+ row_filter: Optional[str] = None,
+ quota_project_id: Optional[str] = None) -> None:
self._table = table
self._project = project
self._change_function = change_function
@@ -580,9 +584,11 @@ class _ExecuteQueryFn(beam.DoFn):
self._change_timestamp_column = change_timestamp_column
self._columns = columns
self._row_filter = row_filter
+ self._quota_project_id = quota_project_id
def setup(self) -> None:
- self._bq_wrapper = bigquery_tools.BigQueryWrapper()
+ self._bq_wrapper = bigquery_tools.BigQueryWrapper(
+ quota_project_id=self._quota_project_id)
if self._location is None:
table_ref = bigquery_tools.parse_table_reference(
self._table, project=self._project)
@@ -711,11 +717,13 @@ class _ReadStorageStreamsSDF(beam.DoFn,
batch_arrow_read: bool = True,
change_timestamp_column: str = 'change_timestamp',
max_split_rounds: int = 1,
- emit_raw_batches: bool = False) -> None:
+ emit_raw_batches: bool = False,
+ quota_project_id: Optional[str] = None) -> None:
self._batch_arrow_read = batch_arrow_read
self._change_timestamp_column = change_timestamp_column
self._max_split_rounds = max_split_rounds
self._emit_raw_batches = emit_raw_batches
+ self._quota_project_id = quota_project_id
self._storage_client = None
def _ensure_client(self) -> None:
@@ -727,7 +735,8 @@ class _ReadStorageStreamsSDF(beam.DoFn,
"""
if self._storage_client is None:
_LOGGER.info('[Read] creating BigQueryReadClient')
- self._storage_client = bq_storage.BigQueryReadClient()
+ self._storage_client = bigquery_tools.create_bq_storage_client(
+ self._quota_project_id)
def setup(self) -> None:
self._ensure_client()
@@ -1072,9 +1081,13 @@ class _CleanupTempTablesFn(beam.DoFn):
STREAMS_READ = beam.transforms.userstate.CombiningValueStateSpec(
'streams_read', sum)
+ def __init__(self, quota_project_id: Optional[str] = None) -> None:
+ self._quota_project_id = quota_project_id
+
def setup(self) -> None:
_LOGGER.info('[Cleanup] setup: creating BigQueryWrapper')
- self._bq_wrapper = bigquery_tools.BigQueryWrapper()
+ self._bq_wrapper = bigquery_tools.BigQueryWrapper(
+ quota_project_id=self._quota_project_id)
def process(
self,
@@ -1247,6 +1260,9 @@ class ReadBigQueryChangeHistory(beam.PTransform):
'project must be specified either in ReadBigQueryChangeHistory '
'or in pipeline options (--project)')
+ quota_project_id = bigquery_tools._quota_project_id_from_options(
+ pbegin.pipeline.options)
+
start_time = Timestamp(self._start_time or time.time())
stop_time = (
Timestamp(self._stop_time)
@@ -1290,7 +1306,8 @@ class ReadBigQueryChangeHistory(beam.PTransform):
start_time=start_time,
stop_time=stop_time,
poll_interval=poll_interval,
- location=self._location)))
+ location=self._location,
+ quota_project_id=quota_project_id)))
# CommitQueryResults: Reshuffle commits _QueryResult (temp table ref)
# so that if the Read SDF retries, it re-reads the existing temp table
@@ -1311,7 +1328,8 @@ class ReadBigQueryChangeHistory(beam.PTransform):
change_type_column=self._change_type_column,
change_timestamp_column=self._change_timestamp_column,
columns=self._columns,
- row_filter=self._row_filter))
+ row_filter=self._row_filter,
+ quota_project_id=quota_project_id))
| 'CommitQueryResults' >> beam.Reshuffle())
emit_raw = self._reshuffle_decompress
@@ -1321,7 +1339,8 @@ class ReadBigQueryChangeHistory(beam.PTransform):
batch_arrow_read=self._batch_arrow_read,
change_timestamp_column=self._change_timestamp_column,
max_split_rounds=self._max_split_rounds,
- emit_raw_batches=emit_raw))
+ emit_raw_batches=emit_raw,
+ quota_project_id=quota_project_id))
if emit_raw:
read_sdf = read_sdf.with_output_types(tuple[bytes, bytes])
else:
@@ -1334,7 +1353,8 @@ class ReadBigQueryChangeHistory(beam.PTransform):
_ = (
read_outputs[_CLEANUP_TAG]
- | 'CleanupTempTables' >> beam.ParDo(_CleanupTempTablesFn()))
+ | 'CleanupTempTables' >> beam.ParDo(
+ _CleanupTempTablesFn(quota_project_id=quota_project_id)))
if emit_raw:
# Reshuffle raw Arrow batches for fan-out, then decompress and
diff --git a/sdks/python/apache_beam/io/gcp/bigquery_change_history_test.py
b/sdks/python/apache_beam/io/gcp/bigquery_change_history_test.py
index 11b467f26d4..e1b1eb69ae3 100644
--- a/sdks/python/apache_beam/io/gcp/bigquery_change_history_test.py
+++ b/sdks/python/apache_beam/io/gcp/bigquery_change_history_test.py
@@ -23,12 +23,18 @@ Tests:
- ReadBigQueryChangeHistory validation
"""
+import contextlib
import datetime
import unittest
+from unittest import mock
+from apache_beam.io.gcp import bigquery_change_history as bqch
+from apache_beam.io.gcp import bigquery_tools
from apache_beam.io.gcp.bigquery_change_history import
ReadBigQueryChangeHistory
from apache_beam.io.gcp.bigquery_change_history import build_changes_query
from apache_beam.io.gcp.bigquery_change_history import compute_ranges
+from apache_beam.options.pipeline_options import PipelineOptions
+from apache_beam.testing.test_pipeline import TestPipeline
from apache_beam.utils.timestamp import Duration
from apache_beam.utils.timestamp import Timestamp
@@ -205,5 +211,74 @@ class ValidationTest(unittest.TestCase):
self.assertEqual(t._buffer_sec, 10)
[email protected](HttpError is None, 'GCP dependencies are not installed')
+class QuotaProjectTest(unittest.TestCase):
+ """Tests that the quota project reaches every BigQuery client."""
+ _QUOTA = 'my-billing-project'
+
+ def _assert_wrapper_quota(self, fn):
+ with mock.patch.object(bigquery_tools, 'BigQueryWrapper') as mock_wrapper:
+ fn.setup()
+ self.assertEqual(
+ mock_wrapper.call_args.kwargs['quota_project_id'], self._QUOTA)
+
+ def test_poll_fn_setup_passes_quota_project(self):
+ fn = bqch._PollChangeHistoryFn(
+ table='proj:dataset.table',
+ project='p',
+ change_function='APPENDS',
+ buffer=Duration(seconds=10),
+ start_time=Timestamp(0),
+ stop_time=Timestamp(100),
+ poll_interval=Duration(seconds=60),
+ quota_project_id=self._QUOTA)
+ self._assert_wrapper_quota(fn)
+
+ def test_execute_query_fn_setup_passes_quota_project(self):
+ fn = bqch._ExecuteQueryFn(
+ table='proj:dataset.table',
+ project='p',
+ change_function='APPENDS',
+ temp_dataset='tmp',
+ location='US',
+ quota_project_id=self._QUOTA)
+ self._assert_wrapper_quota(fn)
+
+ def test_cleanup_fn_setup_passes_quota_project(self):
+ self._assert_wrapper_quota(bqch._CleanupTempTablesFn(self._QUOTA))
+
+ def test_read_sdf_creates_storage_client_with_quota_project(self):
+ fn = bqch._ReadStorageStreamsSDF(quota_project_id=self._QUOTA)
+ with mock.patch.object(bigquery_tools,
+ 'create_bq_storage_client') as mock_create:
+ fn.setup()
+ mock_create.assert_called_once_with(self._QUOTA)
+
+ def test_expand_resolves_quota_project_from_options(self):
+ options = PipelineOptions(
+ ['--quota_project_id=' + self._QUOTA, '--no_pipeline_type_check'])
+ classes = [
+ '_PollChangeHistoryFn',
+ '_ExecuteQueryFn',
+ '_ReadStorageStreamsSDF',
+ '_CleanupTempTablesFn'
+ ]
+ p = TestPipeline(options=options)
+ with contextlib.ExitStack() as stack:
+ mocks = [
+ stack.enter_context(
+ mock.patch.object(
+ getattr(bqch, name),
+ '__init__',
+ autospec=True,
+ side_effect=getattr(bqch, name).__init__)) for name in
classes
+ ]
+ _ = p | ReadBigQueryChangeHistory(
+ table='p:d.t', project='p', start_time=0, temp_dataset='tmp')
+ for name, m in zip(classes, mocks):
+ self.assertEqual(
+ m.call_args.kwargs['quota_project_id'], self._QUOTA, name)
+
+
if __name__ == '__main__':
unittest.main()
diff --git a/sdks/python/apache_beam/io/gcp/bigquery_file_loads.py
b/sdks/python/apache_beam/io/gcp/bigquery_file_loads.py
index dabe80c5c27..79940d109f0 100644
--- a/sdks/python/apache_beam/io/gcp/bigquery_file_loads.py
+++ b/sdks/python/apache_beam/io/gcp/bigquery_file_loads.py
@@ -383,16 +383,19 @@ class UpdateDestinationSchema(beam.DoFn):
test_client=None,
additional_bq_parameters=None,
step_name=None,
- load_job_project_id=None):
+ load_job_project_id=None,
+ quota_project_id=None):
self.project = project
self._test_client = test_client
+ self._quota_project_id = quota_project_id
self._write_disposition = write_disposition
self._additional_bq_parameters = additional_bq_parameters or {}
self._step_name = step_name
self._load_job_project_id = load_job_project_id
def start_bundle(self):
- self.bq_wrapper = bigquery_tools.BigQueryWrapper(client=self._test_client)
+ self.bq_wrapper = bigquery_tools.BigQueryWrapper(
+ client=self._test_client, quota_project_id=self._quota_project_id)
self._bq_io_metadata = create_bigquery_io_metadata(self._step_name)
self.pending_jobs = []
@@ -525,11 +528,13 @@ class TriggerCopyJobs(beam.DoFn):
write_disposition=None,
test_client=None,
step_name=None,
- load_job_project_id=None):
+ load_job_project_id=None,
+ quota_project_id=None):
self.project = project
self.create_disposition = create_disposition
self.write_disposition = write_disposition
self.test_client = test_client
+ self.quota_project_id = quota_project_id
self._observed_tables = set()
self.bq_io_metadata = None
self._step_name = step_name
@@ -545,7 +550,8 @@ class TriggerCopyJobs(beam.DoFn):
self._observed_tables = set()
def start_bundle(self):
- self.bq_wrapper = bigquery_tools.BigQueryWrapper(client=self.test_client)
+ self.bq_wrapper = bigquery_tools.BigQueryWrapper(
+ client=self.test_client, quota_project_id=self.quota_project_id)
if not self.bq_io_metadata:
self.bq_io_metadata = create_bigquery_io_metadata(self._step_name)
self.pending_jobs = []
@@ -677,10 +683,12 @@ class TriggerLoadJobs(beam.DoFn):
additional_bq_parameters=None,
source_format=None,
step_name=None,
- load_job_project_id=None):
+ load_job_project_id=None,
+ quota_project_id=None):
self.schema = schema
self.project = project
self.test_client = test_client
+ self.quota_project_id = quota_project_id
self.temporary_tables = temporary_tables
self.additional_bq_parameters = additional_bq_parameters or {}
self.source_format = source_format
@@ -709,7 +717,8 @@ class TriggerLoadJobs(beam.DoFn):
return result
def start_bundle(self):
- self.bq_wrapper = bigquery_tools.BigQueryWrapper(client=self.test_client)
+ self.bq_wrapper = bigquery_tools.BigQueryWrapper(
+ client=self.test_client, quota_project_id=self.quota_project_id)
if not self.bq_io_metadata:
self.bq_io_metadata = create_bigquery_io_metadata(self._step_name)
self.pending_jobs = []
@@ -928,11 +937,13 @@ class PartitionFiles(beam.DoFn):
class DeleteTablesFn(beam.DoFn):
- def __init__(self, test_client=None):
+ def __init__(self, test_client=None, quota_project_id=None):
self.test_client = test_client
+ self.quota_project_id = quota_project_id
def start_bundle(self):
- self.bq_wrapper = bigquery_tools.BigQueryWrapper(client=self.test_client)
+ self.bq_wrapper = bigquery_tools.BigQueryWrapper(
+ client=self.test_client, quota_project_id=self.quota_project_id)
def process(self, table_reference):
_LOGGER.info("Deleting table %s", table_reference)
@@ -998,6 +1009,7 @@ class BigQueryBatchFileLoads(beam.PTransform):
self.test_client = test_client
self.schema = schema
+ self.quota_project_id = None
self._temp_file_format = temp_file_format or bigquery_tools.FileFormat.JSON
# If we have multiple destinations, then we will have multiple load jobs,
@@ -1207,7 +1219,8 @@ class BigQueryBatchFileLoads(beam.PTransform):
additional_bq_parameters=self.additional_bq_parameters,
source_format=self._temp_file_format,
step_name=step_name,
- load_job_project_id=self.load_job_project_id),
+ load_job_project_id=self.load_job_project_id,
+ quota_project_id=self.quota_project_id),
load_job_name_pcv,
*self.schema_side_inputs).with_outputs(
TriggerLoadJobs.TEMP_TABLES,
@@ -1228,7 +1241,8 @@ class BigQueryBatchFileLoads(beam.PTransform):
test_client=self.test_client,
additional_bq_parameters=self.additional_bq_parameters,
step_name=step_name,
- load_job_project_id=self.load_job_project_id),
+ load_job_project_id=self.load_job_project_id,
+ quota_project_id=self.quota_project_id),
schema_mod_job_name_pcv))
if self.write_disposition in ('WRITE_EMPTY', 'WRITE_TRUNCATE'):
@@ -1259,7 +1273,8 @@ class BigQueryBatchFileLoads(beam.PTransform):
write_disposition=self.write_disposition,
test_client=self.test_client,
step_name=step_name,
- load_job_project_id=self.load_job_project_id),
+ load_job_project_id=self.load_job_project_id,
+ quota_project_id=self.quota_project_id),
copy_job_name_pcv,
pvalue.AsIter(schema_mod_job_ids_pc)).with_outputs(
TriggerCopyJobs.TRIGGER_DELETE_TEMP_TABLES, main='main'))
@@ -1275,7 +1290,9 @@ class BigQueryBatchFileLoads(beam.PTransform):
| "RemoveTempTables/DeduplicateTables" >> beam.GroupByKey()
| "RemoveTempTables/GetTableNames" >> beam.Keys()
| "RemoveTempTables/Delete" >> beam.ParDo(
- DeleteTablesFn(self.test_client)))
+ DeleteTablesFn(
+ test_client=self.test_client,
+ quota_project_id=self.quota_project_id)))
# Load data directly to destination table
destination_load_job_ids_pc = (
@@ -1290,7 +1307,8 @@ class BigQueryBatchFileLoads(beam.PTransform):
additional_bq_parameters=self.additional_bq_parameters,
source_format=self._temp_file_format,
step_name=step_name,
- load_job_project_id=self.load_job_project_id),
+ load_job_project_id=self.load_job_project_id,
+ quota_project_id=self.quota_project_id),
load_job_name_pcv,
*self.schema_side_inputs).with_outputs(
TriggerLoadJobs.ONGOING_JOBS, main='main')
@@ -1305,6 +1323,8 @@ class BigQueryBatchFileLoads(beam.PTransform):
def expand(self, pcoll):
p = pcoll.pipeline
self.project = self.project or
p.options.view_as(GoogleCloudOptions).project
+ self.quota_project_id = bigquery_tools._quota_project_id_from_options(
+ p.options)
try:
step_name = self.label
except AttributeError:
diff --git a/sdks/python/apache_beam/io/gcp/bigquery_file_loads_test.py
b/sdks/python/apache_beam/io/gcp/bigquery_file_loads_test.py
index 7ae0ccc8e32..7ccda7ecad8 100644
--- a/sdks/python/apache_beam/io/gcp/bigquery_file_loads_test.py
+++ b/sdks/python/apache_beam/io/gcp/bigquery_file_loads_test.py
@@ -19,6 +19,7 @@
# pytype: skip-file
+import contextlib
import logging
import os
import secrets
@@ -1259,6 +1260,52 @@ class TestBigQueryFileLoads(_TestCaseWithTempDirCleanUp):
assert_that(jobs, equal_to(expected_jobs), label='CheckJobs')
+class TestFileLoadsQuotaProject(_TestCaseWithTempDirCleanUp):
+ """Tests that the file loads DoFns attribute quota to the quota project."""
+ @parameterized.expand([
+ param(dofn=bqfl.UpdateDestinationSchema, quota_attr='_quota_project_id'),
+ param(dofn=bqfl.TriggerCopyJobs, quota_attr='quota_project_id'),
+ param(dofn=bqfl.TriggerLoadJobs, quota_attr='quota_project_id'),
+ param(dofn=bqfl.DeleteTablesFn, quota_attr='quota_project_id'),
+ ])
+ def test_start_bundle_passes_quota_project(self, dofn, quota_attr):
+ fn = dofn(quota_project_id='my-billing-project')
+ self.assertEqual(getattr(fn, quota_attr), 'my-billing-project')
+ with mock.patch.object(bigquery_tools, 'BigQueryWrapper') as mock_wrapper:
+ fn.start_bundle()
+ self.assertEqual(
+ mock_wrapper.call_args.kwargs['quota_project_id'],
'my-billing-project')
+
+ def test_expand_resolves_quota_project_from_options(self):
+ options = PipelineOptions(['--quota_project_id=my-billing-project'])
+ transform = bqfl.BigQueryBatchFileLoads(
+ 'project:dataset.table',
+ custom_gcs_temp_location=self._new_tempdir(),
+ test_client=mock.Mock(),
+ validate=False,
+ temp_file_format=bigquery_tools.FileFormat.JSON)
+ p = TestPipeline(options=options)
+ dofns = [
+ bqfl.UpdateDestinationSchema,
+ bqfl.TriggerCopyJobs,
+ bqfl.TriggerLoadJobs,
+ bqfl.DeleteTablesFn
+ ]
+ with contextlib.ExitStack() as stack:
+ mocks = [
+ stack.enter_context(
+ mock.patch.object(
+ dofn, '__init__', autospec=True, side_effect=dofn.__init__))
+ for dofn in dofns
+ ]
+ _ = p | beam.Create(_ELEMENTS) | transform
+ self.assertEqual(transform.quota_project_id, 'my-billing-project')
+ for dofn, m in zip(dofns, mocks):
+ self.assertTrue(m.call_args_list, dofn.__name__)
+ for c in m.call_args_list:
+ self.assertEqual(c.kwargs['quota_project_id'], 'my-billing-project')
+
+
class BigQueryFileLoadsIT(unittest.TestCase):
BIG_QUERY_SCHEMA = (
diff --git a/sdks/python/apache_beam/io/gcp/bigquery_read_internal.py
b/sdks/python/apache_beam/io/gcp/bigquery_read_internal.py
index 136b3cc56b7..2fa95a3d179 100644
--- a/sdks/python/apache_beam/io/gcp/bigquery_read_internal.py
+++ b/sdks/python/apache_beam/io/gcp/bigquery_read_internal.py
@@ -142,11 +142,13 @@ class _PassThroughThenCleanupTempDatasets(PTransform):
Utilizes readiness of PCollection to trigger DoFn.
"""
- def __init__(self, side_input=None):
+ def __init__(self, side_input=None, quota_project_id=None):
self.side_input = side_input
+ self.quota_project_id = quota_project_id
def expand(self, input):
pipeline_options = input.pipeline.options
+ quota_project_id = self.quota_project_id
class PassThrough(beam.DoFn):
def process(self, element):
@@ -154,8 +156,9 @@ class _PassThroughThenCleanupTempDatasets(PTransform):
class CleanUpProjects(beam.DoFn):
def process(self, unused_element, unused_signal, pipeline_details):
- bq = bigquery_tools.BigQueryWrapper.from_pipeline_options(
- pipeline_options)
+ bq = bigquery_tools.BigQueryWrapper(
+ pipeline_options=pipeline_options,
+ quota_project_id=quota_project_id)
pipeline_details = pipeline_details[0]
if 'temp_table_ref' in pipeline_details.keys():
temp_table_ref = pipeline_details['temp_table_ref']
@@ -251,7 +254,7 @@ class _BigQueryReadSplit(beam.transforms.DoFn):
def start_bundle(self):
self.bq = bigquery_tools.BigQueryWrapper(
temp_dataset_id=self._get_temp_dataset_id(),
- client=bigquery_tools.BigQueryWrapper._bigquery_client(self.options))
+ pipeline_options=self.options)
def process(self,
element: 'ReadFromBigQueryRequest') -> Iterable[BoundedSource]:
diff --git a/sdks/python/apache_beam/io/gcp/bigquery_test.py
b/sdks/python/apache_beam/io/gcp/bigquery_test.py
index 22b617fa874..ae03104612c 100644
--- a/sdks/python/apache_beam/io/gcp/bigquery_test.py
+++ b/sdks/python/apache_beam/io/gcp/bigquery_test.py
@@ -53,6 +53,7 @@ from apache_beam.io.gcp.bigquery import WriteToBigQuery
from apache_beam.io.gcp.bigquery import _StreamToBigQuery
from apache_beam.io.gcp.bigquery_read_internal import _BigQueryReadSplit
from apache_beam.io.gcp.bigquery_read_internal import _JsonToDictCoder
+from apache_beam.io.gcp.bigquery_read_internal import
_PassThroughThenCleanupTempDatasets
from apache_beam.io.gcp.bigquery_read_internal import
bigquery_export_destination_uri
from apache_beam.io.gcp.bigquery_tools import JSON_COMPLIANCE_ERROR
from apache_beam.io.gcp.bigquery_tools import BigQueryWrapper
@@ -65,6 +66,7 @@ from apache_beam.io.gcp.tests.bigquery_matcher import
BigqueryFullResultStreamin
from apache_beam.io.gcp.tests.bigquery_matcher import BigQueryTableMatcher
from apache_beam.metrics.metric import Lineage
from apache_beam.options import value_provider
+from apache_beam.options.pipeline_options import GoogleCloudOptions
from apache_beam.options.pipeline_options import PipelineOptions
from apache_beam.options.pipeline_options import StandardOptions
from apache_beam.options.value_provider import RuntimeValueProvider
@@ -338,6 +340,10 @@ class TestJsonToDictCoder(unittest.TestCase):
class TestReadFromBigQuery(unittest.TestCase):
@classmethod
def setUpClass(cls):
+ cls.env_patch = mock.patch.dict(
+ os.environ, {'GOOGLE_CLOUD_PROJECT': 'test-project'})
+ cls.env_patch.start()
+
class UserDefinedOptions(PipelineOptions):
@classmethod
def _add_argparse_args(cls, parser):
@@ -351,6 +357,7 @@ class TestReadFromBigQuery(unittest.TestCase):
@classmethod
def tearDownClass(cls):
+ cls.env_patch.stop()
# Unset the option added in setupClass to avoid interfere with other tests.
# Force a gc so PipelineOptions.__subclass__() no longer contains it.
del cls.UserDefinedOptions
@@ -879,6 +886,343 @@ class TestReadFromBigQuery(unittest.TestCase):
mock_convert.assert_called_once_with(schema, None)
[email protected](
+ HttpError is None or gcp_bigquery is None,
+ 'GCP dependencies are not installed')
+class TestReadFromBigQueryQuotaProject(unittest.TestCase):
+ """Tests for quota_project_id in ReadFromBigQuery sources."""
+ def test_quota_project_id_from_pipeline_options(self):
+ """Test that quota_project_id is read from GoogleCloudOptions."""
+ options = PipelineOptions(['--quota_project_id=my-billing-project'])
+ gcp_options = options.view_as(GoogleCloudOptions)
+ self.assertEqual(gcp_options.quota_project_id, 'my-billing-project')
+
+ def test_quota_project_id_none_by_default_in_options(self):
+ """Test that quota_project_id is None by default in options."""
+ options = PipelineOptions([])
+ gcp_options = options.view_as(GoogleCloudOptions)
+ self.assertIsNone(gcp_options.quota_project_id)
+
+ def test_export_source_explicit_quota_project(self):
+ """Test that explicit quota_project_id is stored in
+ _CustomBigQuerySource."""
+ source = beam_bq._CustomBigQuerySource(
+ method=ReadFromBigQuery.Method.EXPORT,
+ table='project:dataset.table',
+ quota_project_id='my-billing-project')
+ self.assertEqual(source.quota_project_id, 'my-billing-project')
+ self.assertEqual(source._get_quota_project_id(), 'my-billing-project')
+
+ def test_export_source_gets_quota_from_options(self):
+ """Test that _CustomBigQuerySource falls back to options for
+ quota_project_id."""
+ options = PipelineOptions(['--quota_project_id=my-billing-project'])
+ source = beam_bq._CustomBigQuerySource(
+ method=ReadFromBigQuery.Method.EXPORT,
+ table='project:dataset.table',
+ pipeline_options=options)
+ self.assertEqual(source._get_quota_project_id(), 'my-billing-project')
+
+ def test_export_source_explicit_overrides_options(self):
+ """Test that explicit quota_project_id overrides options."""
+ options = PipelineOptions(['--quota_project_id=options-project'])
+ source = beam_bq._CustomBigQuerySource(
+ method=ReadFromBigQuery.Method.EXPORT,
+ table='project:dataset.table',
+ pipeline_options=options,
+ quota_project_id='explicit-project')
+ self.assertEqual(source._get_quota_project_id(), 'explicit-project')
+
+ def test_storage_source_explicit_quota_project(self):
+ """Test that explicit quota_project_id is stored in
+ _CustomBigQueryStorageSource."""
+ source = beam_bq._CustomBigQueryStorageSource(
+ method=ReadFromBigQuery.Method.DIRECT_READ,
+ table='project:dataset.table',
+ quota_project_id='my-billing-project')
+ self.assertEqual(source.quota_project_id, 'my-billing-project')
+ self.assertEqual(source._get_quota_project_id(), 'my-billing-project')
+
+ def test_storage_source_gets_quota_from_options(self):
+ """Test that _CustomBigQueryStorageSource falls back to options."""
+ options = PipelineOptions(['--quota_project_id=my-billing-project'])
+ source = beam_bq._CustomBigQueryStorageSource(
+ method=ReadFromBigQuery.Method.DIRECT_READ,
+ table='project:dataset.table',
+ pipeline_options=options)
+ self.assertEqual(source._get_quota_project_id(), 'my-billing-project')
+
+ @mock.patch.object(bigquery_tools.BigQueryWrapper, '_bigquery_client')
+ def test_export_source_split_passes_quota_to_client(self, mock_bq_client):
+ """split() must pass the transform-level quota_project_id to the BigQuery
+ client, not only the --quota_project_id pipeline option."""
+ mock_bq_client.side_effect = RuntimeError('stop')
+ source = beam_bq._CustomBigQuerySource(
+ method=ReadFromBigQuery.Method.EXPORT,
+ table='project:dataset.table',
+ quota_project_id='my-billing-project')
+ with self.assertRaises(RuntimeError):
+ list(source.split(desired_bundle_size=0))
+ mock_bq_client.assert_called_once()
+ self.assertIsInstance(mock_bq_client.call_args.args[0], PipelineOptions)
+ self.assertEqual(
+ mock_bq_client.call_args.kwargs['quota_project_id'],
+ 'my-billing-project')
+
+ @mock.patch.object(bigquery_tools.BigQueryWrapper, '_bigquery_client')
+ def test_storage_source_split_passes_quota_to_client(self, mock_bq_client):
+ """split() must pass the transform-level quota_project_id to the BigQuery
+ client, not only the --quota_project_id pipeline option."""
+ mock_bq_client.side_effect = RuntimeError('stop')
+ source = beam_bq._CustomBigQueryStorageSource(
+ method=ReadFromBigQuery.Method.DIRECT_READ,
+ table='project:dataset.table',
+ quota_project_id='my-billing-project')
+ with self.assertRaises(RuntimeError):
+ list(source.split(desired_bundle_size=0))
+ mock_bq_client.assert_called_once()
+ self.assertIsInstance(mock_bq_client.call_args.args[0], PipelineOptions)
+ self.assertEqual(
+ mock_bq_client.call_args.kwargs['quota_project_id'],
+ 'my-billing-project')
+
+ @mock.patch.object(bigquery_tools, 'create_bq_storage_client')
+ @mock.patch.object(bigquery_tools.BigQueryWrapper, 'get_table')
+ @mock.patch.object(bigquery_tools.BigQueryWrapper, '_bigquery_client')
+ def test_storage_source_split_passes_quota_to_storage_client(
+ self, unused_mock_bq_client, mock_get_table, mock_create_storage_client):
+ """split() must create the Storage Read client with the quota project
+ and pass it on to each stream source."""
+ mock_get_table.return_value = mock.Mock(numBytes=5)
+ stream = mock.Mock()
+ stream.name = 'projects/p/locations/l/sessions/s/streams/stream1'
+ mock_create_storage_client.return_value.create_read_session.return_value =
(
+ mock.Mock(streams=[stream]))
+ source = beam_bq._CustomBigQueryStorageSource(
+ method=ReadFromBigQuery.Method.DIRECT_READ,
+ table='project:dataset.table',
+ quota_project_id='my-billing-project')
+
+ bundles = list(source.split(desired_bundle_size=1 << 20))
+
+ mock_create_storage_client.assert_called_once_with('my-billing-project')
+ self.assertEqual(len(bundles), 1)
+ self.assertEqual(bundles[0].source.quota_project_id, 'my-billing-project')
+
+ @mock.patch.object(bigquery_tools.BigQueryWrapper, '_bigquery_client')
+ def test_export_source_estimate_size_passes_quota_to_client(
+ self, mock_bq_client):
+ """estimate_size() must attribute quota too, not just split()."""
+ mock_bq_client.side_effect = RuntimeError('stop')
+ source = beam_bq._CustomBigQuerySource(
+ method=ReadFromBigQuery.Method.EXPORT,
+ table='project:dataset.table',
+ quota_project_id='my-billing-project')
+ with self.assertRaises(RuntimeError):
+ source.estimate_size()
+ self.assertEqual(
+ mock_bq_client.call_args.kwargs['quota_project_id'],
+ 'my-billing-project')
+
+ @mock.patch.object(bigquery_tools.BigQueryWrapper, '_bigquery_client')
+ def test_storage_source_estimate_size_passes_quota_to_client(
+ self, mock_bq_client):
+ """estimate_size() must attribute quota too, not just split()."""
+ mock_bq_client.side_effect = RuntimeError('stop')
+ source = beam_bq._CustomBigQueryStorageSource(
+ method=ReadFromBigQuery.Method.DIRECT_READ,
+ table='project:dataset.table',
+ quota_project_id='my-billing-project')
+ with self.assertRaises(RuntimeError):
+ source.estimate_size()
+ self.assertEqual(
+ mock_bq_client.call_args.kwargs['quota_project_id'],
+ 'my-billing-project')
+
+ @mock.patch(
+ 'apache_beam.io.gcp.bigquery.bigquery_schema_tools'
+ '.convert_to_usertype')
+ @mock.patch.object(bigquery_tools, 'BigQueryWrapper')
+ def test_output_type_schema_fetch_passes_quota_to_client(
+ self, mock_wrapper, mock_convert):
+ """The BEAM_ROW schema lookup is a BigQuery API call of this transform, so
+ it must be attributed to the quota project as well."""
+ mock_convert.return_value = beam.Map(lambda x: x)
+ transform = ReadFromBigQuery(
+ table='project:dataset.table',
+ output_type='BEAM_ROW',
+ quota_project_id='my-billing-project')
+ transform._expand_output_type(mock.MagicMock())
+
+ self.assertEqual(
+ mock_wrapper.call_args.kwargs['quota_project_id'],
'my-billing-project')
+
+ @mock.patch.object(bigquery_tools, 'BigQueryWrapper')
+ def test_direct_read_cleanup_passes_quota_to_client(self, mock_wrapper):
+ """The DIRECT_READ temp-dataset cleanup is a BigQuery API call of this
+ transform, so it must be attributed to the quota project as well."""
+ pipeline_details = {
+ 'project_id': 'project', 'bigquery_dataset_labels': {
+ 'k': 'v'
+ }
+ }
+ with TestPipeline() as p:
+ side_input = beam.pvalue.AsList(
+ p | 'Details' >> beam.Create([pipeline_details]))
+ _ = (
+ p | beam.Create([1])
+ | _PassThroughThenCleanupTempDatasets(
+ side_input, quota_project_id='my-billing-project'))
+
+ self.assertEqual(
+ mock_wrapper.call_args.kwargs['quota_project_id'],
'my-billing-project')
+ self.assertIsNotNone(mock_wrapper.call_args.kwargs['pipeline_options'])
+
+
@mock.patch('apache_beam.io.gcp.bigquery._PassThroughThenCleanupTempDatasets')
+ def test_direct_read_wires_quota_project_into_cleanup(self, mock_cleanup):
+ mock_cleanup.return_value = beam.Map(lambda x: x)
+ transform = ReadFromBigQuery(
+ table='project:dataset.table',
+ method=ReadFromBigQuery.Method.DIRECT_READ,
+ quota_project_id='my-billing-project')
+ p = TestPipeline()
+ transform._expand_direct_read(beam.pvalue.PBegin(p))
+
+ self.assertEqual(
+ mock_cleanup.call_args.kwargs['quota_project_id'],
'my-billing-project')
+
+ def test_quota_project_id_in_export_source_display_data(self):
+ """Test that quota_project_id appears in display data for export source."""
+ source = beam_bq._CustomBigQuerySource(
+ method=ReadFromBigQuery.Method.EXPORT,
+ table='project:dataset.table',
+ quota_project_id='my-billing-project')
+ items = {i.key: i.value for i in DisplayData.create_from(source).items}
+ self.assertEqual(items['quota_project_id'], 'my-billing-project')
+
+ def test_quota_project_id_in_direct_read_source_display_data(self):
+ """Test that quota_project_id appears in display data for DIRECT_READ."""
+ source = beam_bq._CustomBigQueryStorageSource(
+ method=ReadFromBigQuery.Method.DIRECT_READ,
+ table='project:dataset.table',
+ quota_project_id='my-billing-project')
+ items = {i.key: i.value for i in DisplayData.create_from(source).items}
+ self.assertEqual(items['quota_project_id'], 'my-billing-project')
+
+ def test_quota_project_id_absent_from_display_data_when_not_set(self):
+ """Test that quota_project_id is dropped from display data when not set."""
+ options = PipelineOptions([])
+ sources = [
+ beam_bq._CustomBigQuerySource(
+ method=ReadFromBigQuery.Method.EXPORT,
+ table='project:dataset.table',
+ pipeline_options=options),
+ beam_bq._CustomBigQueryStorageSource(
+ method=ReadFromBigQuery.Method.DIRECT_READ,
+ table='project:dataset.table',
+ pipeline_options=options)
+ ]
+ for source in sources:
+ keys = {i.key for i in DisplayData.create_from(source).items}
+ self.assertNotIn('quota_project_id', keys)
+
+ def test_stream_source_stores_quota_project_id(self):
+ """Test that quota_project_id is stored in
+ _CustomBigQueryStorageStreamSource."""
+ stream_source = beam_bq._CustomBigQueryStorageStreamSource(
+ read_stream_name='projects/p/locations/l/sessions/s/streams/stream1',
+ use_native_datetime=True,
+ timeout=30.0,
+ quota_project_id='my-billing-project')
+ self.assertEqual(stream_source.quota_project_id, 'my-billing-project')
+
+ def test_stream_source_quota_project_id_none_by_default(self):
+ """Test that quota_project_id is None by default in stream source."""
+ stream_source = beam_bq._CustomBigQueryStorageStreamSource(
+ read_stream_name='projects/p/locations/l/sessions/s/streams/stream1')
+ self.assertIsNone(stream_source.quota_project_id)
+
+ def test_stream_source_split_preserves_quota_project_id(self):
+ """Test that split() preserves quota_project_id."""
+ stream_source = beam_bq._CustomBigQueryStorageStreamSource(
+ read_stream_name='projects/p/locations/l/sessions/s/streams/stream1',
+ use_native_datetime=True,
+ timeout=30.0,
+ quota_project_id='my-billing-project')
+ bundle = stream_source.split(desired_bundle_size=0)
+ self.assertEqual(bundle.source.quota_project_id, 'my-billing-project')
+ self.assertEqual(bundle.source.timeout, 30.0)
+ self.assertEqual(bundle.source.use_native_datetime, True)
+
+ @mock.patch('apache_beam.io.gcp.bigquery_tools.create_bq_storage_client')
+ def test_stream_source_read_arrow_uses_quota_project(
+ self, mock_create_client):
+ """Test that read_arrow() uses create_bq_storage_client
+ with quota_project_id."""
+ mock_client = mock.MagicMock()
+ mock_create_client.return_value = mock_client
+ # Mock read_rows to return empty iterator
+ mock_client.read_rows.return_value.rows.return_value = iter([])
+
+ stream_source = beam_bq._CustomBigQueryStorageStreamSource(
+ read_stream_name='projects/p/locations/l/sessions/s/streams/stream1',
+ use_native_datetime=True,
+ quota_project_id='my-billing-project')
+ # Consume the iterator
+ list(stream_source.read_arrow())
+
+ mock_create_client.assert_called_once_with('my-billing-project')
+
+ @mock.patch('apache_beam.io.gcp.bigquery_tools.create_bq_storage_client')
+ def test_stream_source_read_avro_uses_quota_project(self,
mock_create_client):
+ """Test that read_avro() uses create_bq_storage_client
+ with quota_project_id."""
+ mock_client = mock.MagicMock()
+ mock_create_client.return_value = mock_client
+ # Mock read_rows to return empty iterator
+ mock_client.read_rows.return_value = iter([])
+
+ stream_source = beam_bq._CustomBigQueryStorageStreamSource(
+ read_stream_name='projects/p/locations/l/sessions/s/streams/stream1',
+ use_native_datetime=False,
+ quota_project_id='my-billing-project')
+ # Consume the iterator
+ list(stream_source.read_avro())
+
+ mock_create_client.assert_called_once_with('my-billing-project')
+
+
[email protected](HttpError is None, 'GCP dependencies are not installed')
+class TestWriteToBigQueryQuotaProject(unittest.TestCase):
+ """Tests for quota_project_id in the BigQuery streaming insert path."""
+ @mock.patch.object(bigquery_tools, 'BigQueryWrapper')
+ def test_write_fn_start_bundle_passes_quota_project(self, mock_wrapper):
+ fn = beam_bq.BigQueryWriteFn(
+ batch_size=1, quota_project_id='my-billing-project')
+ fn.start_bundle()
+ self.assertEqual(
+ mock_wrapper.call_args.kwargs['quota_project_id'],
'my-billing-project')
+
+ def test_stream_to_bigquery_resolves_quota_project_from_options(self):
+ options = PipelineOptions(['--quota_project_id=my-billing-project'])
+ p = TestPipeline(options=options)
+ with mock.patch.object(
+ beam_bq.BigQueryWriteFn,
+ '__init__',
+ autospec=True,
+ side_effect=beam_bq.BigQueryWriteFn.__init__) as mock_write_fn:
+ _ = (
+ p
+ | beam.Create([])
+ | WriteToBigQuery(
+ 'project:dataset.table',
+ schema='a:STRING',
+ method=WriteToBigQuery.Method.STREAMING_INSERTS))
+ self.assertEqual(
+ mock_write_fn.call_args.kwargs['quota_project_id'],
+ 'my-billing-project')
+
+
@unittest.skipIf(HttpError is None, 'GCP dependencies are not installed')
class TestBigQuerySink(unittest.TestCase):
def test_table_spec_display_data(self):
@@ -920,10 +1264,14 @@ class TestWriteToBigQuery(unittest.TestCase):
os.remove('insert_calls2')
def setUp(self):
+ self.env_patch = mock.patch.dict(
+ os.environ, {'GOOGLE_CLOUD_PROJECT': 'test-project'})
+ self.env_patch.start()
self._cleanup_files()
def tearDown(self):
self._cleanup_files()
+ self.env_patch.stop()
def test_noop_schema_parsing(self):
expected_table_schema = None
@@ -1481,6 +1829,13 @@ class TestWriteToBigQuery(unittest.TestCase):
HttpError is None or exceptions is None,
'GCP dependencies are not installed')
class BigQueryStreamingInsertsErrorHandling(unittest.TestCase):
+ def setUp(self):
+ self.env_patch = mock.patch.dict(
+ os.environ, {'GOOGLE_CLOUD_PROJECT': 'test-project'})
+ self.env_patch.start()
+
+ def tearDown(self):
+ self.env_patch.stop()
# Running tests with a variety of exceptions from https://googleapis.dev
#
/python/google-api-core/latest/_modules/google/api_core/exceptions.html.
diff --git a/sdks/python/apache_beam/io/gcp/bigquery_tools.py
b/sdks/python/apache_beam/io/gcp/bigquery_tools.py
index cc1a16a659f..d1412d0f855 100644
--- a/sdks/python/apache_beam/io/gcp/bigquery_tools.py
+++ b/sdks/python/apache_beam/io/gcp/bigquery_tools.py
@@ -36,6 +36,7 @@ import logging
import re
import sys
import time
+import urllib.parse
import uuid
from json.decoder import JSONDecodeError
from typing import Optional
@@ -59,6 +60,7 @@ from apache_beam.io.gcp.internal.clients import bigquery
from apache_beam.metrics import monitoring_infos
from apache_beam.metrics.metric import Metrics
from apache_beam.options import value_provider
+from apache_beam.options.pipeline_options import GoogleCloudOptions
from apache_beam.options.pipeline_options import PipelineOptions
from apache_beam.transforms import DoFn
from apache_beam.typehints.row_type import RowTypeConstraint
@@ -80,6 +82,11 @@ except Exception:
gcp_bigquery = None
pass
+try:
+ from google.cloud import bigquery_storage_v1 as bq_storage
+except ImportError:
+ bq_storage = None
+
try:
from orjson import dumps as fast_json_dumps
from orjson import loads as fast_json_loads
@@ -419,6 +426,58 @@ def _build_dataset_encryption_config(kms_key):
return bigquery.EncryptionConfiguration(kmsKeyName=kms_key)
+def _quota_project_id_from_options(pipeline_options):
+ """Returns --quota_project_id from pipeline options, if any."""
+ if pipeline_options is None:
+ return None
+ return pipeline_options.view_as(GoogleCloudOptions).quota_project_id
+
+
+def create_bq_storage_client(quota_project_id=None):
+ """Create a BigQueryReadClient with optional quota project.
+
+ Args:
+ quota_project_id: Optional GCP project ID to use for quota and billing.
+ The client applies it to the credentials it resolves itself.
+
+ Returns:
+ A BigQueryReadClient instance.
+ """
+ if not quota_project_id:
+ return bq_storage.BigQueryReadClient()
+ return bq_storage.BigQueryReadClient(
+ client_options={'quota_project_id': quota_project_id})
+
+
+class _HttpWithHeaders(object):
+ """Wraps an httplib2.Http to add fixed headers to requests to given hosts.
+
+ Sits below the credentials layer, so the headers are applied on every
+ attempt and take precedence over headers set by the credentials. The
+ credentials refresh their token through the same Http, so requests to other
+ hosts (token endpoint, metadata server, IAM Credentials API) are passed
+ through unchanged.
+ """
+ def __init__(self, http, headers, hosts):
+ object.__setattr__(self, '_http', http)
+ object.__setattr__(self, '_headers', headers)
+ object.__setattr__(self, '_hosts', frozenset(hosts))
+
+ def request(self, uri, method='GET', body=None, headers=None, **kwargs):
+ if urllib.parse.urlsplit(uri).hostname in self._hosts:
+ headers = dict(headers or {})
+ headers.update(self._headers)
+ return self._http.request(uri, method, body=body, headers=headers,
**kwargs)
+
+ def __getattr__(self, name):
+ if name == '_http': # object created without __init__ (copy/unpickle)
+ raise AttributeError(name)
+ return getattr(self._http, name)
+
+ def __setattr__(self, name, value):
+ setattr(self._http, name, value)
+
+
class BigQueryWrapper(object):
"""BigQuery client wrapper with utilities for querying.
@@ -427,8 +486,10 @@ class BigQueryWrapper(object):
In addition, it offers various functions used both in sources and sinks
(e.g., find and create tables, query a table, etc.).
- Note that client parameter in constructor is only for testing purposes and
- should not be used in production code.
+ Note that the client and gcp_client parameters in the constructor are only
+ for testing purposes and should not be used in production code. Production
+ code should pass pipeline_options (and optionally quota_project_id) so that
+ every client is built here with the right quota attribution.
"""
# If updating following names, also update the corresponding pydocs in
@@ -438,16 +499,29 @@ class BigQueryWrapper(object):
HISTOGRAM_METRIC_LOGGER = MetricLogger()
- def __init__(self, client=None, temp_dataset_id=None, temp_table_ref=None):
- self.client = client or BigQueryWrapper._bigquery_client(PipelineOptions())
- self.gcp_bq_client = client or gcp_bigquery.Client(
- client_info=ClientInfo(
- user_agent="apache-beam-%s" % apache_beam.__version__))
+ def __init__(
+ self,
+ client=None,
+ temp_dataset_id=None,
+ temp_table_ref=None,
+ quota_project_id=None,
+ *,
+ gcp_client=None,
+ pipeline_options=None):
+ self.quota_project_id = (
+ quota_project_id or _quota_project_id_from_options(pipeline_options))
+ self.client = client or BigQueryWrapper._bigquery_client(
+ pipeline_options or PipelineOptions(),
+ quota_project_id=self.quota_project_id)
+ # Test-only overrides. `client` also stands in for the
google-cloud-bigquery
+ # client, which several tests rely on; pass `gcp_client` as well as
`client`
+ # to override only that one.
+ self._gcp_bq_client = gcp_client or client
self._unique_row_id = 0
# For testing scenarios where we pass in a client we do not want a
# randomized prefix for row IDs.
- self._row_id_prefix = '' if client else uuid.uuid4()
+ self._row_id_prefix = '' if (client or gcp_client) else uuid.uuid4()
self._latency_histogram_metric = Metrics.histogram(
self.__class__,
'latency_histogram_ms',
@@ -474,6 +548,17 @@ class BigQueryWrapper(object):
self.created_temp_dataset = False
+ @property
+ def gcp_bq_client(self):
+ """The google-cloud-bigquery client, created on first use.
+
+ Only the streaming insert path needs it, so it is not built eagerly.
+ """
+ if self._gcp_bq_client is None:
+ self._gcp_bq_client = BigQueryWrapper._gcp_bigquery_client(
+ quota_project_id=self.quota_project_id)
+ return self._gcp_bq_client
+
@property
def unique_row_id(self):
"""Returns a unique row ID (str) used to avoid multiple insertions.
@@ -1490,19 +1575,70 @@ class BigQueryWrapper(object):
@staticmethod
def from_pipeline_options(pipeline_options: PipelineOptions):
- return BigQueryWrapper(
- client=BigQueryWrapper._bigquery_client(pipeline_options))
+ """Create a BigQueryWrapper from pipeline options.
+
+ Args:
+ pipeline_options: Pipeline options containing GCP configuration.
+ The quota_project_id is read from GoogleCloudOptions if set.
+ """
+ return BigQueryWrapper(pipeline_options=pipeline_options)
@staticmethod
- def _bigquery_client(pipeline_options: PipelineOptions):
+ def _bigquery_client(
+ pipeline_options: PipelineOptions, quota_project_id: str = None):
+ """Create a BigQuery API client from pipeline options.
+
+ Args:
+ pipeline_options: Pipeline options for credentials.
+ quota_project_id: Optional quota project ID. If not provided, will be
+ extracted from pipeline_options.
+ """
+ credentials = auth.get_service_credentials(pipeline_options)
+ # Use explicit quota_project_id if provided, otherwise get from options
+ quota_project_id = quota_project_id or _quota_project_id_from_options(
+ pipeline_options)
+ http = get_new_http()
+ if quota_project_id:
+ if credentials is None:
+ # Ignoring the request would silently bill a different project.
+ raise ValueError(
+ 'quota_project_id was set to %r, but no credentials were found to '
+ 'apply it to.' % quota_project_id)
+ # Send the quota project as a request header rather than deriving new
+ # credentials: credentials.with_quota_project() returns a copy without
+ # the cached token, and this keeps the shared credentials untouched.
+ # The header is set below the credentials layer so that it is applied on
+ # every attempt and wins over a quota project the credentials carry.
+ http = _HttpWithHeaders(
+ http, {'x-goog-user-project': quota_project_id},
+ hosts=[
+ urllib.parse.urlsplit(url).hostname for url in (
+ bigquery.BigqueryV2.BASE_URL,
+ bigquery.BigqueryV2.MTLS_BASE_URL)
+ ])
return bigquery.BigqueryV2(
- http=get_new_http(),
- credentials=auth.get_service_credentials(pipeline_options),
+ http=http,
+ credentials=credentials,
response_encoding='utf8',
additional_http_headers={
"user-agent": "apache-beam-%s" % apache_beam.__version__
})
+ @staticmethod
+ def _gcp_bigquery_client(quota_project_id: str = None):
+ """Create a google-cloud-bigquery Client with optional quota project.
+
+ The client applies the quota project to the credentials it resolves
+ itself, so no credentials are derived here.
+ """
+ client_options = None
+ if quota_project_id:
+ client_options = {'quota_project_id': quota_project_id}
+ return gcp_bigquery.Client(
+ client_options=client_options,
+ client_info=ClientInfo(
+ user_agent="apache-beam-%s" % apache_beam.__version__))
+
class RowAsDictJsonCoder(coders.Coder):
"""A coder for a table row (represented as a dict) to/from a JSON string.
diff --git a/sdks/python/apache_beam/io/gcp/bigquery_tools_test.py
b/sdks/python/apache_beam/io/gcp/bigquery_tools_test.py
index 9f8b78ee5c3..de4f3e4066a 100644
--- a/sdks/python/apache_beam/io/gcp/bigquery_tools_test.py
+++ b/sdks/python/apache_beam/io/gcp/bigquery_tools_test.py
@@ -23,6 +23,7 @@ import io
import json
import logging
import math
+import os
import re
import unittest
from typing import Optional
@@ -35,10 +36,12 @@ import pytz
from parameterized import parameterized
import apache_beam as beam
+from apache_beam.io.gcp import bigquery_tools
from apache_beam.io.gcp import resource_identifiers
from apache_beam.io.gcp.bigquery_tools import JSON_COMPLIANCE_ERROR
from apache_beam.io.gcp.bigquery_tools import AvroRowWriter
from apache_beam.io.gcp.bigquery_tools import BigQueryJobTypes
+from apache_beam.io.gcp.bigquery_tools import BigQueryWrapper
from apache_beam.io.gcp.bigquery_tools import JsonRowWriter
from apache_beam.io.gcp.bigquery_tools import RowAsDictJsonCoder
from apache_beam.io.gcp.bigquery_tools import beam_row_from_dict
@@ -299,14 +302,16 @@ class TestBigQueryWrapper(unittest.TestCase):
@mock.patch('google.cloud._http.JSONConnection.http')
def test_user_agent_insert_all(
self, http_mock, patched_skip_get_credentials, patched_sleep):
- wrapper = beam.io.gcp.bigquery_tools.BigQueryWrapper()
- try:
- wrapper._insert_all_rows('p', 'd', 't', [{'name': 'any'}], None)
- except: # pylint: disable=bare-except
- # Ignore errors. The errors come from the fact that we did not mock
- # the response from the API, so the overall insert_all_rows call fails
- # soon after the BQ API is called.
- pass
+ # Set GOOGLE_CLOUD_PROJECT to ensure Client creation succeeds in test env
+ with mock.patch.dict(os.environ, {'GOOGLE_CLOUD_PROJECT': 'test-project'}):
+ wrapper = beam.io.gcp.bigquery_tools.BigQueryWrapper()
+ try:
+ wrapper._insert_all_rows('p', 'd', 't', [{'name': 'any'}], None)
+ except: # pylint: disable=bare-except
+ # Ignore errors. The errors come from the fact that we did not mock
+ # the response from the API, so the overall insert_all_rows call fails
+ # soon after the BQ API is called.
+ pass
call = http_mock.request.mock_calls[-2]
self.assertIn('apache-beam-', call[2]['headers']['User-Agent'])
@@ -1199,8 +1204,6 @@ class TestGeographyTypeSupport(unittest.TestCase):
def test_geography_field_conversion(self):
"""Test that GEOGRAPHY fields are converted correctly."""
- from apache_beam.io.gcp.bigquery_tools import BigQueryWrapper
-
# Create a mock field with GEOGRAPHY type
field = bigquery.TableFieldSchema()
field.type = 'GEOGRAPHY'
@@ -1322,8 +1325,6 @@ class TestGeographyTypeSupport(unittest.TestCase):
def test_geography_with_special_characters(self):
"""Test GEOGRAPHY values with special characters and geometries."""
- from apache_beam.io.gcp.bigquery_tools import BigQueryWrapper
-
field = bigquery.TableFieldSchema()
field.type = 'GEOGRAPHY'
field.name = 'complex_geo'
@@ -1502,6 +1503,213 @@ class TestTypeOverrides(unittest.TestCase):
self.assertEqual(typehints_dict, [("data", Optional[dict])])
[email protected](HttpError is None, 'GCP dependencies are not installed')
+class TestBigQueryWrapperQuotaProject(unittest.TestCase):
+ """Tests for quota_project_id in BigQueryWrapper."""
+ @mock.patch(
+ 'apache_beam.io.gcp.bigquery_tools.BigQueryWrapper._bigquery_client')
+ @mock.patch(
+ 'apache_beam.io.gcp.bigquery_tools.BigQueryWrapper._gcp_bigquery_client')
+ def test_quota_project_id_stored(self, mock_gcp_client, mock_bq_client):
+ """Test that quota_project_id is stored in BigQueryWrapper."""
+ mock_bq_client.return_value = mock.Mock()
+ mock_gcp_client.return_value = mock.Mock()
+
+ wrapper = BigQueryWrapper(quota_project_id='my-billing-project')
+ self.assertEqual(wrapper.quota_project_id, 'my-billing-project')
+
+ @mock.patch(
+ 'apache_beam.io.gcp.bigquery_tools.BigQueryWrapper._bigquery_client')
+ @mock.patch(
+ 'apache_beam.io.gcp.bigquery_tools.BigQueryWrapper._gcp_bigquery_client')
+ def test_from_pipeline_options_reads_quota_from_options(
+ self, mock_gcp_client, mock_bq_client):
+ """Test from_pipeline_options reads quota_project_id from
+ GoogleCloudOptions."""
+ from apache_beam.options.pipeline_options import PipelineOptions
+
+ mock_bq_client.return_value = mock.Mock()
+ mock_gcp_client.return_value = mock.Mock()
+
+ options = PipelineOptions(['--quota_project_id=my-billing-project'])
+ wrapper = BigQueryWrapper.from_pipeline_options(options)
+
+ self.assertEqual(wrapper.quota_project_id, 'my-billing-project')
+
+ @mock.patch(
+ 'apache_beam.io.gcp.bigquery_tools.BigQueryWrapper._bigquery_client')
+ @mock.patch(
+ 'apache_beam.io.gcp.bigquery_tools.BigQueryWrapper._gcp_bigquery_client')
+ def test_from_pipeline_options_none_when_not_set(
+ self, mock_gcp_client, mock_bq_client):
+ """Test from_pipeline_options returns None when quota_project_id not
set."""
+ from apache_beam.options.pipeline_options import PipelineOptions
+
+ mock_bq_client.return_value = mock.Mock()
+ mock_gcp_client.return_value = mock.Mock()
+
+ options = PipelineOptions([])
+ wrapper = BigQueryWrapper.from_pipeline_options(options)
+
+ self.assertIsNone(wrapper.quota_project_id)
+
+ def test_http_with_headers_sets_headers_and_delegates(self):
+ from apache_beam.io.gcp.bigquery_tools import _HttpWithHeaders
+ inner = mock.MagicMock()
+ inner.request.return_value = ('response', b'content')
+ inner.redirect_codes = {301, 302}
+ http = _HttpWithHeaders(
+ inner, {'x-goog-user-project': 'p'}, hosts=['example.com'])
+
+ result = http.request(
+ 'https://example.com',
+ method='GET',
+ headers={
+ 'x-goog-user-project': 'other', 'accept': 'json'
+ },
+ redirections=3)
+
+ self.assertEqual(result, ('response', b'content'))
+ inner.request.assert_called_once_with(
+ 'https://example.com',
+ 'GET',
+ body=None,
+ headers={
+ 'x-goog-user-project': 'p', 'accept': 'json'
+ },
+ redirections=3)
+ # Attribute reads and writes reach the wrapped http.
+ self.assertEqual(http.redirect_codes, {301, 302})
+ http.redirect_codes = {301}
+ self.assertEqual(inner.redirect_codes, {301})
+
+ def test_http_with_headers_survives_copy(self):
+ import copy
+
+ from apache_beam.io.gcp.bigquery_tools import _HttpWithHeaders
+ inner = mock.MagicMock()
+ http = copy.copy(
+ _HttpWithHeaders(
+ inner, {'x-goog-user-project': 'p'}, hosts=['example.com']))
+
+ http.request('https://example.com')
+ inner.request.assert_called_once_with(
+ 'https://example.com',
+ 'GET',
+ body=None,
+ headers={'x-goog-user-project': 'p'})
+
+ def test_http_with_headers_skips_other_hosts(self):
+ from apache_beam.io.gcp.bigquery_tools import _HttpWithHeaders
+ inner = mock.MagicMock()
+ http = _HttpWithHeaders(
+ inner, {'x-goog-user-project': 'p'}, hosts=['example.com'])
+
+ http.request('https://oauth2.googleapis.com/token', headers={'a': 'b'})
+ inner.request.assert_called_once_with(
+ 'https://oauth2.googleapis.com/token',
+ 'GET',
+ body=None,
+ headers={'a': 'b'})
+
+ @mock.patch('apache_beam.io.gcp.bigquery_tools.bq_storage')
+ def test_create_bq_storage_client_with_quota_project(self, mock_bq_storage):
+ """The quota project is passed as a client option; the client applies it
+ to the credentials it resolves itself."""
+ bigquery_tools.create_bq_storage_client('my-billing-project')
+ mock_bq_storage.BigQueryReadClient.assert_called_once_with(
+ client_options={'quota_project_id': 'my-billing-project'})
+
+ @mock.patch('apache_beam.io.gcp.bigquery_tools.bq_storage')
+ def test_create_bq_storage_client_without_quota_project(
+ self, mock_bq_storage):
+ """Without a quota project the client is created with defaults."""
+ bigquery_tools.create_bq_storage_client(None)
+ mock_bq_storage.BigQueryReadClient.assert_called_once_with()
+
+ @mock.patch('apache_beam.io.gcp.bigquery_tools.auth.get_service_credentials')
+ def test_bigquery_client_raises_without_credentials(
+ self, mock_get_credentials):
+ """An explicit quota project must fail loudly rather than fall back to
+ apitools' own credential discovery, which would bill another project."""
+ from apache_beam.options.pipeline_options import PipelineOptions
+ mock_get_credentials.return_value = None
+
+ with self.assertRaisesRegex(ValueError, 'no credentials were found'):
+ BigQueryWrapper._bigquery_client(
+ PipelineOptions(), quota_project_id='requested-project')
+
+ @mock.patch('apache_beam.io.gcp.bigquery_tools.get_new_http')
+ @mock.patch('apache_beam.io.gcp.bigquery_tools.auth.get_service_credentials')
+ def test_bigquery_client_quota_project_wins_over_credentials(
+ self, mock_get_credentials, mock_get_new_http):
+ """The requested quota project must reach every BigQuery request,
+ including the retry after a 401 refresh, even when the credentials carry
+ their own, but not the token refresh requests."""
+ import httplib2
+ from google.auth import credentials as ga_credentials
+
+ from apache_beam.internal.gcp import auth
+ from apache_beam.options.pipeline_options import PipelineOptions
+
+ class CredentialsWithQuotaProject(ga_credentials.Credentials):
+ def __init__(self):
+ super().__init__()
+ self._quota_project_id = 'adc-project'
+
+ def refresh(self, request):
+ # Token refreshes go through the same Http as the API requests.
+ request('https://oauth2.googleapis.com/token', method='POST')
+ self.token = 'token'
+
+ seen_headers = []
+ token_request_headers = []
+ statuses = iter([401, 200])
+
+ class FakeHttp(object):
+ connections = {}
+ redirect_codes = set()
+
+ def request(self, uri, method='GET', body=None, headers=None, **kwargs):
+ if uri.startswith('https://oauth2.googleapis.com/'):
+ token_request_headers.append(dict(headers or {}))
+ return httplib2.Response({'status': 200}), b'{}'
+ seen_headers.append(dict(headers))
+ response = httplib2.Response({
+ 'status': next(statuses), 'content-type': 'application/json'
+ })
+ return response, b'{}'
+
+ mock_get_new_http.return_value = FakeHttp()
+ mock_get_credentials.return_value = auth._ApitoolsCredentialsAdapter(
+ CredentialsWithQuotaProject())
+
+ client = BigQueryWrapper._bigquery_client(
+ PipelineOptions(), quota_project_id='requested-project')
+ client.projects.List(bigquery.BigqueryProjectsListRequest())
+
+ self.assertEqual(len(seen_headers), 2)
+ for headers in seen_headers:
+ self.assertEqual(headers['x-goog-user-project'], 'requested-project')
+ self.assertEqual(headers['authorization'], 'Bearer token')
+ # The quota project must not be sent to the token endpoint.
+ self.assertTrue(token_request_headers)
+ for headers in token_request_headers:
+ self.assertNotIn('x-goog-user-project', headers)
+
+ @mock.patch('apache_beam.io.gcp.bigquery_tools.gcp_bigquery.Client')
+ def test_gcp_bigquery_client_passes_quota_project(self, mock_client):
+ BigQueryWrapper._gcp_bigquery_client(quota_project_id='my-billing-project')
+ self.assertEqual(
+ mock_client.call_args.kwargs['client_options'],
+ {'quota_project_id': 'my-billing-project'})
+
+ @mock.patch('apache_beam.io.gcp.bigquery_tools.gcp_bigquery.Client')
+ def test_gcp_bigquery_client_no_quota_project_by_default(self, mock_client):
+ BigQueryWrapper._gcp_bigquery_client()
+ self.assertIsNone(mock_client.call_args.kwargs['client_options'])
+
+
if __name__ == '__main__':
logging.getLogger().setLevel(logging.INFO)
unittest.main()
diff --git a/sdks/python/apache_beam/options/pipeline_options.py
b/sdks/python/apache_beam/options/pipeline_options.py
index 08ba34dabee..f0406e2505b 100644
--- a/sdks/python/apache_beam/options/pipeline_options.py
+++ b/sdks/python/apache_beam/options/pipeline_options.py
@@ -1197,6 +1197,17 @@ class GoogleCloudOptions(PipelineOptions):
action='store_true',
help='Throttling counter in GcsIO is enabled by default. Set '
'--no_gcsio_throttling_counter to avoid it.')
+ parser.add_argument(
+ '--quota_project_id',
+ default=None,
+ help='GCP project ID used for quota and billing attribution of Google '
+ 'Cloud API requests (sets the x-goog-user-project header), if '
+ 'different from the project the data resides in. Currently honored '
+ 'by BigQuery IO transforms (ReadFromBigQuery, ReadAllFromBigQuery, '
+ 'WriteToBigQuery with FILE_LOADS or STREAMING_INSERTS, '
+ 'ReadBigQueryChangeHistory); STORAGE_WRITE_API (Java-backed) and '
+ 'other GCP IOs do not use it yet. The credentials used must have the '
+ 'serviceusage.services.use permission on this project.')
parser.add_argument(
'--enable_gcsio_blob_generation',
default=False,
