From d260a308bdf05be80f0fc374ba6c9e33e91d268a Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Sybren=20A=2E=20St=C3=BCvel?= <sybren@blender.org>
Date: Mon, 27 Nov 2023 12:26:18 +0700
Subject: [PATCH] Worker: enable write-ahead logging on the database

Now the Worker and the Manager share the same database initialisation
code (enabling foreign key constraints + write-ahead logging).

The foreign key constraints were already enabled before, but now it's done
with (a copy of) the same code as the Manager.
---
 internal/worker/persistence/db.go | 87 ++++++++++++++++++++++++++-----
 1 file changed, 74 insertions(+), 13 deletions(-)

diff --git a/internal/worker/persistence/db.go b/internal/worker/persistence/db.go
index ecf8838c..07519e89 100644
--- a/internal/worker/persistence/db.go
+++ b/internal/worker/persistence/db.go
@@ -71,34 +71,52 @@ func openDBWithConfig(dsn string, config *gorm.Config) (*DB, error) {
 		return nil, err
 	}
 
+	db := &DB{
+		gormDB: gormDB,
+	}
+
+	// Close the database connection if there was some error. This prevents
+	// leaking database connections & should remove any write-ahead-log files.
+	closeConnOnReturn := true
+	defer func() {
+		if !closeConnOnReturn {
+			return
+		}
+		if err := db.Close(); err != nil {
+			log.Debug().AnErr("cause", err).Msg("cannot close database connection")
+		}
+	}()
+
 	// Use the generic sql.DB interface to set some connection pool options.
 	sqlDB, err := gormDB.DB()
 	if err != nil {
 		return nil, err
 	}
+
 	// Only allow a single database connection, to avoid SQLITE_BUSY errors.
 	// It's not certain that this'll improve the situation, but it's worth a try.
 	sqlDB.SetMaxIdleConns(1) // Max num of connections in the idle connection pool.
 	sqlDB.SetMaxOpenConns(1) // Max num of open connections to the database.
 
-	// Enable foreign key checks.
+	// Always enable foreign key checks, to make SQLite behave like a real database.
 	log.Trace().Msg("enabling SQLite foreign key checks")
-	if tx := gormDB.Exec("PRAGMA foreign_keys = 1"); tx.Error != nil {
-		return nil, fmt.Errorf("enabling foreign keys: %w", tx.Error)
-	}
-	var fkEnabled int
-	if tx := gormDB.Raw("PRAGMA foreign_keys").Scan(&fkEnabled); tx.Error != nil {
-		return nil, fmt.Errorf("checking whether the database has foreign key checks enabled: %w", tx.Error)
-	}
-	if fkEnabled == 0 {
-		log.Error().Msg("SQLite database does not want to enable foreign keys, this may cause data loss")
+	if err := db.pragmaForeignKeys(true); err != nil {
+		return nil, err
 	}
 
-	db := DB{
-		gormDB: gormDB,
+	// Write-ahead-log journal may improve writing speed.
+	log.Trace().Msg("enabling SQLite write-ahead-log journal mode")
+	if tx := gormDB.Exec("PRAGMA journal_mode = WAL"); tx.Error != nil {
+		return nil, fmt.Errorf("enabling SQLite write-ahead-log journal mode: %w", tx.Error)
+	}
+	// Switching from 'full' (default) to 'normal' sync may improve writing speed.
+	log.Trace().Msg("enabling SQLite 'normal' synchronisation")
+	if tx := gormDB.Exec("PRAGMA synchronous = normal"); tx.Error != nil {
+		return nil, fmt.Errorf("enabling SQLite 'normal' sync mode: %w", tx.Error)
 	}
 
-	return &db, nil
+	closeConnOnReturn = false
+	return db, nil
 }
 
 // nowFunc returns 'now' in UTC, so that GORM-managed times (createdAt,
@@ -126,3 +144,46 @@ func (db *DB) Close() error {
 	}
 	return nil
 }
+
+func (db *DB) pragmaForeignKeys(enabled bool) error {
+	var (
+		value int
+		noun  string
+	)
+	switch enabled {
+	case false:
+		value = 0
+		noun = "disabl"
+	case true:
+		value = 1
+		noun = "enabl"
+	}
+
+	log.Trace().Msgf("%sing SQLite foreign key checks", noun)
+
+	// SQLite doesn't seem to like SQL parameters for `PRAGMA`, so `PRAGMA foreign_keys = ?` doesn't work.
+	sql := fmt.Sprintf("PRAGMA foreign_keys = %d", value)
+
+	if tx := db.gormDB.Exec(sql); tx.Error != nil {
+		return fmt.Errorf("%sing foreign keys: %w", noun, tx.Error)
+	}
+	fkEnabled, err := db.areForeignKeysEnabled()
+	if err != nil {
+		return err
+	}
+	if fkEnabled != enabled {
+		return fmt.Errorf("SQLite database does not want to %se foreign keys, this may cause data loss", noun)
+	}
+
+	return nil
+}
+
+func (db *DB) areForeignKeysEnabled() (bool, error) {
+	log.Trace().Msg("checking whether SQLite foreign key checks are enabled")
+
+	var fkEnabled int
+	if tx := db.gormDB.Raw("PRAGMA foreign_keys").Scan(&fkEnabled); tx.Error != nil {
+		return false, fmt.Errorf("checking whether the database has foreign key checks are enabled: %w", tx.Error)
+	}
+	return fkEnabled != 0, nil
+}