diff --git a/CHANGES.md b/CHANGES.md index f0d5d06b9d9f..db85c8e3c0b8 100644 --- a/CHANGES.md +++ b/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 diff --git a/sdks/go/pkg/beam/io/bigqueryio/bigquery.go b/sdks/go/pkg/beam/io/bigqueryio/bigquery.go index bbe951969b92..81414a785a19 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,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. 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. 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 } diff --git a/sdks/go/pkg/beam/io/bigqueryio/bigquery_test.go b/sdks/go/pkg/beam/io/bigqueryio/bigquery_test.go index 9f83eeb847f7..07cedbe04e2a 100644 --- a/sdks/go/pkg/beam/io/bigqueryio/bigquery_test.go +++ b/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{ 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 face2ef5841a..9a0b61170808 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 3c6ee776b67e..d5cda9d61b11 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 @@ -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 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 3902fb1fca33..b076e45b72fb 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.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); diff --git a/sdks/python/apache_beam/io/gcp/bigquery.py b/sdks/python/apache_beam/io/gcp/bigquery.py index 38acd29da7d9..6b40614cd56a 100644 --- a/sdks/python/apache_beam/io/gcp/bigquery.py +++ b/sdks/python/apache_beam/io/gcp/bigquery.py @@ -716,7 +716,8 @@ def __init__( 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.' @@ -750,6 +751,7 @@ def __init__( 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 @@ -769,12 +771,13 @@ def display_data(self): 'use_legacy_sql': self.use_legacy_sql, 'bigquery_job_labels': json.dumps(self.bigquery_job_labels), 'export_file_format': export_format, + 'quota_project_id': self._get_quota_project_id() or '', '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 @@ -836,6 +839,29 @@ def _get_project(self): project = self.project return project + def _get_quota_project_id(self): + """Returns the quota project ID for API calls. + + Prefers the explicit quota_project_id parameter, falls back to + quota_project_id from GoogleCloudOptions. + """ + if self.quota_project_id: + return self.quota_project_id + if self.options is not None: + return self.options.view_as(GoogleCloudOptions).quota_project_id + return None + + 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) @@ -850,10 +876,9 @@ def _create_source(self, path, coder): 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) @@ -986,6 +1011,22 @@ def _export_files(self, bq): return table.schema, metadata_list +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 _CustomBigQueryStorageSource(BoundedSource): """A base class for BoundedSource implementations which read from BigQuery using the BigQuery Storage API. @@ -1043,7 +1084,8 @@ def __init__( 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( @@ -1082,6 +1124,7 @@ def __init__( 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.""" @@ -1093,6 +1136,29 @@ def _get_project(self): return project return self.project + def _get_quota_project_id(self): + """Returns the quota project ID for API calls. + + Prefers the explicit quota_project_id parameter, falls back to + quota_project_id from GoogleCloudOptions. + """ + if self.quota_project_id: + return self.quota_project_id + if self.pipeline_options is not None: + return self.pipeline_options.view_as(GoogleCloudOptions).quota_project_id + return None + + 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: @@ -1171,8 +1237,7 @@ def display_data(self): 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 @@ -1219,10 +1284,8 @@ def estimate_size(self): 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) @@ -1255,7 +1318,7 @@ def split(self, desired_bundle_size, start_position=None, stop_position=None): if self.row_restriction is not None: requested_session.read_options.row_restriction = self.row_restriction - storage_client = bq_storage.BigQueryReadClient() + storage_client = _create_bq_storage_client(self._get_quota_project_id()) stream_count = 0 if desired_bundle_size > 0: table_size = self._get_table_size(bq, self.table_reference) @@ -1286,8 +1349,10 @@ def split(self, desired_bundle_size, start_position=None, stop_position=None): 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: @@ -1321,10 +1386,12 @@ def __init__( 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 { @@ -1347,7 +1414,10 @@ def split(self, desired_bundle_size, start_position=None, stop_position=None): 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) @@ -1383,7 +1453,7 @@ def retry_delay_callback(delay): def read_arrow(self): - storage_client = bq_storage.BigQueryReadClient() + storage_client = _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 @@ -1402,7 +1472,7 @@ def read_arrow(self): yield py_row def read_avro(self): - storage_client = bq_storage.BigQueryReadClient() + storage_client = _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 @@ -3024,6 +3094,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 @@ -3124,10 +3200,12 @@ def _expand_output_type(self, output_pcollection): '%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( @@ -3209,7 +3287,9 @@ def _get_pipeline_details(unused_elm): 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_read_internal.py b/sdks/python/apache_beam/io/gcp/bigquery_read_internal.py index 136b3cc56b7e..2fa95a3d1794 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 @@ def process(self, element): 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 @@ def _get_temp_dataset_project(self): 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 dcadee7f6a1a..a91221ca3cd4 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 _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 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 @@ def test_repeatable_field_is_properly_converted(self): 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 @@ def tearDown(self): @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 @@ -827,6 +834,289 @@ def test_expand_output_type_uses_query_schema(self): mock_convert.assert_called_once_with(schema, None) +@unittest.skipIf( + 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.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') + display_data = source.display_data() + self.assertEqual(display_data['quota_project_id'], 'my-billing-project') + + def test_quota_project_id_empty_in_display_data_when_not_set(self): + """Test that quota_project_id is empty string in display data + when not set.""" + options = PipelineOptions([]) + source = beam_bq._CustomBigQuerySource( + method=ReadFromBigQuery.Method.EXPORT, + table='project:dataset.table', + pipeline_options=options) + display_data = source.display_data() + self.assertEqual(display_data['quota_project_id'], '') + + 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._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._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') + + @mock.patch('apache_beam.io.gcp.bigquery.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.""" + beam_bq._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.bq_storage') + def test_create_bq_storage_client_without_quota_project( + self, mock_bq_storage): + """Test _create_bq_storage_client without quota project uses default.""" + beam_bq._create_bq_storage_client(None) + mock_bq_storage.BigQueryReadClient.assert_called_once_with() + + @unittest.skipIf(HttpError is None, 'GCP dependencies are not installed') class TestBigQuerySink(unittest.TestCase): def test_table_spec_display_data(self): @@ -868,10 +1158,14 @@ def _cleanup_files(self): 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 @@ -1429,6 +1723,13 @@ def test_copy_load_job_exception(self, exception_type, error_message): 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 0d62ec5233c1..d10cfccefee5 100644 --- a/sdks/python/apache_beam/io/gcp/bigquery_tools.py +++ b/sdks/python/apache_beam/io/gcp/bigquery_tools.py @@ -59,6 +59,7 @@ 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 @@ -346,6 +347,35 @@ 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 + + +class _HttpWithHeaders(object): + """Wraps an httplib2.Http to add fixed headers to every request. + + Sits below the credentials layer, so the headers are applied on every + attempt and take precedence over headers set by the credentials. + """ + def __init__(self, http, headers): + object.__setattr__(self, '_http', http) + object.__setattr__(self, '_headers', headers) + + def request(self, uri, method='GET', body=None, headers=None, **kwargs): + headers = dict(headers or {}) + headers.update(self._headers) + return self._http.request(uri, method, body=body, headers=headers, **kwargs) + + def __getattr__(self, 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. @@ -354,8 +384,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 @@ -365,16 +397,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', @@ -401,6 +446,17 @@ def __init__(self, client=None, temp_dataset_id=None, temp_table_ref=None): 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. @@ -1414,19 +1470,64 @@ def convert_row_to_dict(self, row, schema): @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}) 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 078c42160941..8b9e265f939b 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 json import logging import math +import os import re import unittest from typing import Optional @@ -39,6 +40,7 @@ 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 @@ -236,14 +238,16 @@ def test_delete_dataset_retries_for_timeouts(self, patched_time_sleep): @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']) @@ -1106,8 +1110,6 @@ def test_geography_in_bigquery_type_mapping(self): 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' @@ -1229,8 +1231,6 @@ def test_geography_json_encoding(self): 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' @@ -1409,6 +1409,157 @@ def test_type_overrides_json_to_dict(self): self.assertEqual(typehints_dict, [("data", Optional[dict])]) +@unittest.skipIf(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'}) + + 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}) + + @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 request, including the + retry after a 401 refresh, even when the credentials carry their own.""" + 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): + self.token = 'token' + + seen_headers = [] + statuses = iter([401, 200]) + + class FakeHttp(object): + connections = {} + redirect_codes = set() + + def request(self, uri, method='GET', body=None, headers=None, **kwargs): + 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') + + @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 9c460f7678e4..12928bab5133 100644 --- a/sdks/python/apache_beam/options/pipeline_options.py +++ b/sdks/python/apache_beam/options/pipeline_options.py @@ -1189,6 +1189,14 @@ def _add_argparse_args(cls, parser): 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 to use for quota and billing purposes. ' + 'If not specified, the project associated with the credentials ' + 'will be used for quota. This is useful when running pipelines ' + 'that access resources in a different project than the one ' + 'associated with the credentials.') parser.add_argument( '--enable_gcsio_blob_generation', default=False,