Repository: G-Research/fasttrackml
Pull request #485 (2023-10-18 12:45:30)

Implement (again) database reset for SQLite
The infamous database refactor removed the ability to reset the SQLite database. This PR restores this ability, whilst also finalizing the intended refactoring (based on TODOs). The `--database-pool-max` is also now used in the SQLite database provider.

The SQLite reset is not the most efficient now that it has to be done _after_ opening the connection, as we cannot simply delete the files anymore. We could potentially change that by having `database.Reset(dsn string)` be responsible for calling the provider's implementation, and letting the provider do its logic. Thoughts?


Patch:
diff --git a/pkg/cmd/import.go b/pkg/cmd/import.go
index 0049e55..7be865c 100644
--- a/pkg/cmd/import.go
+++ b/pkg/cmd/import.go
@@ -43,7 +43,6 @@ func initDBs() (input, output database.DBProvider, err error) {
 		viper.GetString("input-database-uri"),
 		time.Second*1,
 		20,
-		false,
 	)
 	if err != nil {
 		return input, output, fmt.Errorf("error connecting to input DB: %w", err)
@@ -53,7 +52,6 @@ func initDBs() (input, output database.DBProvider, err error) {
 		viper.GetString("output-database-uri"),
 		time.Second*1,
 		20,
-		false,
 	)
 	if err != nil {
 		return input, output, fmt.Errorf("error connecting to output DB: %w", err)
diff --git a/pkg/cmd/server.go b/pkg/cmd/server.go
index e159e60..1214a34 100644
--- a/pkg/cmd/server.go
+++ b/pkg/cmd/server.go
@@ -151,12 +151,17 @@ func initDB(config *mlflowConfig.ServiceConfig) (database.DBProvider, error) {
 		config.DatabaseURI,
 		config.DatabaseSlowThreshold,
 		config.DatabasePoolMax,
-		config.DatabaseReset,
 	)
 	if err != nil {
 		return nil, fmt.Errorf("error connecting to DB: %w", err)
 	}
 
+	if config.DatabaseReset {
+		if err := db.Reset(); err != nil {
+			return nil, eris.Wrap(err, "error resetting database")
+		}
+	}
+
 	if err := database.CheckAndMigrateDB(config.DatabaseMigrate, db.GormDB()); err != nil {
 		return nil, eris.Wrap(err, "error running database migration")
 	}
diff --git a/pkg/database/db_factory.go b/pkg/database/db_factory.go
index 08d8bfe..3a11e1c 100644
--- a/pkg/database/db_factory.go
+++ b/pkg/database/db_factory.go
@@ -5,12 +5,11 @@ import (
 	"time"
 
 	"github.com/rotisserie/eris"
-	log "github.com/sirupsen/logrus"
 )
 
 // NewDBProvider creates a DBProvider of the correct type from the parameters.
 func NewDBProvider(
-	dsn string, slowThreshold time.Duration, poolMax int, reset bool,
+	dsn string, slowThreshold time.Duration, poolMax int,
 ) (db DBProvider, err error) {
 	dsnURL, err := url.Parse(dsn)
 	if err != nil {
@@ -22,7 +21,6 @@ func NewDBProvider(
 			*dsnURL,
 			slowThreshold,
 			poolMax,
-			reset,
 		)
 		if err != nil {
 			return nil, eris.Wrap(err, "error creating sqlite provider")
@@ -32,7 +30,6 @@ func NewDBProvider(
 			*dsnURL,
 			slowThreshold,
 			poolMax,
-			reset,
 		)
 		if err != nil {
 			return nil, eris.Wrap(err, "error creating postgres provider")
@@ -41,15 +38,5 @@ func NewDBProvider(
 		return nil, eris.New("unsupported database type")
 	}
 
-	// TODO:DSuhinin - it shouldn't be there. NewDBProvider has to only create an instance without any hidden logic.
-	if reset {
-		log.Infof("reseting database")
-		if err := db.Reset(); err != nil {
-			//nolint:errcheck,gosec
-			db.Close()
-			return nil, eris.Wrap(err, "error resetting database")
-		}
-	}
-
 	return db, nil
 }
diff --git a/pkg/database/postgres_db_instance.go b/pkg/database/postgres_db_instance.go
index fd41dc5..06da994 100644
--- a/pkg/database/postgres_db_instance.go
+++ b/pkg/database/postgres_db_instance.go
@@ -18,7 +18,7 @@ type PostgresDBInstance struct {
 
 // NewPostgresDBInstance constructs a Postgres DbInstance.
 func NewPostgresDBInstance(
-	dsnURL url.URL, slowThreshold time.Duration, poolMax int, reset bool,
+	dsnURL url.URL, slowThreshold time.Duration, poolMax int,
 ) (*PostgresDBInstance, error) {
 	db := PostgresDBInstance{
 		DBInstance: DBInstance{dsn: dsnURL.String()},
diff --git a/pkg/database/sqlite_db_instance.go b/pkg/database/sqlite_db_instance.go
index 738c1f3..cb06e32 100644
--- a/pkg/database/sqlite_db_instance.go
+++ b/pkg/database/sqlite_db_instance.go
@@ -2,9 +2,8 @@ package database
 
 import (
 	"database/sql"
-	"errors"
+	"fmt"
 	"net/url"
-	"os"
 	"regexp"
 	"slices"
 	"strings"
@@ -30,7 +29,7 @@ type SqliteDBInstance struct {
 
 // NewSqliteDBInstance creates a SqliteDBInstance.
 func NewSqliteDBInstance(
-	dsnURL url.URL, slowThreshold time.Duration, poolMax int, reset bool,
+	dsnURL url.URL, slowThreshold time.Duration, poolMax int,
 ) (*SqliteDBInstance, error) {
 	db := SqliteDBInstance{
 		DBInstance: DBInstance{dsn: dsnURL.String()},
@@ -47,17 +46,6 @@ func NewSqliteDBInstance(
 	sourceURL := dsnURL
 	sourceURL.RawQuery = query.Encode()
 
-	if reset && query.Get("mode") != "memory" {
-		file := sourceURL.Host
-		if file == "" {
-			file = sourceURL.Path
-		}
-		log.Infof("Removing database file %s", file)
-		if err := os.Remove(file); err != nil && !errors.Is(err, os.ErrNotExist) {
-			return nil, eris.Wrap(err, "failed to remove database file")
-		}
-	}
-
 	if !slices.Contains(sql.Drivers(), SQLiteCustomDriverName) {
 		sql.Register(SQLiteCustomDriverName, &sqlite3.SQLiteDriver{
 			ConnectHook: func(conn *sqlite3.SQLiteConn) error {
@@ -104,6 +92,7 @@ func NewSqliteDBInstance(
 		return nil, eris.Wrap(err, "failed to connect to database")
 	}
 	db.closers = append(db.closers, replicaDB)
+	replicaDB.SetMaxOpenConns(poolMax)
 	replicaConn = sqlite.Dialector{
 		Conn: replicaDB,
 	}
@@ -142,6 +131,40 @@ func NewSqliteDBInstance(
 }
 
 // Reset resets database.
-func (f SqliteDBInstance) Reset() error {
-	return eris.New("reset for sqlite database not supported")
+func (db SqliteDBInstance) Reset() error {
+	log.Info("Resetting database schema")
+
+	var tables []string
+	if err := db.
+		Table("sqlite_schema").
+		Select("name").
+		Where("type = ?", "table").
+		Find(&tables).
+		Error; err != nil {
+		return eris.Wrap(err, "error getting tables")
+	}
+
+	if err := db.Transaction(func(tx *gorm.DB) error {
+		tx.Exec("PRAGMA foreign_keys = OFF")
+		defer tx.Exec("PRAGMA foreign_keys = ON")
+
+		for _, table := range tables {
+			if err := tx.Exec(fmt.Sprintf("DROP TABLE %s", table)).Error; err != nil {
+				return eris.Wrapf(err, "error dropping table %q", table)
+			}
+		}
+
+		return nil
+	}); err != nil {
+		return eris.Wrap(err, "error dropping tables")
+	}
+
+	if err := db.Exec("VACUUM").Error; err != nil {
+		return eris.Wrap(err, "error vacuuming database")
+	}
+	if err := db.Exec("PRAGMA wal_checkpoint(TRUNCATE)").Error; err != nil {
+		return eris.Wrap(err, "error checkpointing database")
+	}
+
+	return nil
 }


Test patch:
diff --git a/pkg/database/db_factory_test.go b/pkg/database/db_factory_test.go
index 05b526a..51cb27c 100644
--- a/pkg/database/db_factory_test.go
+++ b/pkg/database/db_factory_test.go
@@ -27,7 +27,6 @@ func TestMakeDBProvider(t *testing.T) {
 				tt.dsn,
 				time.Second*2,
 				2,
-				false,
 			)
 			assert.Nil(t, err)
 			assert.NotNil(t, db)
diff --git a/tests/integration/golang/database/import_test.go b/tests/integration/golang/database/import_test.go
index 8c44ac2..56224da 100644
--- a/tests/integration/golang/database/import_test.go
+++ b/tests/integration/golang/database/import_test.go
@@ -51,7 +51,6 @@ func (s *ImportTestSuite) SetupTest() {
 		helpers.GetInputDatabaseUri(),
 		1*time.Second,
 		20,
-		false,
 	)
 	assert.Nil(s.T(), err)
 	assert.Nil(s.T(), database.CheckAndMigrateDB(true, db.GormDB()))
@@ -119,7 +118,6 @@ func (s *ImportTestSuite) SetupTest() {
 		helpers.GetOutputDatabaseUri(),
 		1*time.Second,
 		20,
-		false,
 	)
 	assert.Nil(s.T(), err)
 	assert.Nil(s.T(), database.CheckAndMigrateDB(true, db.GormDB()))
diff --git a/tests/integration/golang/helpers/test.go b/tests/integration/golang/helpers/test.go
index d7dda5b..258f868 100644
--- a/tests/integration/golang/helpers/test.go
+++ b/tests/integration/golang/helpers/test.go
@@ -34,7 +34,6 @@ func (s *BaseTestSuite) SetupTest(t *testing.T) {
 			GetDatabaseUri(),
 			1*time.Second,
 			20,
-			false,
 		)
 		assert.Nil(t, err)
 		db = instance.GormDB()