diff --git a/go.mod b/go.mod index 4d27ca437..3ff2ce4f8 100644 --- a/go.mod +++ b/go.mod @@ -13,6 +13,7 @@ require ( github.com/go-chi/chi/v5 v5.1.0 github.com/go-chi/httplog/v2 v2.1.1 github.com/go-chi/render v1.0.3 + github.com/go-sql-driver/mysql v1.8.1 github.com/goccy/go-yaml v1.15.13 github.com/google/go-cmp v0.6.0 github.com/google/uuid v1.6.0 @@ -55,7 +56,6 @@ require ( github.com/felixge/httpsnoop v1.0.4 // indirect github.com/go-logr/logr v1.4.2 // indirect github.com/go-logr/stdr v1.2.2 // indirect - github.com/go-sql-driver/mysql v1.8.1 // indirect github.com/golang-sql/civil v0.0.0-20220223132316-b832511892a9 // indirect github.com/golang-sql/sqlexp v0.1.0 // indirect github.com/golang/groupcache v0.0.0-20241129210726-2c02b8208cf8 // indirect diff --git a/internal/server/config.go b/internal/server/config.go index d54683225..34a12f0a2 100644 --- a/internal/server/config.go +++ b/internal/server/config.go @@ -25,6 +25,7 @@ import ( cloudsqlmssqlsrc "github.com/googleapis/genai-toolbox/internal/sources/cloudsqlmssql" cloudsqlmysqlsrc "github.com/googleapis/genai-toolbox/internal/sources/cloudsqlmysql" cloudsqlpgsrc "github.com/googleapis/genai-toolbox/internal/sources/cloudsqlpg" + mysqlsrc "github.com/googleapis/genai-toolbox/internal/sources/mysql" neo4jrc "github.com/googleapis/genai-toolbox/internal/sources/neo4j" postgressrc "github.com/googleapis/genai-toolbox/internal/sources/postgres" spannersrc "github.com/googleapis/genai-toolbox/internal/sources/spanner" @@ -163,6 +164,12 @@ func (c *SourceConfigs) UnmarshalYAML(unmarshal func(interface{}) error) error { return fmt.Errorf("unable to parse as %q: %w", k.Kind, err) } (*c)[name] = actual + case mysqlsrc.SourceKind: + actual := mysqlsrc.Config{Name: name} + if err := u.Unmarshal(&actual); err != nil { + return fmt.Errorf("unable to parse as %q: %w", k.Kind, err) + } + (*c)[name] = actual case spannersrc.SourceKind: actual := spannersrc.Config{Name: name, Dialect: "googlesql"} if err := u.Unmarshal(&actual); err != nil { diff --git a/internal/sources/mysql/mysql.go b/internal/sources/mysql/mysql.go new file mode 100644 index 000000000..c1321e807 --- /dev/null +++ b/internal/sources/mysql/mysql.go @@ -0,0 +1,95 @@ +// 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 mysql + +import ( + "context" + "database/sql" + "fmt" + + _ "github.com/go-sql-driver/mysql" + "github.com/googleapis/genai-toolbox/internal/sources" + "go.opentelemetry.io/otel/trace" +) + +const SourceKind string = "mysql" + +// validate interface +var _ sources.SourceConfig = Config{} + +type Config struct { + Name string `yaml:"name"` + Kind string `yaml:"kind"` + Host string `yaml:"host"` + Port string `yaml:"port"` + User string `yaml:"user"` + Password string `yaml:"password"` + Database string `yaml:"database"` +} + +func (r Config) SourceConfigKind() string { + return SourceKind +} + +func (r Config) Initialize(ctx context.Context, tracer trace.Tracer) (sources.Source, error) { + pool, err := initMySQLConnectionPool(ctx, tracer, r.Name, r.Host, r.Port, r.User, r.Password, r.Database) + if err != nil { + return nil, fmt.Errorf("unable to create pool: %w", err) + } + + err = pool.PingContext(context.Background()) + if err != nil { + return nil, fmt.Errorf("unable to connect successfully: %w", err) + } + + s := &Source{ + Name: r.Name, + Kind: SourceKind, + Pool: pool, + } + return s, nil +} + +var _ sources.Source = &Source{} + +type Source struct { + Name string `yaml:"name"` + Kind string `yaml:"kind"` + Pool *sql.DB +} + +func (s *Source) SourceKind() string { + return SourceKind +} + +func (s *Source) MySQLPool() *sql.DB { + return s.Pool +} + +func initMySQLConnectionPool(ctx context.Context, tracer trace.Tracer, name, host, port, user, pass, dbname string) (*sql.DB, error) { + //nolint:all // Reassigned ctx + ctx, span := sources.InitConnectionSpan(ctx, tracer, SourceKind, name) + defer span.End() + + // Configure the driver to connect to the database + dsn := fmt.Sprintf("%s:%s@tcp(%s:%s)/%s?parseTime=true", user, pass, host, port, dbname) + + // Interact with the driver directly as you normally would + pool, err := sql.Open("mysql", dsn) + if err != nil { + return nil, fmt.Errorf("sql.Open: %w", err) + } + return pool, nil +} diff --git a/internal/sources/mysql/mysql_test.go b/internal/sources/mysql/mysql_test.go new file mode 100644 index 000000000..3e312a20e --- /dev/null +++ b/internal/sources/mysql/mysql_test.go @@ -0,0 +1,70 @@ +// 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 mysql_test + +import ( + "testing" + + "github.com/google/go-cmp/cmp" + "github.com/googleapis/genai-toolbox/internal/server" + "github.com/googleapis/genai-toolbox/internal/sources/mysql" + "github.com/googleapis/genai-toolbox/internal/testutils" + "gopkg.in/yaml.v3" +) + +func TestParseFromYamlCloudSQLMySQL(t *testing.T) { + tcs := []struct { + desc string + in string + want server.SourceConfigs + }{ + { + desc: "basic example", + in: ` + sources: + my-mysql-instance: + kind: mysql + host: 0.0.0.0 + port: my-host + database: my_db + `, + want: server.SourceConfigs{ + "my-mysql-instance": mysql.Config{ + Name: "my-mysql-instance", + Kind: mysql.SourceKind, + Host: "0.0.0.0", + Port: "my-host", + Database: "my_db", + }, + }, + }, + } + for _, tc := range tcs { + t.Run(tc.desc, func(t *testing.T) { + got := struct { + Sources server.SourceConfigs `yaml:"sources"` + }{} + // Parse contents + err := yaml.Unmarshal(testutils.FormatYaml(tc.in), &got) + if err != nil { + t.Fatalf("unable to unmarshal: %s", err) + } + if !cmp.Equal(tc.want, got.Sources) { + t.Fatalf("incorrect parse: want %v, got %v", tc.want, got.Sources) + } + }) + } + +} diff --git a/internal/tools/mysqlsql/mysqlsql.go b/internal/tools/mysqlsql/mysqlsql.go index 968f4abdf..4317969bb 100644 --- a/internal/tools/mysqlsql/mysqlsql.go +++ b/internal/tools/mysqlsql/mysqlsql.go @@ -22,6 +22,7 @@ import ( "github.com/googleapis/genai-toolbox/internal/sources" "github.com/googleapis/genai-toolbox/internal/sources/cloudsqlmysql" + "github.com/googleapis/genai-toolbox/internal/sources/mysql" "github.com/googleapis/genai-toolbox/internal/tools" ) @@ -33,8 +34,9 @@ type compatibleSource interface { // validate compatible sources are still compatible var _ compatibleSource = &cloudsqlmysql.Source{} +var _ compatibleSource = &mysql.Source{} -var compatibleSources = [...]string{cloudsqlmysql.SourceKind} +var compatibleSources = [...]string{cloudsqlmysql.SourceKind, mysql.SourceKind} type Config struct { Name string `yaml:"name"`