diff --git a/internal/tools/bigquery/bigqueryexecutesql/bigqueryexecutesql.go b/internal/tools/bigquery/bigqueryexecutesql/bigqueryexecutesql.go index d0b497dd9ab3..16d704d7dda1 100644 --- a/internal/tools/bigquery/bigqueryexecutesql/bigqueryexecutesql.go +++ b/internal/tools/bigquery/bigqueryexecutesql/bigqueryexecutesql.go @@ -135,25 +135,6 @@ func (t Tool) Invoke(ctx context.Context, params tools.ParamValues) (any, error) 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. @@ -177,10 +158,21 @@ func (t Tool) Invoke(ctx context.Context, params tools.ParamValues) (any, error) } out = append(out, vMap) } - if out == nil { + // If the query returned any rows, return them directly. + if len(out) > 0 { + return out, nil + } + + // This handles the standard case for a SELECT query that successfully + // executes but returns zero rows. + if statementType == "SELECT" { return "The query returned 0 rows.", nil } - return out, nil + // This is the fallback for a successful query that doesn't return content. + // In most cases, this will be for DML/DDL statements like INSERT, UPDATE, CREATE, etc. + // However, it is also possible that this was a query that was expected to return rows + // but returned none, a case that we cannot distinguish here. + return "Query executed successfully and returned no content.", nil } func (t Tool) ParseParams(data map[string]any, claims map[string]map[string]any) (tools.ParamValues, error) { diff --git a/tests/bigquery/bigquery_integration_test.go b/tests/bigquery/bigquery_integration_test.go index 13d8f1ad2507..93b8a6c2dd40 100644 --- a/tests/bigquery/bigquery_integration_test.go +++ b/tests/bigquery/bigquery_integration_test.go @@ -161,7 +161,7 @@ func TestBigQueryToolEndpoints(t *testing.T) { ) tests.RunToolInvokeWithTemplateParameters(t, tableNameTemplateParam, templateParamTestConfig) - runBigQueryExecuteSqlToolInvokeTest(t, select1Want, invokeParamWant, tableNameParam) + runBigQueryExecuteSqlToolInvokeTest(t, select1Want, invokeParamWant, tableNameParam, ddlWant) runBigQueryDataTypeTests(t) runBigQueryListDatasetToolInvokeTest(t, datasetName) runBigQueryGetDatasetInfoToolInvokeTest(t, datasetName, datasetInfoWant) @@ -411,7 +411,7 @@ func addBigQuerySqlToolConfig(t *testing.T, config map[string]any, toolStatement return config } -func runBigQueryExecuteSqlToolInvokeTest(t *testing.T, select1Want, invokeParamWant, tableNameParam string) { +func runBigQueryExecuteSqlToolInvokeTest(t *testing.T, select1Want, invokeParamWant, tableNameParam, ddlWant string) { // Get ID token idToken, err := tests.GetGoogleIdToken(tests.ClientId) if err != nil { @@ -447,7 +447,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."`, + want: ddlWant, isErr: true, }, { @@ -471,7 +471,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":"DROP TABLE t"}`)), - want: `"Operation completed successfully."`, + want: ddlWant, isErr: true, }, { @@ -479,7 +479,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: `"Operation completed successfully."`, + want: ddlWant, isErr: false, }, {