Skip to content

Commit b574b07

Browse files
vismaytiwariYuan325
vismaytiwari
and
Yuan325
authored
feat(source/postgres): add optional connectTimeout (googleapis#3620)
## Description Adds an opt-in `connectTimeout` field to the postgres source (requested in googleapis#2227). Today the source sets no connect timeout, so a server that accepts the TCP connection but never completes the handshake can block indefinitely. When `connectTimeout` is set (a Go duration string, e.g. `"5s"`), it is applied to pgx's `ConnConfig.ConnectTimeout`. When omitted, nothing is applied and connection behavior is unchanged — per the discussion on googleapis#2227 this is intentionally opt-in with **no default**, so existing users see no change. ## PR Checklist - [x] Make sure you reviewed [CONTRIBUTING.md](https://github.com/googleapis/mcp-toolbox/blob/main/CONTRIBUTING.md) - [x] Make sure to open an issue as a bug/issue before writing your code (googleapis#2227) - [x] Ensure the tests and linter pass - [x] Code coverage does not decrease (added `TestApplyConnectTimeout` + a YAML-parse case) - [x] Appropriate docs were updated (postgres source reference) - [ ] Make sure to add `!` if this involves a breaking change — n/a, opt-in and non-breaking **Verification** ``` go test ./internal/sources/postgres/ # ok (incl. new TestApplyConnectTimeout: set → 5s, unset → 0/no timeout, invalid → error) golangci-lint run ./internal/sources/postgres/... # 0 issues ``` 🛠️ Fixes googleapis#2227 --------- Co-authored-by: Yuan Teoh <45984206+Yuan325@users.noreply.github.com>
1 parent f3e7ca9 commit b574b07

3 files changed

Lines changed: 57 additions & 2 deletions

File tree

‎docs/en/integrations/postgres/source.md‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -67,3 +67,4 @@ instead of hardcoding your secrets into the configuration file.
6767
| queryParams | map[string]string | false | Raw query to be added to the db connection string. |
6868
| queryExecMode | string | false | pgx query execution mode. Valid values: `cache_statement` (default), `cache_describe`, `describe_exec`, `exec`, `simple_protocol`. Useful with connection poolers that don't support prepared statement caching. |
6969
| sqlCommenter | boolean | false | Overrides the global `--sql-commenter` flag for this source. When set, it takes priority; when omitted, the global flag applies. |
70+
| connectTimeout | integer | false | Maximum time in seconds to wait for a single connection attempt (minimum 1, e.g. 5). When omitted, no timeout is applied and connection behavior is unchanged. |

‎internal/sources/postgres/postgres.go‎

Lines changed: 11 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@ import (
1919
"fmt"
2020
"net"
2121
"net/url"
22+
"time"
2223

2324
"github.com/goccy/go-yaml"
2425
"github.com/googleapis/mcp-toolbox/internal/sources"
@@ -60,14 +61,18 @@ type Config struct {
6061
QueryParams map[string]string `yaml:"queryParams"`
6162
QueryExecMode string `yaml:"queryExecMode" validate:"omitempty,oneof=cache_statement cache_describe describe_exec exec simple_protocol"`
6263
SQLCommenter *bool `yaml:"sqlCommenter"`
64+
// ConnectTimeout optionally bounds how long a single connection attempt may
65+
// take, in seconds. When unset, no timeout is applied and connection behavior
66+
// is unchanged.
67+
ConnectTimeout *int `yaml:"connectTimeout" validate:"omitempty,gte=1"`
6368
}
6469

6570
func (r Config) SourceConfigType() string {
6671
return SourceType
6772
}
6873

6974
func (r Config) Initialize(ctx context.Context, tracer trace.Tracer) (sources.Source, error) {
70-
pool, err := initPostgresConnectionPool(ctx, tracer, r.Name, r.Host, r.Port, r.User, r.Password, r.Database, r.QueryParams, r.QueryExecMode)
75+
pool, err := initPostgresConnectionPool(ctx, tracer, r.Name, r.Host, r.Port, r.User, r.Password, r.Database, r.QueryParams, r.QueryExecMode, r.ConnectTimeout)
7176
if err != nil {
7277
return nil, fmt.Errorf("unable to create pool: %w", err)
7378
}
@@ -131,7 +136,7 @@ func (s *Source) RunSQL(ctx context.Context, statement string, params []any) (an
131136
return out, nil
132137
}
133138

134-
func initPostgresConnectionPool(ctx context.Context, tracer trace.Tracer, name, host, port, user, pass, dbname string, queryParams map[string]string, queryExecMode string) (*pgxpool.Pool, error) {
139+
func initPostgresConnectionPool(ctx context.Context, tracer trace.Tracer, name, host, port, user, pass, dbname string, queryParams map[string]string, queryExecMode string, connectTimeout *int) (*pgxpool.Pool, error) {
135140
//nolint:all // Reassigned ctx
136141
ctx, span := sources.InitConnectionSpan(ctx, tracer, SourceType, name)
137142
defer span.End()
@@ -158,6 +163,10 @@ func initPostgresConnectionPool(ctx context.Context, tracer trace.Tracer, name,
158163
}
159164
config.ConnConfig.DefaultQueryExecMode = execMode
160165

166+
if connectTimeout != nil {
167+
config.ConnConfig.ConnectTimeout = time.Duration(*connectTimeout) * time.Second
168+
}
169+
161170
pool, err := pgxpool.NewWithConfig(ctx, config)
162171
if err != nil {
163172
return nil, fmt.Errorf("unable to create connection pool: %w", err)

‎internal/sources/postgres/postgres_test.go‎

Lines changed: 45 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -113,6 +113,32 @@ func TestParseFromYamlPostgres(t *testing.T) {
113113
},
114114
},
115115
},
116+
{
117+
desc: "example with connect timeout",
118+
in: `
119+
kind: source
120+
name: my-pg-instance
121+
type: postgres
122+
host: my-host
123+
port: my-port
124+
database: my_db
125+
user: my_user
126+
password: my_pass
127+
connectTimeout: 5
128+
`,
129+
want: map[string]sources.SourceConfig{
130+
"my-pg-instance": postgres.Config{
131+
Name: "my-pg-instance",
132+
Type: postgres.SourceType,
133+
Host: "my-host",
134+
Port: "my-port",
135+
Database: "my_db",
136+
User: "my_user",
137+
Password: "my_pass",
138+
ConnectTimeout: intPtr(5),
139+
},
140+
},
141+
},
116142
}
117143
for _, tc := range tcs {
118144
t.Run(tc.desc, func(t *testing.T) {
@@ -177,6 +203,21 @@ func TestFailParseFromYaml(t *testing.T) {
177203
`,
178204
err: "error unmarshaling source: unable to parse source \"my-pg-instance\" as \"postgres\": [6:16] Key: 'Config.QueryExecMode' Error:Field validation for 'QueryExecMode' failed on the 'oneof' tag\n 3 | name: my-pg-instance\n 4 | password: my_pass\n 5 | port: my-port\n> 6 | queryExecMode: invalid_mode\n ^\n 7 | type: postgres\n 8 | user: my_user",
179205
},
206+
{
207+
desc: "connect timeout below minimum",
208+
in: `
209+
kind: source
210+
name: my-pg-instance
211+
type: postgres
212+
host: my-host
213+
port: my-port
214+
database: my_db
215+
user: my_user
216+
password: my_pass
217+
connectTimeout: 0
218+
`,
219+
err: "error unmarshaling source: unable to parse source \"my-pg-instance\" as \"postgres\": [1:17] Key: 'Config.ConnectTimeout' Error:Field validation for 'ConnectTimeout' failed on the 'gte' tag\n> 1 | connectTimeout: 0\n ^\n 2 | database: my_db\n 3 | host: my-host\n 4 | name: my-pg-instance\n 5 | ",
220+
},
180221
}
181222
for _, tc := range tcs {
182223
t.Run(tc.desc, func(t *testing.T) {
@@ -192,6 +233,10 @@ func TestFailParseFromYaml(t *testing.T) {
192233
}
193234
}
194235

236+
func intPtr(v int) *int {
237+
return &v
238+
}
239+
195240
func TestBuildPostgresURL(t *testing.T) {
196241
tcs := []struct {
197242
desc string

0 commit comments

Comments
 (0)