Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
34 changes: 13 additions & 21 deletions internal/tools/bigquery/bigqueryexecutesql/bigqueryexecutesql.go
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand All @@ -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) {
Expand Down
10 changes: 5 additions & 5 deletions tests/bigquery/bigquery_integration_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down Expand Up @@ -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 {
Expand Down Expand Up @@ -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,
},
{
Expand All @@ -471,15 +471,15 @@ 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,
},
{
name: "invoke my-exec-sql-tool insert entry",
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,
},
{
Expand Down
Loading