From 96e73eb9aa7ba849b24eae15477456d8bbb1c9b7 Mon Sep 17 00:00:00 2001
From: Evan Jones <evan.jones@datadoghq.com>
Date: Fri, 15 Apr 2022 14:46:00 -0400
Subject: [PATCH] conn: Implement driver.Validator, SessionResetter for
 cancelation

Commit 8446d16b89 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.
---
 conn.go        | 18 ++++++++++++
 issues_test.go | 77 ++++++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 95 insertions(+)

diff --git a/conn.go b/conn.go
index c47a8b3a..1f7ebc65 100644
--- a/conn.go
+++ b/conn.go
@@ -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
+}
diff --git a/issues_test.go b/issues_test.go
index 26a70282..e2c93925 100644
--- a/issues_test.go
+++ b/issues_test.go
@@ -2,6 +2,7 @@ package pq
 
 import (
 	"context"
+	"database/sql"
 	"errors"
 	"testing"
 	"time"
@@ -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)
+}