Merge pull request #101 from otoolep/new_db_laer

Support restore to in-memory databases

CHANGELOG.md

@@ -1,5 +1,6 @@
 ## 2.2.2 (unreleased)
 - [PR #96](https://github.com/otoolep/rqlite/pull/96): Add build time to status output.
+- [PR #101](https://github.com/otoolep/rqlite/pull/101): Fix restore to in-memory databases.
 
 ## 2.2.1 (April 19th 2016)
 - [PR #95](https://github.com/otoolep/rqlite/pull/95): Correctly set HTTP authentication.

cmd/rqlited/main.go

@@ -23,7 +23,6 @@ import (
 	"runtime/pprof"
 
 	"github.com/otoolep/rqlite/auth"
-	sql "github.com/otoolep/rqlite/db"
 	httpd "github.com/otoolep/rqlite/http"
 	"github.com/otoolep/rqlite/store"
 )
@@ -143,9 +142,7 @@ func main() {
 	if err != nil {
 		log.Fatalf("failed to determine absolute data path: %s", err.Error())
 	}
-	dbConf := sql.NewConfig()
-	dbConf.DSN = dsn
-	dbConf.Memory = inMem
+	dbConf := store.NewDBConfig(dsn, inMem)
 	store := store.New(dbConf, dataPath, raftAddr)
 	if err := store.Open(joinAddr == ""); err != nil {
 		log.Fatalf("failed to open store: %s", err.Error())

db/db.go

@@ -19,29 +19,12 @@ func init() {
 	DBVersion, _, _ = sqlite3.Version()
 }
 
-// Config represents the configuration of the SQLite database.
-type Config struct {
-	DSN    string // Datasource name
-	Memory bool   // In-memory database enabled?
-}
-
-// NewConfig returns an instance of Config for the database at path.
-func NewConfig() *Config {
-	return &Config{}
-}
-
-// FQDSN returns the fully-qualified datasource name.
-func (c *Config) FQDSN(path string) string {
-	if c.DSN != "" {
-		return fmt.Sprintf("file:%s?%s", path, c.DSN)
-	}
-	return path
-}
-
 // DB is the SQL database.
 type DB struct {
 	sqlite3conn *sqlite3.SQLiteConn // Driver connection to database.
 	path        string              // Path to database file.
+	dsn         string              // DSN, if any.
+	memory      bool                // In-memory only.
 }
 
 // Result represents the outcome of an operation that changes rows.
@@ -61,23 +44,64 @@ type Rows struct {
 	Time    float64         `json:"time,omitempty"`
 }
 
-// Open an existing database, creating it if it does not exist.
+// Open opens a file-based database, creating it if it does not exist.
 func Open(dbPath string) (*DB, error) {
-	d := sqlite3.SQLiteDriver{}
-	dbc, err := d.Open(dbPath)
+	return open(fqdsn(dbPath, ""))
+}
+
+// OpenwithDSN opens a file-based database, creating it if it does not exist.
+func OpenWithDSN(dbPath, dsn string) (*DB, error) {
+	return open(fqdsn(dbPath, dsn))
+}
+
+// OpenInMemory opens an in-memory database.
+func OpenInMemory() (*DB, error) {
+	return open(fqdsn(":memory:", ""))
+}
+
+// OpenInMemoryWithDSN opens an in-memory database with a specific DSN.
+func OpenInMemoryWithDSN(dsn string) (*DB, error) {
+	return open(fqdsn(":memory:", dsn))
+}
+
+// LoadInMemoryWithDSN loads an in-memory database with that at the path,
+// with the specified DSN
+func LoadInMemoryWithDSN(dbPath, dsn string) (*DB, error) {
+	db, err := OpenInMemoryWithDSN(dsn)
 	if err != nil {
 		return nil, err
 	}
 
-	return &DB{
-		sqlite3conn: dbc.(*sqlite3.SQLiteConn),
-		path:        dbPath,
-	}, nil
+	srcDB, err := Open(dbPath)
+	if err != nil {
+		return nil, err
+	}
+
+	bk, err := db.sqlite3conn.Backup("main", srcDB.sqlite3conn, "main")
+	if err != nil {
+		return nil, err
+	}
+
+	for {
+		done, err := bk.Step(-1)
+		if err != nil {
+			bk.Finish()
+			return nil, err
+		}
+		if done {
+			break
+		}
+		time.Sleep(bkDelay * time.Millisecond)
+	}
+
+	if err := bk.Finish(); err != nil {
+		return nil, err
+	}
+	if err := srcDB.Close(); err != nil {
+		return nil, err
 	}
 
-// OpenWithConfiguration an existing database, creating it if it does not exist.
-func OpenWithConfiguration(dbPath string, conf *Config) (*DB, error) {
-	return Open(conf.FQDSN(dbPath))
+	return db, nil
 }
 
 // Close closes the underlying database connection.
@@ -85,6 +109,19 @@ func (db *DB) Close() error {
 	return db.sqlite3conn.Close()
 }
 
+func open(dbPath string) (*DB, error) {
+	d := sqlite3.SQLiteDriver{}
+	dbc, err := d.Open(dbPath)
+	if err != nil {
+		return nil, err
+	}
+
+	return &DB{
+		sqlite3conn: dbc.(*sqlite3.SQLiteConn),
+		path:        dbPath,
+	}, nil
+}
+
 // Execute executes queries that modify the database.
 func (db *DB) Execute(queries []string, tx, xTime bool) ([]*Result, error) {
 	type Execer interface {
@@ -298,3 +335,11 @@ func (db *DB) Backup(path string) error {
 
 	return nil
 }
+
+// fqdsn returns the fully-qualified datasource name.
+func fqdsn(path, dsn string) string {
+	if dsn != "" {
+		return fmt.Sprintf("file:%s?%s", path, dsn)
+	}
+	return path
+}

db/db_test.go

@@ -12,14 +12,6 @@ import (
  * Lowest-layer database tests
  */
 
-func Test_Config(t *testing.T) {
-	c := NewConfig()
-	c.DSN = "cache=shared&mode=memory"
-	if c.FQDSN("/foo/bar/db.sqlite") != "file:/foo/bar/db.sqlite?cache=shared&mode=memory" {
-		t.Fatalf("Fully qualified DSN not correct, got: %s", c.FQDSN("/foo/bar/db.sqlite"))
-	}
-}
-
 func Test_DbFileCreation(t *testing.T) {
 	dir, err := ioutil.TempDir("", "rqlite-test-")
 	defer os.RemoveAll(dir)
@@ -56,6 +48,39 @@ func Test_TableCreation(t *testing.T) {
 	}
 }
 
+func Test_LoadInMemory(t *testing.T) {
+	db, path := mustCreateDatabase()
+	defer db.Close()
+	defer os.Remove(path)
+
+	_, err := db.Execute([]string{"CREATE TABLE foo (id INTEGER NOT NULL PRIMARY KEY, name TEXT)"}, false, false)
+	if err != nil {
+		t.Fatalf("failed to create table: %s", err.Error())
+	}
+
+	r, err := db.Query([]string{"SELECT * FROM foo"}, false, false)
+	if err != nil {
+		t.Fatalf("failed to query empty table: %s", err.Error())
+	}
+	if exp, got := `[{"columns":["id","name"],"types":["integer","text"]}]`, asJSON(r); exp != got {
+		t.Fatalf("unexpected results for query, expected %s, got %s", exp, got)
+	}
+
+	inmem, err := LoadInMemoryWithDSN(path, "")
+	if err != nil {
+		t.Fatalf("failed to create loaded in-memory database: %s", err.Error())
+	}
+
+	// Ensure it has been loaded correctly into the database
+	r, err = inmem.Query([]string{"SELECT * FROM foo"}, false, false)
+	if err != nil {
+		t.Fatalf("failed to query empty table: %s", err.Error())
+	}
+	if exp, got := `[{"columns":["id","name"],"types":["integer","text"]}]`, asJSON(r); exp != got {
+		t.Fatalf("unexpected results for query, expected %s, got %s", exp, got)
+	}
+}
+
 func Test_SimpleSingleStatements(t *testing.T) {
 	db, path := mustCreateDatabase()
 	defer db.Close()

store/store.go

@@ -58,6 +58,17 @@ type command struct {
 	Timings bool        `json:"timings,omitempty"`
 }
 
+// DBConfig represents the configuration of the underlying SQLite database.
+type DBConfig struct {
+	DSN    string // Any custom DSN
+	Memory bool   // Whether the database is in-memory only.
+}
+
+// NewDBConfig returns a new DB config instance.
+func NewDBConfig(dsn string, memory bool) *DBConfig {
+	return &DBConfig{DSN: dsn, Memory: memory}
+}
+
 // Store is a SQLite database, where all changes are made via Raft consensus.
 type Store struct {
 	raftDir  string
@@ -67,25 +78,20 @@ type Store struct {
 
 	ln     *networkLayer // Raft network between nodes.
 	raft   *raft.Raft    // The consensus mechanism.
-	dbConf *sql.Config   // SQLite database config.
-	dbPath string        // Path to database file.
+	dbConf *DBConfig     // SQLite database config.
+	dbPath string        // Path to underlying SQLite file, if not in-memory.
 	db     *sql.DB       // The underlying SQLite store.
 
 	logger *log.Logger
 }
 
 // New returns a new Store.
-func New(dbConf *sql.Config, dir, bind string) *Store {
-	dbPath := filepath.Join(dir, sqliteFile)
-	if dbConf.Memory {
-		dbPath = ":memory:"
-	}
-
+func New(dbConf *DBConfig, dir, bind string) *Store {
 	return &Store{
 		raftDir:  dir,
 		raftBind: bind,
 		dbConf:   dbConf,
-		dbPath:   dbPath,
+		dbPath:   filepath.Join(dir, sqliteFile),
 		logger:   log.New(os.Stderr, "[store] ", log.LstdFlags),
 	}
 }
@@ -98,19 +104,26 @@ func (s *Store) Open(enableSingle bool) error {
 	}
 
 	// Create the database. Unless it's a memory-based database, it must be deleted
+	var db *sql.DB
+	var err error
 	if !s.dbConf.Memory {
 		// as it will be rebuilt from (possibly) a snapshot and committed log entries.
 		if err := os.Remove(s.dbPath); err != nil && !os.IsNotExist(err) {
 			return err
 		}
+		db, err = sql.OpenWithDSN(s.dbPath, s.dbConf.DSN)
+		if err != nil {
+			return err
 		}
-
-	db, err := sql.OpenWithConfiguration(s.dbPath, s.dbConf)
+		s.logger.Println("SQLite database opened at", s.dbPath)
+	} else {
+		db, err = sql.OpenInMemoryWithDSN(s.dbConf.DSN)
 		if err != nil {
 			return err
 		}
+		s.logger.Println("SQLite in-memory database opened")
+	}
 	s.db = db
-	s.logger.Println("SQLite database opened at", s.dbConf.FQDSN(s.dbPath))
 
 	// Setup Raft configuration.
 	config := raft.DefaultConfig()
@@ -234,16 +247,18 @@ func (s *Store) WaitForAppliedIndex(idx uint64, timeout time.Duration) error {
 // Stats returns stats for the store.
 func (s *Store) Stats() (map[string]interface{}, error) {
 	dbStatus := map[string]interface{}{
-		"path":    s.dbPath,
-		"dns":     s.dbConf.FQDSN(s.dbPath),
+		"dns":     s.dbConf.DSN,
 		"version": sql.DBVersion,
 	}
 	if !s.dbConf.Memory {
+		dbStatus["path"] = s.dbPath
 		stat, err := os.Stat(s.dbPath)
 		if err != nil {
 			return nil, err
 		}
 		dbStatus["size"] = stat.Size()
+	} else {
+		dbStatus["path"] = ":memory:"
 	}
 
 	status := map[string]interface{}{
@@ -410,7 +425,7 @@ func (s *Store) Snapshot() (raft.FSMSnapshot, error) {
 
 // Restore restores the database to a previous state.
 func (s *Store) Restore(rc io.ReadCloser) error {
-	if err := os.Remove(s.dbPath); err != nil && !os.IsNotExist(err) {
+	if err := s.db.Close(); err != nil {
 		return err
 	}
 
@@ -419,14 +434,37 @@ func (s *Store) Restore(rc io.ReadCloser) error {
 		return err
 	}
 
+	var db *sql.DB
+	if !s.dbConf.Memory {
+		// Write snapshot over any existing database file.
 		if err := ioutil.WriteFile(s.dbPath, b, 0660); err != nil {
 			return err
 		}
 
-	db, err := sql.OpenWithConfiguration(s.dbPath, s.dbConf)
+		// Re-open it.
+		db, err = sql.OpenWithDSN(s.dbPath, s.dbConf.DSN)
 		if err != nil {
 			return err
 		}
+	} else {
+		// In memory. Copy to temporary file, and then load memory from file.
+		f, err := ioutil.TempFile("", "rqlilte-snap-")
+		if err != nil {
+			return err
+		}
+		f.Close()
+		defer os.Remove(f.Name())
+
+		if err := ioutil.WriteFile(f.Name(), b, 0660); err != nil {
+			return err
+		}
+
+		// Load an in-memory database from the snapshot now on disk.
+		db, err = sql.LoadInMemoryWithDSN(f.Name(), s.dbConf.DSN)
+		if err != nil {
+			return err
+		}
+	}
 	s.db = db
 
 	return nil

store/store_test.go

@@ -7,8 +7,6 @@ import (
 	"path/filepath"
 	"testing"
 	"time"
-
-	sql "github.com/otoolep/rqlite/db"
 )
 
 type mockSnapshotSink struct {
@@ -227,7 +225,7 @@ func Test_MultiNodeExecuteQuery(t *testing.T) {
 	}
 }
 
-func Test_SingleNodeSnapshot(t *testing.T) {
+func Test_SingleNodeSnapshotOnDisk(t *testing.T) {
 	s := mustNewStore(false)
 	defer os.RemoveAll(s.Path())
 
@@ -289,12 +287,73 @@ func Test_SingleNodeSnapshot(t *testing.T) {
 	}
 }
 
+func Test_SingleNodeSnapshotInMem(t *testing.T) {
+	s := mustNewStore(true)
+	defer os.RemoveAll(s.Path())
+
+	if err := s.Open(true); err != nil {
+		t.Fatalf("failed to open single-node store: %s", err.Error())
+	}
+	defer s.Close(true)
+	s.WaitForLeader(10 * time.Second)
+
+	queries := []string{
+		`CREATE TABLE foo (id INTEGER NOT NULL PRIMARY KEY, name TEXT)`,
+		`INSERT INTO foo(id, name) VALUES(1, "fiona")`,
+	}
+	_, err := s.Execute(queries, false, false)
+	if err != nil {
+		t.Fatalf("failed to execute on single node: %s", err.Error())
+	}
+	_, err = s.Query([]string{`SELECT * FROM foo`}, false, false, None)
+	if err != nil {
+		t.Fatalf("failed to query single node: %s", err.Error())
+	}
+
+	// Snap the node and write to disk.
+	f, err := s.Snapshot()
+	if err != nil {
+		t.Fatalf("failed to snapshot node: %s", err.Error())
+	}
+
+	snapDir := mustTempDir()
+	defer os.RemoveAll(snapDir)
+	snapFile, err := os.Create(filepath.Join(snapDir, "snapshot"))
+	if err != nil {
+		t.Fatalf("failed to create snapshot file: %s", err.Error())
+	}
+	sink := &mockSnapshotSink{snapFile}
+	if err := f.Persist(sink); err != nil {
+		t.Fatalf("failed to persist snapshot to disk: %s", err.Error())
+	}
+
+	// Check restoration.
+	snapFile, err = os.Open(filepath.Join(snapDir, "snapshot"))
+	if err != nil {
+		t.Fatalf("failed to open snapshot file: %s", err.Error())
+	}
+	if err := s.Restore(snapFile); err != nil {
+		t.Fatalf("failed to restore snapshot from disk: %s", err.Error())
+	}
+
+	// Ensure database is back in the correct state.
+	r, err := s.Query([]string{`SELECT * FROM foo`}, false, false, None)
+	if err != nil {
+		t.Fatalf("failed to query single node: %s", err.Error())
+	}
+	if exp, got := `["id","name"]`, asJSON(r[0].Columns); exp != got {
+		t.Fatalf("unexpected results for query\nexp: %s\ngot: %s", exp, got)
+	}
+	if exp, got := `[[1,"fiona"]]`, asJSON(r[0].Values); exp != got {
+		t.Fatalf("unexpected results for query\nexp: %s\ngot: %s", exp, got)
+	}
+}
+
 func mustNewStore(inmem bool) *Store {
 	path := mustTempDir()
 	defer os.RemoveAll(path)
 
-	cfg := sql.NewConfig()
-	cfg.Memory = inmem
+	cfg := NewDBConfig("", inmem)
 	s := New(cfg, path, "localhost:0")
 	if s == nil {
 		panic("failed to create new store")

system_test/helpers.go

@@ -11,7 +11,6 @@ import (
 	"strings"
 	"time"
 
-	sql "github.com/otoolep/rqlite/db"
 	httpd "github.com/otoolep/rqlite/http"
 	"github.com/otoolep/rqlite/store"
 )
@@ -172,7 +171,7 @@ func mustNewNode(enableSingle bool) *Node {
 		Dir: mustTempDir(),
 	}
 
-	dbConf := sql.NewConfig()
+	dbConf := store.NewDBConfig("", false)
 	node.Store = store.New(dbConf, node.Dir, "localhost:0")
 	if err := node.Store.Open(enableSingle); err != nil {
 		node.Deprovision()