Skip to content

Commit 28f86d0

Browse files
authored
Prevent usage of failed connections and improve connection management (#1156)
1 parent d26f1cd commit 28f86d0

File tree

9 files changed

+278
-81
lines changed

9 files changed

+278
-81
lines changed

pkg/plugin/driver.go

Lines changed: 42 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -73,7 +73,6 @@ func getPDCDialContext(settings Settings) (func(context.Context, string) (net.Co
7373
return nil, errors.New("unable to cast SOCKS proxy dialer to context proxy dialer")
7474
}
7575

76-
// Return a function matching the expected signature
7776
return func(ctx context.Context, addr string) (net.Conn, error) {
7877
return contextDialer.DialContext(ctx, "tcp", addr)
7978
}, nil
@@ -180,30 +179,26 @@ func (h *Clickhouse) Connect(ctx context.Context, config backend.DataSourceInsta
180179
httpHeaders[k] = v
181180
}
182181

183-
timeout := time.Duration(t) * time.Second
184-
ctx, cancel := context.WithTimeout(ctx, timeout)
185-
defer cancel()
186-
187182
opts := &clickhouse.Options{
188-
ClientInfo: clickhouse.ClientInfo{
189-
Products: getClientInfoProducts(ctx),
190-
},
191-
TLS: tlsConfig,
192-
Addr: []string{fmt.Sprintf("%s:%d", settings.Host, settings.Port)},
193-
HttpUrlPath: settings.Path,
194-
HttpHeaders: httpHeaders,
183+
Addr: []string{fmt.Sprintf("%s:%d", settings.Host, settings.Port)},
195184
Auth: clickhouse.Auth{
196-
Username: settings.Username,
197-
Password: settings.Password,
198185
Database: settings.DefaultDatabase,
186+
Password: settings.Password,
187+
Username: settings.Username,
188+
},
189+
ClientInfo: clickhouse.ClientInfo{
190+
Products: getClientInfoProducts(ctx),
199191
},
200192
Compression: &clickhouse.Compression{
201193
Method: compression,
202194
},
203195
DialTimeout: time.Duration(t) * time.Second,
204-
ReadTimeout: time.Duration(qt) * time.Second,
196+
HttpHeaders: httpHeaders,
197+
HttpUrlPath: settings.Path,
205198
Protocol: protocol,
199+
ReadTimeout: time.Duration(qt) * time.Second,
206200
Settings: customSettings,
201+
TLS: tlsConfig,
207202
}
208203

209204
// dialCtx is used to create a connection to PDC, if it is enabled
@@ -215,21 +210,44 @@ func (h *Clickhouse) Connect(ctx context.Context, config backend.DataSourceInsta
215210
opts.DialContext = dialCtx
216211
}
217212

213+
ctx, cancel := context.WithTimeout(ctx, time.Duration(t)*time.Second)
214+
defer cancel()
215+
218216
db := clickhouse.OpenDB(opts)
219217

218+
// Set connection pool settings
219+
if i, err := strconv.Atoi(settings.ConnMaxLifetime); err == nil {
220+
db.SetConnMaxLifetime(time.Duration(i) * time.Minute)
221+
}
222+
if i, err := strconv.Atoi(settings.MaxIdleConns); err == nil {
223+
db.SetMaxIdleConns(i)
224+
}
225+
if i, err := strconv.Atoi(settings.MaxOpenConns); err == nil {
226+
db.SetMaxOpenConns(i)
227+
}
228+
220229
select {
221230
case <-ctx.Done():
222-
return db, fmt.Errorf("the operation was cancelled: %w", ctx.Err())
231+
return nil, fmt.Errorf("the operation was cancelled before starting: %w", ctx.Err())
223232
default:
224-
err := db.PingContext(ctx)
225-
if err != nil {
226-
if exception, ok := err.(*clickhouse.Exception); ok {
227-
log.DefaultLogger.Error("[%d] %s \n%s\n", exception.Code, exception.Message, exception.StackTrace)
228-
} else {
229-
log.DefaultLogger.Error(err.Error())
230-
}
231-
return db, nil
233+
// proceed
234+
}
235+
236+
// `sqlds` normally calls `db.PingContext()` to check if the connection is alive,
237+
// however, as ClickHouse returns its own non-standard `Exception` type, we need
238+
// to handle it here so that we can log the error code, message and stack trace
239+
if err := db.PingContext(ctx); err != nil {
240+
if ctx.Err() != nil {
241+
return nil, fmt.Errorf("the operation was cancelled during execution: %w", ctx.Err())
242+
}
243+
244+
if exception, ok := err.(*clickhouse.Exception); ok {
245+
log.DefaultLogger.Error("[%d] %s \n%s\n", exception.Code, exception.Message, exception.StackTrace)
246+
} else {
247+
log.DefaultLogger.Error(err.Error())
232248
}
249+
250+
return nil, err
233251
}
234252

235253
return db, settings.isValid()

pkg/plugin/driver_test.go

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -160,12 +160,11 @@ func TestHTTPConnect(t *testing.T) {
160160
username := getEnv("CLICKHOUSE_USERNAME", "default")
161161
password := getEnv("CLICKHOUSE_PASSWORD", "")
162162
ssl := getEnv("CLICKHOUSE_SSL", "false")
163-
path := "custom-path"
164163
clickhouse := plugin.Clickhouse{}
165164
t.Run("should not error when valid settings passed", func(t *testing.T) {
166165
secure := map[string]string{}
167166
secure["password"] = password
168-
settings := backend.DataSourceInstanceSettings{JSONData: []byte(fmt.Sprintf(`{ "server": "%s", "port": %s, "path": "%s", "username": "%s", "secure": %s, "protocol": "http" }`, host, port, path, username, ssl)), DecryptedSecureJSONData: secure}
167+
settings := backend.DataSourceInstanceSettings{JSONData: []byte(fmt.Sprintf(`{ "server": "%s", "port": %s, "username": "%s", "secure": %s, "protocol": "http" }`, host, port, username, ssl)), DecryptedSecureJSONData: secure}
169168
_, err := clickhouse.Connect(context.Background(), settings, json.RawMessage{})
170169
assert.Equal(t, nil, err)
171170
})

pkg/plugin/settings.go

Lines changed: 17 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -34,8 +34,11 @@ type Settings struct {
3434

3535
DefaultDatabase string `json:"defaultDatabase,omitempty"`
3636

37-
DialTimeout string `json:"dialTimeout,omitempty"`
38-
QueryTimeout string `json:"queryTimeout,omitempty"`
37+
ConnMaxLifetime string `json:"connMaxLifetime,omitempty"`
38+
DialTimeout string `json:"dialTimeout,omitempty"`
39+
QueryTimeout string `json:"queryTimeout,omitempty"`
40+
MaxIdleConns string `json:"maxIdleConns,omitempty"`
41+
MaxOpenConns string `json:"maxOpenConns,omitempty"`
3942

4043
HttpHeaders map[string]string `json:"-"`
4144
ForwardGrafanaHeaders bool `json:"forwardGrafanaHeaders,omitempty"`
@@ -181,12 +184,24 @@ func LoadSettings(ctx context.Context, config backend.DataSourceInstanceSettings
181184
}
182185
}
183186

187+
// Set default values
184188
if strings.TrimSpace(settings.DialTimeout) == "" {
185189
settings.DialTimeout = "10"
186190
}
187191
if strings.TrimSpace(settings.QueryTimeout) == "" {
188192
settings.QueryTimeout = "60"
189193
}
194+
if strings.TrimSpace(settings.ConnMaxLifetime) == "" {
195+
settings.ConnMaxLifetime = "5"
196+
}
197+
if strings.TrimSpace(settings.MaxIdleConns) == "" {
198+
settings.MaxIdleConns = "25"
199+
}
200+
if strings.TrimSpace(settings.MaxOpenConns) == "" {
201+
settings.MaxOpenConns = "50"
202+
}
203+
204+
// Load secure settings
190205
password, ok := config.DecryptedSecureJSONData["password"]
191206
if ok {
192207
settings.Password = password

pkg/plugin/settings_test.go

Lines changed: 13 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -63,7 +63,10 @@ func TestLoadSettings(t *testing.T) {
6363
TlsCACert: "caCert",
6464
TlsClientCert: "clientCert",
6565
TlsClientKey: "clientKey",
66+
ConnMaxLifetime: "5",
6667
DialTimeout: "10",
68+
MaxIdleConns: "25",
69+
MaxOpenConns: "50",
6770
QueryTimeout: "60",
6871
HttpHeaders: map[string]string{
6972
"test-plain-1": "value-1",
@@ -100,7 +103,10 @@ func TestLoadSettings(t *testing.T) {
100103
InsecureSkipVerify: true,
101104
TlsClientAuth: true,
102105
TlsAuthWithCACert: true,
106+
ConnMaxLifetime: "5",
103107
DialTimeout: "10",
108+
MaxIdleConns: "25",
109+
MaxOpenConns: "50",
104110
QueryTimeout: "60",
105111
ProxyOptions: nil,
106112
},
@@ -115,10 +121,13 @@ func TestLoadSettings(t *testing.T) {
115121
},
116122
},
117123
wantSettings: Settings{
118-
Host: "test",
119-
Port: 443,
120-
DialTimeout: "10",
121-
QueryTimeout: "60",
124+
Host: "test",
125+
Port: 443,
126+
ConnMaxLifetime: "5",
127+
DialTimeout: "10",
128+
MaxIdleConns: "25",
129+
MaxOpenConns: "50",
130+
QueryTimeout: "60",
122131
},
123132
wantErr: nil,
124133
},

src/components/configEditor/QuerySettingsConfig.test.tsx

Lines changed: 90 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -7,9 +7,15 @@ describe('QuerySettingsConfig', () => {
77
it('should render', () => {
88
const result = render(
99
<QuerySettingsConfig
10-
dialTimeout='10'
11-
queryTimeout='10'
12-
validateSql
10+
connMaxLifetime={'5'}
11+
dialTimeout={'5'}
12+
maxIdleConns={'5'}
13+
maxOpenConns={'5'}
14+
queryTimeout={'5'}
15+
validateSql={true}
16+
onConnMaxIdleConnsChange={() => {}}
17+
onConnMaxLifetimeChange={() => {}}
18+
onConnMaxOpenConnsChange={() => {}}
1319
onDialTimeoutChange={() => {}}
1420
onQueryTimeoutChange={() => {}}
1521
onValidateSqlChange={() => {}}
@@ -22,6 +28,9 @@ describe('QuerySettingsConfig', () => {
2228
const onDialTimeout = jest.fn();
2329
const result = render(
2430
<QuerySettingsConfig
31+
onConnMaxIdleConnsChange={() => {}}
32+
onConnMaxLifetimeChange={() => {}}
33+
onConnMaxOpenConnsChange={() => {}}
2534
onDialTimeoutChange={onDialTimeout}
2635
onQueryTimeoutChange={() => {}}
2736
onValidateSqlChange={() => {}}
@@ -33,14 +42,17 @@ describe('QuerySettingsConfig', () => {
3342
expect(input).toBeInTheDocument();
3443
fireEvent.change(input, { target: { value: '10' } });
3544
fireEvent.blur(input);
36-
expect(onDialTimeout).toBeCalledTimes(1);
37-
expect(onDialTimeout).toBeCalledWith(expect.any(Object));
45+
expect(onDialTimeout).toHaveBeenCalledTimes(1);
46+
expect(onDialTimeout).toHaveBeenCalledWith(expect.any(Object));
3847
});
3948

4049
it('should call onQueryTimeout when changed', () => {
4150
const onQueryTimeout = jest.fn();
4251
const result = render(
4352
<QuerySettingsConfig
53+
onConnMaxIdleConnsChange={() => {}}
54+
onConnMaxLifetimeChange={() => {}}
55+
onConnMaxOpenConnsChange={() => {}}
4456
onDialTimeoutChange={() => {}}
4557
onQueryTimeoutChange={onQueryTimeout}
4658
onValidateSqlChange={() => {}}
@@ -52,14 +64,17 @@ describe('QuerySettingsConfig', () => {
5264
expect(input).toBeInTheDocument();
5365
fireEvent.change(input, { target: { value: '10' } });
5466
fireEvent.blur(input);
55-
expect(onQueryTimeout).toBeCalledTimes(1);
56-
expect(onQueryTimeout).toBeCalledWith(expect.any(Object));
67+
expect(onQueryTimeout).toHaveBeenCalledTimes(1);
68+
expect(onQueryTimeout).toHaveBeenCalledWith(expect.any(Object));
5769
});
5870

5971
it('should call onValidateSqlChange when changed', () => {
6072
const onValidateSqlChange = jest.fn();
6173
const result = render(
6274
<QuerySettingsConfig
75+
onConnMaxIdleConnsChange={() => {}}
76+
onConnMaxLifetimeChange={() => {}}
77+
onConnMaxOpenConnsChange={() => {}}
6378
onDialTimeoutChange={() => {}}
6479
onQueryTimeoutChange={() => {}}
6580
onValidateSqlChange={onValidateSqlChange}
@@ -70,7 +85,73 @@ describe('QuerySettingsConfig', () => {
7085
const input = result.getByRole('checkbox');
7186
expect(input).toBeInTheDocument();
7287
fireEvent.click(input);
73-
expect(onValidateSqlChange).toBeCalledTimes(1);
74-
expect(onValidateSqlChange).toBeCalledWith(expect.any(Object));
88+
expect(onValidateSqlChange).toHaveBeenCalledTimes(1);
89+
expect(onValidateSqlChange).toHaveBeenCalledWith(expect.any(Object));
90+
});
91+
92+
it('should call onConnMaxIdleConnsChange when changed', () => {
93+
const onConnMaxIdleConnsChange = jest.fn();
94+
const result = render(
95+
<QuerySettingsConfig
96+
onConnMaxIdleConnsChange={onConnMaxIdleConnsChange}
97+
onConnMaxLifetimeChange={() => {}}
98+
onConnMaxOpenConnsChange={() => {}}
99+
onDialTimeoutChange={() => {}}
100+
onQueryTimeoutChange={() => {}}
101+
onValidateSqlChange={() => {}}
102+
/>
103+
);
104+
expect(result.container.firstChild).not.toBeNull();
105+
106+
const input = result.getByPlaceholderText(allLabels.components.Config.QuerySettingsConfig.maxIdleConns.placeholder);
107+
expect(input).toBeInTheDocument();
108+
fireEvent.change(input, { target: { value: '10' } });
109+
fireEvent.blur(input);
110+
expect(onConnMaxIdleConnsChange).toHaveBeenCalledTimes(1);
111+
expect(onConnMaxIdleConnsChange).toHaveBeenCalledWith(expect.any(Object));
112+
});
113+
114+
it('should call onConnMaxLifetimeChange when changed', () => {
115+
const onConnMaxLifetimeChange = jest.fn();
116+
const result = render(
117+
<QuerySettingsConfig
118+
onConnMaxIdleConnsChange={() => {}}
119+
onConnMaxLifetimeChange={onConnMaxLifetimeChange}
120+
onConnMaxOpenConnsChange={() => {}}
121+
onDialTimeoutChange={() => {}}
122+
onQueryTimeoutChange={() => {}}
123+
onValidateSqlChange={() => {}}
124+
/>
125+
);
126+
expect(result.container.firstChild).not.toBeNull();
127+
128+
const input = result.getByPlaceholderText(allLabels.components.Config.QuerySettingsConfig.connMaxLifetime.placeholder);
129+
expect(input).toBeInTheDocument();
130+
fireEvent.change(input, { target: { value: '10' } });
131+
fireEvent.blur(input);
132+
expect(onConnMaxLifetimeChange).toHaveBeenCalledTimes(1);
133+
expect(onConnMaxLifetimeChange).toHaveBeenCalledWith(expect.any(Object));
134+
});
135+
136+
it('should call onConnMaxOpenConnsChange when changed', () => {
137+
const onConnMaxOpenConnsChange = jest.fn();
138+
const result = render(
139+
<QuerySettingsConfig
140+
onConnMaxIdleConnsChange={() => {}}
141+
onConnMaxLifetimeChange={() => {}}
142+
onConnMaxOpenConnsChange={onConnMaxOpenConnsChange}
143+
onDialTimeoutChange={() => {}}
144+
onQueryTimeoutChange={() => {}}
145+
onValidateSqlChange={() => {}}
146+
/>
147+
);
148+
expect(result.container.firstChild).not.toBeNull();
149+
150+
const input = result.getByPlaceholderText(allLabels.components.Config.QuerySettingsConfig.maxOpenConns.placeholder);
151+
expect(input).toBeInTheDocument();
152+
fireEvent.change(input, { target: { value: '10' } });
153+
fireEvent.blur(input);
154+
expect(onConnMaxOpenConnsChange).toHaveBeenCalledTimes(1);
155+
expect(onConnMaxOpenConnsChange).toHaveBeenCalledWith(expect.any(Object));
75156
});
76157
});

0 commit comments

Comments
 (0)