github/genai-toolbox
1
0
Fork 0
mirror of https://github.com/googleapis/genai-toolbox.git synced 2026-01-09 15:38:08 -05:00
Code Issues Packages Projects Releases 14 Wiki Activity

feat(tools/bigquery,mssql,mysql,postgres,spanner,tidb): Add query logging to execute-sql tools (#1069)

Browse Source 
fix: https://github.com/googleapis/genai-toolbox/issues/1052
This commit is contained in:
Wenxin Du
2025-08-04 23:01:19 -04:00
committed by GitHub
parent 8d0fa6783a
commit 0527532bd7
6 changed files with 48 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

8
internal/tools/bigquery/bigqueryexecutesql/bigqueryexecutesql.go
View File

@@ -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"
"github.com/googleapis/genai-toolbox/internal/util"
bigqueryrestapi "google.golang.org/api/bigquery/v2"
"google.golang.org/api/iterator"
)
@@ -135,6 +136,13 @@ func (t Tool) Invoke(ctx context.Context, params tools.ParamValues) (any, error)
query := t.Client.Query(sql)
query.Location = t.Client.Location
// Log the query executed for debugging.
logger, err := util.LoggerFromContext(ctx)
if err != nil {
return nil, fmt.Errorf("error getting logger: %s", err)
}
logger.DebugContext(ctx, "executing `%s` tool query: %s", kind, sql)
// 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.

9
internal/tools/mssql/mssqlexecutesql/mssqlexecutesql.go
View File

@@ -24,6 +24,7 @@ import (
"github.com/googleapis/genai-toolbox/internal/sources/cloudsqlmssql"
"github.com/googleapis/genai-toolbox/internal/sources/mssql"
"github.com/googleapis/genai-toolbox/internal/tools"
"github.com/googleapis/genai-toolbox/internal/util"
)
const kind string = "mssql-execute-sql"
@@ -122,6 +123,14 @@ func (t Tool) Invoke(ctx context.Context, params tools.ParamValues) (any, error)
if !ok {
return nil, fmt.Errorf("unable to get cast %s", sliceParams[0])
}
// Log the query executed for debugging.
logger, err := util.LoggerFromContext(ctx)
if err != nil {
return nil, fmt.Errorf("error getting logger: %s", err)
}
logger.DebugContext(ctx, "executing `%s` tool query: %s", kind, sql)
results, err := t.Pool.QueryContext(ctx, sql)
if err != nil {
return nil, fmt.Errorf("unable to execute query: %w", err)

8
internal/tools/mysql/mysqlexecutesql/mysqlexecutesql.go
View File

@@ -24,6 +24,7 @@ import (
"github.com/googleapis/genai-toolbox/internal/sources/cloudsqlmysql"
"github.com/googleapis/genai-toolbox/internal/sources/mysql"
"github.com/googleapis/genai-toolbox/internal/tools"
"github.com/googleapis/genai-toolbox/internal/util"
)
const kind string = "mysql-execute-sql"
@@ -123,6 +124,13 @@ func (t Tool) Invoke(ctx context.Context, params tools.ParamValues) (any, error)
return nil, fmt.Errorf("unable to get cast %s", sliceParams[0])
}
// Log the query executed for debugging.
logger, err := util.LoggerFromContext(ctx)
if err != nil {
return nil, fmt.Errorf("error getting logger: %s", err)
}
logger.DebugContext(ctx, "executing `%s` tool query: %s", kind, sql)
results, err := t.Pool.QueryContext(ctx, sql)
if err != nil {
return nil, fmt.Errorf("unable to execute query: %w", err)

7
internal/tools/postgres/postgresexecutesql/postgresexecutesql.go
View File

@@ -24,6 +24,7 @@ import (
"github.com/googleapis/genai-toolbox/internal/sources/cloudsqlpg"
"github.com/googleapis/genai-toolbox/internal/sources/postgres"
"github.com/googleapis/genai-toolbox/internal/tools"
"github.com/googleapis/genai-toolbox/internal/util"
"github.com/jackc/pgx/v5/pgxpool"
)
@@ -124,6 +125,12 @@ func (t Tool) Invoke(ctx context.Context, params tools.ParamValues) (any, error)
if !ok {
return nil, fmt.Errorf("unable to get cast %s", sliceParams[0])
}
// Log the query executed for debugging.
logger, err := util.LoggerFromContext(ctx)
if err != nil {
return nil, fmt.Errorf("error getting logger: %s", err)
}
logger.DebugContext(ctx, "executing `%s` tool query: %s", kind, sql)
results, err := t.Pool.Query(ctx, sql)
if err != nil {

8
internal/tools/spanner/spannerexecutesql/spannerexecutesql.go
View File

@@ -23,6 +23,7 @@ import (
"github.com/googleapis/genai-toolbox/internal/sources"
spannerdb "github.com/googleapis/genai-toolbox/internal/sources/spanner"
"github.com/googleapis/genai-toolbox/internal/tools"
"github.com/googleapis/genai-toolbox/internal/util"
"google.golang.org/api/iterator"
)
@@ -151,6 +152,13 @@ func (t Tool) Invoke(ctx context.Context, params tools.ParamValues) (any, error)
return nil, fmt.Errorf("unable to get cast %s", sliceParams[0])
}
// Log the query executed for debugging.
logger, err := util.LoggerFromContext(ctx)
if err != nil {
return nil, fmt.Errorf("error getting logger: %s", err)
}
logger.DebugContext(ctx, "executing `%s` tool query: %s", kind, sql)
var results []any
var opErr error
stmt := spanner.Statement{SQL: sql}

8
internal/tools/tidb/tidbexecutesql/tidbexecutesql.go
View File

@@ -23,6 +23,7 @@ import (
"github.com/googleapis/genai-toolbox/internal/sources"
"github.com/googleapis/genai-toolbox/internal/sources/tidb"
"github.com/googleapis/genai-toolbox/internal/tools"
"github.com/googleapis/genai-toolbox/internal/util"
)
const kind string = "tidb-execute-sql"
@@ -121,6 +122,13 @@ func (t Tool) Invoke(ctx context.Context, params tools.ParamValues) (any, error)
return nil, fmt.Errorf("unable to get cast %s", sliceParams[0])
}
// Log the query executed for debugging.
logger, err := util.LoggerFromContext(ctx)
if err != nil {
return nil, fmt.Errorf("error getting logger: %s", err)
}
logger.DebugContext(ctx, "executing `%s` tool query: %s", kind, sql)
results, err := t.Pool.QueryContext(ctx, sql)
if err != nil {
return nil, fmt.Errorf("unable to execute query: %w", err)