github/genai-toolbox
1
0
Fork 0
mirror of https://github.com/googleapis/genai-toolbox.git synced 2026-02-13 08:35:15 -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"
)
const SourceKind string = "cockroachdb"
const SourceType string = "cockroachdb"
var _ sources.SourceConfig = Config{}
func init() {
if !sources.Register(SourceKind, newConfig) {
panic(fmt.Sprintf("source kind %q already registered", SourceKind))
if !sources.Register(SourceType, newConfig) {
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
}
func (r Config) SourceConfigKind() string {
return SourceKind
}
func (r Config) SourceConfigType() string {
return SourceType
}
@@ -127,10 +122,6 @@ type Source struct {
Pool *pgxpool.Pool
}
func (s *Source) SourceKind() string {
return SourceKind
}
func (s *Source) SourceType() string {
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) {
//nolint:all
ctx, span := sources.InitConnectionSpan(ctx, tracer, SourceKind, name)
ctx, span := sources.InitConnectionSpan(ctx, tracer, SourceType, name)
defer span.End()
userAgent, err := util.UserAgentFromContext(ctx)

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

@@ -30,11 +30,11 @@ import (
"github.com/jackc/pgx/v5"
)
const kind string = "cockroachdb-execute-sql"
const resourceType string = "cockroachdb-execute-sql"
func init() {
if !tools.Register(kind, newConfig) {
panic(fmt.Sprintf("tool kind %q already registered", kind))
if !tools.Register(resourceType, newConfig) {
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)
}
var compatibleSources = [...]string{cockroachdb.SourceKind}
var compatibleSources = [...]string{cockroachdb.SourceType}
type Config struct {
Name string `yaml:"name" validate:"required"`
@@ -62,12 +62,8 @@ type Config struct {
var _ tools.ToolConfig = Config{}
func (cfg Config) ToolConfigKind() string {
return kind
}
func (cfg Config) ToolConfigType() string {
return kind
return resourceType
}
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)
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.")

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

@@ -29,7 +29,7 @@ import (
"github.com/jackc/pgx/v5"
)
const kind string = "cockroachdb-list-schemas"
const resourceType string = "cockroachdb-list-schemas"
const listSchemasStatement = `
SELECT
@@ -44,8 +44,8 @@ const listSchemasStatement = `
`
func init() {
if !tools.Register(kind, newConfig) {
panic(fmt.Sprintf("tool kind %q already registered", kind))
if !tools.Register(resourceType, newConfig) {
panic(fmt.Sprintf("tool type %q already registered", resourceType))
}
}
@@ -63,7 +63,7 @@ type compatibleSource interface {
var _ compatibleSource = &cockroachdb.Source{}
var compatibleSources = [...]string{cockroachdb.SourceKind}
var compatibleSources = [...]string{cockroachdb.SourceType}
type Config struct {
Name string `yaml:"name" validate:"required"`
@@ -75,12 +75,8 @@ type Config struct {
var _ tools.ToolConfig = Config{}
func (cfg Config) ToolConfigKind() string {
return kind
}
func (cfg Config) ToolConfigType() string {
return kind
return resourceType
}
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)
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{}

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

@@ -29,7 +29,7 @@ import (
"github.com/jackc/pgx/v5"
)
const kind string = "cockroachdb-list-tables"
const resourceType string = "cockroachdb-list-tables"
const listTablesStatement = `
WITH desired_relkinds AS (
@@ -104,8 +104,8 @@ const listTablesStatement = `
`
func init() {
if !tools.Register(kind, newConfig) {
panic(fmt.Sprintf("tool kind %q already registered", kind))
if !tools.Register(resourceType, newConfig) {
panic(fmt.Sprintf("tool type %q already registered", resourceType))
}
}
@@ -123,7 +123,7 @@ type compatibleSource interface {
var _ compatibleSource = &cockroachdb.Source{}
var compatibleSources = [...]string{cockroachdb.SourceKind}
var compatibleSources = [...]string{cockroachdb.SourceType}
type Config struct {
Name string `yaml:"name" validate:"required"`
@@ -135,12 +135,8 @@ type Config struct {
var _ tools.ToolConfig = Config{}
func (cfg Config) ToolConfigKind() string {
return kind
}
func (cfg Config) ToolConfigType() string {
return kind
return resourceType
}
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)
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{

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

@@ -30,11 +30,11 @@ import (
"github.com/jackc/pgx/v5"
)
const kind string = "cockroachdb-sql"
const resourceType string = "cockroachdb-sql"
func init() {
if !tools.Register(kind, newConfig) {
panic(fmt.Sprintf("tool kind %q already registered", kind))
if !tools.Register(resourceType, newConfig) {
panic(fmt.Sprintf("tool type %q already registered", resourceType))
}
}
@@ -52,7 +52,7 @@ type compatibleSource interface {
var _ compatibleSource = &cockroachdb.Source{}
var compatibleSources = [...]string{cockroachdb.SourceKind}
var compatibleSources = [...]string{cockroachdb.SourceType}
type Config struct {
Name string `yaml:"name" validate:"required"`
@@ -67,12 +67,8 @@ type Config struct {
var _ tools.ToolConfig = Config{}
func (cfg Config) ToolConfigKind() string {
return kind
}
func (cfg Config) ToolConfigType() string {
return kind
return resourceType
}
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)
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)

10
tests/cockroachdb/cockroachdb_integration_test.go
View File

@@ -31,8 +31,8 @@ import (
)
var (
CockroachDBSourceKind = "cockroachdb"
CockroachDBToolKind = "cockroachdb-sql"
CockroachDBSourceType = "cockroachdb"
CockroachDBToolType = "cockroachdb-sql"
CockroachDBDatabase = getEnvOrDefault("COCKROACHDB_DATABASE", "defaultdb")
CockroachDBHost = getEnvOrDefault("COCKROACHDB_HOST", "localhost")
CockroachDBPort = getEnvOrDefault("COCKROACHDB_PORT", "26257")
@@ -53,7 +53,7 @@ func getCockroachDBVars(t *testing.T) map[string]any {
}
return map[string]any{
"type": CockroachDBSourceKind,
"type": CockroachDBSourceType,
"host": CockroachDBHost,
"port": CockroachDBPort,
"database": CockroachDBDatabase,
@@ -128,13 +128,13 @@ func TestCockroachDB(t *testing.T) {
defer teardownTable2(t)
// 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)
toolsFile = addCockroachDBExecuteSqlConfig(t, toolsFile, sourceConfig)
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...)
if err != nil {