sqlite: add optional ConnLogger

ConnLogger can be used to log executed statements for a connection. The
interface captures events for Begin, Exec, Commit, and Rollback calls.

Updates tailscale/corp#33577

Author: alisdair

sqlite.go

@@ -127,9 +127,19 @@ func Connector(sqliteURI string, connInitFunc ConnInitFunc, tracer sqliteh.Trace
 	}
 }
 
+func ConnectorWithLogger(sqliteURI string, connInitFunc ConnInitFunc, tracer sqliteh.Tracer, makeLogger func() ConnLogger) driver.Connector {
+	return &connector{
+		name:         sqliteURI,
+		tracer:       tracer,
+		makeLogger:   makeLogger,
+		connInitFunc: connInitFunc,
+	}
+}
+
 type connector struct {
 	name         string
 	tracer       sqliteh.Tracer
+	makeLogger   func() ConnLogger
 	connInitFunc ConnInitFunc
 }
 
@@ -152,12 +162,14 @@ func (p *connector) Connect(ctx context.Context) (driver.Conn, error) {
 		}
 		return nil, err
 	}
-
 	c := &conn{
 		db:     db,
 		tracer: p.tracer,
 		id:     sqliteh.TraceConnID(maxConnID.Add(1)),
 	}
+	if p.makeLogger != nil {
+		c.logger = p.makeLogger()
+	}
 	if p.connInitFunc != nil {
 		if err := p.connInitFunc(ctx, c); err != nil {
 			db.Close()
@@ -179,6 +191,7 @@ type conn struct {
 	db       sqliteh.DB
 	id       sqliteh.TraceConnID
 	tracer   sqliteh.Tracer
+	logger   ConnLogger
 	stmts    map[string]*stmt // persisted statements
 	txState  txState
 	readOnly bool
@@ -341,6 +354,9 @@ func (c *conn) txInit(ctx context.Context) error {
 			return err
 		}
 	} else {
+		if c.logger != nil {
+			c.logger.Begin()
+		}
 		// TODO(crawshaw): offer BEGIN DEFERRED (and BEGIN CONCURRENT?)
 		// semantics via a context annotation function.
 		if err := c.execInternal(ctx, "BEGIN IMMEDIATE"); err != nil {
@@ -381,6 +397,9 @@ func (tx *connTx) Commit() error {
 	if tx.conn.tracer != nil {
 		tx.conn.tracer.Commit(tx.conn.id, err)
 	}
+	if tx.conn.logger != nil && err == nil {
+		tx.conn.logger.Commit()
+	}
 	return err
 }
 
@@ -394,6 +413,9 @@ func (tx *connTx) Rollback() error {
 	if tx.conn.tracer != nil {
 		tx.conn.tracer.Rollback(tx.conn.id, err)
 	}
+	if tx.conn.logger != nil {
+		tx.conn.logger.Rollback()
+	}
 	return err
 }
 
@@ -490,6 +512,9 @@ func (s *stmt) ExecContext(ctx context.Context, args []driver.NamedValue) (drive
 	if err := s.bindAll(args); err != nil {
 		return nil, s.reserr("Stmt.Exec(Bind)", err)
 	}
+	if s.conn.logger != nil {
+		s.conn.logger.Statement(s.stmt.ExpandedSQL())
+	}
 
 	if ctx.Value(queryCancelKey{}) != nil {
 		done := make(chan struct{})
@@ -1068,3 +1093,20 @@ func WithQueryCancel(ctx context.Context) context.Context {
 
 // queryCancelKey is a context key for query context enforcement.
 type queryCancelKey struct{}
+
+// ConnLogger is implemented by the caller to support statement-level logging for
+// write transactions. Only Exec calls are logged, not Query calls, as this is
+// intended as a mechanism to replay failed transactions.
+type ConnLogger interface {
+	// Begin is called when a writable transaction is opened.
+	Begin()
+
+	// Statement is called with evaluated SQL when a statement is executed.
+	Statement(sql string)
+
+	// Commit is called when a transaction successfully commits.
+	Commit()
+
+	// Rollback is called when a transaction is rolled back.
+	Rollback()
+}

sqlite_test.go

@@ -13,6 +13,7 @@ import (
 	"os"
 	"reflect"
 	"runtime"
+	"slices"
 	"strconv"
 	"strings"
 	"sync"
@@ -1354,3 +1355,118 @@ func TestDisableFunction(t *testing.T) {
 		t.Fatal("Attempting to use the LOWER function after disabling should have failed")
 	}
 }
+
+type connLogger struct {
+	ch         chan []string
+	statements []string
+}
+
+func (cl *connLogger) Begin() {
+	cl.statements = nil
+}
+
+func (cl *connLogger) Statement(s string) {
+	cl.statements = append(cl.statements, s)
+}
+
+func (cl *connLogger) Commit() {
+	cl.ch <- cl.statements
+}
+
+func (cl *connLogger) Rollback() {
+	cl.statements = nil
+}
+
+func TestConnLogger_writable(t *testing.T) {
+	for _, commit := range []bool{true, false} {
+		doneStatement := "ROLLBACK"
+		if commit {
+			doneStatement = "COMMIT"
+		}
+		t.Run(doneStatement, func(t *testing.T) {
+			ctx := context.Background()
+			ch := make(chan []string, 1)
+			txl := connLogger{ch: ch}
+			makeLogger := func() ConnLogger { return &txl }
+			db := sql.OpenDB(ConnectorWithLogger("file:"+t.TempDir()+"/test.db", nil, nil, makeLogger))
+			configDB(t, db)
+
+			tx, err := db.BeginTx(ctx, nil)
+			if err != nil {
+				t.Fatal(err)
+			}
+			if _, err := tx.Exec("CREATE TABLE T (x INTEGER)"); err != nil {
+				t.Fatal(err)
+			}
+			if _, err := tx.Exec("INSERT INTO T VALUES (?)", 1); err != nil {
+				t.Fatal(err)
+			}
+			done := tx.Rollback
+			if commit {
+				done = tx.Commit
+			}
+			if err := done(); err != nil {
+				t.Fatal(err)
+			}
+			if !commit {
+				select {
+				case got := <-ch:
+					t.Errorf("unexpectedly logged statements for rollback:\n%s", strings.Join(got, "\n"))
+				default:
+					return
+				}
+			}
+
+			want := []string{
+				"BEGIN IMMEDIATE",
+				"CREATE TABLE T (x INTEGER)",
+				"INSERT INTO T VALUES (1)",
+				doneStatement,
+			}
+			select {
+			case got := <-ch:
+				if !slices.Equal(got, want) {
+					t.Errorf("unexpected log statements. got:\n%s\n\nwant:\n%s", strings.Join(got, "\n"), strings.Join(want, "\n"))
+				}
+			default:
+				t.Fatal("no logged statements after commit")
+			}
+		})
+	}
+}
+
+func TestConnLogger_commit_error_retry(t *testing.T) {
+	ctx := context.Background()
+	ch := make(chan []string, 1)
+	txl := connLogger{ch: ch}
+	makeLogger := func() ConnLogger { return &txl }
+	db := sql.OpenDB(ConnectorWithLogger("file:"+t.TempDir()+"/test.db", nil, nil, makeLogger))
+	configDB(t, db)
+
+	if _, err := db.Exec("PRAGMA foreign_keys = ON"); err != nil {
+		t.Fatal(err)
+	}
+	if _, err := db.Exec("CREATE TABLE A (x INTEGER PRIMARY KEY)"); err != nil {
+		t.Fatal(err)
+	}
+	if _, err := db.Exec("CREATE TABLE B (x INTEGER REFERENCES A(X) DEFERRABLE INITIALLY DEFERRED)"); err != nil {
+		t.Fatal(err)
+	}
+
+	tx, err := db.BeginTx(ctx, nil)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if _, err := tx.Exec("INSERT INTO B VALUES (?)", 1); err != nil {
+		t.Fatal(err)
+	}
+	if err := tx.Commit(); err == nil {
+		t.Fatal("expected Commit to error, but didn't")
+	}
+	select {
+	case got := <-ch:
+		t.Errorf("unexpectedly logged statements for errored commit:\n%s", strings.Join(got, "\n"))
+	default:
+		return
+	}
+}