Skip to content

Commit

Permalink
Fix connection leak on query retry reconnect (#79)
Browse files Browse the repository at this point in the history
* Close released conn on a reconnect attempt

* fix err message

* tests: do actual DB connection on each driver connect
  • Loading branch information
jkaflik committed Dec 6, 2022
1 parent cd68d26 commit e30c0ed
Show file tree
Hide file tree
Showing 2 changed files with 28 additions and 13 deletions.
19 changes: 15 additions & 4 deletions datasource.go
Expand Up @@ -227,11 +227,10 @@ func (ds *SQLDatasource) handleQuery(ctx context.Context, req backend.DataQuery,
if errors.Is(err, ErrorQuery) && !errors.Is(err, context.DeadlineExceeded) {
for i := 0; i < ds.driverSettings.Retries; i++ {
backend.Logger.Warn(fmt.Sprintf("query failed. retrying %d times", i))
db, err := ds.c.Connect(dbConn.settings, q.ConnectionArgs)
db, err := ds.dbReconnect(dbConn, q, cacheKey)
if err != nil {
return nil, err
}
ds.storeDBConnection(cacheKey, dbConnection{db, dbConn.settings})

if ds.driverSettings.Pause > 0 {
time.Sleep(time.Duration(ds.driverSettings.Pause * int(time.Second)))
Expand All @@ -247,11 +246,10 @@ func (ds *SQLDatasource) handleQuery(ctx context.Context, req backend.DataQuery,
if errors.Is(err, context.DeadlineExceeded) {
for i := 0; i < ds.driverSettings.Retries; i++ {
backend.Logger.Warn(fmt.Sprintf("connection timed out. retrying %d times", i))
db, err := ds.c.Connect(dbConn.settings, q.ConnectionArgs)
db, err := ds.dbReconnect(dbConn, q, cacheKey)
if err != nil {
continue
}
ds.storeDBConnection(cacheKey, dbConnection{db, dbConn.settings})

res, err = QueryDB(ctx, db, ds.c.Converters(), fillMode, q)
if err == nil {
Expand All @@ -263,6 +261,19 @@ func (ds *SQLDatasource) handleQuery(ctx context.Context, req backend.DataQuery,
return nil, err
}

func (ds *SQLDatasource) dbReconnect(dbConn dbConnection, q *Query, cacheKey string) (*sql.DB, error) {
if err := dbConn.db.Close(); err != nil {
backend.Logger.Warn(fmt.Sprintf("closing existing connection failed: %s", err.Error()))
}

db, err := ds.c.Connect(dbConn.settings, q.ConnectionArgs)
if err != nil {
return nil, err
}
ds.storeDBConnection(cacheKey, dbConnection{db, dbConn.settings})
return db, nil
}

// CheckHealth pings the connected SQL database
func (ds *SQLDatasource) CheckHealth(ctx context.Context, req *backend.CheckHealthRequest) (*backend.CheckHealthResult, error) {
key := defaultKey(getDatasourceUID(*req.PluginContext.DataSourceInstanceSettings))
Expand Down
22 changes: 13 additions & 9 deletions datasource_test.go
Expand Up @@ -18,13 +18,13 @@ import (
)

type fakeDriver struct {
db *sql.DB
openDBfn func() (*sql.DB, error)

Driver
}

func (d fakeDriver) Connect(backend.DataSourceInstanceSettings, json.RawMessage) (db *sql.DB, err error) {
return d.db, nil
return d.openDBfn()
}

func (d fakeDriver) Macros() Macros {
Expand All @@ -41,7 +41,7 @@ func Test_getDBConnectionFromQuery(t *testing.T) {
db := &sql.DB{}
db2 := &sql.DB{}
db3 := &sql.DB{}
d := &fakeDriver{db: db3}
d := &fakeDriver{openDBfn: func() (*sql.DB, error) { return db3, nil }}
tests := []struct {
desc string
dsUID string
Expand Down Expand Up @@ -144,7 +144,7 @@ func Test_timeout_retries(t *testing.T) {
t.Errorf("failed to connect to mock driver: %v", err)
}
timeoutDriver := fakeDriver{
db: db,
openDBfn: func() (*sql.DB, error) { return db, nil },
}
retries := 5
max := time.Duration(testTimeout) * time.Second
Expand Down Expand Up @@ -178,12 +178,15 @@ func Test_error_retries(t *testing.T) {
}
mockDriver := "sqlmock-error"
mock.RegisterDriver(mockDriver, handler)
db, err := sql.Open(mockDriver, "")
if err != nil {
t.Errorf("failed to connect to mock driver: %v", err)
}

timeoutDriver := fakeDriver{
db: db,
openDBfn: func() (*sql.DB, error) {
db, err := sql.Open(mockDriver, "")
if err != nil {
t.Errorf("failed to connect to mock driver: %v", err)
}
return db, nil
},
}
retries := 5
max := time.Duration(10) * time.Second
Expand All @@ -192,6 +195,7 @@ func Test_error_retries(t *testing.T) {

key := defaultKey(dsUID)
// Add the mandatory default db
db, _ := timeoutDriver.Connect(settings, nil)
ds.storeDBConnection(key, dbConnection{db, settings})
ctx := context.Background()

Expand Down

0 comments on commit e30c0ed

Please sign in to comment.