From 801ca16cd7c76094763f255fe75395bb7541a5ed Mon Sep 17 00:00:00 2001
From: RebeccaMahany <rebeccamahany@gmail.com>
Date: Mon, 27 Jan 2025 12:44:28 -0500
Subject: [PATCH] Force migration rather than halting launcher startup

---
 .../storage/sqlite/keyvalue_store_sqlite.go   | 16 ++++++++++-
 .../sqlite/keyvalue_store_sqlite_test.go      | 28 +++++++++++++++++++
 2 files changed, 43 insertions(+), 1 deletion(-)

diff --git a/ee/agent/storage/sqlite/keyvalue_store_sqlite.go b/ee/agent/storage/sqlite/keyvalue_store_sqlite.go
index 61b943384..99d5f6199 100644
--- a/ee/agent/storage/sqlite/keyvalue_store_sqlite.go
+++ b/ee/agent/storage/sqlite/keyvalue_store_sqlite.go
@@ -187,7 +187,21 @@ func (s *sqliteStore) migrate() error {
 
 	// don't prevent DB access for a missing migration, this is the result of a downgrade after previously
 	// running a migration
-	if err := m.Up(); err != nil && !errors.Is(err, migrate.ErrNoChange) && !isMissingMigrationError(err) {
+	if err := m.Up(); err != nil {
+		// Not actually errors for us -- we're in a successful state
+		if errors.Is(err, migrate.ErrNoChange) || isMissingMigrationError(err) {
+			return nil
+		}
+
+		// If we need to force, do that
+		if errDirty, ok := err.(migrate.ErrDirty); ok {
+			if err := m.Force(errDirty.Version); err != nil {
+				return fmt.Errorf("forcing migration version %d: %w", errDirty.Version, err)
+			}
+			return nil
+		}
+
+		// Some other error -- return it
 		return fmt.Errorf("running migrations: %w", err)
 	}
 
diff --git a/ee/agent/storage/sqlite/keyvalue_store_sqlite_test.go b/ee/agent/storage/sqlite/keyvalue_store_sqlite_test.go
index 966206ffd..3c123f6d9 100644
--- a/ee/agent/storage/sqlite/keyvalue_store_sqlite_test.go
+++ b/ee/agent/storage/sqlite/keyvalue_store_sqlite_test.go
@@ -7,6 +7,7 @@ import (
 	"testing"
 
 	"github.com/golang-migrate/migrate/v4"
+	"github.com/golang-migrate/migrate/v4/database/sqlite"
 	"github.com/golang-migrate/migrate/v4/source/iofs"
 	"github.com/kolide/launcher/ee/agent/flags/keys"
 	"github.com/stretchr/testify/require"
@@ -68,6 +69,33 @@ func TestOpenRW_DatabaseIsCorrupt(t *testing.T) {
 	require.NoError(t, s.Close(), "closing test store")
 }
 
+func TestOpenRW_DatabaseIsDirty(t *testing.T) {
+	t.Parallel()
+
+	testRootDir := t.TempDir()
+
+	// Create the database
+	s, err := OpenRW(context.TODO(), testRootDir, StartupSettingsStore)
+	require.NoError(t, err, "expected no error creating test store")
+
+	// Mark the migration as dirty
+	_, err = s.conn.Exec(fmt.Sprintf(`UPDATE %s SET dirty = 1;`, sqlite.DefaultMigrationsTable))
+	require.NoError(t, err, "marking migration as dirty")
+
+	// Close the connection
+	require.NoError(t, s.Close(), "expected no error closing test store")
+
+	// Open a new connection and expect that it succeeds, forcing the dirty migration
+	s2, err := OpenRW(context.TODO(), testRootDir, StartupSettingsStore)
+	require.NoError(t, err, "expected no error opening test store with dirty migration")
+	require.NoError(t, s2.Close(), "expected no error closing test store")
+
+	// Open and close again successfully just to be sure
+	s3, err := OpenRW(context.TODO(), testRootDir, StartupSettingsStore)
+	require.NoError(t, err, "expected no error opening test store with dirty migration")
+	require.NoError(t, s3.Close(), "expected no error closing test store")
+}
+
 func TestOpenRW_InvalidTable(t *testing.T) {
 	t.Parallel()