Restore the SQLite metastore from its backup when the local database is empty

cli/cmd/runtime/start.go

@@ -62,9 +62,18 @@ import (
 // Env var keys must be prefixed with RILL_RUNTIME_ and are converted from snake_case to CamelCase.
 // For example RILL_RUNTIME_HTTP_PORT is mapped to Config.HTTPPort.
 type Config struct {
-	MetastoreDriver         string                 `default:"sqlite" split_words:"true"`
-	MetastoreURL            string                 `default:"file:rill?mode=memory&cache=shared" split_words:"true"`
-	MetastoreID             string                 `split_words:"true"`
+	// MetastoreDriver specifies the database driver for the metastore.
+	MetastoreDriver string `default:"sqlite" split_words:"true"`
+	// MetastoreURL specifies the connection string for the metastore database.
+	// It defaults to an in-memory SQLite database.
+	MetastoreURL string `default:"file:rill?mode=memory&cache=shared" split_words:"true"`
+	// MetastoreID is an optional globally unique ID for the metastore.
+	// It is currently used to identify backups in object storage.
+	MetastoreID string `split_words:"true"`
+	// MetastoreBackupsEnable enables periodic backups of the metastore to object storage.
+	// It also enables restoring from the latest backup if the metastore is missing and a valid backup exists in object storage.
+	// It requires MetastoreID and DataBucket to be set.
+	MetastoreBackupsEnable  bool                   `default:"true" split_words:"true"`
 	RedisURL                string                 `default:"" split_words:"true"`
 	MetricsExporter         observability.Exporter `default:"prometheus" split_words:"true"`
 	TracesExporter          observability.Exporter `default:"" split_words:"true"`
@@ -229,11 +238,12 @@ func StartCmd(ch *cmdutil.Helper) *cobra.Command {
 			ctx := graceful.WithCancelOnTerminate(context.Background())
 			// Init runtime
 			metastoreConfig, err := structpb.NewStruct(map[string]any{
-				"dsn": conf.MetastoreURL,
-				"id":  conf.MetastoreID,
+				"dsn":            conf.MetastoreURL,
+				"id":             conf.MetastoreID,
+				"backups_enable": conf.MetastoreBackupsEnable,
 			})
 			if err != nil {
-				logger.Fatal("error: could not creat metastore metastore config", zap.Error(err))
+				logger.Fatal("could not create metastore config", zap.Error(err))
 			}
 			opts := &runtime.Options{
 				ConnectionCacheSize:          conf.ConnectionCacheSize,

runtime/drivers/sqlite/backups.go

@@ -22,9 +22,13 @@ import (
 	_ "modernc.org/sqlite"
 )
 
+// Name of the SQLite snapshot file in the backup directory.
+// It is the only file in a backup that can be restored from; the Parquet files are for analytics.
+const backupSnapshotName = "snapshot.db"
+
 var (
 	// Maximum size of the SQLite snapshot for backup.
-	backupMaxSizeBytes int64 = 5 * 1024 * 1024 * 1024 // 5 GB
+	backupMaxSizeBytes int64 = 15 * 1024 * 1024 * 1024 // 15 GB
 
 	// Max time a backup may run for.
 	backupMaxDuration = 10 * time.Minute
@@ -56,17 +60,19 @@ var (
 //
 // It is a no-op unless the following pre-requisites are in place:
 // 1. An external bucket is configured on the storage client.
-// 2. A backup ID is provided in the connection config (through the "id" config parameter, currently propagates from RILL_RUNTIME_METASTORE_ID).
-// 3. The SQLite database is file-based and doesn't exceed backupMaxSizeBytes in size.
+// 2. Backups are enabled in the connection config (through the "backups_enable" config parameter, currently propagates from RILL_RUNTIME_METASTORE_BACKUPS_ENABLE).
+// 3. A backup ID is provided in the connection config (through the "id" config parameter, currently propagates from RILL_RUNTIME_METASTORE_ID).
+// 4. The SQLite database is file-based and doesn't exceed backupMaxSizeBytes in size.
 //
-// It is a best-effort backup used for analytics. There are currently no guarantees on backups and no restore functionality.
-// Backups are performed at midnight UTC every day if the runtime is running at that time.
+// It is a best-effort backup. Backups are performed at midnight UTC every day if the runtime is running at that time.
+// The snapshot.db file can be restored on startup; see restoreBackupIfEmpty() for details.
+// The Parquet files are only used for downstream analytics.
 //
 // Backups are stored in the external bucket under the path "shared/metastore/{backupID}/" (the "shared/metastore" prefix is not applied here, but where the connection is opened).
 // The directory will contain a snapshot.db SQLite file and Parquet files for each of the tables defined in parquetBackupQueries.
 func (c *connection) startBackups() {
-	// It's a no-op if no backup ID is provided.
-	if c.backupID == "" {
+	// It's a no-op unless backups are configured.
+	if c.backupID == "" || !c.backupsEnable {
 		return
 	}
 
@@ -150,7 +156,7 @@ func (c *connection) backup(ctx context.Context, bucket *blob.Bucket) error {
 	defer os.RemoveAll(tmpDir)
 
 	// Capture a snapshot of the SQLite database
-	snapshotPath := filepath.Join(tmpDir, "snapshot.db")
+	snapshotPath := filepath.Join(tmpDir, backupSnapshotName)
 	_, err = c.db.ExecContext(ctx, fmt.Sprintf("VACUUM INTO '%s'", snapshotPath))
 	if err != nil {
 		return fmt.Errorf("failed to create SQLite snapshot: %w", err)
@@ -162,7 +168,7 @@ func (c *connection) backup(ctx context.Context, bucket *blob.Bucket) error {
 		return fmt.Errorf("failed to open SQLite snapshot for upload: %w", err)
 	}
 	defer f.Close()
-	err = bucket.Upload(ctx, "snapshot.db", f, &blob.WriterOptions{
+	err = bucket.Upload(ctx, backupSnapshotName, f, &blob.WriterOptions{
 		ContentType: "application/octet-stream",
 	})
 	if err != nil {

runtime/drivers/sqlite/backups_test.go

@@ -30,8 +30,9 @@ func TestBackup(t *testing.T) {
 
 	// Create sqlite handle
 	cfg := map[string]any{
-		"dsn": dbPath,
-		"id":  "test-backup",
+		"dsn":            dbPath,
+		"id":             "test-backup",
+		"backups_enable": true,
 	}
 	h, err := driver{}.Open("", "", cfg, storage.MustNew(storageDir, nil), activity.NewNoopClient(), zap.NewNop())
 	require.NoError(t, err)

runtime/drivers/sqlite/restore.go

@@ -0,0 +1,188 @@
+package sqlite
+
+import (
+	"context"
+	"database/sql"
+	"errors"
+	"fmt"
+	"io/fs"
+	"os"
+	"time"
+
+	"github.com/jmoiron/sqlx"
+	"github.com/rilldata/rill/runtime/storage"
+	"go.uber.org/zap"
+	"gocloud.dev/blob"
+	"gocloud.dev/gcerrors"
+)
+
+// Max time a restore may run for.
+// It blocks runtime startup, so we bound it instead of letting a slow download hang the process.
+// It needs to accommodate downloading a snapshot of up to backupMaxSizeBytes, hence the generous bound.
+var restoreMaxDuration = 30 * time.Minute
+
+// restoreBackupIfEmpty restores the SQLite database at the given DSN from the latest backup in object storage,
+// but only if the database has no data yet. It is called from driver.Open() before the connection handle is created,
+// since it replaces the database file. See startBackups() for how backups are produced and where they are stored.
+//
+// It is a no-op (returning nil) if the database is in-memory, if it already has data,
+// if no bucket is configured on the storage client, or if the backup directory doesn't contain a snapshot.
+// Any other failure is returned as an error, which fails runtime startup.
+// Starting with an empty database would be worse: the next backup would overwrite the snapshot we failed to restore.
+// If a restore fails persistently and the runtime needs to start anyway, set RILL_RUNTIME_METASTORE_BACKUPS_ENABLE=false.
+// That skips the restore, but it also disables backups, so the existing snapshot is left untouched.
+func restoreBackupIfEmpty(ctx context.Context, st *storage.Client, backupID, dsn string, logger *zap.Logger) error {
+	ctx, cancel := context.WithTimeout(ctx, restoreMaxDuration)
+	defer cancel()
+
+	// Check if the database is a restore candidate before touching object storage.
+	// Opening a bucket resolves cloud credentials, which we shouldn't require on a startup that doesn't need it.
+	dbPath, ok, err := shouldRestoreBackup(ctx, dsn)
+	if err != nil {
+		return err
+	}
+	if !ok {
+		return nil
+	}
+
+	// Open bucket scoped to the backup directory.
+	// Return early (no-op) if a bucket isn't available.
+	bucket, ok, err := st.OpenBucket(ctx, backupID)
+	if err != nil {
+		return fmt.Errorf("could not open backup bucket: %w", err)
+	}
+	if !ok {
+		return nil
+	}
+	defer bucket.Close()
+
+	return restoreBackup(ctx, bucket, dbPath, logger.With(zap.String("backup_id", backupID)))
+}
+
+// shouldRestoreBackup reports whether the SQLite database at the given DSN is empty and should be restored from a backup,
+// along with the path of the database file to restore into.
+//
+// "Empty" means migrations have never run on the database.
+// That is a wider condition than the database file not existing, and deliberately so:
+// it also covers a zero-length file and a previous startup that created the file and then crashed before migrating.
+// Without those cases, a single failed startup would permanently disable restores,
+// and the next backup would overwrite the good snapshot with an empty database.
+//
+// It closes its connection before returning, which is what makes it safe for the caller to replace the database file:
+// an open SQLite connection keeps referencing the file it originally opened.
+func shouldRestoreBackup(ctx context.Context, dsn string) (dbPath string, ok bool, err error) {
+	db, err := sqlx.Open("sqlite", dsn)
+	if err != nil {
+		return "", false, fmt.Errorf("failed to open database: %w", err)
+	}
+	db.SetMaxOpenConns(1)
+	defer db.Close()
+
+	// Ask SQLite for the file path instead of parsing the DSN, which can take several forms.
+	// An empty path means the database is in-memory and can't be restored into.
+	err = db.QueryRowContext(ctx, `SELECT file FROM pragma_database_list WHERE name = 'main';`).Scan(&dbPath)
+	if err != nil {
+		return "", false, fmt.Errorf("failed to find database file path: %w", err)
+	}
+	if dbPath == "" || dbPath == ":memory:" || dbPath == "file::memory:" {
+		return "", false, nil
+	}
+
+	// The migration version table doesn't exist on a database that has never been migrated.
+	// We check sqlite_master instead of matching on the "no such table" error string.
+	var tables int
+	err = db.QueryRowContext(ctx, `SELECT count(*) FROM sqlite_master WHERE type = 'table' AND name = ?`, migrationVersionTable).Scan(&tables)
+	if err != nil {
+		return "", false, fmt.Errorf("failed to check for migration version table: %w", err)
+	}
+	if tables != 0 {
+		// The table is created with version 0 before the first migration is applied, so it may exist on an empty database.
+		// It may also exist with no rows at all, since Migrate() creates the table and inserts the row as two separate statements.
+		var version int
+		err = db.QueryRowContext(ctx, fmt.Sprintf("SELECT version FROM %s", migrationVersionTable)).Scan(&version)
+		if err != nil && !errors.Is(err, sql.ErrNoRows) {
+			return "", false, fmt.Errorf("failed to read migration version: %w", err)
+		}
+		if version > 0 {
+			return "", false, nil
+		}
+	}
+
+	return dbPath, true, nil
+}
+
+// restoreBackup downloads the snapshot from the backup bucket and moves it into place at dbPath.
+// It assumes the bucket is already scoped to the correct backup directory, that dbPath is not currently open,
+// and that the directory containing dbPath exists (it does for any dbPath that shouldRestoreBackup has connected to).
+// It is a no-op (returning nil) if the backup directory doesn't contain a snapshot,
+// which is the normal case for a new deployment.
+func restoreBackup(ctx context.Context, bucket *blob.Bucket, dbPath string, logger *zap.Logger) error {
+	attrs, err := bucket.Attributes(ctx, backupSnapshotName)
+	if err != nil {
+		if gcerrors.Code(err) == gcerrors.NotFound {
+			logger.Info("sqlite: no backup found, starting with an empty database")
+			return nil
+		}
+		return fmt.Errorf("failed to check for backup snapshot: %w", err)
+	}
+
+	// Download the snapshot to a temporary file in the same directory as the database, so the rename below is atomic.
+	// The name is deterministic and the file is truncated on create, so a crashed restore doesn't leak a file per attempt.
+	tmpPath := dbPath + ".restore"
+	defer os.Remove(tmpPath)
+	f, err := os.Create(tmpPath)
+	if err != nil {
+		return fmt.Errorf("failed to create temporary file for snapshot: %w", err)
+	}
+	defer f.Close() // Best-effort; we also close it explicitly below.
+	err = bucket.Download(ctx, backupSnapshotName, f, nil)
+	if err != nil {
+		return fmt.Errorf("failed to download snapshot: %w", err)
+	}
+	err = f.Close()
+	if err != nil {
+		return fmt.Errorf("failed to write snapshot: %w", err)
+	}
+
+	// Verify the download is a well-formed SQLite database containing a metastore schema.
+	// This catches a truncated or corrupted download before we overwrite anything.
+	// We deliberately don't run PRAGMA integrity_check or quick_check: snapshots can be several GB,
+	// and the underlying object storage already verifies transfer integrity.
+	snapshotDB, err := sqlx.Open("sqlite", tmpPath)
+	if err != nil {
+		return fmt.Errorf("failed to open downloaded snapshot: %w", err)
+	}
+	snapshotDB.SetMaxOpenConns(1)
+	defer snapshotDB.Close() // Idempotent; we also close it explicitly below.
+	var snapshotVersion int
+	err = snapshotDB.QueryRowContext(ctx, fmt.Sprintf("SELECT version FROM %s", migrationVersionTable)).Scan(&snapshotVersion)
+	if err != nil {
+		return fmt.Errorf("downloaded snapshot is not a valid metastore database: %w", err)
+	}
+	err = snapshotDB.Close()
+	if err != nil {
+		return fmt.Errorf("failed to close downloaded snapshot: %w", err)
+	}
+
+	// Remove journal files left behind by the database we're replacing.
+	// They describe a different database and would corrupt the restored one.
+	for _, suffix := range []string{"-wal", "-shm", "-journal"} {
+		err := os.Remove(dbPath + suffix)
+		if err != nil && !errors.Is(err, fs.ErrNotExist) {
+			return fmt.Errorf("failed to remove stale journal file: %w", err)
+		}
+	}
+
+	err = os.Rename(tmpPath, dbPath)
+	if err != nil {
+		return fmt.Errorf("failed to move snapshot into place: %w", err)
+	}
+
+	// Log at warn level: this only happens after data loss, and the snapshot may be up to a day old.
+	logger.Warn("sqlite: restored database from backup",
+		zap.Time("snapshot_time", attrs.ModTime),
+		zap.Int64("snapshot_size_bytes", attrs.Size),
+		zap.Int("snapshot_migration_version", snapshotVersion),
+	)
+	return nil
+}