Support attributing BigQuery quota and billing to a custom GCP project (Python, Java, Go)

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/apache/beam/issues/37431)).
 
 ## New Features / Improvements
 

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,22 +161,38 @@ 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 {
 	s = s.Scope("bigquery.Query")
 	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
 	}

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"))
+	})
+}

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();

sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryServicesImpl.java

@@ -1800,7 +1800,8 @@ private static boolean nextBackOff(Sleeper sleeper, BackOff backoff) throws Inte
   }
 
   /** 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 @@ private static Bigquery.Builder newBigQueryClient(BigQueryOptions options) {
     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 @@ private static BigQueryWriteClient newBigQueryWriteClient(BigQueryOptions option
       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 void onRetryAttempt(Status status, Metadata metadata) {
       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();