github/genai-toolbox
1
0
Fork 0
mirror of https://github.com/googleapis/genai-toolbox.git synced 2026-02-17 02:22:30 -05:00
Code Issues Packages Projects Releases 14 Wiki Activity

fix(sources/cockroachdb): update kind to type (#2465)

Browse Source 
Fix failing integration test, clean up source code from `kind` to
`type`.
This commit is contained in:
Yuan Teoh
2026-02-12 15:59:55 -08:00
committed by GitHub
parent f032389a07
commit 2d341acaa6
6 changed files with 32 additions and 57 deletions
Download Patch File Download Diff File Expand all files Collapse all files

15
internal/sources/cockroachdb/cockroachdb.go
View File

@@ -34,14 +34,13 @@ import (
"go.opentelemetry.io/otel/trace" "go.opentelemetry.io/otel/trace"
) )
const SourceKind string = "cockroachdb"
const SourceType string = "cockroachdb" const SourceType string = "cockroachdb"
var _ sources.SourceConfig = Config{} var _ sources.SourceConfig = Config{}
func init() { func init() {
if !sources.Register(SourceKind, newConfig) { if !sources.Register(SourceType, newConfig) {
panic(fmt.Sprintf("source kind %q already registered", SourceKind)) panic(fmt.Sprintf("source type %q already registered", SourceType))
} }
} }
@@ -94,10 +93,6 @@ type Config struct {
ClusterID string `yaml:"clusterID"` // Optional cluster identifier for telemetry ClusterID string `yaml:"clusterID"` // Optional cluster identifier for telemetry
} }
func (r Config) SourceConfigKind() string {
return SourceKind
}
func (r Config) SourceConfigType() string { func (r Config) SourceConfigType() string {
return SourceType return SourceType
} }
@@ -127,10 +122,6 @@ type Source struct {
Pool *pgxpool.Pool Pool *pgxpool.Pool
} }
func (s *Source) SourceKind() string {
return SourceKind
}
func (s *Source) SourceType() string { func (s *Source) SourceType() string {
return SourceType return SourceType
} }
@@ -379,7 +370,7 @@ func (s *Source) EmitTelemetry(ctx context.Context, event TelemetryEvent) {
func initCockroachDBConnectionPoolWithRetry(ctx context.Context, tracer trace.Tracer, name, host, port, user, pass, dbname string, queryParams map[string]string, maxRetries int, baseDelay time.Duration) (*pgxpool.Pool, error) { func initCockroachDBConnectionPoolWithRetry(ctx context.Context, tracer trace.Tracer, name, host, port, user, pass, dbname string, queryParams map[string]string, maxRetries int, baseDelay time.Duration) (*pgxpool.Pool, error) {
//nolint:all //nolint:all
ctx, span := sources.InitConnectionSpan(ctx, tracer, SourceKind, name) ctx, span := sources.InitConnectionSpan(ctx, tracer, SourceType, name)
defer span.End() defer span.End()
userAgent, err := util.UserAgentFromContext(ctx) userAgent, err := util.UserAgentFromContext(ctx)

16
internal/tools/cockroachdb/cockroachdbexecutesql/cockroachdbexecutesql.go
View File

@@ -30,11 +30,11 @@ import (
"github.com/jackc/pgx/v5" "github.com/jackc/pgx/v5"
) )
const kind string = "cockroachdb-execute-sql" const resourceType string = "cockroachdb-execute-sql"
func init() { func init() {
if !tools.Register(kind, newConfig) { if !tools.Register(resourceType, newConfig) {
panic(fmt.Sprintf("tool kind %q already registered", kind)) panic(fmt.Sprintf("tool type %q already registered", resourceType))
} }
} }
@@ -50,7 +50,7 @@ type compatibleSource interface {
Query(ctx context.Context, sql string, args ...interface{}) (pgx.Rows, error) Query(ctx context.Context, sql string, args ...interface{}) (pgx.Rows, error)
} }
var compatibleSources = [...]string{cockroachdb.SourceKind} var compatibleSources = [...]string{cockroachdb.SourceType}
type Config struct { type Config struct {
Name string `yaml:"name" validate:"required"` Name string `yaml:"name" validate:"required"`
@@ -62,12 +62,8 @@ type Config struct {
var _ tools.ToolConfig = Config{} var _ tools.ToolConfig = Config{}
func (cfg Config) ToolConfigKind() string {
return kind
}
func (cfg Config) ToolConfigType() string { func (cfg Config) ToolConfigType() string {
return kind return resourceType
} }
func (cfg Config) Initialize(srcs map[string]sources.Source) (tools.Tool, error) { func (cfg Config) Initialize(srcs map[string]sources.Source) (tools.Tool, error) {
@@ -78,7 +74,7 @@ func (cfg Config) Initialize(srcs map[string]sources.Source) (tools.Tool, error)
_, ok = rawS.(compatibleSource) _, ok = rawS.(compatibleSource)
if !ok { if !ok {
return nil, fmt.Errorf("invalid source for %q tool: source kind must be one of %q", kind, compatibleSources) return nil, fmt.Errorf("invalid source for %q tool: source type must be one of %q", resourceType, compatibleSources)
} }
sqlParameter := parameters.NewStringParameter("sql", "The sql to execute.") sqlParameter := parameters.NewStringParameter("sql", "The sql to execute.")

16
internal/tools/cockroachdb/cockroachdblistschemas/cockroachdblistschemas.go
View File

@@ -29,7 +29,7 @@ import (
"github.com/jackc/pgx/v5" "github.com/jackc/pgx/v5"
) )
const kind string = "cockroachdb-list-schemas" const resourceType string = "cockroachdb-list-schemas"
const listSchemasStatement = ` const listSchemasStatement = `
SELECT SELECT
@@ -44,8 +44,8 @@ const listSchemasStatement = `
` `
func init() { func init() {
if !tools.Register(kind, newConfig) { if !tools.Register(resourceType, newConfig) {
panic(fmt.Sprintf("tool kind %q already registered", kind)) panic(fmt.Sprintf("tool type %q already registered", resourceType))
} }
} }
@@ -63,7 +63,7 @@ type compatibleSource interface {
var _ compatibleSource = &cockroachdb.Source{} var _ compatibleSource = &cockroachdb.Source{}
var compatibleSources = [...]string{cockroachdb.SourceKind} var compatibleSources = [...]string{cockroachdb.SourceType}
type Config struct { type Config struct {
Name string `yaml:"name" validate:"required"` Name string `yaml:"name" validate:"required"`
@@ -75,12 +75,8 @@ type Config struct {
var _ tools.ToolConfig = Config{} var _ tools.ToolConfig = Config{}
func (cfg Config) ToolConfigKind() string {
return kind
}
func (cfg Config) ToolConfigType() string { func (cfg Config) ToolConfigType() string {
return kind return resourceType
} }
func (cfg Config) Initialize(srcs map[string]sources.Source) (tools.Tool, error) { func (cfg Config) Initialize(srcs map[string]sources.Source) (tools.Tool, error) {
@@ -91,7 +87,7 @@ func (cfg Config) Initialize(srcs map[string]sources.Source) (tools.Tool, error)
_, ok = rawS.(compatibleSource) _, ok = rawS.(compatibleSource)
if !ok { if !ok {
return nil, fmt.Errorf("invalid source for %q tool: source kind must be one of %q", kind, compatibleSources) return nil, fmt.Errorf("invalid source for %q tool: source type must be one of %q", resourceType, compatibleSources)
} }
allParameters := parameters.Parameters{} allParameters := parameters.Parameters{}

16
internal/tools/cockroachdb/cockroachdblisttables/cockroachdblisttables.go
View File

@@ -29,7 +29,7 @@ import (
"github.com/jackc/pgx/v5" "github.com/jackc/pgx/v5"
) )
const kind string = "cockroachdb-list-tables" const resourceType string = "cockroachdb-list-tables"
const listTablesStatement = ` const listTablesStatement = `
WITH desired_relkinds AS ( WITH desired_relkinds AS (
@@ -104,8 +104,8 @@ const listTablesStatement = `
` `
func init() { func init() {
if !tools.Register(kind, newConfig) { if !tools.Register(resourceType, newConfig) {
panic(fmt.Sprintf("tool kind %q already registered", kind)) panic(fmt.Sprintf("tool type %q already registered", resourceType))
} }
} }
@@ -123,7 +123,7 @@ type compatibleSource interface {
var _ compatibleSource = &cockroachdb.Source{} var _ compatibleSource = &cockroachdb.Source{}
var compatibleSources = [...]string{cockroachdb.SourceKind} var compatibleSources = [...]string{cockroachdb.SourceType}
type Config struct { type Config struct {
Name string `yaml:"name" validate:"required"` Name string `yaml:"name" validate:"required"`
@@ -135,12 +135,8 @@ type Config struct {
var _ tools.ToolConfig = Config{} var _ tools.ToolConfig = Config{}
func (cfg Config) ToolConfigKind() string {
return kind
}
func (cfg Config) ToolConfigType() string { func (cfg Config) ToolConfigType() string {
return kind return resourceType
} }
func (cfg Config) Initialize(srcs map[string]sources.Source) (tools.Tool, error) { func (cfg Config) Initialize(srcs map[string]sources.Source) (tools.Tool, error) {
@@ -151,7 +147,7 @@ func (cfg Config) Initialize(srcs map[string]sources.Source) (tools.Tool, error)
_, ok = rawS.(compatibleSource) _, ok = rawS.(compatibleSource)
if !ok { if !ok {
return nil, fmt.Errorf("invalid source for %q tool: source kind must be one of %q", kind, compatibleSources) return nil, fmt.Errorf("invalid source for %q tool: source type must be one of %q", resourceType, compatibleSources)
} }
allParameters := parameters.Parameters{ allParameters := parameters.Parameters{

16
internal/tools/cockroachdb/cockroachdbsql/cockroachdbsql.go
View File

@@ -30,11 +30,11 @@ import (
"github.com/jackc/pgx/v5" "github.com/jackc/pgx/v5"
) )
const kind string = "cockroachdb-sql" const resourceType string = "cockroachdb-sql"
func init() { func init() {
if !tools.Register(kind, newConfig) { if !tools.Register(resourceType, newConfig) {
panic(fmt.Sprintf("tool kind %q already registered", kind)) panic(fmt.Sprintf("tool type %q already registered", resourceType))
} }
} }
@@ -52,7 +52,7 @@ type compatibleSource interface {
var _ compatibleSource = &cockroachdb.Source{} var _ compatibleSource = &cockroachdb.Source{}
var compatibleSources = [...]string{cockroachdb.SourceKind} var compatibleSources = [...]string{cockroachdb.SourceType}
type Config struct { type Config struct {
Name string `yaml:"name" validate:"required"` Name string `yaml:"name" validate:"required"`
@@ -67,12 +67,8 @@ type Config struct {
var _ tools.ToolConfig = Config{} var _ tools.ToolConfig = Config{}
func (cfg Config) ToolConfigKind() string {
return kind
}
func (cfg Config) ToolConfigType() string { func (cfg Config) ToolConfigType() string {
return kind return resourceType
} }
func (cfg Config) Initialize(srcs map[string]sources.Source) (tools.Tool, error) { func (cfg Config) Initialize(srcs map[string]sources.Source) (tools.Tool, error) {
@@ -83,7 +79,7 @@ func (cfg Config) Initialize(srcs map[string]sources.Source) (tools.Tool, error)
_, ok = rawS.(compatibleSource) _, ok = rawS.(compatibleSource)
if !ok { if !ok {
return nil, fmt.Errorf("invalid source for %q tool: source kind must be one of %q", kind, compatibleSources) return nil, fmt.Errorf("invalid source for %q tool: source type must be one of %q", resourceType, compatibleSources)
} }
allParameters, paramManifest, err := parameters.ProcessParameters(cfg.TemplateParameters, cfg.Parameters) allParameters, paramManifest, err := parameters.ProcessParameters(cfg.TemplateParameters, cfg.Parameters)

10
tests/cockroachdb/cockroachdb_integration_test.go
View File

@@ -31,8 +31,8 @@ import (
) )
var ( var (
CockroachDBSourceKind = "cockroachdb" CockroachDBSourceType = "cockroachdb"
CockroachDBToolKind = "cockroachdb-sql" CockroachDBToolType = "cockroachdb-sql"
CockroachDBDatabase = getEnvOrDefault("COCKROACHDB_DATABASE", "defaultdb") CockroachDBDatabase = getEnvOrDefault("COCKROACHDB_DATABASE", "defaultdb")
CockroachDBHost = getEnvOrDefault("COCKROACHDB_HOST", "localhost") CockroachDBHost = getEnvOrDefault("COCKROACHDB_HOST", "localhost")
CockroachDBPort = getEnvOrDefault("COCKROACHDB_PORT", "26257") CockroachDBPort = getEnvOrDefault("COCKROACHDB_PORT", "26257")
@@ -53,7 +53,7 @@ func getCockroachDBVars(t *testing.T) map[string]any {
} }
return map[string]any{ return map[string]any{
"type": CockroachDBSourceKind, "type": CockroachDBSourceType,
"host": CockroachDBHost, "host": CockroachDBHost,
"port": CockroachDBPort, "port": CockroachDBPort,
"database": CockroachDBDatabase, "database": CockroachDBDatabase,
@@ -128,13 +128,13 @@ func TestCockroachDB(t *testing.T) {
defer teardownTable2(t) defer teardownTable2(t)
// Write config into a file and pass it to command // Write config into a file and pass it to command
toolsFile := tests.GetToolsConfig(sourceConfig, CockroachDBToolKind, paramToolStmt, idParamToolStmt, nameParamToolStmt, arrayToolStmt, authToolStmt) toolsFile := tests.GetToolsConfig(sourceConfig, CockroachDBToolType, paramToolStmt, idParamToolStmt, nameParamToolStmt, arrayToolStmt, authToolStmt)
// Add execute-sql tool with write-enabled source (CockroachDB MCP security requires explicit opt-in) // Add execute-sql tool with write-enabled source (CockroachDB MCP security requires explicit opt-in)
toolsFile = addCockroachDBExecuteSqlConfig(t, toolsFile, sourceConfig) toolsFile = addCockroachDBExecuteSqlConfig(t, toolsFile, sourceConfig)
tmplSelectCombined, tmplSelectFilterCombined := tests.GetPostgresSQLTmplToolStatement() tmplSelectCombined, tmplSelectFilterCombined := tests.GetPostgresSQLTmplToolStatement()
toolsFile = tests.AddTemplateParamConfig(t, toolsFile, CockroachDBToolKind, tmplSelectCombined, tmplSelectFilterCombined, "") toolsFile = tests.AddTemplateParamConfig(t, toolsFile, CockroachDBToolType, tmplSelectCombined, tmplSelectFilterCombined, "")
cmd, cleanup, err := tests.StartCmd(ctx, toolsFile, args...) cmd, cleanup, err := tests.StartCmd(ctx, toolsFile, args...)
if err != nil { if err != nil {