From ce9416058027e4e7e29b4a924dc1e816ab020d18 Mon Sep 17 00:00:00 2001
From: James Tucker <james@tailscale.com>
Date: Fri, 26 May 2023 16:45:45 -0700
Subject: [PATCH] sqlite: implement driver.SessionResetter and driver.Validator

These implementations enable us to maintain connections in database/sql
pools.

Fixes #70
---
 sqlite.go      | 26 ++++++++++++++++++++++++++
 sqlite_test.go |  7 +++++++
 2 files changed, 33 insertions(+)

diff --git a/sqlite.go b/sqlite.go
index e049e3c..fe77515 100644
--- a/sqlite.go
+++ b/sqlite.go
@@ -275,6 +275,32 @@ func (c *conn) BeginTx(ctx context.Context, opts driver.TxOptions) (driver.Tx, e
 // Raw is so ConnInitFunc can cast to SQLConn.
 func (c *conn) Raw(fn func(any) error) error { return fn(c) }
 
+// ResetSession is called prior to executing a query on the connection
+// if the connection has been used before. If the driver returns ErrBadConn
+// the connection is discarded.
+func (c *conn) ResetSession(ctx context.Context) error {
+	if c.txState != txStateNone {
+		return errors.New("cannot ResetSession on a connection with an active transaction")
+	}
+
+	// ensure that all prepared statements are reset and unbound
+	if c.stmts != nil {
+		for _, s := range c.stmts {
+			if err := s.resetAndClear(); err != nil {
+				return err
+			}
+		}
+	}
+
+	return nil
+}
+
+// IsValid is called prior to placing the connection into the
+// connection pool. The connection will be discarded if false is returned.
+func (c *conn) IsValid() bool {
+	return true
+}
+
 type readOnlyKey struct{}
 
 // ReadOnly applies the query_only pragma to the connection.
diff --git a/sqlite_test.go b/sqlite_test.go
index 2070b39..5206bee 100644
--- a/sqlite_test.go
+++ b/sqlite_test.go
@@ -20,6 +20,13 @@ import (
 	"github.com/tailscale/sqlite/sqliteh"
 )
 
+// SessionResetter and Validator are required in order for conns to be returned
+// to database/sql connection pools.
+var (
+	_ driver.SessionResetter = (*conn)(nil)
+	_ driver.Validator       = (*conn)(nil)
+)
+
 func TestOpenDB(t *testing.T) {
 	db := openTestDB(t)
 	var journalMode string