diff --git a/CHANGES.md b/CHANGES.md index 537c3a32447c..d7328eda2933 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/apache/beam/issues/37431)). ## New Features / Improvements diff --git a/sdks/go/pkg/beam/io/bigqueryio/bigquery.go b/sdks/go/pkg/beam/io/bigqueryio/bigquery.go index bbe951969b92..98818ccf1af1 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. 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. 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 9f83eeb847f7..4cf0c46b9f62 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 1826f4a1fd21..64382d64bb51 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 7a924c6aa021..a921be3700ec 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 @@ 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 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 90b2e1fb8723..642d35b43905 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; @@ -2311,6 +2313,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 2a52b31edf7f..f3d64bfcfa6c 100644 --- a/sdks/python/apache_beam/io/gcp/bigquery.py +++ b/sdks/python/apache_beam/io/gcp/bigquery.py @@ -720,7 +720,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.' @@ -754,6 +755,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 @@ -773,12 +775,15 @@ 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': 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 @@ def _get_project(self): 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 @@ 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) @@ -1051,7 +1072,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( @@ -1090,6 +1112,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.""" @@ -1101,6 +1124,23 @@ def _get_project(self): 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 @@ def display_data(self): '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 @@ 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) @@ -1278,7 +1318,8 @@ 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 = 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 @@ 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: @@ -1337,10 +1380,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 { @@ -1363,7 +1408,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) @@ -1399,7 +1447,8 @@ def retry_delay_callback(delay): 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 @@ def read_arrow(self): 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 @@ def __init__( 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 @@ def __init__( 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 @@ def start_bundle(self): 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 @@ def process(self, element): 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 @@ def expand(self, input): 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 @@ 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( @@ -3228,7 +3293,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_change_history.py b/sdks/python/apache_beam/io/gcp/bigquery_change_history.py index 1e90c2974778..9f836190d9ca 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 @@ def __init__( 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 @@ def __init__( 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 @@ def __init__( 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 @@ def __init__( 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 @@ def __init__( 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 @@ def _ensure_client(self) -> None: """ 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 @@ def expand(self, pbegin: beam.pvalue.PBegin) -> beam.PCollection: '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 @@ def expand(self, pbegin: beam.pvalue.PBegin) -> beam.PCollection: 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 @@ def expand(self, pbegin: beam.pvalue.PBegin) -> beam.PCollection: 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 @@ def expand(self, pbegin: beam.pvalue.PBegin) -> beam.PCollection: 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 @@ def expand(self, pbegin: beam.pvalue.PBegin) -> beam.PCollection: _ = ( 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 11b467f26d49..e1b1eb69ae38 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 @@ - 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 @@ def test_default_buffer(self): self.assertEqual(t._buffer_sec, 10) +@unittest.skipIf(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 dabe80c5c27e..79940d109f07 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 @@ def __init__( 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 @@ def __init__( 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 @@ def setup(self): 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 @@ def __init__( 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 @@ def display_data(self): 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 @@ def process(self, element): 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 @@ def __init__( 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 @@ def _load_data( 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 @@ def _load_data( 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 @@ def _load_data( 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 @@ def _load_data( | "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 @@ def _load_data( 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 @@ def _load_data( 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 7ae0ccc8e32f..7ccda7ecad8d 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 @@ def __call__(self): 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 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 22b617fa8741..ae03104612c9 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 @@ -879,6 +886,343 @@ 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, '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') + + +@unittest.skipIf(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 @@ 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 @@ -1481,6 +1829,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 cc1a16a659fe..d1412d0f855e 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 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.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 @@ 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 @@ 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. @@ -1490,19 +1575,70 @@ 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}, + 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 9f8b78ee5c3b..de4f3e4066af 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 @@ -35,10 +36,12 @@ 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 @@ 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']) @@ -1199,8 +1204,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' @@ -1322,8 +1325,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' @@ -1502,6 +1503,213 @@ 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'}, 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 08ba34dabee4..f0406e2505bc 100644 --- a/sdks/python/apache_beam/options/pipeline_options.py +++ b/sdks/python/apache_beam/options/pipeline_options.py @@ -1197,6 +1197,17 @@ 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 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,