From 6a223b82a466540ea5de1e9f8a7bfea31d77e07d Mon Sep 17 00:00:00 2001
From: Daniel Theophanes <kardianos@gmail.com>
Date: Mon, 2 Oct 2017 11:16:53 -0700
Subject: database/sql: add driver.ResetSessioner and add pool support

A single database connection ususally maps to a single session.
A connection pool is logically also a session pool. Most
sessions have a way to reset the session state which is desirable
to prevent one bad query from poisoning another later query with
temp table name conflicts or other persistent session resources.

It also lets drivers provide users with better error messages from
queryies when the underlying transport or query method fails.
Internally the driver connection should now be marked as bad, but
return the actual connection. When ResetSession is called on the
connection it should return driver.ErrBadConn to remove it from
the connection pool. Previously drivers had to choose between
meaningful error messages or poisoning the connection pool.

Lastly update TestPoolExhaustOnCancel from relying on a
WAIT query fixing a flaky timeout issue exposed by this
change.

Fixes #22049
Fixes #20807

Change-Id: I2b5df6d954a38d0ad93bf1922ec16e74c827274c
Reviewed-on: https://go-review.googlesource.com/73033
Run-TryBot: Daniel Theophanes <kardianos@gmail.com>
TryBot-Result: Gobot Gobot <gobot@golang.org>
Reviewed-by: Ian Lance Taylor <iant@golang.org>
---
 src/database/sql/fakedb_test.go | 62 +++++++++++++++++++++++++++++++++++++++++
 1 file changed, 62 insertions(+)

(limited to 'src/database/sql/fakedb_test.go')

diff --git a/src/database/sql/fakedb_test.go b/src/database/sql/fakedb_test.go
index 4dcd096ca4..070b783453 100644
--- a/src/database/sql/fakedb_test.go
+++ b/src/database/sql/fakedb_test.go
@@ -55,6 +55,22 @@ type fakeDriver struct {
 	dbs        map[string]*fakeDB
 }
 
+type fakeConnector struct {
+	name string
+
+	waiter func(context.Context)
+}
+
+func (c *fakeConnector) Connect(context.Context) (driver.Conn, error) {
+	conn, err := fdriver.Open(c.name)
+	conn.(*fakeConn).waiter = c.waiter
+	return conn, err
+}
+
+func (c *fakeConnector) Driver() driver.Driver {
+	return fdriver
+}
+
 type fakeDB struct {
 	name string
 
@@ -107,6 +123,16 @@ type fakeConn struct {
 	// bad connection tests; see isBad()
 	bad       bool
 	stickyBad bool
+
+	skipDirtySession bool // tests that use Conn should set this to true.
+
+	// dirtySession tests ResetSession, true if a query has executed
+	// until ResetSession is called.
+	dirtySession bool
+
+	// The waiter is called before each query. May be used in place of the "WAIT"
+	// directive.
+	waiter func(context.Context)
 }
 
 func (c *fakeConn) touchMem() {
@@ -298,6 +324,9 @@ func (c *fakeConn) isBad() bool {
 	if c.stickyBad {
 		return true
 	} else if c.bad {
+		if c.db == nil {
+			return false
+		}
 		// alternate between bad conn and not bad conn
 		c.db.badConn = !c.db.badConn
 		return c.db.badConn
@@ -306,6 +335,21 @@ func (c *fakeConn) isBad() bool {
 	}
 }
 
+func (c *fakeConn) isDirtyAndMark() bool {
+	if c.skipDirtySession {
+		return false
+	}
+	if c.currTx != nil {
+		c.dirtySession = true
+		return false
+	}
+	if c.dirtySession {
+		return true
+	}
+	c.dirtySession = true
+	return false
+}
+
 func (c *fakeConn) Begin() (driver.Tx, error) {
 	if c.isBad() {
 		return nil, driver.ErrBadConn
@@ -337,6 +381,14 @@ func setStrictFakeConnClose(t *testing.T) {
 	testStrictClose = t
 }
 
+func (c *fakeConn) ResetSession(ctx context.Context) error {
+	c.dirtySession = false
+	if c.isBad() {
+		return driver.ErrBadConn
+	}
+	return nil
+}
+
 func (c *fakeConn) Close() (err error) {
 	drv := fdriver.(*fakeDriver)
 	defer func() {
@@ -572,6 +624,10 @@ func (c *fakeConn) PrepareContext(ctx context.Context, query string) (driver.Stm
 		stmt.cmd = cmd
 		parts = parts[1:]
 
+		if c.waiter != nil {
+			c.waiter(ctx)
+		}
+
 		if stmt.wait > 0 {
 			wait := time.NewTimer(stmt.wait)
 			select {
@@ -662,6 +718,9 @@ func (s *fakeStmt) ExecContext(ctx context.Context, args []driver.NamedValue) (d
 	if s.c.stickyBad || (hookExecBadConn != nil && hookExecBadConn()) {
 		return nil, driver.ErrBadConn
 	}
+	if s.c.isDirtyAndMark() {
+		return nil, errors.New("session is dirty")
+	}
 
 	err := checkSubsetTypes(s.c.db.allowAny, args)
 	if err != nil {
@@ -774,6 +833,9 @@ func (s *fakeStmt) QueryContext(ctx context.Context, args []driver.NamedValue) (
 	if s.c.stickyBad || (hookQueryBadConn != nil && hookQueryBadConn()) {
 		return nil, driver.ErrBadConn
 	}
+	if s.c.isDirtyAndMark() {
+		return nil, errors.New("session is dirty")
+	}
 
 	err := checkSubsetTypes(s.c.db.allowAny, args)
 	if err != nil {
-- 
cgit v1.3-5-g9baa