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

CHANGES.md

@@ -195,6 +195,7 @@
 * (YAML) Added WriteToDatadog transform ([#38362](https://github.com/apache/beam/issues/38362)).
 * (Java) Flink 2.1 and 2.2 support is added ([#38947](https://github.com/apache/beam/issues/38947)) ([#38978](https://github.com/apache/beam/issues/38978)); Flink 1.17 and 1.18 support is dropped.
 * (Python) MqttIO is now supported in Python via cross-language ([#21060](https://github.com/apache/beam/issues/21060)).
+* 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 (Python), `--bigQueryQuotaProjectId` pipeline option (Java), and `bigqueryio.WithQuotaProject` read/query option (Go) ([#37431](https://github.com/apache/beam/issues/37431)).
 
 ## Breaking Changes
 

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,14 @@ 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 {
+func Read(s beam.Scope, project, table string, t reflect.Type, options ...func(*QueryOptions) error) beam.PCollection {
 	mustParseTable(table)
 
 	s = s.Scope("bigquery.Read")
 
 	stmt := constructSelectStatement(t, bigQueryTag, table)
 
-	return query(s, project, stmt, t)
+	return query(s, project, stmt, t, options...)
 }
 
 func constructSelectStatement(t reflect.Type, tagKey string, table string) string {
@@ -114,6 +115,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 +133,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 +153,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 {
@@ -182,7 +206,7 @@ type queryFn struct {
 }
 
 func (f *queryFn) ProcessElement(ctx context.Context, _ []byte, emit func(beam.X)) error {
-	client, err := bigquery.NewClient(ctx, f.Project)
+	client, err := bigquery.NewClient(ctx, f.Project, clientOptions(f.Options)...)
 	if err != nil {
 		return err
 	}

sdks/go/pkg/beam/io/bigqueryio/bigquery_test.go

@@ -149,6 +149,25 @@ 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) = %d options, want 0", len(got))
+	}
+	if got := clientOptions(QueryOptions{QuotaProject: "quota-project"}); len(got) != 1 {
+		t.Errorf("clientOptions(with quota) = %d options, want 1", len(got))
+	}
+}
+
 func TestWithQueryParameters(t *testing.T) {
 	t.Run("WithQueryParameters sets parameters correctly", func(t *testing.T) {
 		params := []bigquery.QueryParameter{

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

@@ -1753,7 +1753,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 =
@@ -1771,6 +1772,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));
@@ -1801,6 +1814,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())
@@ -1926,6 +1943,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();

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.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;
@@ -2173,6 +2175,58 @@ public RetryInfo parseBytes(byte[] serialized) {
     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);