Use a transaction to run data migrations

internal/server/data/data.go

@@ -37,6 +37,10 @@ func NewDB(connection gorm.Dialector, dbOpts NewDBOptions) (*DB, error) {
 		return nil, fmt.Errorf("db conn: %w", err)
 	}
 	dataDB := &DB{DB: db}
+	tx, err := dataDB.Begin(context.TODO())
+	if err != nil {
+		return nil, err
+	}
 
 	opts := migrator.Options{
 		InitSchema: initializeSchema,
@@ -47,10 +51,17 @@ func NewDB(connection gorm.Dialector, dbOpts NewDBOptions) (*DB, error) {
 			return loadDBKey(tx, dbOpts.EncryptionKeyProvider, dbOpts.RootKeyID)
 		},
 	}
-	m := migrator.New(dataDB, opts, migrations())
+	m := migrator.New(tx, opts, migrations())
 	if err := m.Migrate(); err != nil {
+		if err := tx.Rollback(); err != nil {
+			logging.L.Warn().Err(err).Msg("failed to rollback")
+		}
 		return nil, fmt.Errorf("migration failed: %w", err)
 	}
+	if err := tx.Commit(); err != nil {
+		return nil, fmt.Errorf("commit migrations: %w", err)
+	}
+
 	if err := initialize(dataDB); err != nil {
 		return nil, fmt.Errorf("initialize database: %w", err)
 	}

internal/server/data/migrations.go

@@ -507,6 +507,15 @@ func addDefaultOrganization() *migrator.Migration {
 	return &migrator.Migration{
 		ID: "2022-08-10T13:35",
 		Migrate: func(tx migrator.DB) error {
+			row := tx.QueryRow(`SELECT count(id) from organizations where name='Default'`)
+			var count int
+			if err := row.Scan(&count); err != nil {
+				return err
+			}
+			if count > 0 {
+				return nil
+			}
+
 			stmt := `
 INSERT INTO organizations(id, name, created_at, updated_at)
 VALUES (?, ?, ?, ?);

internal/server/data/migrations_test.go

@@ -47,8 +47,9 @@ func TestMigrations(t *testing.T) {
 			t.Fatalf("there are more test cases than migrations")
 		}
 		mgs := allMigrations[:index+1]
+		currentMigration := mgs[len(mgs)-1]
 
-		if mID := mgs[len(mgs)-1].ID; mID != tc.label.Name {
+		if mID := currentMigration.ID; mID != tc.label.Name {
 			t.Error("the list of test cases is not in the same order as the list of migrations")
 			t.Fatalf("test case %v was run with migration ID %v", tc.label.Name, mID)
 		}
@@ -79,6 +80,10 @@ func TestMigrations(t *testing.T) {
 		err := m.Migrate()
 		assert.NilError(t, err)
 
+		t.Run("run again to check idempotency", func(t *testing.T) {
+			err := currentMigration.Migrate(db)
+			assert.NilError(t, err)
+		})
 		tc.expected(t, db)
 	}
 

internal/server/data/migrator/helpers.go

@@ -1,6 +1,10 @@
 package migrator
 
-import "github.com/infrahq/infra/internal/logging"
+import (
+	"strings"
+
+	"github.com/infrahq/infra/internal/logging"
+)
 
 // HasTable returns true if the database has a table with name. Returns
 // false if the table does not exist, or if there was a failure querying the
@@ -81,3 +85,22 @@ func HasConstraint(tx DB, table string, constraint string) bool {
 	}
 	return count != 0
 }
+
+// HasFunction returns true if the database already has the function. Returns
+// false if the database table does not have the function, or if there was a
+// failure querying the database.
+func HasFunction(tx DB, funcName string) bool {
+	stmt := `
+		SELECT count(*)
+		FROM pg_proc
+		INNER JOIN pg_namespace ON pg_proc.pronamespace = pg_namespace.oid
+		WHERE proname = ? AND nspname = CURRENT_SCHEMA()`
+
+	var count int
+	// function names are stored in lowercase, so convert to lowercase
+	if err := tx.QueryRow(stmt, strings.ToLower(funcName)).Scan(&count); err != nil {
+		logging.L.Warn().Err(err).Msg("failed to check if function exists")
+		return false
+	}
+	return count != 0
+}

internal/server/data/sqlfunc.go

@@ -8,6 +8,10 @@ func sqlFunctionsMigration() *migrator.Migration {
 	return &migrator.Migration{
 		ID: "2022-08-22T14:58:00Z",
 		Migrate: func(tx migrator.DB) error {
+			if migrator.HasFunction(tx, "uidIntToStr") {
+				return nil
+			}
+
 			sql := `
 			CREATE FUNCTION uidIntToStr(id BIGINT) RETURNS text
 			LANGUAGE PLPGSQL
@@ -81,7 +85,6 @@ func sqlFunctionsMigration() *migrator.Migration {
 			`
 			_, err := tx.Exec(sql)
 			return err
-
 		},
 	}
 }