sqlite: implement backup api

Signed-off-by: David Crawshaw <crawshaw@tailscale.com>

Author: crawshaw

cgosqlite/cgosqlite.go

@@ -159,6 +159,53 @@ func (db *DB) TxnState(schema string) sqliteh.TxnState {
 	return sqliteh.TxnState(C.sqlite3_txn_state(db.db, cSchema))
 }
 
+func (db *DB) BackupInit(dstSchema string, src sqliteh.DB, srcSchema string) (sqliteh.Backup, error) {
+	var cDstSchema, cSrcSchema *C.char
+	if dstSchema != "" {
+		cDstSchema = C.CString(dstSchema)
+		defer C.free(unsafe.Pointer(cDstSchema))
+	}
+	if srcSchema != "" {
+		cSrcSchema = C.CString(srcSchema)
+		defer C.free(unsafe.Pointer(cSrcSchema))
+	}
+
+	b := C.sqlite3_backup_init(db.db, cDstSchema, src.(*DB).db, cSrcSchema)
+	if b == nil {
+		// sqlite3_backup_init docs tell us the error is on the dst DB.
+		return nil, sqliteh.ErrCode(db.ExtendedErrCode())
+	}
+	return &backup{backup: b}, nil
+}
+
+type backup struct {
+	backup *C.sqlite3_backup
+}
+
+func (b *backup) Step(numPages int) (more bool, err error) {
+	res := C.sqlite3_backup_step(b.backup, C.int(numPages))
+	if res == C.SQLITE_DONE {
+		return false, nil
+	}
+	if res == C.SQLITE_OK {
+		return true, nil
+	}
+	return false, errCode(res)
+}
+
+func (b *backup) Finish() error {
+	res := C.sqlite3_backup_finish(b.backup)
+	return errCode(res)
+}
+
+func (b *backup) Remaining() int {
+	return int(C.sqlite3_backup_remaining(b.backup))
+}
+
+func (b *backup) PageCount() int {
+	return int(C.sqlite3_backup_pagecount(b.backup))
+}
+
 func (db *DB) Prepare(query string, prepFlags sqliteh.PrepareFlags) (stmt sqliteh.Stmt, remainingQuery string, err error) {
 	csql := C.CString(query)
 	defer C.free(unsafe.Pointer(csql))

sqlite.go

@@ -123,20 +123,25 @@ type connector struct {
 	connInitFunc ConnInitFunc
 }
 
+func errWithMsg(db sqliteh.DB, err error, loc string) error {
+	if ec, ok := err.(sqliteh.ErrCode); ok {
+		e := &Error{
+			Code: sqliteh.Code(ec),
+			Loc:  loc,
+		}
+		if db != nil {
+			e.Msg = db.ErrMsg()
+		}
+		return e
+	}
+	return err
+}
+
 func (p *connector) Driver() driver.Driver { return drv{} }
 func (p *connector) Connect(ctx context.Context) (driver.Conn, error) {
 	db, err := Open(p.name, sqliteh.OpenFlagsDefault, "")
 	if err != nil {
-		if ec, ok := err.(sqliteh.ErrCode); ok {
-			e := &Error{
-				Code: sqliteh.Code(ec),
-				Loc:  "Open",
-			}
-			if db != nil {
-				e.Msg = db.ErrMsg()
-			}
-			err = e
-		}
+		err = errWithMsg(db, err, "Open")
 		if db != nil {
 			db.Close()
 		}
@@ -838,3 +843,35 @@ func WithPersist(ctx context.Context) context.Context {
 
 // persistQuery is used as a context value.
 type persistQuery struct{}
+
+// DB executes fn with the sqliteh.DB underlying sqlconn.
+func DB(sqlconn SQLConn, fn func(sqliteh.DB) error) error {
+	return sqlconn.Raw(func(driverConn interface{}) error {
+		c, ok := driverConn.(*conn)
+		if !ok {
+			return fmt.Errorf("sqlite.Checkpoint: sql.Conn is not the sqlite driver: %T", driverConn)
+		}
+		return fn(c.db)
+	})
+}
+
+// Backup backups the specified database from srcConn to dstConn.
+func Backup(dstConn SQLConn, dstSchema string, srcConn SQLConn, srcSchema string) error {
+	return DB(dstConn, func(dst sqliteh.DB) error {
+		return DB(srcConn, func(src sqliteh.DB) error {
+			b, err := dst.BackupInit(dstSchema, src, srcSchema)
+			if err != nil {
+				return errWithMsg(dst, err, "Backup")
+			}
+			more := true
+			for more {
+				more, err = b.Step(1024)
+				if err != nil {
+					b.Finish()
+					return errWithMsg(dst, err, "Step")
+				}
+			}
+			return errWithMsg(dst, b.Finish(), "Finish")
+		})
+	})
+}

sqlite_test.go

@@ -779,6 +779,59 @@ func TestAttachOrderingDeadlock(t *testing.T) {
 	}
 }
 
+func TestBackup(t *testing.T) {
+	src, err := sql.Open("sqlite3", "file:src?mode=memory")
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer src.Close()
+	dst, err := sql.Open("sqlite3", "file:dst?mode=memory")
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer dst.Close()
+	ctx := context.Background()
+
+	srcConn, err := src.Conn(ctx)
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer srcConn.Close()
+	err = ExecScript(srcConn, `
+		ATTACH 'file:src2?mode=memory' AS src2;
+		CREATE TABLE t1 (c);
+		INSERT INTO t1 VALUES ('a');
+		CREATE TABLE src2.t2 (c);
+		INSERT INTO src2.t2 VALUES ('b');
+	`)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	dstConn, err := src.Conn(ctx)
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer dstConn.Close()
+	if err := Backup(dstConn, "main", srcConn, "main"); err != nil {
+		t.Fatal(err)
+	}
+	if _, err := dstConn.ExecContext(ctx, "ATTACH 'file:dst2?mode=memory' AS dst2;"); err != nil {
+		t.Fatal(err)
+	}
+	var count int
+	if err := dstConn.QueryRowContext(ctx, "SELECT count(*) FROM t1").Scan(&count); err != nil || count != 1 {
+		t.Fatalf("err=%v, count=%d", err, count)
+	}
+	if err := Backup(dstConn, "dst2", srcConn, "src2"); err != nil {
+		t.Fatal(err)
+	}
+	count = 0
+	if err := dstConn.QueryRowContext(ctx, "SELECT count(*) FROM dst2.t2").Scan(&count); err != nil || count != 1 {
+		t.Fatalf("err=%v, count=%d", err, count)
+	}
+}
+
 func BenchmarkPersist(b *testing.B) {
 	ctx := context.Background()
 	db := openTestDB(b)

sqliteh/sqliteh.go

@@ -53,6 +53,9 @@ type DB interface {
 	AutoCheckpoint(n int) error
 	// TxnState is sqlite3_txn_state.
 	TxnState(schema string) TxnState
+	// BackupInit is sqlite3_backup_init, this DB is the destination.
+	// https://www.sqlite.org/c3ref/backup_finish.html#sqlite3backupinit
+	BackupInit(dstSchema string, src DB, srcSchema string) (Backup, error)
 }
 
 // Stmt is an sqlite3_stmt* database connection object.
@@ -165,6 +168,17 @@ type Stmt interface {
 	ColumnTableName(col int) string
 }
 
+// Backup is an sqlite3_backup object.
+// https://www.sqlite.org/c3ref/backup_finish.html
+type Backup interface {
+	// Step is called repeatedly to transfer data between the two DBs.
+	Step(numPages int) (more bool, err error)
+	// Finish releases all resources associated with the Backup.
+	Finish() error
+	Remaining() int
+	PageCount() int
+}
+
 // ColumnType are constants for each of the SQLite datatypes.
 // https://www.sqlite.org/c3ref/c_blob.html
 type ColumnType int