Skip to content

Commit

Permalink
conn: Implement driver.Validator, SessionResetter for cancelation
Browse files Browse the repository at this point in the history
Commit 8446d16 released in 1.10.4 changed how some cancelled query
errors were returned. This has caused a lib/pq application I work on
to start returning "driver: bad connection". This is because we were
cancelling a query, after looking at some of the rows. This causes a
"bad" connection to be returned to the connection pool.

To prevent this, implement the driver.Validator and
driver.SessionResetter interfaces. The database/sql/driver package
recommends implementing them:

"All Conn implementations should implement the following interfaces:
Pinger, SessionResetter, and Validator"

Add two tests for this behaviour. One of these tests passed with
1.10.3 but fails with newer versions. The other never passed, but
does after this change.
  • Loading branch information
evanj authored and rafiss committed Apr 14, 2023
1 parent 922c00e commit 96e73eb
Show file tree
Hide file tree
Showing 2 changed files with 95 additions and 0 deletions.
18 changes: 18 additions & 0 deletions conn.go
Original file line number Diff line number Diff line change
Expand Up @@ -2081,3 +2081,21 @@ func alnumLowerASCII(ch rune) rune {
}
return -1 // discard
}

// The database/sql/driver package says:
// All Conn implementations should implement the following interfaces: Pinger, SessionResetter, and Validator.
var _ driver.Pinger = &conn{}
var _ driver.SessionResetter = &conn{}
var _ driver.Validator = &conn{}

func (cn *conn) ResetSession(ctx context.Context) error {
// Ensure bad connections are reported: From database/sql/driver:
// If a connection is never returned to the connection pool but immediately reused, then
// ResetSession is called prior to reuse but IsValid is not called.
return cn.err.get()
}

func (cn *conn) IsValid() bool {
// panic("TODO IsValid")
return cn.err.get() == nil
}
77 changes: 77 additions & 0 deletions issues_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@ package pq

import (
"context"
"database/sql"
"errors"
"testing"
"time"
Expand Down Expand Up @@ -79,3 +80,79 @@ func TestIssue1062(t *testing.T) {
}
}
}

func connIsValid(t *testing.T, db *sql.DB) {
t.Helper()

ctx := context.Background()
conn, err := db.Conn(ctx)
if err != nil {
t.Fatal(err)
}
defer conn.Close()

// the connection must be valid
err = conn.PingContext(ctx)
if err != nil {
t.Errorf("PingContext err=%#v", err)
}
// close must not return an error
err = conn.Close()
if err != nil {
t.Errorf("Close err=%#v", err)
}
}

func TestQueryCancelRace(t *testing.T) {
db := openTestConn(t)
defer db.Close()

// cancel a query while executing on Postgres: must return the cancelled error code
ctx, cancel := context.WithCancel(context.Background())
go func() {
time.Sleep(10 * time.Millisecond)
cancel()
}()
row := db.QueryRowContext(ctx, "select pg_sleep(0.5)")
var pgSleepVoid string
err := row.Scan(&pgSleepVoid)
if pgErr := (*Error)(nil); !(errors.As(err, &pgErr) && pgErr.Code == cancelErrorCode) {
t.Fatalf("expected cancelled error; err=%#v", err)
}

// get a connection: it must be a valid
connIsValid(t, db)
}

// Test cancelling a scan after it is started. This broke with 1.10.4.
func TestQueryCancelledReused(t *testing.T) {
db := openTestConn(t)
defer db.Close()

ctx, cancel := context.WithCancel(context.Background())
// run a query that returns a lot of data
rows, err := db.QueryContext(ctx, "select generate_series(1, 10000)")
if err != nil {
t.Fatal(err)
}

// scan the first value
if !rows.Next() {
t.Error("expected rows.Next() to return true")
}
var i int
err = rows.Scan(&i)
if err != nil {
t.Fatal(err)
}
if i != 1 {
t.Error(i)
}

// cancel the context and close rows, ignoring errors
cancel()
rows.Close()

// get a connection: it must be valid
connIsValid(t, db)
}

0 comments on commit 96e73eb

Please sign in to comment.