fix(tools/bigquery-execute-sql): ensure invoke always returns a non-null value (#925)

- Added a dry run step to identify the query type (e.g., SELECT, DML),
which allows the tool to correctly handle the query's output.
- The recommended high-level client, cloud.google.com/go/bigquery, does
not expose the statement type from a dry run. To circumvent this
limitation, the low-level BigQuery REST API client
(google.golang.org/api/bigquery/v2) was added to gain access to these
necessary details.

fixes: #915

---------

Co-authored-by: Yuan Teoh <45984206+Yuan325@users.noreply.github.com>

internal/sources/bigquery/bigquery.go

@@ -24,6 +24,7 @@ import (
 	"github.com/googleapis/genai-toolbox/internal/util"
 	"go.opentelemetry.io/otel/trace"
 	"golang.org/x/oauth2/google"
+	bigqueryrestapi "google.golang.org/api/bigquery/v2"
 	"google.golang.org/api/option"
 )
 
@@ -61,15 +62,17 @@ func (r Config) SourceConfigKind() string {
 
 func (r Config) Initialize(ctx context.Context, tracer trace.Tracer) (sources.Source, error) {
 	// Initializes a BigQuery Google SQL source
-	client, err := initBigQueryConnection(ctx, tracer, r.Name, r.Project, r.Location)
+	client, restService, err := initBigQueryConnection(ctx, tracer, r.Name, r.Project, r.Location)
 	if err != nil {
 		return nil, err
 	}
+
 	s := &Source{
-		Name:     r.Name,
-		Kind:     SourceKind,
-		Client:   client,
-		Location: r.Location,
+		Name:        r.Name,
+		Kind:        SourceKind,
+		Client:      client,
+		RestService: restService,
+		Location:    r.Location,
 	}
 	return s, nil
 
@@ -79,10 +82,11 @@ var _ sources.Source = &Source{}
 
 type Source struct {
 	// BigQuery Google SQL struct with client
-	Name     string `yaml:"name"`
-	Kind     string `yaml:"kind"`
-	Client   *bigqueryapi.Client
-	Location string `yaml:"location"`
+	Name        string `yaml:"name"`
+	Kind        string `yaml:"kind"`
+	Client      *bigqueryapi.Client
+	RestService *bigqueryrestapi.Service
+	Location    string `yaml:"location"`
 }
 
 func (s *Source) SourceKind() string {
@@ -94,30 +98,42 @@ func (s *Source) BigQueryClient() *bigqueryapi.Client {
 	return s.Client
 }
 
+func (s *Source) BigQueryRestService() *bigqueryrestapi.Service {
+	return s.RestService
+}
+
 func initBigQueryConnection(
 	ctx context.Context,
 	tracer trace.Tracer,
 	name string,
 	project string,
 	location string,
-) (*bigqueryapi.Client, error) {
+) (*bigqueryapi.Client, *bigqueryrestapi.Service, error) {
 	ctx, span := sources.InitConnectionSpan(ctx, tracer, SourceKind, name)
 	defer span.End()
 
 	cred, err := google.FindDefaultCredentials(ctx, bigqueryapi.Scope)
 	if err != nil {
-		return nil, fmt.Errorf("failed to find default Google Cloud credentials with scope %q: %w", bigqueryapi.Scope, err)
+		return nil, nil, fmt.Errorf("failed to find default Google Cloud credentials with scope %q: %w", bigqueryapi.Scope, err)
 	}
 
 	userAgent, err := util.UserAgentFromContext(ctx)
 	if err != nil {
-		return nil, err
+		return nil, nil, err
 	}
 
+	// Initialize the high-level BigQuery client
 	client, err := bigqueryapi.NewClient(ctx, project, option.WithUserAgent(userAgent), option.WithCredentials(cred))
-	client.Location = location
 	if err != nil {
-		return nil, fmt.Errorf("failed to create BigQuery client for project %q: %w", project, err)
+		return nil, nil, fmt.Errorf("failed to create BigQuery client for project %q: %w", project, err)
 	}
-	return client, nil
+	client.Location = location
+
+	// Initialize the low-level BigQuery REST service using the same credentials
+	restService, err := bigqueryrestapi.NewService(ctx, option.WithUserAgent(userAgent), option.WithCredentials(cred))
+	if err != nil {
+		return nil, nil, fmt.Errorf("failed to create BigQuery v2 service: %w", err)
+	}
+
+	return client, restService, nil
 }

internal/tools/bigquery/bigqueryexecutesql/bigqueryexecutesql.go

@@ -23,6 +23,7 @@ import (
 	"github.com/googleapis/genai-toolbox/internal/sources"
 	bigqueryds "github.com/googleapis/genai-toolbox/internal/sources/bigquery"
 	"github.com/googleapis/genai-toolbox/internal/tools"
+	bigqueryrestapi "google.golang.org/api/bigquery/v2"
 	"google.golang.org/api/iterator"
 )
 
@@ -44,6 +45,7 @@ func newConfig(ctx context.Context, name string, decoder *yaml.Decoder) (tools.T
 
 type compatibleSource interface {
 	BigQueryClient() *bigqueryapi.Client
+	BigQueryRestService() *bigqueryrestapi.Service
 }
 
 // validate compatible sources are still compatible
@@ -95,6 +97,7 @@ func (cfg Config) Initialize(srcs map[string]sources.Source) (tools.Tool, error)
 		Parameters:   parameters,
 		AuthRequired: cfg.AuthRequired,
 		Client:       s.BigQueryClient(),
+		RestService:  s.BigQueryRestService(),
 		manifest:     tools.Manifest{Description: cfg.Description, Parameters: parameters.Manifest(), AuthRequired: cfg.AuthRequired},
 		mcpManifest:  mcpManifest,
 	}
@@ -110,6 +113,7 @@ type Tool struct {
 	AuthRequired []string         `yaml:"authRequired"`
 	Parameters   tools.Parameters `yaml:"parameters"`
 	Client       *bigqueryapi.Client
+	RestService  *bigqueryrestapi.Service
 	manifest     tools.Manifest
 	mcpManifest  tools.McpManifest
 }
@@ -121,18 +125,46 @@ func (t Tool) Invoke(ctx context.Context, params tools.ParamValues) (any, error)
 		return nil, fmt.Errorf("unable to get cast %s", sliceParams[0])
 	}
 
+	dryRunJob, err := dryRunQuery(ctx, t.RestService, t.Client.Project(), sql)
+	if err != nil {
+		return nil, fmt.Errorf("query validation failed during dry run: %w", err)
+	}
+
+	statementType := dryRunJob.Statistics.Query.StatementType
+	// JobStatistics.QueryStatistics.StatementType
 	query := t.Client.Query(sql)
 	query.Location = t.Client.Location
 
+	// This block handles Data Manipulation Language (DML) and Data Definition Language (DDL) statements.
+	// These statements (e.g., INSERT, UPDATE, CREATE TABLE) do not return a row set.
+	// Instead, we execute them as a job, wait for completion, and return a success
+	// message, including the number of affected rows for DML operations.
+	if statementType != "SELECT" {
+		job, err := query.Run(ctx)
+		if err != nil {
+			return nil, fmt.Errorf("failed to start DML/DDL job: %w", err)
+		}
+		status, err := job.Wait(ctx)
+		if err != nil {
+			return nil, fmt.Errorf("failed to wait for DML/DDL job to complete: %w", err)
+		}
+		if err := status.Err(); err != nil {
+			return nil, fmt.Errorf("DML/DDL job failed with error: %w", err)
+		}
+		return "Operation completed successfully.", nil
+	}
+
+	// This block handles SELECT statements, which return a row set.
+	// We iterate through the results, convert each row into a map of
+	// column names to values, and return the collection of rows.
+	var out []any
 	it, err := query.Read(ctx)
 	if err != nil {
 		return nil, fmt.Errorf("unable to execute query: %w", err)
 	}
-
-	var out []any
 	for {
 		var row map[string]bigqueryapi.Value
-		err := it.Next(&row)
+		err = it.Next(&row)
 		if err == iterator.Done {
 			break
 		}
@@ -145,7 +177,9 @@ func (t Tool) Invoke(ctx context.Context, params tools.ParamValues) (any, error)
 		}
 		out = append(out, vMap)
 	}
-
+	if out == nil {
+		return "The query returned 0 rows.", nil
+	}
 	return out, nil
 }
 
@@ -164,3 +198,23 @@ func (t Tool) McpManifest() tools.McpManifest {
 func (t Tool) Authorized(verifiedAuthServices []string) bool {
 	return tools.IsAuthorized(t.AuthRequired, verifiedAuthServices)
 }
+
+// dryRunQuery performs a dry run of the SQL query to validate it and get metadata.
+func dryRunQuery(ctx context.Context, restService *bigqueryrestapi.Service, projectID string, sql string) (*bigqueryrestapi.Job, error) {
+	useLegacySql := false
+	jobToInsert := &bigqueryrestapi.Job{
+		Configuration: &bigqueryrestapi.JobConfiguration{
+			DryRun: true,
+			Query: &bigqueryrestapi.JobConfigurationQuery{
+				Query:        sql,
+				UseLegacySql: &useLegacySql,
+			},
+		},
+	}
+
+	insertResponse, err := restService.Jobs.Insert(projectID, jobToInsert).Context(ctx).Do()
+	if err != nil {
+		return nil, fmt.Errorf("failed to insert dry run job: %w", err)
+	}
+	return insertResponse, nil
+}

tests/bigquery/bigquery_integration_test.go

@@ -380,6 +380,7 @@ func runBigQueryExecuteSqlToolInvokeTest(t *testing.T, select1Want, invokeParamW
 			api:           "http://127.0.0.1:5000/api/tool/my-exec-sql-tool/invoke",
 			requestHeader: map[string]string{},
 			requestBody:   bytes.NewBuffer([]byte(`{"sql":"CREATE TABLE t (id SERIAL PRIMARY KEY, name TEXT)"}`)),
+			want:          `"Operation completed successfully."`,
 			isErr:         true,
 		},
 		{
@@ -390,11 +391,20 @@ func runBigQueryExecuteSqlToolInvokeTest(t *testing.T, select1Want, invokeParamW
 			want:          invokeParamWant,
 			isErr:         false,
 		},
+		{
+			name:          "invoke my-exec-sql-tool with no matching rows",
+			api:           "http://127.0.0.1:5000/api/tool/my-exec-sql-tool/invoke",
+			requestHeader: map[string]string{},
+			requestBody:   bytes.NewBuffer([]byte(fmt.Sprintf("{\"sql\":\"SELECT * FROM %s WHERE id = 999\"}", tableNameParam))),
+			want:          `"The query returned 0 rows."`,
+			isErr:         false,
+		},
 		{
 			name:          "invoke my-exec-sql-tool drop table",
 			api:           "http://127.0.0.1:5000/api/tool/my-exec-sql-tool/invoke",
 			requestHeader: map[string]string{},
 			requestBody:   bytes.NewBuffer([]byte(`{"sql":"DROP TABLE t"}`)),
+			want:          `"Operation completed successfully."`,
 			isErr:         true,
 		},
 		{
@@ -402,7 +412,7 @@ func runBigQueryExecuteSqlToolInvokeTest(t *testing.T, select1Want, invokeParamW
 			api:           "http://127.0.0.1:5000/api/tool/my-exec-sql-tool/invoke",
 			requestHeader: map[string]string{},
 			requestBody:   bytes.NewBuffer([]byte(fmt.Sprintf("{\"sql\":\"INSERT INTO %s (id, name) VALUES (4, 'test_name')\"}", tableNameParam))),
-			want:          "null",
+			want:          `"Operation completed successfully."`,
 			isErr:         false,
 		},
 		{