feat(bigquery): job creation mode GA (#12225)

shollyman · web-flow · commit 1d8990dbf256 · 2025-05-27T09:00:51.000-07:00
* feat(bigquery): job creation mode GA

Previously, use of the QUERY_PREVIEW_ENABLED environmental variable
governed whether this client would attempt to execute queries without
creating jobs.  Now that this feature is GA, this client has been
updated so that it defaults to using optional job mode, and exposes
a new ClientOption (WithJobCreationMode) to override this behavior.
diff --git a/bigquery/bigquery.go b/bigquery/bigquery.go
@@ -21,7 +21,6 @@ import (
 	"io"
 	"net/http"
 	"net/url"
-	"os"
 	"strings"
 	"time"
 
@@ -61,8 +60,8 @@ type Client struct {
 	bqs       *bq.Service
 	rc        *readClient
 
-	// governs use of preview query features.
-	enableQueryPreview bool
+	// container for custom client options
+	customConfig *customClientConfig
 }
 
 // DetectProjectID is a sentinel value that instructs [NewClient] to detect the
@@ -79,12 +78,6 @@ const DetectProjectID = "*detect-project-id*"
 //
 // If the project ID is set to [DetectProjectID], NewClient will attempt to detect
 // the project ID from credentials.
-//
-// This client supports enabling query-related preview features via environmental
-// variables.  By setting the environment variable QUERY_PREVIEW_ENABLED to the string
-// "TRUE", the client will enable preview features, though behavior may still be
-// controlled via the bigquery service as well.  Currently, the feature(s) in scope
-// include: short mode queries (query execution without corresponding job metadata).
 func NewClient(ctx context.Context, projectID string, opts ...option.ClientOption) (*Client, error) {
 	o := []option.ClientOption{
 		option.WithScopes(Scope),
@@ -102,17 +95,13 @@ func NewClient(ctx context.Context, projectID string, opts ...option.ClientOptio
 		return nil, err
 	}
 
-	var preview bool
-	if v, ok := os.LookupEnv("QUERY_PREVIEW_ENABLED"); ok {
-		if strings.ToUpper(v) == "TRUE" {
-			preview = true
-		}
-	}
+	// gather any custom client options
+	custom := newCustomClientConfig(opts...)
 
 	c := &Client{
-		projectID:          projectID,
-		bqs:                bqs,
-		enableQueryPreview: preview,
+		projectID:    projectID,
+		bqs:          bqs,
+		customConfig: custom,
 	}
 	return c, nil
 }
diff --git a/bigquery/integration_test.go b/bigquery/integration_test.go
@@ -801,10 +801,10 @@ func TestIntegration_SimpleRowResults(t *testing.T) {
 	if client == nil {
 		t.Skip("Integration tests skipped")
 	}
-	beforePreview := client.enableQueryPreview
+	beforeMode := client.customConfig.jobCreationMode
 	// ensure we restore the preview setting on test exit
 	defer func() {
-		client.enableQueryPreview = beforePreview
+		client.customConfig.jobCreationMode = beforeMode
 	}()
 	ctx := context.Background()
 
@@ -845,8 +845,8 @@ func TestIntegration_SimpleRowResults(t *testing.T) {
 		},
 	}
 
-	t.Run("nopreview_group", func(t *testing.T) {
-		client.enableQueryPreview = false
+	t.Run("jobrequired_group", func(t *testing.T) {
+		client.customConfig.jobCreationMode = JobCreationModeRequired
 		for _, tc := range testCases {
 			curCase := tc
 			t.Run(curCase.description, func(t *testing.T) {
@@ -860,8 +860,8 @@ func TestIntegration_SimpleRowResults(t *testing.T) {
 			})
 		}
 	})
-	t.Run("preview_group", func(t *testing.T) {
-		client.enableQueryPreview = true
+	t.Run("joboptional_group", func(t *testing.T) {
+		client.customConfig.jobCreationMode = JobCreationModeOptional
 		for _, tc := range testCases {
 			curCase := tc
 			t.Run(curCase.description, func(t *testing.T) {
diff --git a/bigquery/options.go b/bigquery/options.go
@@ -0,0 +1,73 @@
+// Copyright 2025 Google LLC
+//
+// Licensed 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 bigquery
+
+import (
+	"google.golang.org/api/option"
+	"google.golang.org/api/option/internaloption"
+)
+
+// type for collecting custom ClientOption values.
+type customClientConfig struct {
+	jobCreationMode JobCreationMode
+}
+
+type customClientOption interface {
+	option.ClientOption
+	ApplyCustomClientOpt(*customClientConfig)
+}
+
+func newCustomClientConfig(opts ...option.ClientOption) *customClientConfig {
+	conf := &customClientConfig{}
+	for _, opt := range opts {
+		if cOpt, ok := opt.(customClientOption); ok {
+			cOpt.ApplyCustomClientOpt(conf)
+		}
+	}
+	return conf
+}
+
+// JobCreationMode controls how job creation is handled.  Some queries may
+// be run without creating a job to expedite fetching results.
+type JobCreationMode string
+
+var (
+	// JobCreationModeUnspecified is the default (unspecified) option.
+	JobCreationModeUnspecified JobCreationMode = "JOB_CREATION_MODE_UNSPECIFIED"
+	// JobCreationModeRequired indicates job creation is required.
+	JobCreationModeRequired JobCreationMode = "JOB_CREATION_REQUIRED"
+	// JobCreationModeOptional indicates job creation is optional, and returning
+	// results immediately is prioritized.  The conditions under which BigQuery
+	// can choose to avoid job creation are internal and subject to change.
+	JobCreationModeOptional JobCreationMode = "JOB_CREATION_OPTIONAL"
+)
+
+// WithDefaultJobCreationMode is a ClientOption that governs the job creation
+// mode used when executing queries that can be accelerated via the jobs.Query
+// API.  Users may experience performance improvements by leveraging the
+// JobCreationModeOptional mode.
+func WithDefaultJobCreationMode(mode JobCreationMode) option.ClientOption {
+	return &applierJobCreationMode{mode: mode}
+}
+
+// applier for propagating the custom client option to the config object
+type applierJobCreationMode struct {
+	internaloption.EmbeddableAdapter
+	mode JobCreationMode
+}
+
+func (s *applierJobCreationMode) ApplyCustomClientOpt(c *customClientConfig) {
+	c.jobCreationMode = s.mode
+}
diff --git a/bigquery/options_test.go b/bigquery/options_test.go
@@ -0,0 +1,64 @@
+// Copyright 2025 Google LLC
+//
+// Licensed 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
+//
+//     https://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 bigquery
+
+import (
+	"testing"
+
+	"github.com/google/go-cmp/cmp"
+	"google.golang.org/api/option"
+)
+
+func TestCustomClientOptions(t *testing.T) {
+	testCases := []struct {
+		desc    string
+		options []option.ClientOption
+		want    *customClientConfig
+	}{
+		{
+			desc: "no options",
+			want: &customClientConfig{
+				jobCreationMode: "",
+			},
+		},
+		{
+			desc: "jobmode required",
+			options: []option.ClientOption{
+				WithDefaultJobCreationMode(JobCreationModeRequired),
+			},
+			want: &customClientConfig{
+				jobCreationMode: JobCreationModeRequired,
+			},
+		},
+		{
+			desc: "jobmode optional",
+			options: []option.ClientOption{
+				WithDefaultJobCreationMode(JobCreationModeOptional),
+			},
+			want: &customClientConfig{
+				jobCreationMode: JobCreationModeOptional,
+			},
+		},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.desc, func(t *testing.T) {
+			gotCfg := newCustomClientConfig(tc.options...)
+			if diff := cmp.Diff(gotCfg, tc.want, cmp.AllowUnexported(customClientConfig{})); diff != "" {
+				t.Errorf("diff in case (%s):\n%v", tc.desc, diff)
+			}
+		})
+	}
+}
diff --git a/bigquery/query.go b/bigquery/query.go
@@ -502,8 +502,11 @@ func (q *Query) probeFastPath() (*bq.QueryRequest, error) {
 			DatasetId: q.QueryConfig.DefaultDatasetID,
 		}
 	}
-	if q.client.enableQueryPreview {
-		qRequest.JobCreationMode = "JOB_CREATION_OPTIONAL"
+
+	if custCfg := q.client.customConfig; custCfg != nil {
+		if custCfg.jobCreationMode != "" {
+			qRequest.JobCreationMode = string(custCfg.jobCreationMode)
+		}
 	}
 	return qRequest, nil
 }
