This is an automated email from the ASF dual-hosted git repository.
jrmccluskey pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/beam.git
The following commit(s) were added to refs/heads/master by this push:
new d11883d91a6 Add biquery query parameter option to support parametrize
query (#39389)
d11883d91a6 is described below
commit d11883d91a66277ed0a29b24daa998412a2fbd45
Author: Makoto Nagai <[email protected]>
AuthorDate: Wed Sep 9 23:26:55 2026 +0900
Add biquery query parameter option to support parametrize query (#39389)
* Add biquery query parameters to support parametrize query
* 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
* Update coder.go
Fix typo
---------
Co-authored-by: Jack McCluskey
<[email protected]>
---
sdks/go.mod | 2 +-
sdks/go/pkg/beam/io/bigqueryio/bigquery.go | 32 +-
sdks/go/pkg/beam/io/bigqueryio/bigquery_test.go | 35 ++
sdks/go/pkg/beam/io/bigqueryio/coder.go | 88 ++++
sdks/go/pkg/beam/io/bigqueryio/coder_test.go | 495 +++++++++++++++++++++
sdks/go/test/integration/integration.go | 4 +-
.../integration/io/bigqueryio/bigqueryio_test.go | 182 +++++++-
.../test/integration/io/bigqueryio/helper_test.go | 73 ++-
8 files changed, 878 insertions(+), 33 deletions(-)
diff --git a/sdks/go.mod b/sdks/go.mod
index 222aa328b1b..293bd34354e 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.56.1
github.com/golang-cz/devslog v0.0.17
@@ -137,7 +138,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.12.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 a80661e22d3..bbe951969b9 100644
--- a/sdks/go/pkg/beam/io/bigqueryio/bigquery.go
+++ b/sdks/go/pkg/beam/io/bigqueryio/bigquery.go
@@ -114,6 +114,11 @@ 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.
+ // 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
}
// UseStandardSQL enables BigQuery's Standard SQL dialect when executing a
query.
@@ -124,6 +129,15 @@ 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
+ return nil
+ }
+}
+
// Query executes a query. The output must have a schema compatible with the
given
// type, t. It returns a PCollection<t>.
func Query(s beam.Scope, project, q string, t reflect.Type, options
...func(*QueryOptions) error) beam.PCollection {
@@ -142,7 +156,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 +175,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 +192,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 b955136e833..9f83eeb847f 100644
--- a/sdks/go/pkg/beam/io/bigqueryio/bigquery_test.go
+++ b/sdks/go/pkg/beam/io/bigqueryio/bigquery_test.go
@@ -148,3 +148,38 @@ 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[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
new file mode 100644
index 00000000000..baa82ca473d
--- /dev/null
+++ b/sdks/go/pkg/beam/io/bigqueryio/coder.go
@@ -0,0 +1,88 @@
+// 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 (
+ "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) {
+ if params == nil {
+ return []byte{}, nil
+ }
+ // 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(
+ "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(data []byte) ([]bigquery.QueryParameter, error) {
+ if len(data) == 0 {
+ return []bigquery.QueryParameter{}, 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
new file mode 100644
index 00000000000..b569e9400b6
--- /dev/null
+++ b/sdks/go/pkg/beam/io/bigqueryio/coder_test.go
@@ -0,0 +1,495 @@
+// 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 (
+ "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: "string",
+ val: bigquery.QueryParameter{
+ Name: "key",
+ Value: "value",
+ },
+ },
+ {
+ 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) {
+ 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)
+ }
+
+ 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 a0eef7d1018..535a7fb8413 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_[^WQ].*",
"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 337083745c5..47701cb3d51 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
+ }
})
}
}
@@ -161,18 +309,18 @@ 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()
// 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 e39ed87d7ff..d55da8fbdd0 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"
@@ -31,11 +32,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 +46,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 +67,64 @@ 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)
+ }
+}
+
+// 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)
}
}
