From 412e221e98a312351c6d902149c098986f857237 Mon Sep 17 00:00:00 2001 From: i05nagai Date: Sun, 8 Feb 2026 15:58:41 +0900 Subject: [PATCH] Add biquery query parameters to support parametrize query --- sdks/go/pkg/beam/io/bigqueryio/bigquery.go | 29 ++++++++- .../pkg/beam/io/bigqueryio/bigquery_test.go | 27 ++++++++ sdks/go/pkg/beam/io/bigqueryio/coder.go | 41 ++++++++++++ sdks/go/pkg/beam/io/bigqueryio/coder_test.go | 62 +++++++++++++++++++ sdks/go/test/integration/integration.go | 4 +- .../io/bigqueryio/bigqueryio_test.go | 4 +- .../integration/io/bigqueryio/helper_test.go | 28 +++++---- 7 files changed, 178 insertions(+), 17 deletions(-) create mode 100644 sdks/go/pkg/beam/io/bigqueryio/coder.go create mode 100644 sdks/go/pkg/beam/io/bigqueryio/coder_test.go diff --git a/sdks/go/pkg/beam/io/bigqueryio/bigquery.go b/sdks/go/pkg/beam/io/bigqueryio/bigquery.go index a80661e22d33..8a7c8f1d7ae6 100644 --- a/sdks/go/pkg/beam/io/bigqueryio/bigquery.go +++ b/sdks/go/pkg/beam/io/bigqueryio/bigquery.go @@ -114,6 +114,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 + // Parameters are the query parameters for parameterized queries. + // bigquery.QueryParameter cannot be encoded/decoded by beam coder. + parameters []bigquery.QueryParameter } // UseStandardSQL enables BigQuery's Standard SQL dialect when executing a query. @@ -124,6 +127,14 @@ func UseStandardSQL() func(qo *QueryOptions) error { } } +// WithQueryParameters sets the query parameters for parameterized queries. +func WithQueryParameters(params ...bigquery.QueryParameter) func(qo *QueryOptions) error { + return func(qo *QueryOptions) error { + qo.parameters = params + return nil + } +} + // 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 { @@ -142,7 +153,16 @@ func query(s beam.Scope, project, query string, t reflect.Type, options ...func( } imp := beam.Impulse(s) - return beam.ParDo(s, &queryFn{Project: project, Query: query, Type: beam.EncodedType{T: t}, Options: queryOptions}, imp, beam.TypeDefinition{Var: beam.XType, T: t}) + queryParameters, err := encodeQueryParameters(queryOptions.parameters) + if err != nil { + panic(errors.Wrapf(err, "bigqueryio.Query: failed to encode query parameters")) + } + return beam.ParDo( + s, + &queryFn{Project: project, Query: query, Type: beam.EncodedType{T: t}, QueryParameters: queryParameters, Options: queryOptions}, + imp, + beam.TypeDefinition{Var: beam.XType, T: t}, + ) } type queryFn struct { @@ -152,6 +172,8 @@ type queryFn struct { Query string `json:"query"` // Type is the encoded schema type. Type beam.EncodedType `json:"type"` + // QueryParameters are serialized query parameters for parameterized queries. + QueryParameters []byte `json:"query_parameters"` // Options specifies additional query execution options. Options QueryOptions `json:"options"` } @@ -167,6 +189,11 @@ func (f *queryFn) ProcessElement(ctx context.Context, _ []byte, emit func(beam.X if !f.Options.UseStandardSQL { q.UseLegacySQL = true } + parameters, err := decodeQueryParameters(f.QueryParameters) + if err != nil { + return errors.Wrapf(err, "bigqueryio.queryFn: failed to decode query parameters") + } + q.Parameters = parameters it, err := q.Read(ctx) if err != nil { diff --git a/sdks/go/pkg/beam/io/bigqueryio/bigquery_test.go b/sdks/go/pkg/beam/io/bigqueryio/bigquery_test.go index b955136e8335..aa0169e5b23a 100644 --- a/sdks/go/pkg/beam/io/bigqueryio/bigquery_test.go +++ b/sdks/go/pkg/beam/io/bigqueryio/bigquery_test.go @@ -148,3 +148,30 @@ func Test_mustInferSchema(t *testing.T) { }) } } + +func TestWithQueryParameters(t *testing.T) { + t.Run("WithQueryParameters sets parameters correctly", func(t *testing.T) { + params := []bigquery.QueryParameter{ + {Name: "param1", Value: "value1"}, + {Name: "param2", Value: 42}, + } + + opt := WithQueryParameters(params...) + queryOpts := &QueryOptions{} + if err := opt(queryOpts); err != nil { + t.Fatalf("WithQueryParameters() failed: %v", err) + } + + if len(queryOpts.parameters) != 2 { + t.Errorf("Expected 2 parameters, got %d", len(queryOpts.parameters)) + } + + if queryOpts.parameters[0].Name != "param1" { + t.Errorf("Expected param1, got %s", queryOpts.parameters[0].Name) + } + + if queryOpts.parameters[1].Name != "param2" { + t.Errorf("Expected param2, got %s", queryOpts.parameters[1].Name) + } + }) +} diff --git a/sdks/go/pkg/beam/io/bigqueryio/coder.go b/sdks/go/pkg/beam/io/bigqueryio/coder.go new file mode 100644 index 000000000000..883ecb69ccbb --- /dev/null +++ b/sdks/go/pkg/beam/io/bigqueryio/coder.go @@ -0,0 +1,41 @@ +// Licensed to the Apache Software Foundation (ASF) under one or more +// contributor license agreements. See the NOTICE file distributed with +// this work for additional information regarding copyright ownership. +// The ASF licenses this file to You under the Apache License, Version 2.0 +// (the "License"); you may not use this file except in compliance with +// the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package bigqueryio + +import ( + "encoding/json" + "fmt" + + "cloud.google.com/go/bigquery" +) + +func encodeQueryParameters(params []bigquery.QueryParameter) ([]byte, error) { + out, err := json.Marshal(params) + if err != nil { + return nil, fmt.Errorf("error encoding JSON: %w", err) + } + + return out, nil +} + +func decodeQueryParameters(bytes []byte) ([]bigquery.QueryParameter, error) { + var out []bigquery.QueryParameter + if err := json.Unmarshal(bytes, &out); err != nil { + return nil, fmt.Errorf("error decoding JSON: %w", err) + } + + return out, nil +} diff --git a/sdks/go/pkg/beam/io/bigqueryio/coder_test.go b/sdks/go/pkg/beam/io/bigqueryio/coder_test.go new file mode 100644 index 000000000000..b64a39828eab --- /dev/null +++ b/sdks/go/pkg/beam/io/bigqueryio/coder_test.go @@ -0,0 +1,62 @@ +// Licensed to the Apache Software Foundation (ASF) under one or more +// contributor license agreements. See the NOTICE file distributed with +// this work for additional information regarding copyright ownership. +// The ASF licenses this file to You under the Apache License, Version 2.0 +// (the "License"); you may not use this file except in compliance with +// the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package bigqueryio + +import ( + "testing" + + "cloud.google.com/go/bigquery" + "github.com/google/go-cmp/cmp" +) + +func Test_encodeDecodeQueryOptions(t *testing.T) { + tests := []struct { + name string + val bigquery.QueryParameter + }{ + { + name: "Encode/decode QueryOptions with parameters", + val: bigquery.QueryParameter{ + Name: "key", + Value: "value", + }, + }, + { + name: "Encode/decode QueryOptions with nil parameter", + val: bigquery.QueryParameter{ + Name: "key", + Value: nil, + }, + }, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + encoded, err := encodeQueryParameters([]bigquery.QueryParameter{tt.val}) + if err != nil { + t.Fatalf("encodeQueryParameters() error = %v", err) + } + + decoded, err := decodeQueryParameters(encoded) + if err != nil { + t.Fatalf("decodeQueryParameters() error = %v", err) + } + + if diff := cmp.Diff(tt.val, decoded[0]); diff != "" { + t.Errorf("encode/decode mismatch (-want +got):\n%s", diff) + } + }) + } +} diff --git a/sdks/go/test/integration/integration.go b/sdks/go/test/integration/integration.go index a0eef7d10189..9a2cfa8f90e3 100644 --- a/sdks/go/test/integration/integration.go +++ b/sdks/go/test/integration/integration.go @@ -67,7 +67,7 @@ var directFilters = []string{ // The direct runner does not yet support cross-language. "TestXLang.*", "TestKafkaIO.*", - "TestBigQueryIO.*", + "TestBigQueryIO_[^W].*", "TestBigtableIO.*", "TestSpannerIO.*", "TestDebeziumIO_BasicRead", @@ -200,7 +200,7 @@ var flinkFilters = []string{ "TestTestStreamTimersEventTime", "TestTimers_EventTime_WithNoOutputTimestamp", // Encounter error: TimestampCombiner moved element from TIMESTAMP_MAX_VALUE to earlier time (end of global window) for window GlobalWindow - "TestTimers_ProcessingTime.*", // Flink doesn't support processing time timers. + "TestTimers_ProcessingTime.*", // Flink doesn't support processing time timers. // no support for BundleFinalizer "TestParDoBundleFinalizer.*", diff --git a/sdks/go/test/integration/io/bigqueryio/bigqueryio_test.go b/sdks/go/test/integration/io/bigqueryio/bigqueryio_test.go index 337083745c50..7cb8051be4b8 100644 --- a/sdks/go/test/integration/io/bigqueryio/bigqueryio_test.go +++ b/sdks/go/test/integration/io/bigqueryio/bigqueryio_test.go @@ -161,10 +161,10 @@ func TestBigQueryIO_Write(t *testing.T) { tableID := fmt.Sprintf("%s_temp_%v", "go_bqio_it", time.Now().UnixNano()) tableName := fmt.Sprintf("%s.%s", *integration.BigQueryDataset, tableID) if tt.preCreate { - newTempTable(t, tableName, ddlTestRowSchema) + newTempTable(t, tableName, ddlTestRowSchema, project) } t.Cleanup(func() { - deleteTempTable(t, tableName) + deleteTempTable(t, tableName, project) }) createTestRows := &CreateTestRowsFn{seed: time.Now().UnixNano()} p, s := beam.NewPipelineWithRoot() diff --git a/sdks/go/test/integration/io/bigqueryio/helper_test.go b/sdks/go/test/integration/io/bigqueryio/helper_test.go index e39ed87d7fff..c8404410800d 100644 --- a/sdks/go/test/integration/io/bigqueryio/helper_test.go +++ b/sdks/go/test/integration/io/bigqueryio/helper_test.go @@ -31,11 +31,11 @@ import ( // // newTable takes the name of a BigQuery dataset and a DDL schema for the data, // and generates that table with a unique suffix and an expiration time of a day later. -func newTempTable(t *testing.T, name string, schema string) { +func newTempTable(t *testing.T, name string, schema string, projectID string) { t.Helper() query := fmt.Sprintf("CREATE TABLE `%s`(%s) OPTIONS(expiration_timestamp=TIMESTAMP_ADD(CURRENT_TIMESTAMP(), INTERVAL 1 DAY))", name, schema) - cmd := exec.Command("bq", "query", "--use_legacy_sql=false", query) + cmd := exec.Command("bq", "query", fmt.Sprintf("--project_id=%s", projectID), "--use_legacy_sql=false", query) _, err := cmd.CombinedOutput() if err != nil { t.Fatalf("error creating BigQuery table: %v", err) @@ -45,11 +45,11 @@ func newTempTable(t *testing.T, name string, schema string) { // deleteTable deletes a BigQuery table using BigQuery's Data Definition Language (DDL) and the // "bq query" console command. Reference: https://cloud.google.com/bigquery/docs/reference/standard-sql/data-definition-language -func deleteTempTable(t *testing.T, table string) { +func deleteTempTable(t *testing.T, table string, projectID string) { t.Helper() query := fmt.Sprintf("DROP TABLE IF EXISTS `%s`", table) - cmd := exec.Command("bq", "query", "--use_legacy_sql=false", query) + cmd := exec.Command("bq", "query", fmt.Sprintf("--project_id=%s", projectID), "--use_legacy_sql=false", query) t.Logf("Deleting BigQuery table %v", table) _, err := cmd.CombinedOutput() if err != nil { @@ -66,16 +66,20 @@ func checkTableExistsAndNonEmpty(ctx context.Context, t *testing.T, project, tab } defer client.Close() - tableRef := client.Dataset(*integration.BigQueryDataset).Table(tableID) - metadata, err := tableRef.Metadata(ctx) + q := client.Query(fmt.Sprintf("SELECT COUNT(*) FROM `%s.%s.%s`", project, *integration.BigQueryDataset, tableID)) + it, err := q.Read(ctx) if err != nil { - t.Fatalf("unable to find table: %v", err) + t.Fatalf("error querying table row count: %v", err) } - streamingBuffer := metadata.StreamingBuffer - if streamingBuffer == nil { - t.Fatalf("there's no streaming buffer for the table") + var row []bigquery.Value + if err := it.Next(&row); err != nil { + t.Fatalf("error reading row count result: %v", err) } - if streamingBuffer.EstimatedRows != inputSize { - t.Fatalf("streamingBuffer.EstimatedRows = %v, want %v", streamingBuffer.EstimatedRows, inputSize) + count, ok := row[0].(int64) + if !ok { + t.Fatalf("unexpected type for row count: %T", row[0]) + } + if count != inputSize { + t.Fatalf("row count = %v, want %v", count, inputSize) } }