feat(chmigrate): add WithReplicated option

ch/query_select.go

@@ -52,7 +52,7 @@ func (q *SelectQuery) Err(err error) *SelectQuery {
 	return q
 }
 
-func (q *SelectQuery) Apply(fn func(*SelectQuery) *SelectQuery) *SelectQuery {
+func (q *SelectQuery) WithQuery(fn func(*SelectQuery) *SelectQuery) *SelectQuery {
 	return fn(q)
 }
 

ch/query_table_create.go

@@ -33,6 +33,10 @@ func (q *CreateTableQuery) Model(model any) *CreateTableQuery {
 	return q
 }
 
+func (q *CreateTableQuery) WithQuery(fn func(*CreateTableQuery) *CreateTableQuery) *CreateTableQuery {
+	return fn(q)
+}
+
 // ------------------------------------------------------------------------------
 
 func (q *CreateTableQuery) Table(tables ...string) *CreateTableQuery {

chmigrate/migration.go

@@ -5,6 +5,7 @@ import (
 	"bytes"
 	"context"
 	"fmt"
+	"io"
 	"io/fs"
 	"sort"
 	"strings"
@@ -14,8 +15,6 @@ import (
 )
 
 type Migration struct {
-	ch.CHModel `ch:"engine:CollapsingMergeTree(sign)"`
-
 	Name       string `ch:",pk"`
 	Comment    string `ch:"-"`
 	GroupID    int64
@@ -42,7 +41,11 @@ func NewSQLMigrationFunc(fsys fs.FS, name string) MigrationFunc {
 		if err != nil {
 			return err
 		}
+		return Exec(ctx, db, f)
+	}
+}
 
+func Exec(ctx context.Context, db *ch.DB, f io.Reader) error {
 	scanner := bufio.NewScanner(f)
 	var queries []string
 
@@ -73,15 +76,13 @@ func NewSQLMigrationFunc(fsys fs.FS, name string) MigrationFunc {
 	}
 
 	for _, q := range queries {
-			_, err = db.ExecContext(ctx, q)
+		if _, err := db.ExecContext(ctx, q); err != nil {
-			if err != nil {
 			return err
 		}
 	}
 
 	return nil
 }
-}
 
 const goTemplate = `package %s
 

chmigrate/migrator.go

@@ -27,6 +27,12 @@ func WithLocksTableName(table string) MigratorOption {
 	}
 }
 
+func WithReplicated(on bool) MigratorOption {
+	return func(m *Migrator) {
+		m.replicated = on
+	}
+}
+
 // WithMarkAppliedOnSuccess sets the migrator to only mark migrations as applied/unapplied
 // when their up/down is successful
 func WithMarkAppliedOnSuccess(enabled bool) MigratorOption {
@@ -43,6 +49,7 @@ type Migrator struct {
 
 	table                string
 	locksTable           string
+	replicated           bool
 	markAppliedOnSuccess bool
 }
 
@@ -95,6 +102,12 @@ func (m *Migrator) migrationsWithStatus(ctx context.Context) (MigrationSlice, in
 func (m *Migrator) Init(ctx context.Context) error {
 	if _, err := m.db.NewCreateTable().
 		Model((*Migration)(nil)).
+		WithQuery(func(q *ch.CreateTableQuery) *ch.CreateTableQuery {
+			if m.replicated {
+				return q.Engine("ReplicatedCollapsingMergeTree(sign)")
+			}
+			return q.Engine("CollapsingMergeTree(sign)")
+		}).
 		ModelTableExpr(m.table).
 		IfNotExists().
 		Exec(ctx); err != nil {
@@ -102,6 +115,12 @@ func (m *Migrator) Init(ctx context.Context) error {
 	}
 	if _, err := m.db.NewCreateTable().
 		Model((*migrationLock)(nil)).
+		WithQuery(func(q *ch.CreateTableQuery) *ch.CreateTableQuery {
+			if m.replicated {
+				return q.Engine("ReplicatedMergeTree")
+			}
+			return q.Engine("MergeTree")
+		}).
 		ModelTableExpr(m.locksTable).
 		IfNotExists().
 		Exec(ctx); err != nil {