support TursoDB database imports

Author: pedrocarlo

internal/cmd/db_create.go

@@ -155,7 +155,12 @@ func createDatabase(client *turso.Client, name, location, groupName string, seed
 	if sizeLimitFlag != "" {
 		return createDatabaseV2(client, name, location, groupName, seed, spinner)
 	}
-	if seed != nil && seed.Type != "database" && seed.Type != "upload" {
+	// Only fork seeds ("database") can go through the v3 API. File seeds
+	// ("database_upload") must use the v2 flow: the v3 branch never uploads
+	// the file after creating the database. (This used to compare against
+	// "upload", a seed type that doesn't exist, which routed uploads to v2
+	// by accident.)
+	if seed != nil && seed.Type != "database" {
 		return createDatabaseV2(client, name, location, groupName, seed, spinner)
 	}
 	orgID, err := tryResolveOrgID(client)

internal/cmd/db_import.go

@@ -9,11 +9,19 @@ import (
 	"github.com/spf13/cobra"
 )
 
+// Import defaults to TursoDB (MVCC) databases — the only mode the cloud
+// supports for new databases. The flag is separate from createCmd's
+// --tursodb (which shares the tursoDBFlag global with a false default) so
+// the two commands can have different defaults; RunE copies it over before
+// delegating to CreateDatabase.
+var importTursoDBFlag bool
+
 func init() {
 	dbCmd.AddCommand(importCmd)
 	addGroupFlag(importCmd)
 	addRemoteEncryptionKeyFlag(importCmd)
 	addRemoteEncryptionCipherFlag(importCmd)
+	importCmd.Flags().BoolVar(&importTursoDBFlag, "tursodb", true, "Import into a TursoDB (MVCC) database.")
 }
 
 var importCmd = &cobra.Command{
@@ -43,6 +51,7 @@ var importCmd = &cobra.Command{
 		}
 
 		fromFileFlag = filename
+		tursoDBFlag = importTursoDBFlag
 		name := sanitizeDatabaseName(filename)
 		return CreateDatabase(name)
 	},

internal/cmd/group_flag.go

@@ -8,6 +8,7 @@ import (
 	"log"
 	"os"
 	"os/exec"
+	"path/filepath"
 	"strconv"
 	"strings"
 	"time"
@@ -340,11 +341,110 @@ func runQuickCheck(file string) error {
 	return nil
 }
 
+// checkpointWALBeforeUpload folds any pending WAL frames into the main
+// database file and verifies no data is left behind in sidecar files. The
+// upload ships only the main database file, so frames still sitting in
+// <file>-wal (or a hot rollback journal) would silently be missing from the
+// imported database.
+func checkpointWALBeforeUpload(file string) error {
+	if out, err := exec.Command("sqlite3", "-list", file, "PRAGMA wal_checkpoint(TRUNCATE);").CombinedOutput(); err != nil {
+		return fmt.Errorf("could not checkpoint database %s: %w: %s", file, err, out)
+	}
+	for _, sidecar := range []struct{ suffix, hint string }{
+		{"-wal", "close all connections to the database and retry the import"},
+		{"-journal", "the database has a leftover rollback journal; open and cleanly close it with sqlite3 first"},
+	} {
+		if err := checkSidecarEmpty(file+sidecar.suffix, sidecar.hint); err != nil {
+			return err
+		}
+	}
+	return nil
+}
+
+func checkSidecarEmpty(sidecarPath, hint string) error {
+	info, err := os.Stat(sidecarPath)
+	if errors.Is(err, os.ErrNotExist) {
+		return nil
+	}
+	if err != nil {
+		return fmt.Errorf("could not check %s: %w", sidecarPath, err)
+	}
+	if info.Size() > 0 {
+		return fmt.Errorf("%s is not empty, importing would lose the data it holds: %s", sidecarPath, hint)
+	}
+	return nil
+}
+
+// tursodbLogPath returns the logical-log sidecar path tursodb uses for a
+// database file, mirroring turso_core's `with_extension("db-log")`: the
+// file's extension (if any) is replaced with "db-log".
+func tursodbLogPath(file string) string {
+	return strings.TrimSuffix(file, filepath.Ext(file)) + ".db-log"
+}
+
 func handleDBFileAWS(file string, cipher string) (*turso.DBSeed, error) {
-	if err := sqliteFileIntegrityChecks(file, cipher); err != nil {
+	format, err := sniffSQLiteFileFormat(file)
+	if err != nil {
 		return nil, err
 	}
 
+	if format == fileFormatRollback {
+		// The server only accepts WAL or MVCC format files. Converting to WAL
+		// is exactly the remediation the error message used to instruct users
+		// to run themselves, and sqlite3 is already a hard requirement here.
+		fmt.Printf("File %s uses a rollback journal; converting it to WAL mode for import.\n", file)
+		if out, err := exec.Command("sqlite3", file, "PRAGMA journal_mode=WAL;").CombinedOutput(); err != nil {
+			return nil, fmt.Errorf("could not convert %s to WAL mode: %w: %s", file, err, out)
+		}
+		if format, err = sniffSQLiteFileFormat(file); err != nil {
+			return nil, err
+		}
+	}
+
+	switch format {
+	case fileFormatWAL:
+		if err := checkpointWALBeforeUpload(file); err != nil {
+			return nil, err
+		}
+		if err := sqliteFileIntegrityChecks(file, cipher); err != nil {
+			return nil, err
+		}
+	case fileFormatMVCC:
+		// sqlite3-based checks can't run on MVCC (tursodb format) files: the
+		// sqlite3 binary reports them as not-a-database. The server verifies
+		// the file with the tursodb engine after upload.
+		if !tursoDBFlag {
+			return nil, fmt.Errorf("%s is in tursodb (MVCC) format and can only be imported into a tursodb database", file)
+		}
+		if cipher != "" {
+			return nil, errors.New("remote encryption is not supported when importing tursodb (MVCC) format files")
+		}
+		// The upload ships only the main database file: a non-empty tursodb
+		// logical log next to it means the file is not a fully checkpointed
+		// snapshot and importing it would lose the log's data.
+		if err := checkSidecarEmpty(tursodbLogPath(file), "checkpoint the database with tursodb before importing"); err != nil {
+			return nil, err
+		}
+		fileInfo, err := os.Stat(file)
+		if err != nil {
+			return nil, fmt.Errorf("failed to get file info: %w", err)
+		}
+		if fileInfo.Size() > MaxAWSDBSizeBytes {
+			return nil, errors.New("database file size exceeds maximum allowed size of 20 GB")
+		}
+	case fileFormatNotSQLite:
+		isDump, err := checkIfDump(file)
+		if err != nil {
+			return nil, fmt.Errorf("failed to get file header: %w", err)
+		}
+		if isDump {
+			return nil, fmt.Errorf("%s is a sqlite3 dump, not a sqlite3 database. Please import a sqlite database", file)
+		}
+		return nil, fmt.Errorf("file %s is not a valid SQLite database file", file)
+	default:
+		return nil, fmt.Errorf("file %s has an unsupported SQLite file format", file)
+	}
+
 	seed := &turso.DBSeed{
 		Type:     "database_upload",
 		Filepath: file,
@@ -365,6 +465,16 @@ func handleDBFile(client *turso.Client, file string, isAWS bool, cipher string)
 		return handleDBFileAWS(file, cipher)
 	}
 
+	format, err := sniffSQLiteFileFormat(file)
+	if err != nil {
+		return nil, err
+	}
+	if format == fileFormatMVCC {
+		// non-AWS groups are seeded by replaying a .dump, which sqlite3 cannot
+		// produce from an MVCC (tursodb format) file
+		return nil, fmt.Errorf("%s is in tursodb (MVCC) format and can only be imported into AWS groups", file)
+	}
+
 	if err := checkSQLiteFile(file); err != nil {
 		return nil, err
 	}

internal/cmd/group_flag_test.go

@@ -74,3 +74,49 @@ func TestRunQuickCheck(t *testing.T) {
 		require.Error(t, err)
 	})
 }
+
+func TestTursodbLogPath(t *testing.T) {
+	require.Equal(t, "data.db-log", tursodbLogPath("data.db"))
+	require.Equal(t, "data.db-log", tursodbLogPath("data.sqlite"))
+	require.Equal(t, "/some/dir/mydb.db-log", tursodbLogPath("/some/dir/mydb.db"))
+	require.Equal(t, "noext.db-log", tursodbLogPath("noext"))
+}
+
+func TestCheckSidecarEmpty(t *testing.T) {
+	dir := t.TempDir()
+
+	t.Run("missing sidecar is fine", func(t *testing.T) {
+		require.NoError(t, checkSidecarEmpty(filepath.Join(dir, "missing-wal"), "hint"))
+	})
+
+	t.Run("empty sidecar is fine", func(t *testing.T) {
+		path := filepath.Join(dir, "empty-wal")
+		require.NoError(t, os.WriteFile(path, nil, 0644))
+		require.NoError(t, checkSidecarEmpty(path, "hint"))
+	})
+
+	t.Run("non-empty sidecar errors with hint", func(t *testing.T) {
+		path := filepath.Join(dir, "full-wal")
+		require.NoError(t, os.WriteFile(path, []byte("frames"), 0644))
+		err := checkSidecarEmpty(path, "close all connections")
+		require.Error(t, err)
+		require.Contains(t, err.Error(), "importing would lose the data it holds")
+		require.Contains(t, err.Error(), "close all connections")
+	})
+}
+
+func TestCheckpointWALBeforeUpload(t *testing.T) {
+	if _, err := exec.LookPath("sqlite3"); err != nil {
+		t.Skip("sqlite3 not available, skipping test")
+	}
+
+	dbPath := createTestDatabase(t, 10*1024)
+	require.NoError(t, checkpointWALBeforeUpload(dbPath))
+
+	// after the checkpoint no data-bearing sidecars may remain
+	for _, suffix := range []string{"-wal", "-journal"} {
+		if info, err := os.Stat(dbPath + suffix); err == nil {
+			require.Zero(t, info.Size(), "%s must be empty after checkpoint", dbPath+suffix)
+		}
+	}
+}

internal/cmd/sqlite_header.go

@@ -0,0 +1,52 @@
+package cmd
+
+import (
+	"io"
+	"os"
+)
+
+type sqliteFileFormat int
+
+const (
+	fileFormatNotSQLite sqliteFileFormat = iota
+	fileFormatRollback
+	fileFormatWAL
+	fileFormatMVCC
+	fileFormatUnknown
+)
+
+const sqliteMagic = "SQLite format 3\x00"
+
+// sniffSQLiteFileFormat classifies a database file by its SQLite header: the
+// 16-byte magic string plus the read/write format version bytes at offsets
+// 18/19 (1 = rollback journal, 2 = WAL, 255 = MVCC, i.e. tursodb format).
+//
+// This must run before any sqlite3 shellout: the sqlite3 binary reports
+// MVCC-format files as not-a-database, so they have to be routed around the
+// sqlite3-based checks entirely.
+func sniffSQLiteFileFormat(path string) (sqliteFileFormat, error) {
+	file, err := os.Open(path)
+	if err != nil {
+		return fileFormatNotSQLite, err
+	}
+	defer file.Close()
+
+	header := make([]byte, 20)
+	if _, err := io.ReadFull(file, header); err != nil {
+		// too short to hold a SQLite header
+		return fileFormatNotSQLite, nil
+	}
+	if string(header[:len(sqliteMagic)]) != sqliteMagic {
+		return fileFormatNotSQLite, nil
+	}
+	readVersion, writeVersion := header[18], header[19]
+	switch {
+	case readVersion == 1 && writeVersion == 1:
+		return fileFormatRollback, nil
+	case readVersion == 2 && writeVersion == 2:
+		return fileFormatWAL, nil
+	case readVersion == 255 && writeVersion == 255:
+		return fileFormatMVCC, nil
+	}
+	return fileFormatUnknown, nil
+}

internal/cmd/sqlite_header_test.go

@@ -0,0 +1,90 @@
+package cmd
+
+import (
+	"os"
+	"os/exec"
+	"path/filepath"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+func writeHeaderFile(t *testing.T, magic string, readVersion, writeVersion byte) string {
+	t.Helper()
+	header := make([]byte, 100)
+	copy(header, magic)
+	header[18] = readVersion
+	header[19] = writeVersion
+	path := filepath.Join(t.TempDir(), "test.db")
+	require.NoError(t, os.WriteFile(path, header, 0644))
+	return path
+}
+
+func TestSniffSQLiteFileFormat(t *testing.T) {
+	for _, tc := range []struct {
+		name                      string
+		magic                     string
+		readVersion, writeVersion byte
+		expected                  sqliteFileFormat
+	}{
+		{"rollback journal", sqliteMagic, 1, 1, fileFormatRollback},
+		{"wal", sqliteMagic, 2, 2, fileFormatWAL},
+		{"mvcc", sqliteMagic, 255, 255, fileFormatMVCC},
+		{"mixed version bytes", sqliteMagic, 2, 255, fileFormatUnknown},
+		{"unknown version bytes", sqliteMagic, 42, 42, fileFormatUnknown},
+		{"wrong magic", "Not a SQLite db\x00", 2, 2, fileFormatNotSQLite},
+	} {
+		t.Run(tc.name, func(t *testing.T) {
+			path := writeHeaderFile(t, tc.magic, tc.readVersion, tc.writeVersion)
+			format, err := sniffSQLiteFileFormat(path)
+			require.NoError(t, err)
+			require.Equal(t, tc.expected, format)
+		})
+	}
+
+	t.Run("file shorter than header", func(t *testing.T) {
+		path := filepath.Join(t.TempDir(), "short.db")
+		require.NoError(t, os.WriteFile(path, []byte("SQLite"), 0644))
+		format, err := sniffSQLiteFileFormat(path)
+		require.NoError(t, err)
+		require.Equal(t, fileFormatNotSQLite, format)
+	})
+
+	t.Run("missing file", func(t *testing.T) {
+		_, err := sniffSQLiteFileFormat(filepath.Join(t.TempDir(), "missing.db"))
+		require.Error(t, err)
+	})
+}
+
+func TestSniffSQLiteFileFormatOnRealDatabases(t *testing.T) {
+	if _, err := exec.LookPath("sqlite3"); err != nil {
+		t.Skip("sqlite3 not available, skipping test")
+	}
+
+	t.Run("wal database", func(t *testing.T) {
+		dbPath := createTestDatabase(t, 10*1024)
+		format, err := sniffSQLiteFileFormat(dbPath)
+		require.NoError(t, err)
+		require.Equal(t, fileFormatWAL, format)
+	})
+
+	t.Run("rollback database converts to wal", func(t *testing.T) {
+		dbPath := filepath.Join(t.TempDir(), "rollback.db")
+		cmd := exec.Command("sqlite3", "-list", dbPath,
+			"PRAGMA page_size=4096;",
+			"CREATE TABLE data (id INTEGER PRIMARY KEY);")
+		require.NoError(t, cmd.Run(), "failed to create test database")
+
+		format, err := sniffSQLiteFileFormat(dbPath)
+		require.NoError(t, err)
+		require.Equal(t, fileFormatRollback, format)
+
+		// the conversion handleDBFileAWS performs for rollback files
+		out, err := exec.Command("sqlite3", dbPath, "PRAGMA journal_mode=WAL;").CombinedOutput()
+		require.NoError(t, err, "convert to WAL: %s", out)
+
+		format, err = sniffSQLiteFileFormat(dbPath)
+		require.NoError(t, err)
+		require.Equal(t, fileFormatWAL, format)
+	})
+}