From 412e221e98a312351c6d902149c098986f857237 Mon Sep 17 00:00:00 2001 From: i05nagai Date: Sun, 8 Feb 2026 15:58:41 +0900 Subject: [PATCH 1/3] 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) } } From c4675c6ee8b0b7543a351265890cc058be5f823d Mon Sep 17 00:00:00 2001 From: i05nagai Date: Wed, 22 Jul 2026 07:52:54 +0900 Subject: [PATCH 2/3] Use encoding/gob for encode/decode to preserve type information QueryParameter and QueryParameterValue contains filed with any type. Converting go types to encoding/json lose type informations some of go types. The registered types are ones supported in the implementation in paramValue method. https://github.com/googleapis/google-cloud-go/blob/bigquery/v1.79.0/bigquery/params.go#L434 --- sdks/go.mod | 2 +- sdks/go/pkg/beam/io/bigqueryio/bigquery.go | 5 +- .../pkg/beam/io/bigqueryio/bigquery_test.go | 8 + sdks/go/pkg/beam/io/bigqueryio/coder.go | 73 ++- sdks/go/pkg/beam/io/bigqueryio/coder_test.go | 439 +++++++++++++++++- sdks/go/test/integration/integration.go | 2 +- .../io/bigqueryio/bigqueryio_test.go | 178 ++++++- .../integration/io/bigqueryio/helper_test.go | 45 ++ 8 files changed, 718 insertions(+), 34 deletions(-) diff --git a/sdks/go.mod b/sdks/go.mod index d0660291836e..87468c092fae 100644 --- a/sdks/go.mod +++ b/sdks/go.mod @@ -69,6 +69,7 @@ require ( ) require ( + cloud.google.com/go v0.123.0 github.com/avast/retry-go/v4 v4.7.0 github.com/fsouza/fake-gcs-server v1.52.3 github.com/golang-cz/devslog v0.0.17 @@ -139,7 +140,6 @@ require ( ) require ( - cloud.google.com/go v0.123.0 // indirect cloud.google.com/go/compute/metadata v0.9.0 // indirect cloud.google.com/go/iam v1.11.0 // indirect cloud.google.com/go/longrunning v1.2.0 // indirect diff --git a/sdks/go/pkg/beam/io/bigqueryio/bigquery.go b/sdks/go/pkg/beam/io/bigqueryio/bigquery.go index 8a7c8f1d7ae6..bbe951969b92 100644 --- a/sdks/go/pkg/beam/io/bigqueryio/bigquery.go +++ b/sdks/go/pkg/beam/io/bigqueryio/bigquery.go @@ -115,7 +115,9 @@ 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. + // In the current implementation, user-defines types are not supported in Value field. + // Use *bigquery.QueryParameterValue to build STRUCT/ARRAY parameters + // or use go primitive types explicitly. parameters []bigquery.QueryParameter } @@ -128,6 +130,7 @@ func UseStandardSQL() func(qo *QueryOptions) error { } // 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 { return func(qo *QueryOptions) error { qo.parameters = params diff --git a/sdks/go/pkg/beam/io/bigqueryio/bigquery_test.go b/sdks/go/pkg/beam/io/bigqueryio/bigquery_test.go index aa0169e5b23a..9f83eeb847f7 100644 --- a/sdks/go/pkg/beam/io/bigqueryio/bigquery_test.go +++ b/sdks/go/pkg/beam/io/bigqueryio/bigquery_test.go @@ -170,8 +170,16 @@ func TestWithQueryParameters(t *testing.T) { t.Errorf("Expected param1, got %s", queryOpts.parameters[0].Name) } + if queryOpts.parameters[0].Value != "value1" { + t.Errorf("Expected value1, got %s", queryOpts.parameters[0].Value) + } + if queryOpts.parameters[1].Name != "param2" { t.Errorf("Expected param2, got %s", queryOpts.parameters[1].Name) } + + if queryOpts.parameters[1].Value != 42 { + t.Errorf("Expected 42, got %v", queryOpts.parameters[1].Value) + } }) } diff --git a/sdks/go/pkg/beam/io/bigqueryio/coder.go b/sdks/go/pkg/beam/io/bigqueryio/coder.go index 883ecb69ccbb..cabd6e0dab45 100644 --- a/sdks/go/pkg/beam/io/bigqueryio/coder.go +++ b/sdks/go/pkg/beam/io/bigqueryio/coder.go @@ -16,26 +16,73 @@ package bigqueryio import ( - "encoding/json" - "fmt" + "bytes" + "encoding/gob" + "io" + "math/big" + "time" "cloud.google.com/go/bigquery" + "cloud.google.com/go/civil" + + "github.com/apache/beam/sdks/v2/go/pkg/beam/internal/errors" ) +// init registers the concrete types that bigquery.QueryParameter.Value may hold +// so gob can encode/decode them through the interface{} field. +func init() { + // bigquery's paramType/paramValue only recognize + // *bigquery.QueryParameterValue as an explicitly-typed parameter. + gob.Register(&bigquery.QueryParameterValue{}) + gob.Register(bigquery.NullInt64{}) + gob.Register(bigquery.NullFloat64{}) + gob.Register(bigquery.NullString{}) + gob.Register(bigquery.NullBool{}) + gob.Register(bigquery.NullTimestamp{}) + gob.Register(bigquery.NullDate{}) + gob.Register(bigquery.NullTime{}) + gob.Register(bigquery.NullDateTime{}) + gob.Register(bigquery.NullGeography{}) + gob.Register(bigquery.NullJSON{}) + gob.Register(civil.Date{}) + gob.Register(civil.Time{}) + gob.Register(civil.DateTime{}) + gob.Register(time.Time{}) + gob.Register(&big.Rat{}) + gob.Register(&bigquery.IntervalValue{}) + gob.Register(&bigquery.RangeValue{}) +} + func encodeQueryParameters(params []bigquery.QueryParameter) ([]byte, error) { - out, err := json.Marshal(params) - if err != nil { - return nil, fmt.Errorf("error encoding JSON: %w", err) + if params == nil { + return []byte{}, nil } - - return out, nil + // validate each element to tell which paramaeter is unsupported. + for _, p := range params { + if err := gob.NewEncoder(io.Discard).Encode([]bigquery.QueryParameter{p}); err != nil { + return nil, errors.Errorf( + "bigqueryio: query parameter %q has unsupported value type %T (%v). "+ + "WithQueryParameters only supports bool, string, numeric types, []byte, "+ + "time.Time, civil.Date/Time/DateTime, *big.Rat, bigquery.Null* types, "+ + "*bigquery.IntervalValue, *bigquery.RangeValue, and *bigquery.QueryParameterValue. "+ + "For STRUCT/ARRAY parameters or other custom types, build a "+ + "*bigquery.QueryParameterValue explicitly instead", p.Name, p.Value, err) + } + } + var buf bytes.Buffer + if err := gob.NewEncoder(&buf).Encode(params); err != nil { + return nil, err + } + return buf.Bytes(), 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) +func decodeQueryParameters(data []byte) ([]bigquery.QueryParameter, error) { + if len(data) == 0 { + return []bigquery.QueryParameter{}, nil } - - return out, nil + var params []bigquery.QueryParameter + if err := gob.NewDecoder(bytes.NewReader(data)).Decode(¶ms); err != nil { + return nil, err + } + return params, nil } diff --git a/sdks/go/pkg/beam/io/bigqueryio/coder_test.go b/sdks/go/pkg/beam/io/bigqueryio/coder_test.go index b64a39828eab..b569e9400b6d 100644 --- a/sdks/go/pkg/beam/io/bigqueryio/coder_test.go +++ b/sdks/go/pkg/beam/io/bigqueryio/coder_test.go @@ -16,31 +16,422 @@ package bigqueryio import ( + "math/big" + "reflect" + "strings" "testing" + "time" "cloud.google.com/go/bigquery" + "cloud.google.com/go/civil" "github.com/google/go-cmp/cmp" ) +// bigRatComparer lets cmp.Diff compare *big.Rat values. +var bigRatComparer = cmp.Comparer(func(x, y *big.Rat) bool { + if x == nil || y == nil { + return x == y + } + return x.Cmp(y) == 0 +}) + func Test_encodeDecodeQueryOptions(t *testing.T) { tests := []struct { name string val bigquery.QueryParameter }{ { - name: "Encode/decode QueryOptions with parameters", + name: "string", val: bigquery.QueryParameter{ Name: "key", Value: "value", }, }, { - name: "Encode/decode QueryOptions with nil parameter", + name: "nil", val: bigquery.QueryParameter{ Name: "key", Value: nil, }, }, + { + name: "int", + val: bigquery.QueryParameter{ + Name: "key", + Value: 50, + }, + }, + { + name: "int8", + val: bigquery.QueryParameter{ + Name: "key", + Value: int8(50), + }, + }, + { + name: "int16", + val: bigquery.QueryParameter{ + Name: "key", + Value: int16(50), + }, + }, + { + name: "int32", + val: bigquery.QueryParameter{ + Name: "key", + Value: int32(50), + }, + }, + { + name: "int64", + val: bigquery.QueryParameter{ + Name: "key", + Value: int64(50), + }, + }, + { + name: "uint8", + val: bigquery.QueryParameter{ + Name: "key", + Value: uint8(50), + }, + }, + { + name: "uint16", + val: bigquery.QueryParameter{ + Name: "key", + Value: uint16(50), + }, + }, + { + name: "uint32", + val: bigquery.QueryParameter{ + Name: "key", + Value: uint32(50), + }, + }, + { + name: "float32", + val: bigquery.QueryParameter{ + Name: "key", + Value: float32(50.5), + }, + }, + { + name: "float64", + val: bigquery.QueryParameter{ + Name: "key", + Value: 50.0, + }, + }, + { + name: "bool", + val: bigquery.QueryParameter{ + Name: "key", + Value: true, + }, + }, + { + name: "[]byte", + val: bigquery.QueryParameter{ + Name: "key", + Value: []byte("value"), + }, + }, + { + name: "[]int slice", + val: bigquery.QueryParameter{ + Name: "key", + Value: []int{1, 2, 3}, + }, + }, + { + name: "[]string slice", + val: bigquery.QueryParameter{ + Name: "key", + Value: []string{"a", "b", "c"}, + }, + }, + { + name: "time.Time", + val: bigquery.QueryParameter{ + Name: "key", + Value: time.Date(2024, 1, 2, 3, 4, 5, 0, time.UTC), + }, + }, + { + name: "civil.Date", + val: bigquery.QueryParameter{ + Name: "key", + Value: civil.Date{Year: 2024, Month: 1, Day: 2}, + }, + }, + { + name: "civil.Time", + val: bigquery.QueryParameter{ + Name: "key", + Value: civil.Time{Hour: 3, Minute: 4, Second: 5}, + }, + }, + { + name: "civil.DateTime", + val: bigquery.QueryParameter{ + Name: "key", + Value: civil.DateTime{ + Date: civil.Date{Year: 2024, Month: 1, Day: 2}, + Time: civil.Time{Hour: 3, Minute: 4, Second: 5}, + }, + }, + }, + { + name: "*big.Rat", + val: bigquery.QueryParameter{ + Name: "key", + Value: big.NewRat(22, 7), + }, + }, + { + name: "*bigquery.IntervalValue", + val: bigquery.QueryParameter{ + Name: "key", + Value: &bigquery.IntervalValue{ + Years: 1, + Months: 2, + Days: 3, + Hours: 4, + Minutes: 5, + Seconds: 6, + SubSecondNanos: 7, + }, + }, + }, + { + name: "*bigquery.RangeValue", + val: bigquery.QueryParameter{ + Name: "key", + Value: &bigquery.RangeValue{ + Start: civil.Date{Year: 2024, Month: 1, Day: 1}, + End: civil.Date{Year: 2024, Month: 12, Day: 31}, + }, + }, + }, + { + name: "NullInt64 valid", + val: bigquery.QueryParameter{ + Name: "key", + Value: bigquery.NullInt64{Int64: 50, Valid: true}, + }, + }, + { + name: "NullInt64 null", + val: bigquery.QueryParameter{ + Name: "key", + Value: bigquery.NullInt64{Valid: false}, + }, + }, + { + name: "NullFloat64 valid", + val: bigquery.QueryParameter{ + Name: "key", + Value: bigquery.NullFloat64{Float64: 50.5, Valid: true}, + }, + }, + { + name: "NullFloat64 null", + val: bigquery.QueryParameter{ + Name: "key", + Value: bigquery.NullFloat64{Valid: false}, + }, + }, + { + name: "NullString valid", + val: bigquery.QueryParameter{ + Name: "key", + Value: bigquery.NullString{StringVal: "value", Valid: true}, + }, + }, + { + name: "NullString null", + val: bigquery.QueryParameter{ + Name: "key", + Value: bigquery.NullString{Valid: false}, + }, + }, + { + name: "NullBool valid", + val: bigquery.QueryParameter{ + Name: "key", + Value: bigquery.NullBool{Bool: true, Valid: true}, + }, + }, + { + name: "NullBool null", + val: bigquery.QueryParameter{ + Name: "key", + Value: bigquery.NullBool{Valid: false}, + }, + }, + { + name: "NullTimestamp valid", + val: bigquery.QueryParameter{ + Name: "key", + Value: bigquery.NullTimestamp{Timestamp: time.Date(2024, 1, 2, 3, 4, 5, 0, time.UTC), Valid: true}, + }, + }, + { + name: "NullTimestamp null", + val: bigquery.QueryParameter{ + Name: "key", + Value: bigquery.NullTimestamp{Valid: false}, + }, + }, + { + name: "NullDate valid", + val: bigquery.QueryParameter{ + Name: "key", + Value: bigquery.NullDate{Date: civil.Date{Year: 2024, Month: 1, Day: 2}, Valid: true}, + }, + }, + { + name: "NullDate null", + val: bigquery.QueryParameter{ + Name: "key", + Value: bigquery.NullDate{Valid: false}, + }, + }, + { + name: "NullTime valid", + val: bigquery.QueryParameter{ + Name: "key", + Value: bigquery.NullTime{Time: civil.Time{Hour: 3, Minute: 4, Second: 5}, Valid: true}, + }, + }, + { + name: "NullTime null", + val: bigquery.QueryParameter{ + Name: "key", + Value: bigquery.NullTime{Valid: false}, + }, + }, + { + name: "NullDateTime valid", + val: bigquery.QueryParameter{ + Name: "key", + Value: bigquery.NullDateTime{ + DateTime: civil.DateTime{ + Date: civil.Date{Year: 2024, Month: 1, Day: 2}, + Time: civil.Time{Hour: 3, Minute: 4, Second: 5}, + }, + Valid: true, + }, + }, + }, + { + name: "NullDateTime null", + val: bigquery.QueryParameter{ + Name: "key", + Value: bigquery.NullDateTime{Valid: false}, + }, + }, + { + name: "NullGeography valid", + val: bigquery.QueryParameter{ + Name: "key", + Value: bigquery.NullGeography{GeographyVal: "POINT(1 2)", Valid: true}, + }, + }, + { + name: "NullGeography null", + val: bigquery.QueryParameter{ + Name: "key", + Value: bigquery.NullGeography{Valid: false}, + }, + }, + { + name: "NullJSON valid", + val: bigquery.QueryParameter{ + Name: "key", + Value: bigquery.NullJSON{JSONVal: `{"a":1}`, Valid: true}, + }, + }, + { + name: "NullJSON null", + val: bigquery.QueryParameter{ + Name: "key", + Value: bigquery.NullJSON{Valid: false}, + }, + }, + { + name: "QueryParameterValue BIGNUMERIC", + val: bigquery.QueryParameter{ + Name: "key", + Value: &bigquery.QueryParameterValue{ + Type: bigquery.StandardSQLDataType{ + TypeKind: "BIGNUMERIC", + }, + Value: "12345678901234567890123456789012345678901234567890.12345678901234567890123456789012345678901234567890", + }, + }, + }, + { + name: "QueryParameterValue ARRAY of STRUCT", + val: bigquery.QueryParameter{ + Name: "key", + Value: &bigquery.QueryParameterValue{ + Type: bigquery.StandardSQLDataType{ + ArrayElementType: &bigquery.StandardSQLDataType{ + StructType: &bigquery.StandardSQLStructType{ + Fields: []*bigquery.StandardSQLField{ + { + Name: "NumberField", + Type: &bigquery.StandardSQLDataType{ + TypeKind: "INT64", + }, + }, + }, + }, + }, + }, + ArrayValue: []bigquery.QueryParameterValue{ + {StructValue: map[string]bigquery.QueryParameterValue{ + "NumberField": { + Value: int64(42), + }, + }}, + {StructValue: map[string]bigquery.QueryParameterValue{ + "NumberField": { + Value: int64(43), + }, + }}, + }, + }, + }, + }, + { + name: "QueryParameterValue STRUCT", + val: bigquery.QueryParameter{ + Name: "key", + Value: &bigquery.QueryParameterValue{ + Type: bigquery.StandardSQLDataType{ + StructType: &bigquery.StandardSQLStructType{ + Fields: []*bigquery.StandardSQLField{ + { + Name: "NumberField", + Type: &bigquery.StandardSQLDataType{ + TypeKind: "INT64", + }, + }, + }, + }, + }, + StructValue: map[string]bigquery.QueryParameterValue{ + "NumberField": { + Value: int64(42), + }, + }, + }, + }, + }, } for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { @@ -54,9 +445,51 @@ func Test_encodeDecodeQueryOptions(t *testing.T) { t.Fatalf("decodeQueryParameters() error = %v", err) } - if diff := cmp.Diff(tt.val, decoded[0]); diff != "" { + gotType := reflect.TypeOf(decoded[0].Value) + wantType := reflect.TypeOf(tt.val.Value) + if gotType != wantType { + t.Errorf("type not preserved: want %v, got %v", wantType, gotType) + } + + if diff := cmp.Diff(tt.val, decoded[0], bigRatComparer); diff != "" { t.Errorf("encode/decode mismatch (-want +got):\n%s", diff) } }) } } + +// customParam is a struct type deliberately not registered with gob. +type customParam struct { + Foo string + Bar int64 +} + +func TestEncodeQueryParameters_UnsupportedType(t *testing.T) { + t.Parallel() + _, err := encodeQueryParameters([]bigquery.QueryParameter{ + {Name: "custom", Value: customParam{Foo: "a", Bar: 1}}, + }) + if err == nil { + t.Fatal("encodeQueryParameters() error = nil, want error for unsupported type") + } + if !strings.Contains(err.Error(), `"custom"`) { + t.Errorf("encodeQueryParameters() error = %q, want it to name the parameter %q", err.Error(), "custom") + } + if !strings.Contains(err.Error(), "bigqueryio.customParam") { + t.Errorf("encodeQueryParameters() error = %q, want it to name the type %q", err.Error(), "bigqueryio.customParam") + } +} + +func TestEncodeQueryParameters_UnsupportedTypeAmongSupported(t *testing.T) { + t.Parallel() + _, err := encodeQueryParameters([]bigquery.QueryParameter{ + {Name: "ok", Value: "value1"}, + {Name: "bad", Value: customParam{Foo: "a", Bar: 1}}, + }) + if err == nil { + t.Fatal("encodeQueryParameters() error = nil, want error for unsupported type") + } + if !strings.Contains(err.Error(), `"bad"`) { + t.Errorf("encodeQueryParameters() error = %q, want it to name the offending parameter %q, not the supported one", err.Error(), "bad") + } +} diff --git a/sdks/go/test/integration/integration.go b/sdks/go/test/integration/integration.go index 9a2cfa8f90e3..535a7fb8413c 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_[^W].*", + "TestBigQueryIO_[^WQ].*", "TestBigtableIO.*", "TestSpannerIO.*", "TestDebeziumIO_BasicRead", diff --git a/sdks/go/test/integration/io/bigqueryio/bigqueryio_test.go b/sdks/go/test/integration/io/bigqueryio/bigqueryio_test.go index 7cb8051be4b8..47701cb3d510 100644 --- a/sdks/go/test/integration/io/bigqueryio/bigqueryio_test.go +++ b/sdks/go/test/integration/io/bigqueryio/bigqueryio_test.go @@ -20,6 +20,7 @@ import ( "flag" "fmt" "math/rand" + "reflect" "strings" "testing" "time" @@ -30,6 +31,7 @@ import ( "github.com/apache/beam/sdks/v2/go/pkg/beam/options/gcpopts" "github.com/apache/beam/sdks/v2/go/pkg/beam/register" _ "github.com/apache/beam/sdks/v2/go/pkg/beam/runners/dataflow" + "github.com/apache/beam/sdks/v2/go/pkg/beam/testing/passert" "github.com/apache/beam/sdks/v2/go/pkg/beam/testing/ptest" "github.com/apache/beam/sdks/v2/go/test/integration" ) @@ -79,11 +81,29 @@ func shuffleText() []string { return words } +func generateTestRows(seed int64, size int) []TestRow { + rand.Seed(seed) + words := shuffleText() + rows := make([]TestRow, size) + for i := 0; i < size; i++ { + rows[i] = TestRow{ + Counter: int64(i), + RandData: RandData{ + Flip: rand.Int63n(2) != 0, + Num: rand.Int63(), + Word: words[i], + }, + } + } + return rows +} + // RandData is a struct of various types of random data. type RandData struct { - Flip bool `bigquery:"flip"` // Flip is a bool with a random chance of either result (a coin flip). - Num int64 `bigquery:"num"` // Num is a random int64. - Word string `bigquery:"word"` // Word is a randomly selected word from a sample text. + Flip bool `bigquery:"flip"` // Flip is a bool with a random chance of either result (a coin flip). + Num int64 `bigquery:"num"` // Num is a random int64. + Word string `bigquery:"word"` // Word is a randomly selected word from a sample text. + NullableWord bigquery.NullString `bigquery:"nullable_word"` // Word is a randomly selected word from a sample text. } // ddlSchema is a string for BigQuery data definition language that corresponds to TestRow. @@ -91,7 +111,8 @@ const ddlTestRowSchema = "counter INT64 NOT NULL, " + "rand_data STRUCT<" + "flip BOOL NOT NULL," + "num INT64 NOT NULL," + - "word STRING NOT NULL" + + "word STRING NOT NULL," + + "nullable_word STRING" + "> NOT NULL" // CreateTestRowsFn is a DoFn that creates randomized TestRows based on a seed. @@ -101,16 +122,143 @@ type CreateTestRowsFn struct { // ProcessElement creates a number of TestRows, populating the randomized data. func (fn *CreateTestRowsFn) ProcessElement(_ []byte, emit func(TestRow)) { - rand.Seed(fn.seed) - words := shuffleText() - for i := 0; i < inputSize; i++ { - emit(TestRow{ - Counter: int64(i), - RandData: RandData{ - Flip: rand.Int63n(2) != 0, - Num: rand.Int63(), - Word: words[i], + rows := generateTestRows(fn.seed, inputSize) + for _, row := range rows { + emit(row) + } +} + +func TestBigQueryIO_Query(t *testing.T) { + integration.CheckFilters(t) + checkFlags(t) + + ctx := context.Background() + // Get the GCP project + // this assumes dataflow is running in the same project as the project in which the bigquery dataset + // is located + project := gcpopts.GetProject(ctx) + bigqueryDataset := *integration.BigQueryDataset + + tests := []struct { + name string + query string + queryOptions []func(*bigqueryio.QueryOptions) error + preCreate bool + insertData bool + expectedCount int + wantErr bool + }{ + { + name: "Query without parameters", + query: "SELECT * FROM `%s`", + queryOptions: []func(*bigqueryio.QueryOptions) error{ + bigqueryio.UseStandardSQL(), + }, + preCreate: true, + insertData: true, + expectedCount: inputSize, + wantErr: false, + }, + { + name: "Query with parameters", + query: "SELECT * FROM `%s` WHERE counter = @key1 and rand_data = @key2", + queryOptions: []func(*bigqueryio.QueryOptions) error{ + bigqueryio.UseStandardSQL(), + bigqueryio.WithQueryParameters([]bigquery.QueryParameter{ + bigquery.QueryParameter{ + Name: "key1", + Value: -1, // to ensure expectedCount is 0 + }, + bigquery.QueryParameter{ + Name: "key2", + Value: &bigquery.QueryParameterValue{ + Type: bigquery.StandardSQLDataType{ + StructType: &bigquery.StandardSQLStructType{ + Fields: []*bigquery.StandardSQLField{ + { + Name: "flip", + Type: &bigquery.StandardSQLDataType{ + TypeKind: "BOOL", + }, + }, + { + Name: "num", + Type: &bigquery.StandardSQLDataType{ + TypeKind: "INT64", + }, + }, + { + Name: "word", + Type: &bigquery.StandardSQLDataType{ + TypeKind: "STRING", + }, + }, + { + Name: "nullable_word", + Type: &bigquery.StandardSQLDataType{ + TypeKind: "STRING", + }, + }, + }, + }, + }, + StructValue: map[string]bigquery.QueryParameterValue{ + "flip": { + Value: true, + }, + "num": { + Value: int64(1), + }, + "word": { + Value: "ipsum", + }, + "nullable_word": { + Value: bigquery.NullString{ + StringVal: "nullable", + Valid: true, + }, + }, + }, + }, + }, + }...), }, + preCreate: true, + insertData: true, + expectedCount: 0, + wantErr: false, + }, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + tableID := fmt.Sprintf("%s_temp_%v", "go_bqio_it", time.Now().UnixNano()) + tableName := fmt.Sprintf("%s.%s", bigqueryDataset, tableID) + fullyQualifiedTableName := fmt.Sprintf("%s.%s.%s", project, bigqueryDataset, tableID) + if tt.preCreate { + newTempTable(t, tableName, ddlTestRowSchema, project) + } + testRows := generateTestRows(time.Now().UnixNano(), inputSize) + if tt.insertData { + insertTestRows(ctx, t, project, bigqueryDataset, tableID, testRows) + waitForRows(ctx, t, project, bigqueryDataset, tableID, len(testRows)) + } + t.Cleanup(func() { + deleteTempTable(t, tableName, project) + }) + + p, s := beam.NewPipelineWithRoot() + + q := fmt.Sprintf(tt.query, fullyQualifiedTableName) + queryType := reflect.TypeOf((*TestRow)(nil)).Elem() + pcol := bigqueryio.Query(s, project, q, queryType, tt.queryOptions...) + + passert.Count(s, pcol, fmt.Sprintf("Should have %d rows", tt.expectedCount), tt.expectedCount) + if err := ptest.Run(p); (err != nil) != tt.wantErr { + t.Fatalf("ptest.Run() err = %v, wantErr %v", err, tt.wantErr) + } else if err != nil { + // Pipeline failed as expected, return early + return + } }) } } @@ -171,8 +319,8 @@ func TestBigQueryIO_Write(t *testing.T) { // Generate elements and write to table. rows := beam.ParDo(s, createTestRows, beam.Impulse(s)) - bigqueryio.Write(s, project, fmt.Sprintf("%s:%s", project, tableName), rows, - bigqueryio.WithCreateDisposition(tt.createDisposition)) + table := fmt.Sprintf("%s:%s", project, tableName) + bigqueryio.Write(s, project, table, rows, bigqueryio.WithCreateDisposition(tt.createDisposition)) if err := ptest.Run(p); (err != nil) != tt.wantErr { t.Fatalf("ptest.Run() err = %v, wantErr %v", err, tt.wantErr) diff --git a/sdks/go/test/integration/io/bigqueryio/helper_test.go b/sdks/go/test/integration/io/bigqueryio/helper_test.go index c8404410800d..d55da8fbdd04 100644 --- a/sdks/go/test/integration/io/bigqueryio/helper_test.go +++ b/sdks/go/test/integration/io/bigqueryio/helper_test.go @@ -20,6 +20,7 @@ import ( "fmt" "os/exec" "testing" + "time" "cloud.google.com/go/bigquery" "github.com/apache/beam/sdks/v2/go/test/integration" @@ -83,3 +84,47 @@ func checkTableExistsAndNonEmpty(ctx context.Context, t *testing.T, project, tab t.Fatalf("row count = %v, want %v", count, inputSize) } } + +// waitForRows polls until the table has at least wantCount rows visible, or times out. +func waitForRows(ctx context.Context, t *testing.T, project, dataset, tableID string, wantCount int) { + t.Helper() + + client, err := bigquery.NewClient(ctx, project) + if err != nil { + t.Fatalf("error creating BigQuery client: %v", err) + } + defer client.Close() + + deadline := time.Now().Add(2 * time.Minute) + for time.Now().Before(deadline) { + q := client.Query(fmt.Sprintf("SELECT COUNT(*) FROM `%s.%s.%s`", project, dataset, tableID)) + it, err := q.Read(ctx) + if err == nil { + var row []bigquery.Value + if err := it.Next(&row); err == nil { + if count, ok := row[0].(int64); ok && int(count) >= wantCount { + return + } + } + } + time.Sleep(5 * time.Second) + } + t.Fatalf("timed out waiting for %d rows in %s.%s", wantCount, dataset, tableID) +} + +func insertTestRows(ctx context.Context, t *testing.T, project, dataset, tableID string, rows []TestRow) { + t.Helper() + + client, err := bigquery.NewClient(ctx, project) + if err != nil { + t.Fatalf("error creating BigQuery client: %v", err) + } + defer client.Close() + ds := client.Dataset(dataset) + tbl := ds.Table(tableID) + inserter := tbl.Inserter() + + if err := inserter.Put(ctx, rows); err != nil { + t.Fatalf("error inserting test rows: %v", err) + } +} From a1c316fdc29b3e7999446f4b9eac752fc5bd11dd Mon Sep 17 00:00:00 2001 From: Jack McCluskey <34928439+jrmccluskey@users.noreply.github.com> Date: Wed, 9 Sep 2026 09:43:57 -0400 Subject: [PATCH 3/3] Update coder.go Fix typo --- sdks/go/pkg/beam/io/bigqueryio/coder.go | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/sdks/go/pkg/beam/io/bigqueryio/coder.go b/sdks/go/pkg/beam/io/bigqueryio/coder.go index cabd6e0dab45..baa82ca473da 100644 --- a/sdks/go/pkg/beam/io/bigqueryio/coder.go +++ b/sdks/go/pkg/beam/io/bigqueryio/coder.go @@ -57,7 +57,7 @@ func encodeQueryParameters(params []bigquery.QueryParameter) ([]byte, error) { if params == nil { return []byte{}, nil } - // validate each element to tell which paramaeter is unsupported. + // validate each element to tell which parameter is unsupported. for _, p := range params { if err := gob.NewEncoder(io.Discard).Encode([]bigquery.QueryParameter{p}); err != nil { return nil, errors.Errorf(