From a93025479d31376816aad79a5e037a6450c9fa76 Mon Sep 17 00:00:00 2001
From: Jack Christensen <jack@jackchristensen.com>
Date: Sat, 13 Feb 2016 13:03:06 -0600
Subject: [PATCH] Add AfterClose and Conn to Rows

---
 conn_pool.go |  7 ++++++-
 query.go     | 32 +++++++++++++++++++++++++-------
 2 files changed, 31 insertions(+), 8 deletions(-)

diff --git a/conn_pool.go b/conn_pool.go
index f6d5bf74..695dc564 100644
--- a/conn_pool.go
+++ b/conn_pool.go
@@ -243,7 +243,8 @@ func (p *ConnPool) Query(sql string, args ...interface{}) (*Rows, error) {
 		return rows, err
 	}
 
-	rows.pool = p
+	rows.AfterClose(p.rowsAfterClose)
+
 	return rows, nil
 }
 
@@ -295,3 +296,7 @@ func (p *ConnPool) BeginIso(iso string) (*Tx, error) {
 func (p *ConnPool) txAfterClose(tx *Tx) {
 	p.Release(tx.Conn())
 }
+
+func (p *ConnPool) rowsAfterClose(rows *Rows) {
+	p.Release(rows.Conn())
+}
diff --git a/query.go b/query.go
index 9c9aa4a8..49c1e311 100644
--- a/query.go
+++ b/query.go
@@ -39,7 +39,6 @@ func (r *Row) Scan(dest ...interface{}) (err error) {
 // the *Conn can be used again. Rows are closed by explicitly calling Close(),
 // calling Next() until it returns false, or when a fatal error occurs.
 type Rows struct {
-	pool       *ConnPool
 	conn       *Conn
 	mr         *msgReader
 	fields     []FieldDescription
@@ -47,13 +46,14 @@ type Rows struct {
 	rowCount   int
 	columnIdx  int
 	err        error
-	closed     bool
 	startTime  time.Time
 	sql        string
 	args       []interface{}
 	log        func(lvl int, msg string, ctx ...interface{})
 	shouldLog  func(lvl int) bool
+	afterClose func(*Rows)
 	unlockConn bool
+	closed     bool
 }
 
 func (rows *Rows) FieldDescriptions() []FieldDescription {
@@ -70,11 +70,6 @@ func (rows *Rows) close() {
 		rows.unlockConn = false
 	}
 
-	if rows.pool != nil {
-		rows.pool.Release(rows.conn)
-		rows.pool = nil
-	}
-
 	rows.closed = true
 
 	if rows.err == nil {
@@ -85,6 +80,10 @@ func (rows *Rows) close() {
 	} else if rows.shouldLog(LogLevelError) {
 		rows.log(LogLevelError, "Query", "sql", rows.sql, "args", logQueryArgs(rows.args))
 	}
+
+	if rows.afterClose != nil {
+		rows.afterClose(rows)
+	}
 }
 
 func (rows *Rows) readUntilReadyForQuery() {
@@ -195,6 +194,11 @@ func (rows *Rows) Next() bool {
 	}
 }
 
+// Conn returns the *Conn this *Rows is using.
+func (rows *Rows) Conn() *Conn {
+	return rows.conn
+}
+
 func (rows *Rows) nextColumn() (*ValueReader, bool) {
 	if rows.closed {
 		return nil, false
@@ -469,6 +473,20 @@ func (rows *Rows) Values() ([]interface{}, error) {
 	return values, rows.Err()
 }
 
+// AfterClose adds f to a LILO queue of functions that will be called when
+// rows is closed.
+func (rows *Rows) AfterClose(f func(*Rows)) {
+	if rows.afterClose == nil {
+		rows.afterClose = f
+	} else {
+		prevFn := rows.afterClose
+		rows.afterClose = func(rows *Rows) {
+			f(rows)
+			prevFn(rows)
+		}
+	}
+}
+
 // Query executes sql with args. If there is an error the returned *Rows will
 // be returned in an error state. So it is allowed to ignore the error returned
 // from Query and handle it in *Rows.