From 00c02255cc1c1976a0741d6ca8d5d4d333615032 Mon Sep 17 00:00:00 2001
From: Andrey Pshenkin <andrey.pshenkin@gmail.com>
Date: Tue, 29 Oct 2024 17:07:44 +0000
Subject: [PATCH] Implement checksummed migrations

---
 migrate.go      | 32 +++++++++++++++++++++++++-------
 migrate_test.go | 41 +++++++++++++++++++++++++++++++++++++++++
 2 files changed, 66 insertions(+), 7 deletions(-)

diff --git a/migrate.go b/migrate.go
index 7fb56f1a..fcae9c7c 100644
--- a/migrate.go
+++ b/migrate.go
@@ -3,6 +3,7 @@ package migrate
 import (
 	"bytes"
 	"context"
+	"crypto/sha256"
 	"database/sql"
 	"embed"
 	"errors"
@@ -125,9 +126,10 @@ func SetIgnoreUnknown(v bool) {
 }
 
 type Migration struct {
-	Id   string
-	Up   []string
-	Down []string
+	Id       string
+	Checksum string
+	Up       []string
+	Down     []string
 
 	DisableTransactionUp   bool
 	DisableTransactionDown bool
@@ -178,6 +180,7 @@ func (b byId) Less(i, j int) bool { return b[i].Less(b[j]) }
 
 type MigrationRecord struct {
 	Id        string    `db:"id"`
+	Checksum  string    `db:"checksum"`
 	AppliedAt time.Time `db:"applied_at"`
 }
 
@@ -420,6 +423,13 @@ func ParseMigration(id string, r io.ReadSeeker) (*Migration, error) {
 		Id: id,
 	}
 
+	hash := sha256.New()
+	if _, err := io.Copy(hash, r); err != nil {
+		return nil, fmt.Errorf("Error computing migration checksum (%s): %w", id, err)
+	}
+
+	m.Checksum = fmt.Sprintf("%x", hash.Sum(nil))
+
 	parsed, err := sqlparse.ParseMigration(r)
 	if err != nil {
 		return nil, fmt.Errorf("Error parsing migration (%s): %w", id, err)
@@ -565,6 +575,7 @@ func (MigrationSet) applyMigrations(ctx context.Context, dir MigrationDirection,
 		case Up:
 			err = executor.Insert(&MigrationRecord{
 				Id:        migration.Id,
+				Checksum:  migration.Checksum,
 				AppliedAt: time.Now(),
 			})
 			if err != nil {
@@ -643,7 +654,8 @@ func (ms MigrationSet) planMigrationCommon(db *sql.DB, dialect string, m Migrati
 	var existingMigrations []*Migration
 	for _, migrationRecord := range migrationRecords {
 		existingMigrations = append(existingMigrations, &Migration{
-			Id: migrationRecord.Id,
+			Id:       migrationRecord.Id,
+			Checksum: migrationRecord.Checksum,
 		})
 	}
 	sort.Sort(byId(existingMigrations))
@@ -651,14 +663,19 @@ func (ms MigrationSet) planMigrationCommon(db *sql.DB, dialect string, m Migrati
 	// Make sure all migrations in the database are among the found migrations which
 	// are to be applied.
 	if !ms.IgnoreUnknown {
-		migrationsSearch := make(map[string]struct{})
+		migrationsSearch := make(map[string]string)
 		for _, migration := range migrations {
-			migrationsSearch[migration.Id] = struct{}{}
+			migrationsSearch[migration.Id] = migration.Checksum
 		}
 		for _, existingMigration := range existingMigrations {
-			if _, ok := migrationsSearch[existingMigration.Id]; !ok {
+			plannedMigrationChecksum, ok := migrationsSearch[existingMigration.Id]
+			if !ok {
 				return nil, nil, newPlanError(existingMigration, "unknown migration in database")
 			}
+			// Check if the checksums match if exists, otherwise ignore for backward compatibility
+			if existingMigration.Checksum != "" && existingMigration.Checksum != plannedMigrationChecksum {
+				return nil, nil, newPlanError(existingMigration, "wrong migration checksum")
+			}
 		}
 	}
 
@@ -745,6 +762,7 @@ func SkipMax(db *sql.DB, dialect string, m MigrationSource, dir MigrationDirecti
 
 		err = executor.Insert(&MigrationRecord{
 			Id:        migration.Id,
+			Checksum:  migration.Checksum,
 			AppliedAt: time.Now(),
 		})
 		if err != nil {
diff --git a/migrate_test.go b/migrate_test.go
index f1d66d6f..416d7970 100644
--- a/migrate_test.go
+++ b/migrate_test.go
@@ -536,6 +536,47 @@ func (s *SqliteMigrateSuite) TestPlanMigrationWithUnknownDatabaseMigrationApplie
 	c.Assert(err, FitsTypeOf, &PlanError{})
 }
 
+func (s *SqliteMigrateSuite) TestPlanMigrationWithMigrationWithDifferentCheckSumApplied(c *C) {
+	migrations := &MemoryMigrationSource{
+		Migrations: []*Migration{
+			{
+				Id:       "1_create_table.sql",
+				Checksum: "123",
+				Up:       []string{"CREATE TABLE people (id int)"},
+				Down:     []string{"DROP TABLE people"},
+			},
+			{
+				Id:       "2_alter_table.sql",
+				Checksum: "345",
+				Up:       []string{"ALTER TABLE people ADD COLUMN first_name text"},
+				Down:     []string{"SELECT 0"}, // Not really supported
+			},
+			{
+				Id:       "10_add_last_name.sql",
+				Checksum: "567",
+				Up:       []string{"ALTER TABLE people ADD COLUMN last_name text"},
+				Down:     []string{"ALTER TABLE people DROP COLUMN last_name"},
+			},
+		},
+	}
+	n, err := Exec(s.Db, "sqlite3", migrations, Up)
+	c.Assert(err, IsNil)
+	c.Assert(n, Equals, 3)
+
+	// Change the checksum of the migration so that it doesn't match the one in the database
+	migrations.Migrations[1].Checksum = "222"
+
+	_, _, err = PlanMigration(s.Db, "sqlite3", migrations, Up, 0)
+	c.Assert(err, NotNil, Commentf("Up migrations should not have been applied when there "+
+		"is checksum missmatch"))
+	c.Assert(err, FitsTypeOf, &PlanError{})
+
+	_, _, err = PlanMigration(s.Db, "sqlite3", migrations, Down, 0)
+	c.Assert(err, NotNil, Commentf("Down migrations should not have been applied when there "+
+		"is checksum missmatch"))
+	c.Assert(err, FitsTypeOf, &PlanError{})
+}
+
 func (s *SqliteMigrateSuite) TestPlanMigrationWithIgnoredUnknownDatabaseMigrationApplied(c *C) {
 	migrations := &MemoryMigrationSource{
 		Migrations: []*Migration{