Add DB serialization support

serialize.go

@@ -0,0 +1,153 @@
+package sqlite
+
+// #include <sqlite3.h>
+// #include <stdlib.h>
+import "C"
+import (
+	"runtime"
+	"unsafe"
+)
+
+// Serialized contains schema and serialized data for a database.
+type Serialized struct {
+	schema         string
+	data           []byte
+	sqliteOwnsData bool
+	shouldFreeData bool
+}
+
+// NewSerialized creates a new serialized DB from the given schema and data.
+//
+// If copyToSqlite is true, the data will be copied. This should be set to true
+// if this will be used with SQLITE_DESERIALIZE_FREEONCLOSE.
+func NewSerialized(schema string, data []byte, copyToSqlite bool) *Serialized {
+	s := &Serialized{
+		schema: schema,
+		data:   data,
+	}
+	if copyToSqlite {
+		sqliteData := (*[1 << 28]uint8)(unsafe.Pointer(C.sqlite3_malloc(C.int(len(data)))))[:len(data):len(data)]
+		copy(sqliteData, data)
+		s.data = sqliteData
+		s.shouldFreeData = true
+		s.sqliteOwnsData = true
+		runtime.SetFinalizer(s, func(s *Serialized) { s.free() })
+	}
+	return s
+}
+
+// Schema returns the schema for this serialized DB.
+func (s *Serialized) Schema() string {
+	if s.schema == "" {
+		return "main"
+	}
+	return s.schema
+}
+
+// Bytes returns the serialized bytes. Do not mutate this value. This is only
+// valid for the life of its receiver and should be copied for any other
+// longer-term use.
+func (s *Serialized) Bytes() []byte { return s.data }
+
+func (s *Serialized) free() {
+	if len(s.data) > 0 && s.shouldFreeData {
+		s.shouldFreeData = false
+		s.sqliteOwnsData = false
+		C.sqlite3_free(unsafe.Pointer(&s.data[0]))
+	}
+	s.data = nil
+}
+
+// SerializeFlags are flags used for Serialize.
+type SerializeFlags int
+
+const (
+	SQLITE_SERIALIZE_NOCOPY SerializeFlags = C.SQLITE_SERIALIZE_NOCOPY
+)
+
+// Serialize serializes the given schema. Returns nil on error. If
+// SQLITE_SERIALIZE_NOCOPY flag is set, the data may only be valid as long as
+// the database.
+//
+// https://www.sqlite.org/c3ref/serialize.html
+func (conn *Conn) Serialize(schema string, flags ...SerializeFlags) *Serialized {
+	var serializeFlags SerializeFlags
+	for _, f := range flags {
+		serializeFlags |= f
+	}
+
+	cschema := cmain
+	if schema != "" && schema != "main" {
+		cschema = C.CString(schema)
+		defer C.free(unsafe.Pointer(cschema))
+	}
+
+	var csize C.sqlite3_int64
+	res := C.sqlite3_serialize(conn.conn, cschema, &csize, C.uint(serializeFlags))
+	if res == nil {
+		return nil
+	}
+
+	s := &Serialized{
+		schema:         schema,
+		data:           (*[1 << 28]uint8)(unsafe.Pointer(res))[:csize:csize],
+		sqliteOwnsData: true,
+	}
+	// Free the memory only if they didn't specify nocopy
+	if serializeFlags&SQLITE_SERIALIZE_NOCOPY == 0 {
+		s.shouldFreeData = true
+		runtime.SetFinalizer(s, func(s *Serialized) { s.free() })
+	}
+	return s
+}
+
+// DeserializeFlags are flags used for Deserialize.
+type DeserializeFlags int
+
+const (
+	SQLITE_DESERIALIZE_FREEONCLOSE DeserializeFlags = C.SQLITE_DESERIALIZE_FREEONCLOSE
+	SQLITE_DESERIALIZE_RESIZEABLE  DeserializeFlags = C.SQLITE_DESERIALIZE_RESIZEABLE
+	SQLITE_DESERIALIZE_READONLY    DeserializeFlags = C.SQLITE_DESERIALIZE_READONLY
+)
+
+// Reopens the database as in-memory representation of given serialized bytes.
+// The given *Serialized instance should remain referenced (i.e. not GC'd) for
+// the life of the DB since the bytes within are referenced directly.
+//
+// Callers should only use SQLITE_DESERIALIZE_FREEONCLOSE and
+// SQLITE_DESERIALIZE_RESIZEABLE if the param came from Serialize or
+// copyToSqlite was given to NewSerialized.
+//
+// The Serialized parameter should no longer be used after this call.
+//
+// https://www.sqlite.org/c3ref/deserialize.html
+func (conn *Conn) Deserialize(s *Serialized, flags ...DeserializeFlags) error {
+	var deserializeFlags DeserializeFlags
+	for _, f := range flags {
+		deserializeFlags |= f
+	}
+
+	cschema := cmain
+	if s.schema != "" && s.schema != "main" {
+		cschema = C.CString(s.schema)
+		defer C.free(unsafe.Pointer(cschema))
+	}
+
+	// If they set to free on close, remove the free flag from the param
+	if deserializeFlags&SQLITE_DESERIALIZE_FREEONCLOSE == 1 {
+		s.shouldFreeData = false
+	}
+
+	res := C.sqlite3_deserialize(
+		conn.conn,
+		cschema,
+		(*C.uchar)(unsafe.Pointer(&s.data[0])),
+		C.sqlite3_int64(len(s.data)),
+		C.sqlite3_int64(len(s.data)),
+		C.uint(deserializeFlags),
+	)
+	if res != C.SQLITE_OK {
+		return conn.extreserr("Conn.Deserialize", "", res)
+	}
+	return nil
+}

