Add biquery query parameter option to support parametrize query

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

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 {

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)
+		}
+	})
+}

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(&params); err != nil {
+		return nil, err
+	}
+	return params, nil
+}