feat: support multiple queries for clickhouse using up/down-sections

pkg/dbmate/db.go

@@ -385,12 +385,17 @@ func (db *DB) Migrate() error {
 		}
 
 		execMigration := func(tx dbutil.Transaction) error {
-			// run actual migration
-			result, err := tx.Exec(parsed.Up)
-			if err != nil {
-				return drv.QueryError(parsed.Up, err)
-			} else if db.Verbose {
-				db.printVerbose(result)
+			if len(parsed.UpSections) != 0 {
+				for _, section := range parsed.UpSections {
+					if err = db.execQuery(drv, tx, section); err != nil {
+						return err
+					}
+				}
+			} else {
+				// run actual migration
+				if err = db.execQuery(drv, tx, parsed.Up); err != nil {
+					return err
+				}
 			}
 
 			// record migration
@@ -554,12 +559,17 @@ func (db *DB) Rollback() error {
 	}
 
 	execMigration := func(tx dbutil.Transaction) error {
-		// rollback migration
-		result, err := tx.Exec(parsed.Down)
-		if err != nil {
-			return drv.QueryError(parsed.Down, err)
-		} else if db.Verbose {
-			db.printVerbose(result)
+		if len(parsed.DownSections) != 0 {
+			for _, section := range parsed.DownSections {
+				if err = db.execQuery(drv, tx, section); err != nil {
+					return err
+				}
+			}
+		} else {
+			// rollback migration
+			if err = db.execQuery(drv, tx, parsed.Down); err != nil {
+				return err
+			}
 		}
 
 		// remove migration record
@@ -589,6 +599,15 @@ func (db *DB) Rollback() error {
 	return nil
 }
 
+func (db *DB) execQuery(drv Driver, tx dbutil.Transaction, query string) error {
+	if result, err := tx.Exec(query); err != nil {
+		return drv.QueryError(query, err)
+	} else if db.Verbose {
+		db.printVerbose(result)
+	}
+	return nil
+}
+
 // Status shows the status of all migrations
 func (db *DB) Status(quiet bool) (int, error) {
 	results, err := db.FindMigrations()

pkg/dbmate/migration.go

@@ -38,11 +38,17 @@ func (m *Migration) Parse() (*ParsedMigration, error) {
 }
 
 // ParsedMigration contains the migration contents and options
+// It should be only one query pro up- or down-section provided,
+// but it might have multiple sections.
+// If the `UpSections` or `DownSections` are not empty,
+// they are applied instead of `Up` or `Down`.
 type ParsedMigration struct {
-	Up          string
-	UpOptions   ParsedMigrationOptions
-	Down        string
-	DownOptions ParsedMigrationOptions
+	Up           string
+	UpSections   []string
+	UpOptions    ParsedMigrationOptions
+	Down         string
+	DownSections []string
+	DownOptions  ParsedMigrationOptions
 }
 
 // ParsedMigrationOptions is an interface for accessing migration options
@@ -60,7 +66,9 @@ func (m migrationOptions) Transaction() bool {
 
 var (
 	upRegExp              = regexp.MustCompile(`(?m)^--\s*migrate:up(\s*$|\s+\S+)`)
+	upSectionRegExp       = regexp.MustCompile(`(?m)^--\s*migrate:up-section(\s*$|\s+\S+)`)
 	downRegExp            = regexp.MustCompile(`(?m)^--\s*migrate:down(\s*$|\s+\S+)`)
+	downSectionRegExp     = regexp.MustCompile(`(?m)^--\s*migrate:down-section(\s*$|\s+\S+)`)
 	emptyLineRegExp       = regexp.MustCompile(`^\s*$`)
 	commentLineRegExp     = regexp.MustCompile(`^\s*--`)
 	whitespaceRegExp      = regexp.MustCompile(`\s+`)
@@ -80,7 +88,9 @@ var (
 // It will return two Migration objects, the first representing the "up"
 // block and the second representing the "down" block. This function
 // requires that at least an up block was defined and will otherwise
-// return an error.
+// return an error. It returns also UpSections or DownSections (default: empty arrays).
+// If UpSections or DownSections are not empty, they will be executed instead of
+// "up" or "down" blocks.
 func parseMigrationContents(contents string) (*ParsedMigration, error) {
 	upDirectiveStart, hasDefinedUpBlock := getMatchPosition(contents, upRegExp)
 	downDirectiveStart, hasDefinedDownBlock := getMatchPosition(contents, downRegExp)
@@ -101,15 +111,40 @@ func parseMigrationContents(contents string) (*ParsedMigration, error) {
 	upBlock := substring(contents, upDirectiveStart, downDirectiveStart)
 	downBlock := substring(contents, downDirectiveStart, len(contents))
 
+	upSections := parseMigrationSections(upBlock, downDirectiveStart, upSectionRegExp)
+	downSections := parseMigrationSections(downBlock, len(downBlock), downSectionRegExp)
+
 	parsed := ParsedMigration{
-		Up:          upBlock,
-		UpOptions:   parseMigrationOptions(upBlock),
-		Down:        downBlock,
-		DownOptions: parseMigrationOptions(downBlock),
+		Up:           upBlock,
+		UpSections:   upSections,
+		UpOptions:    parseMigrationOptions(upBlock),
+		Down:         downBlock,
+		DownSections: downSections,
+		DownOptions:  parseMigrationOptions(downBlock),
 	}
 	return &parsed, nil
 }
 
+func parseMigrationSections(block string, blockDirectiveEnds int, sectionRegExp *regexp.Regexp) []string {
+	sectionDirectiveStart, hasDefinedSectionBlocks := getMatchPosition(block, sectionRegExp)
+	if !hasDefinedSectionBlocks {
+		return []string{}
+	}
+
+	sectionBlocks := substring(block, sectionDirectiveStart, blockDirectiveEnds)
+
+	sectionIdx := sectionRegExp.FindStringIndex(sectionBlocks)[1]
+	sections := sectionRegExp.Split(sectionBlocks[sectionIdx:], -1)
+	cleanSections := make([]string, 0)
+	for _, section := range sections {
+		section = strings.ReplaceAll(strings.TrimSpace(section), "\n", "")
+		if section != "" {
+			cleanSections = append(cleanSections, section)
+		}
+	}
+	return cleanSections
+}
+
 // parseMigrationOptions parses the migration options out of a block
 // directive into an object that implements the MigrationOptions interface.
 //

pkg/dbmate/migration_test.go

@@ -199,4 +199,40 @@ DROP COLUMN status;
 			require.Equal(t, true, parsed.DownOptions.Transaction())
 		})
 	})
+
+	t.Run("support up and down sections", func(t *testing.T) {
+		migration := `-- migrate:up
+-- migrate:up-section
+create table users (id serial, name text);
+-- migrate:up-section
+select * from users;
+-- migrate:down
+-- migrate:down-section
+drop table users;`
+
+		parsed, err := parseMigrationContents(migration)
+		require.Nil(t, err)
+
+		require.Equal(t, len(parsed.UpSections), 2)
+		require.Equal(t, "create table users (id serial, name text);", parsed.UpSections[0])
+		require.Equal(t, "select * from users;", parsed.UpSections[1])
+
+		require.Equal(t, len(parsed.DownSections), 1)
+		require.Equal(t, "drop table users;", parsed.DownSections[0])
+	})
+
+	t.Run("ignore empty up and down sections", func(t *testing.T) {
+		migration := `-- migrate:up
+-- migrate:up-section
+-- migrate:up-section
+-- migrate:down
+-- migrate:down-section
+`
+
+		parsed, err := parseMigrationContents(migration)
+		require.Nil(t, err)
+
+		require.Equal(t, len(parsed.UpSections), 0)
+		require.Equal(t, len(parsed.DownSections), 0)
+	})
 }

pkg/dbmate/version.go

@@ -1,4 +1,4 @@
 package dbmate
 
 // Version of dbmate
-const Version = "2.21.0"
+const Version = "2.22.0"