serialize_test.go

@@ -0,0 +1,122 @@
+package sqlite_test
+
+import (
+	"reflect"
+	"strconv"
+	"strings"
+	"testing"
+
+	"crawshaw.io/sqlite"
+	"crawshaw.io/sqlite/sqlitex"
+)
+
+func TestSerialize(t *testing.T) {
+	conn, err := sqlite.OpenConn(":memory:", 0)
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer conn.Close()
+
+	// Create table and insert a few records
+	err = sqlitex.Exec(conn, "CREATE TABLE mytable (v1 PRIMARY KEY, v2, v3);", nil)
+	if err != nil {
+		t.Fatal(err)
+	}
+	err = sqlitex.Exec(conn,
+		"INSERT INTO mytable (v1, v2, v3)  VALUES ('foo', 'bar', 'baz'), ('foo2', 'bar2', 'baz2');", nil)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	// Serialize
+	ser := conn.Serialize("")
+	if ser == nil {
+		t.Fatal("unexpected nil")
+	}
+	origLen := len(ser.Bytes())
+	t.Logf("Initial serialized size: %v", origLen)
+
+	// Create new connection, confirm table not there
+	conn, err = sqlite.OpenConn(":memory:", 0)
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer conn.Close()
+	err = sqlitex.Exec(conn, "SELECT * FROM mytable ORDER BY v1;", nil)
+	if err == nil || !strings.Contains(err.Error(), "no such table") {
+		t.Fatalf("expected no-table error, got: %v", err)
+	}
+
+	// Deserialize into connection and allow resizing
+	err = conn.Deserialize(ser, sqlite.SQLITE_DESERIALIZE_FREEONCLOSE|sqlite.SQLITE_DESERIALIZE_RESIZEABLE)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	// Confirm data there
+	data := [][3]string{}
+	err = sqlitex.Exec(conn, "SELECT * FROM mytable ORDER BY v1;", func(stmt *sqlite.Stmt) error {
+		data = append(data, [3]string{stmt.ColumnText(0), stmt.ColumnText(1), stmt.ColumnText(2)})
+		return nil
+	})
+	if err != nil {
+		t.Fatal(err)
+	}
+	expected := [][3]string{{"foo", "bar", "baz"}, {"foo2", "bar2", "baz2"}}
+	if !reflect.DeepEqual(expected, data) {
+		t.Fatalf("expected %v, got %v", expected, data)
+	}
+
+	// Confirm 1000 inserts can be made
+	for i := 0; i < 1000; i++ {
+		toAppend := strconv.Itoa(i + 3)
+		err = sqlitex.Exec(conn, "INSERT INTO mytable (v1, v2, v3)  VALUES ('foo"+
+			toAppend+"', 'bar"+toAppend+"', 'baz3"+toAppend+"')", nil)
+		if err != nil {
+			t.Fatal(err)
+		}
+	}
+
+	// Serialize again, this time with no-copy
+	ser = conn.Serialize("")
+	if ser == nil {
+		t.Fatal("unexpected nil")
+	}
+	newLen := len(ser.Bytes())
+	if newLen <= origLen {
+		t.Fatalf("expected %v > %v", newLen, origLen)
+	}
+	t.Logf("New serialized size: %v", newLen)
+
+	// Copy the serialized bytes but to not let sqlite own them
+	ser = sqlite.NewSerialized(ser.Schema(), ser.Bytes(), false)
+
+	// Create new conn, deserialize read only
+	conn, err = sqlite.OpenConn(":memory:", 0)
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer conn.Close()
+	err = conn.Deserialize(ser, sqlite.SQLITE_DESERIALIZE_READONLY)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	// Count
+	var total int64
+	err = sqlitex.Exec(conn, "SELECT COUNT(1) FROM mytable;", func(stmt *sqlite.Stmt) error {
+		total = stmt.ColumnInt64(0)
+		return nil
+	})
+	if err != nil {
+		t.Fatal(err)
+	} else if total != 1002 {
+		t.Fatalf("expected 1002, got %v", total)
+	}
+
+	// Try to insert again
+	err = sqlitex.Exec(conn, "INSERT INTO mytable (v1, v2, v3)  VALUES ('a', 'b', 'c');", nil)
+	if err == nil || !strings.Contains(err.Error(), "readonly") {
+		t.Fatalf("expected readonly error, got: %v", err)
+	}
+}