added pool.Size() and pool.ExecScript()

sqlitex/pool.go

@@ -208,6 +208,35 @@ func (p *Pool) Close() (err error) {
 	return
 }
 
+// Size returns the maximum number of potential connections in the Pool.
+// Safe to call on a zero Pool
+func (p *Pool) Size() int {
+	return cap(p.free)
+}
+
+// ExecScript executes a script of SQL statements on all connections.
+// Useful for e.g., applying the same pragma to the entire pool.
+// Each invocation of the script is wrapped in a SAVEPOINT transaction,
+// which is rolled back on any error.
+// But if any invocation errors, the pool is closed to avoid inconsistent behaviour
+// between connections
+func (p *Pool) ExecScript(script string) (err error) {
+	// p.mu.RLock()
+	// defer p.mu.Unlock()
+	conns := make([]*sqlite.Conn, p.Size()) //used to drain the pool
+	for i := 0; i < len(conns); i++ {
+		conns[i] = p.Get(nil)
+		if err = ExecScript(conns[i], script); err != nil {
+			p.Close() //ignoring potential Close() errors
+			return err
+		}
+	}
+	for i := 0; i < len(conns); i++ {
+		p.Put(conns[i])
+	}
+	return nil
+}
+
 type strerror struct {
 	msg string
 }

sqlitex/pool_test.go

@@ -243,3 +243,34 @@ func TestPoolPutMatch(t *testing.T) {
 		dbpool1.Put(c)
 	}()
 }
+
+func TestPoolExecScript(t *testing.T) {
+	dbpool := newMemPool(t)
+	defer func() {
+		if err := dbpool.Close(); err != nil {
+			t.Error(err)
+		}
+	}()
+	if err := dbpool.ExecScript("PRAGMA user_version=1;"); err != nil {
+		t.Error(err)
+	}
+	version := 0
+	// check the journal_mode state of each pool
+	fn := func(stmt *sqlite.Stmt) error {
+		version = stmt.ColumnInt(0)
+		return nil
+	}
+	conns := make([]*sqlite.Conn, dbpool.Size()) // to drain the pool
+	for i := 0; i < len(conns); i++ {
+		conns[i] = dbpool.Get(nil)
+		if err := sqlitex.Exec(conns[0], "PRAGMA user_version;", fn); err != nil {
+			t.Fatal(err)
+		}
+		if version != 1 {
+			t.Errorf("version=%d, want 1", version)
+		}
+	}
+	for i := 0; i < len(conns); i++ {
+		dbpool.Put(conns[i])
+	}
+}