Skip to content

Commit 73e2a8c

Browse files
feat(sources/singlestore): Add ConnectionParams to SingleStore Config (#2555)
## Description Similarly to MySQL and Postgres, we are adding `connectionParams` field to the SingleStore config. ## 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: - [x] 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 - [x] Ensure the tests and linter pass - [x] Code coverage does not decrease (if any source code was changed) - [x] Appropriate docs were updated (if necessary) - [x] Make sure to add `!` if this involve a breaking change 🛠️ Fixes #<issue_number_goes_here> Co-authored-by: Wenxin Du <117315983+duwenxin99@users.noreply.github.com>
1 parent 377dc5b commit 73e2a8c

4 files changed

Lines changed: 249 additions & 41 deletions

File tree

docs/en/integrations/singlestore/source.md

Lines changed: 60 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -38,6 +38,11 @@ database user][singlestore-user] to login to the database with.
3838

3939
## Example
4040

41+
### Basic
42+
43+
By default, connections use `tls=preferred`, which enables SSL/TLS if the server
44+
supports it and falls back to unencrypted otherwise.
45+
4146
```yaml
4247
kind: source
4348
name: my-singlestore-source
@@ -50,19 +55,65 @@ password: ${PASSWORD}
5055
queryTimeout: 30s # Optional: query timeout duration
5156
```
5257
58+
### With SSL required
59+
60+
```yaml
61+
kind: sources
62+
name: my-singlestore-source
63+
type: singlestore
64+
host: svc-abc123.svc.singlestore.com
65+
port: 3306
66+
database: my_db
67+
user: ${USER_NAME}
68+
password: ${PASSWORD}
69+
connectionParams:
70+
tls: "true" # Require TLS and verify the server certificate
71+
```
72+
73+
### With SSL verification disabled
74+
75+
```yaml
76+
kind: sources
77+
name: my-singlestore-source
78+
type: singlestore
79+
host: svc-abc123.svc.singlestore.com
80+
port: 3306
81+
database: my_db
82+
user: ${USER_NAME}
83+
password: ${PASSWORD}
84+
connectionParams:
85+
tls: "skip-verify" # Require TLS but skip server certificate verification
86+
```
87+
88+
### Without SSL
89+
90+
```yaml
91+
kind: sources
92+
name: my-singlestore-source
93+
type: singlestore
94+
host: 127.0.0.1
95+
port: 3306
96+
database: my_db
97+
user: ${USER_NAME}
98+
password: ${PASSWORD}
99+
connectionParams:
100+
tls: "false" # Disable TLS
101+
```
102+
53103
{{< notice tip >}}
54104
Use environment variable replacement with the format ${ENV_NAME}
55105
instead of hardcoding your secrets into the configuration file.
56106
{{< /notice >}}
57107
58108
## Reference
59109
60-
| **field** | **type** | **required** | **description** |
61-
|--------------|:--------:|:------------:|-------------------------------------------------------------------------------------------------|
62-
| type | string | true | Must be "singlestore". |
63-
| host | string | true | IP address to connect to (e.g. "127.0.0.1"). |
64-
| port | string | true | Port to connect to (e.g. "3306"). |
65-
| database | string | true | Name of the SingleStore database to connect to (e.g. "my_db"). |
66-
| user | string | true | Name of the SingleStore database user to connect as (e.g. "admin"). |
67-
| password | string | true | Password of the SingleStore database user. |
68-
| queryTimeout | string | false | Maximum time to wait for query execution (e.g. "30s", "2m"). By default, no timeout is applied. |
110+
| **field** | **type** | **required** | **description** |
111+
|------------------|:-----------------:|:------------:|------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|
112+
| type | string | true | Must be "singlestore". |
113+
| host | string | true | IP address or hostname to connect to (e.g. "127.0.0.1"). |
114+
| port | string | true | Port to connect to (e.g. "3306"). |
115+
| database | string | true | Name of the SingleStore database to connect to (e.g. "my_db"). |
116+
| user | string | true | Name of the SingleStore database user to connect as (e.g. "admin"). |
117+
| password | string | true | Password of the SingleStore database user. |
118+
| queryTimeout | string | false | Maximum time to wait for query execution (e.g. "30s", "2m"). By default, no timeout is applied. |
119+
| connectionParams | map[string]string | false | Additional driver parameters appended to the DSN. Supports any [go-sql-driver/mysql DSN parameter](https://github.com/go-sql-driver/mysql#dsn-data-source-name). Commonly used for SSL configuration (see `tls` values below). |

internal/sources/singlestore/singlestore.go

Lines changed: 49 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -19,10 +19,9 @@ import (
1919
"database/sql"
2020
"fmt"
2121
"net/url"
22-
"strings"
2322
"time"
2423

25-
_ "github.com/go-sql-driver/mysql"
24+
"github.com/go-sql-driver/mysql"
2625
"github.com/goccy/go-yaml"
2726
"github.com/googleapis/genai-toolbox/internal/sources"
2827
"github.com/googleapis/genai-toolbox/internal/tools/mysql/mysqlcommon"
@@ -51,14 +50,15 @@ func newConfig(ctx context.Context, name string, decoder *yaml.Decoder) (sources
5150

5251
// Config holds the configuration parameters for connecting to a SingleStore database.
5352
type Config struct {
54-
Name string `yaml:"name" validate:"required"`
55-
Type string `yaml:"type" validate:"required"`
56-
Host string `yaml:"host" validate:"required"`
57-
Port string `yaml:"port" validate:"required"`
58-
User string `yaml:"user" validate:"required"`
59-
Password string `yaml:"password" validate:"required"`
60-
Database string `yaml:"database" validate:"required"`
61-
QueryTimeout string `yaml:"queryTimeout"`
53+
Name string `yaml:"name" validate:"required"`
54+
Type string `yaml:"type" validate:"required"`
55+
Host string `yaml:"host" validate:"required"`
56+
Port string `yaml:"port" validate:"required"`
57+
User string `yaml:"user" validate:"required"`
58+
Password string `yaml:"password" validate:"required"`
59+
Database string `yaml:"database" validate:"required"`
60+
QueryTimeout string `yaml:"queryTimeout"`
61+
ConnectionParams map[string]string `yaml:"connectionParams"`
6262
}
6363

6464
// SourceConfigType returns the type of the source configuration.
@@ -68,7 +68,7 @@ func (r Config) SourceConfigType() string {
6868

6969
// Initialize sets up the SingleStore connection pool and returns a Source.
7070
func (r Config) Initialize(ctx context.Context, tracer trace.Tracer) (sources.Source, error) {
71-
pool, err := initSingleStoreConnectionPool(ctx, tracer, r.Name, r.Host, r.Port, r.User, r.Password, r.Database, r.QueryTimeout)
71+
pool, err := initSingleStoreConnectionPool(ctx, tracer, r)
7272
if err != nil {
7373
return nil, fmt.Errorf("unable to create pool: %w", err)
7474
}
@@ -160,31 +160,52 @@ func (s *Source) RunSQL(ctx context.Context, statement string, params []any) (an
160160
return out, nil
161161
}
162162

163-
func initSingleStoreConnectionPool(ctx context.Context, tracer trace.Tracer, name, host, port, user, pass, dbname, queryTimeout string) (*sql.DB, error) {
163+
func initSingleStoreConnectionPool(ctx context.Context, tracer trace.Tracer, cfg Config) (*sql.DB, error) {
164164
//nolint:all // Reassigned ctx
165-
ctx, span := sources.InitConnectionSpan(ctx, tracer, SourceType, name)
165+
ctx, span := sources.InitConnectionSpan(ctx, tracer, SourceType, cfg.Name)
166166
defer span.End()
167167

168-
// Configure the driver to connect to the database
169-
dsn := fmt.Sprintf("%s:%s@tcp(%s:%s)/%s?parseTime=true&vector_type_project_format=JSON", user, pass, host, port, dbname)
168+
// Build query parameters via url.Values for deterministic order and proper escaping.
169+
connectionParams := url.Values{}
170+
171+
mysqlCfg := mysql.Config{
172+
User: cfg.User,
173+
Passwd: cfg.Password,
174+
Net: "tcp",
175+
Addr: fmt.Sprintf("%s:%s", cfg.Host, cfg.Port),
176+
DBName: cfg.Database,
177+
ParseTime: true,
178+
AllowNativePasswords: true,
179+
CheckConnLiveness: true,
180+
MaxAllowedPacket: 64 << 20,
181+
ConnectionAttributes: "_connector_name:MCP toolbox for Databases",
182+
Params: map[string]string{
183+
"vector_type_project_format": "JSON",
184+
},
185+
}
170186

171-
// Add connection attributes to DSN
172-
customAttrs := []string{"_connector_name"}
173-
customAttrValues := []string{"MCP toolbox for Databases"}
187+
// Default to TLS preferred; can be overridden via connectionParams.
188+
connectionParams.Set("tls", "preferred")
174189

175-
customAttrStrs := make([]string, len(customAttrs))
176-
for i := range customAttrs {
177-
customAttrStrs[i] = fmt.Sprintf("%s:%s", customAttrs[i], customAttrValues[i])
190+
// Derive readTimeout from queryTimeout when provided.
191+
if cfg.QueryTimeout != "" {
192+
timeout, err := time.ParseDuration(cfg.QueryTimeout)
193+
if err != nil {
194+
return nil, fmt.Errorf("invalid queryTimeout %q: %w", cfg.QueryTimeout, err)
195+
}
196+
connectionParams.Set("readTimeout", timeout.String())
178197
}
179-
dsn += "&connectionAttributes=" + url.QueryEscape(strings.Join(customAttrStrs, ","))
180198

181-
// Add query timeout to DSN if specified
182-
if queryTimeout != "" {
183-
timeout, err := time.ParseDuration(queryTimeout)
184-
if err != nil {
185-
return nil, fmt.Errorf("invalid queryTimeout %q: %w", queryTimeout, err)
199+
// Custom user parameters (e.g. tls, compress) — may override defaults above.
200+
for k, v := range cfg.ConnectionParams {
201+
if v == "" {
202+
continue // skip empty values
186203
}
187-
dsn += "&readTimeout=" + timeout.String()
204+
connectionParams.Set(k, v)
205+
}
206+
dsn := mysqlCfg.FormatDSN()
207+
if enc := connectionParams.Encode(); enc != "" {
208+
dsn += "&" + enc
188209
}
189210

190211
// Interact with the driver directly as you normally would

internal/sources/singlestore/singlestore_test.go

Lines changed: 85 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -16,9 +16,12 @@ package singlestore_test
1616

1717
import (
1818
"context"
19+
"strings"
1920
"testing"
2021

2122
"github.com/google/go-cmp/cmp"
23+
"go.opentelemetry.io/otel/trace/noop"
24+
2225
"github.com/googleapis/genai-toolbox/internal/server"
2326
"github.com/googleapis/genai-toolbox/internal/sources"
2427
"github.com/googleapis/genai-toolbox/internal/sources/singlestore"
@@ -81,6 +84,66 @@ func TestParseFromYaml(t *testing.T) {
8184
},
8285
},
8386
},
87+
{
88+
desc: "with connection params",
89+
in: `
90+
kind: source
91+
name: my-s2-instance
92+
type: singlestore
93+
host: 0.0.0.0
94+
port: my-port
95+
database: my_db
96+
user: my_user
97+
password: my_pass
98+
connectionParams:
99+
tls: preferred
100+
compress: true
101+
`,
102+
want: map[string]sources.SourceConfig{
103+
"my-s2-instance": singlestore.Config{
104+
Name: "my-s2-instance",
105+
Type: singlestore.SourceType,
106+
Host: "0.0.0.0",
107+
Port: "my-port",
108+
Database: "my_db",
109+
User: "my_user",
110+
Password: "my_pass",
111+
ConnectionParams: map[string]string{
112+
"tls": "preferred",
113+
"compress": "true",
114+
},
115+
},
116+
},
117+
},
118+
{
119+
desc: "with tls via connection params",
120+
in: `
121+
kind: source
122+
name: my-s2-instance
123+
type: singlestore
124+
host: 0.0.0.0
125+
port: my-port
126+
database: my_db
127+
user: my_user
128+
password: my_pass
129+
connectionParams:
130+
tls: skip-verify
131+
`,
132+
want: map[string]sources.SourceConfig{
133+
"my-s2-instance": singlestore.Config{
134+
Name: "my-s2-instance",
135+
Type: singlestore.SourceType,
136+
Host: "0.0.0.0",
137+
Port: "my-port",
138+
Database: "my_db",
139+
User: "my_user",
140+
Password: "my_pass",
141+
ConnectionParams: map[string]string{
142+
"tls": "skip-verify",
143+
},
144+
},
145+
},
146+
},
84147
}
85148
for _, tc := range tcs {
86149
t.Run(tc.desc, func(t *testing.T) {
@@ -144,3 +207,25 @@ func TestFailParseFromYaml(t *testing.T) {
144207
})
145208
}
146209
}
210+
211+
func TestFailInitialization(t *testing.T) {
212+
t.Parallel()
213+
214+
cfg := singlestore.Config{
215+
Name: "instance",
216+
Type: "singlestore",
217+
Host: "localhost",
218+
Port: "3306",
219+
Database: "db",
220+
User: "user",
221+
Password: "pass",
222+
QueryTimeout: "abc", // invalid duration
223+
}
224+
_, err := cfg.Initialize(context.Background(), noop.NewTracerProvider().Tracer("test"))
225+
if err == nil {
226+
t.Fatalf("expected error for invalid queryTimeout, got nil")
227+
}
228+
if !strings.Contains(err.Error(), "invalid queryTimeout") {
229+
t.Fatalf("unexpected error: %v", err)
230+
}
231+
}

tests/singlestore/singlestore_integration_test.go

Lines changed: 55 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -18,13 +18,16 @@ import (
1818
"context"
1919
"database/sql"
2020
"fmt"
21+
"net/url"
2122
"os"
2223
"regexp"
2324
"strings"
2425
"testing"
2526
"time"
2627

28+
"github.com/go-sql-driver/mysql"
2729
"github.com/google/uuid"
30+
singlestoresrc "github.com/googleapis/genai-toolbox/internal/sources/singlestore"
2831
"github.com/googleapis/genai-toolbox/internal/testutils"
2932
"github.com/googleapis/genai-toolbox/tests"
3033
)
@@ -167,9 +170,50 @@ func addSingleStoreExecuteSQLConfig(t *testing.T, config map[string]any) map[str
167170
return config
168171
}
169172

170-
// Copied over from singlestore.go
171-
func initSingleStoreConnectionPool(host, port, user, pass, dbname string) (*sql.DB, error) {
172-
dsn := fmt.Sprintf("%s:%s@tcp(%s:%s)/%s?parseTime=true", user, pass, host, port, dbname)
173+
// Copied over from singlestore.go, with context and tracer removed
174+
func initSingleStoreConnectionPool(cfg singlestoresrc.Config) (*sql.DB, error) {
175+
// Build query parameters via url.Values for deterministic order and proper escaping.
176+
connectionParams := url.Values{}
177+
178+
mysqlCfg := mysql.Config{
179+
User: cfg.User,
180+
Passwd: cfg.Password,
181+
Net: "tcp",
182+
Addr: fmt.Sprintf("%s:%s", cfg.Host, cfg.Port),
183+
DBName: cfg.Database,
184+
ParseTime: true,
185+
AllowNativePasswords: true,
186+
CheckConnLiveness: true,
187+
MaxAllowedPacket: 64 << 20,
188+
ConnectionAttributes: "_connector_name:MCP toolbox for Databases",
189+
Params: map[string]string{
190+
"vector_type_project_format": "JSON",
191+
},
192+
}
193+
194+
// Default to TLS preferred; can be overridden via connectionParams.
195+
connectionParams.Set("tls", "preferred")
196+
197+
// Derive readTimeout from queryTimeout when provided.
198+
if cfg.QueryTimeout != "" {
199+
timeout, err := time.ParseDuration(cfg.QueryTimeout)
200+
if err != nil {
201+
return nil, fmt.Errorf("invalid queryTimeout %q: %w", cfg.QueryTimeout, err)
202+
}
203+
connectionParams.Set("readTimeout", timeout.String())
204+
}
205+
206+
// Custom user parameters (e.g. tls, compress) — may override defaults above.
207+
for k, v := range cfg.ConnectionParams {
208+
if v == "" {
209+
continue // skip empty values
210+
}
211+
connectionParams.Set(k, v)
212+
}
213+
dsn := mysqlCfg.FormatDSN()
214+
if enc := connectionParams.Encode(); enc != "" {
215+
dsn += "&" + enc
216+
}
173217

174218
// Interact with the driver directly as you normally would
175219
pool, err := sql.Open("mysql", dsn)
@@ -186,7 +230,14 @@ func TestSingleStoreToolEndpoints(t *testing.T) {
186230

187231
args := []string{"--enable-api"}
188232

189-
pool, err := initSingleStoreConnectionPool(SingleStoreHost, SingleStorePort, SingleStoreUser, SingleStorePass, SingleStoreDatabase)
233+
cfg := singlestoresrc.Config{
234+
Host: SingleStoreHost,
235+
Port: SingleStorePort,
236+
User: SingleStoreUser,
237+
Password: SingleStorePass,
238+
Database: SingleStoreDatabase,
239+
}
240+
pool, err := initSingleStoreConnectionPool(cfg)
190241
if err != nil {
191242
t.Fatalf("unable to create SingleStore connection pool: %s", err)
192243
}

0 commit comments

Comments
 (0)