Add biquery query parameter option to support parametrize query

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<t>.
 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 {

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

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

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

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.*",

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()

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