mirror of
https://github.com/googleapis/genai-toolbox.git
synced 2026-04-09 03:02:26 -04:00
chore: introduce orderedmap to preserve column order in SQL results during marshal (#1852)
This commit introduces a new `orderedmap` package to preserve the column order of SQL query results when they are marshaled to JSON. The default Go `json.Marshal` function sorts map keys, which was causing the column order to be lost in the output of the database tools. This commit updates the following tools to use the new `orderedmap` package: - `mysqlexecutesql` - `mssqlexecutesql` - `postgresexecutesql` - `spannerexecutesql` - `sqliteexecutesql` - `bigqueryexecutesql` A new test has been added to the `mysqlexecutesql` tool to verify that the column order is preserved. ## Description > Should include a concise description of the changes (bug or feature), it's > impact, along with a summary of the solution ## PR Checklist > Thank you for opening a Pull Request! Before submitting your PR, there are a > few things you can do to make sure it goes smoothly: - [ ] Make sure you reviewed [CONTRIBUTING.md](https://github.com/googleapis/genai-toolbox/blob/main/CONTRIBUTING.md) - [ ] Make sure to open an issue as a [bug/issue](https://github.com/googleapis/genai-toolbox/issues/new/choose) before writing your code! That way we can discuss the change, evaluate designs, and agree on the general idea - [ ] Ensure the tests and linter pass - [ ] Code coverage does not decrease (if any source code was changed) - [ ] Appropriate docs were updated (if necessary) - [ ] Make sure to add `!` if this involve a breaking change 🛠️ Fixes #1492 --------- Co-authored-by: google-labs-jules[bot] <161369871+google-labs-jules[bot]@users.noreply.github.com> Co-authored-by: Yuan Teoh <yuanteoh@google.com> Co-authored-by: Yuan Teoh <45984206+Yuan325@users.noreply.github.com>
This commit is contained in:
@@ -27,6 +27,7 @@ import (
|
||||
"github.com/googleapis/genai-toolbox/internal/tools"
|
||||
bqutil "github.com/googleapis/genai-toolbox/internal/tools/bigquery/bigquerycommon"
|
||||
"github.com/googleapis/genai-toolbox/internal/util"
|
||||
"github.com/googleapis/genai-toolbox/internal/util/orderedmap"
|
||||
bigqueryrestapi "google.golang.org/api/bigquery/v2"
|
||||
"google.golang.org/api/iterator"
|
||||
)
|
||||
@@ -324,19 +325,19 @@ func (t Tool) Invoke(ctx context.Context, params tools.ParamValues, accessToken
|
||||
return nil, fmt.Errorf("unable to read query results: %w", err)
|
||||
}
|
||||
for {
|
||||
var row map[string]bigqueryapi.Value
|
||||
err = it.Next(&row)
|
||||
var val map[string]bigqueryapi.Value
|
||||
err = it.Next(&val)
|
||||
if err == iterator.Done {
|
||||
break
|
||||
}
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("unable to iterate through query results: %w", err)
|
||||
}
|
||||
vMap := make(map[string]any)
|
||||
for key, value := range row {
|
||||
vMap[key] = value
|
||||
row := orderedmap.Row{}
|
||||
for key, value := range val {
|
||||
row.Add(key, value)
|
||||
}
|
||||
out = append(out, vMap)
|
||||
out = append(out, row)
|
||||
}
|
||||
// If the query returned any rows, return them directly.
|
||||
if len(out) > 0 {
|
||||
|
||||
@@ -25,6 +25,7 @@ import (
|
||||
"github.com/googleapis/genai-toolbox/internal/sources/mssql"
|
||||
"github.com/googleapis/genai-toolbox/internal/tools"
|
||||
"github.com/googleapis/genai-toolbox/internal/util"
|
||||
"github.com/googleapis/genai-toolbox/internal/util/orderedmap"
|
||||
)
|
||||
|
||||
const kind string = "mssql-execute-sql"
|
||||
@@ -152,11 +153,11 @@ func (t Tool) Invoke(ctx context.Context, params tools.ParamValues, accessToken
|
||||
if scanErr != nil {
|
||||
return nil, fmt.Errorf("unable to parse row: %w", scanErr)
|
||||
}
|
||||
vMap := make(map[string]any)
|
||||
row := orderedmap.Row{}
|
||||
for i, name := range cols {
|
||||
vMap[name] = rawValues[i]
|
||||
row.Add(name, rawValues[i])
|
||||
}
|
||||
out = append(out, vMap)
|
||||
out = append(out, row)
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -27,6 +27,7 @@ import (
|
||||
"github.com/googleapis/genai-toolbox/internal/tools"
|
||||
"github.com/googleapis/genai-toolbox/internal/tools/mysql/mysqlcommon"
|
||||
"github.com/googleapis/genai-toolbox/internal/util"
|
||||
"github.com/googleapis/genai-toolbox/internal/util/orderedmap"
|
||||
)
|
||||
|
||||
const kind string = "mysql-execute-sql"
|
||||
@@ -159,20 +160,21 @@ func (t Tool) Invoke(ctx context.Context, params tools.ParamValues, accessToken
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("unable to parse row: %w", err)
|
||||
}
|
||||
vMap := make(map[string]any)
|
||||
row := orderedmap.Row{}
|
||||
for i, name := range cols {
|
||||
val := rawValues[i]
|
||||
if val == nil {
|
||||
vMap[name] = nil
|
||||
row.Add(name, nil)
|
||||
continue
|
||||
}
|
||||
|
||||
vMap[name], err = mysqlcommon.ConvertToType(colTypes[i], val)
|
||||
convertedValue, err := mysqlcommon.ConvertToType(colTypes[i], val)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("errors encountered when converting values: %w", err)
|
||||
}
|
||||
row.Add(name, convertedValue)
|
||||
}
|
||||
out = append(out, vMap)
|
||||
out = append(out, row)
|
||||
}
|
||||
|
||||
if err := results.Err(); err != nil {
|
||||
|
||||
@@ -25,6 +25,7 @@ import (
|
||||
"github.com/googleapis/genai-toolbox/internal/sources/postgres"
|
||||
"github.com/googleapis/genai-toolbox/internal/tools"
|
||||
"github.com/googleapis/genai-toolbox/internal/util"
|
||||
"github.com/googleapis/genai-toolbox/internal/util/orderedmap"
|
||||
"github.com/jackc/pgx/v5/pgxpool"
|
||||
)
|
||||
|
||||
@@ -142,11 +143,11 @@ func (t Tool) Invoke(ctx context.Context, params tools.ParamValues, accessToken
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("unable to parse row: %w", err)
|
||||
}
|
||||
vMap := make(map[string]any)
|
||||
row := orderedmap.Row{}
|
||||
for i, f := range fields {
|
||||
vMap[f.Name] = v[i]
|
||||
row.Add(f.Name, v[i])
|
||||
}
|
||||
out = append(out, vMap)
|
||||
out = append(out, row)
|
||||
}
|
||||
|
||||
if err := results.Err(); err != nil {
|
||||
|
||||
@@ -24,6 +24,7 @@ import (
|
||||
spannerdb "github.com/googleapis/genai-toolbox/internal/sources/spanner"
|
||||
"github.com/googleapis/genai-toolbox/internal/tools"
|
||||
"github.com/googleapis/genai-toolbox/internal/util"
|
||||
"github.com/googleapis/genai-toolbox/internal/util/orderedmap"
|
||||
"google.golang.org/api/iterator"
|
||||
)
|
||||
|
||||
@@ -131,12 +132,12 @@ func processRows(iter *spanner.RowIterator) ([]any, error) {
|
||||
return nil, fmt.Errorf("unable to parse row: %w", err)
|
||||
}
|
||||
|
||||
vMap := make(map[string]any)
|
||||
rowMap := orderedmap.Row{}
|
||||
cols := row.ColumnNames()
|
||||
for i, c := range cols {
|
||||
vMap[c] = row.ColumnValue(i)
|
||||
rowMap.Add(c, row.ColumnValue(i))
|
||||
}
|
||||
out = append(out, vMap)
|
||||
out = append(out, rowMap)
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
@@ -25,6 +25,7 @@ import (
|
||||
"github.com/googleapis/genai-toolbox/internal/sources/sqlite"
|
||||
"github.com/googleapis/genai-toolbox/internal/tools"
|
||||
"github.com/googleapis/genai-toolbox/internal/util"
|
||||
"github.com/googleapis/genai-toolbox/internal/util/orderedmap"
|
||||
)
|
||||
|
||||
const kind string = "sqlite-execute-sql"
|
||||
@@ -155,11 +156,11 @@ func (t Tool) Invoke(ctx context.Context, params tools.ParamValues, accessToken
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("unable to parse row: %w", err)
|
||||
}
|
||||
vMap := make(map[string]any)
|
||||
row := orderedmap.Row{}
|
||||
for i, name := range cols {
|
||||
val := rawValues[i]
|
||||
if val == nil {
|
||||
vMap[name] = nil
|
||||
row.Add(name, nil)
|
||||
continue
|
||||
}
|
||||
|
||||
@@ -167,13 +168,13 @@ func (t Tool) Invoke(ctx context.Context, params tools.ParamValues, accessToken
|
||||
if jsonString, ok := val.(string); ok {
|
||||
var unmarshaledData any
|
||||
if json.Unmarshal([]byte(jsonString), &unmarshaledData) == nil {
|
||||
vMap[name] = unmarshaledData
|
||||
row.Add(name, unmarshaledData)
|
||||
continue
|
||||
}
|
||||
}
|
||||
vMap[name] = val
|
||||
row.Add(name, val)
|
||||
}
|
||||
out = append(out, vMap)
|
||||
out = append(out, row)
|
||||
}
|
||||
|
||||
if err := results.Err(); err != nil {
|
||||
|
||||
@@ -26,6 +26,7 @@ import (
|
||||
"github.com/googleapis/genai-toolbox/internal/testutils"
|
||||
"github.com/googleapis/genai-toolbox/internal/tools"
|
||||
"github.com/googleapis/genai-toolbox/internal/tools/sqlite/sqliteexecutesql"
|
||||
"github.com/googleapis/genai-toolbox/internal/util/orderedmap"
|
||||
_ "modernc.org/sqlite"
|
||||
)
|
||||
|
||||
@@ -159,8 +160,20 @@ func TestTool_Invoke(t *testing.T) {
|
||||
},
|
||||
},
|
||||
want: []any{
|
||||
map[string]any{"id": int64(1), "name": "Alice", "age": int64(30)},
|
||||
map[string]any{"id": int64(2), "name": "Bob", "age": int64(25)},
|
||||
orderedmap.Row{
|
||||
Columns: []orderedmap.Column{
|
||||
{Name: "id", Value: int64(1)},
|
||||
{Name: "name", Value: "Alice"},
|
||||
{Name: "age", Value: int64(30)},
|
||||
},
|
||||
},
|
||||
orderedmap.Row{
|
||||
Columns: []orderedmap.Column{
|
||||
{Name: "id", Value: int64(2)},
|
||||
{Name: "name", Value: "Bob"},
|
||||
{Name: "age", Value: int64(25)},
|
||||
},
|
||||
},
|
||||
},
|
||||
wantErr: false,
|
||||
},
|
||||
@@ -233,7 +246,13 @@ func TestTool_Invoke(t *testing.T) {
|
||||
},
|
||||
},
|
||||
want: []any{
|
||||
map[string]any{"id": int64(1), "null_col": nil, "blob_col": []byte{1, 2, 3}},
|
||||
orderedmap.Row{
|
||||
Columns: []orderedmap.Column{
|
||||
{Name: "id", Value: int64(1)},
|
||||
{Name: "null_col", Value: nil},
|
||||
{Name: "blob_col", Value: []byte{1, 2, 3}},
|
||||
},
|
||||
},
|
||||
},
|
||||
wantErr: false,
|
||||
},
|
||||
@@ -264,8 +283,18 @@ func TestTool_Invoke(t *testing.T) {
|
||||
},
|
||||
},
|
||||
want: []any{
|
||||
map[string]any{"name": "Alice", "item": "Laptop"},
|
||||
map[string]any{"name": "Bob", "item": "Keyboard"},
|
||||
orderedmap.Row{
|
||||
Columns: []orderedmap.Column{
|
||||
{Name: "name", Value: "Alice"},
|
||||
{Name: "item", Value: "Laptop"},
|
||||
},
|
||||
},
|
||||
orderedmap.Row{
|
||||
Columns: []orderedmap.Column{
|
||||
{Name: "name", Value: "Bob"},
|
||||
{Name: "item", Value: "Keyboard"},
|
||||
},
|
||||
},
|
||||
},
|
||||
wantErr: false,
|
||||
},
|
||||
@@ -292,7 +321,7 @@ func TestTool_Invoke(t *testing.T) {
|
||||
}
|
||||
|
||||
if !isEqual {
|
||||
t.Errorf("Tool.Invoke() = %v, want %v", got, tt.want)
|
||||
t.Errorf("Tool.Invoke() = %+v, want %v", got, tt.want)
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
62
internal/util/orderedmap/orderedmap.go
Normal file
62
internal/util/orderedmap/orderedmap.go
Normal file
@@ -0,0 +1,62 @@
|
||||
// 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 orderedmap
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"encoding/json"
|
||||
)
|
||||
|
||||
// Column represents a single column in a row.
|
||||
type Column struct {
|
||||
Name string
|
||||
Value any
|
||||
}
|
||||
|
||||
// Row represents a row of data with columns in a specific order.
|
||||
type Row struct {
|
||||
Columns []Column
|
||||
}
|
||||
|
||||
// Add adds a new column to the row.
|
||||
func (r *Row) Add(name string, value any) {
|
||||
r.Columns = append(r.Columns, Column{Name: name, Value: value})
|
||||
}
|
||||
|
||||
// MarshalJSON implements the json.Marshaler interface for the Row struct.
|
||||
// It marshals the row into a JSON object, preserving the order of the columns.
|
||||
func (r Row) MarshalJSON() ([]byte, error) {
|
||||
var buf bytes.Buffer
|
||||
buf.WriteString("{")
|
||||
for i, col := range r.Columns {
|
||||
if i > 0 {
|
||||
buf.WriteString(",")
|
||||
}
|
||||
// Marshal the key
|
||||
key, err := json.Marshal(col.Name)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
buf.Write(key)
|
||||
buf.WriteString(":")
|
||||
// Marshal the value
|
||||
val, err := json.Marshal(col.Value)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
buf.Write(val)
|
||||
}
|
||||
buf.WriteString("}")
|
||||
return buf.Bytes(), nil
|
||||
}
|
||||
83
internal/util/orderedmap/orderedmap_test.go
Normal file
83
internal/util/orderedmap/orderedmap_test.go
Normal file
@@ -0,0 +1,83 @@
|
||||
// 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 orderedmap
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestRowMarshalJSON(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
row Row
|
||||
want string
|
||||
wantErr bool
|
||||
}{
|
||||
{
|
||||
name: "Simple row",
|
||||
row: Row{
|
||||
Columns: []Column{
|
||||
{Name: "A", Value: 1},
|
||||
{Name: "B", Value: "two"},
|
||||
{Name: "C", Value: true},
|
||||
},
|
||||
},
|
||||
want: `{"A":1,"B":"two","C":true}`,
|
||||
wantErr: false,
|
||||
},
|
||||
{
|
||||
name: "Row with different order",
|
||||
row: Row{
|
||||
Columns: []Column{
|
||||
{Name: "C", Value: true},
|
||||
{Name: "A", Value: 1},
|
||||
{Name: "B", Value: "two"},
|
||||
},
|
||||
},
|
||||
want: `{"C":true,"A":1,"B":"two"}`,
|
||||
wantErr: false,
|
||||
},
|
||||
{
|
||||
name: "Empty row",
|
||||
row: Row{},
|
||||
want: `{}`,
|
||||
wantErr: false,
|
||||
},
|
||||
{
|
||||
name: "Row with nil value",
|
||||
row: Row{
|
||||
Columns: []Column{
|
||||
{Name: "A", Value: 1},
|
||||
{Name: "B", Value: nil},
|
||||
},
|
||||
},
|
||||
want: `{"A":1,"B":null}`,
|
||||
wantErr: false,
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
got, err := json.Marshal(tt.row)
|
||||
if (err != nil) != tt.wantErr {
|
||||
t.Errorf("Row.MarshalJSON() error = %v, wantErr %v", err, tt.wantErr)
|
||||
return
|
||||
}
|
||||
if string(got) != tt.want {
|
||||
t.Errorf("Row.MarshalJSON() = %s, want %s", string(got), tt.want)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user