snapshot of https://github.com/golang-migrate/migrate.git (master, 2026-09-15; blobs > 131072B omitted)

2026-09-15 21:52:57 +05:00
commit 9f3bffc2d3
405 changed files with 27657 additions and 0 deletions
--- a/database/mongodb/mongodb.go
+++ b/database/mongodb/mongodb.go
@@ -0,0 +1,400 @@
+package mongodb
+
+import (
+	"context"
+	"errors"
+	"fmt"
+	"io"
+	"net/url"
+	"os"
+	"strconv"
+	"sync/atomic"
+	"time"
+
+	"github.com/cenkalti/backoff/v4"
+	"github.com/golang-migrate/migrate/v4/database"
+	"go.mongodb.org/mongo-driver/bson"
+	"go.mongodb.org/mongo-driver/mongo"
+	"go.mongodb.org/mongo-driver/mongo/options"
+	"go.mongodb.org/mongo-driver/x/mongo/driver/connstring"
+)
+
+func init() {
+	db := Mongo{}
+	database.Register("mongodb", &db)
+	database.Register("mongodb+srv", &db)
+}
+
+var DefaultMigrationsCollection = "schema_migrations"
+
+const DefaultLockingCollection = "migrate_advisory_lock" // the collection to use for advisory locking by default.
+const lockKeyUniqueValue = 0                             // the unique value to lock on. If multiple clients try to insert the same key, it will fail (locked).
+const DefaultLockTimeout = 15                            // the default maximum time to wait for a lock to be released.
+const DefaultLockTimeoutInterval = 10                    // the default maximum intervals time for the locking timout.
+const DefaultAdvisoryLockingFlag = true                  // the default value for the advisory locking feature flag. Default is true.
+const LockIndexName = "lock_unique_key"                  // the name of the index which adds unique constraint to the locking_key field.
+const contextWaitTimeout = 5 * time.Second               // how long to wait for the request to mongo to block/wait for.
+
+var (
+	ErrNoDatabaseName            = fmt.Errorf("no database name")
+	ErrNilConfig                 = fmt.Errorf("no config")
+	ErrLockTimeoutConfigConflict = fmt.Errorf("both x-advisory-lock-timeout-interval and x-advisory-lock-timout-interval were specified")
+)
+
+type Mongo struct {
+	client   *mongo.Client
+	db       *mongo.Database
+	config   *Config
+	isLocked atomic.Bool
+}
+
+type Locking struct {
+	CollectionName string
+	Timeout        int
+	Enabled        bool
+	Interval       int
+}
+type Config struct {
+	DatabaseName         string
+	MigrationsCollection string
+	TransactionMode      bool
+	Locking              Locking
+}
+type versionInfo struct {
+	Version int  `bson:"version"`
+	Dirty   bool `bson:"dirty"`
+}
+
+type lockObj struct {
+	Key       int       `bson:"locking_key"`
+	Pid       int       `bson:"pid"`
+	Hostname  string    `bson:"hostname"`
+	CreatedAt time.Time `bson:"created_at"`
+}
+type findFilter struct {
+	Key int `bson:"locking_key"`
+}
+
+func WithInstance(instance *mongo.Client, config *Config) (database.Driver, error) {
+	if config == nil {
+		return nil, ErrNilConfig
+	}
+	if len(config.DatabaseName) == 0 {
+		return nil, ErrNoDatabaseName
+	}
+	if len(config.MigrationsCollection) == 0 {
+		config.MigrationsCollection = DefaultMigrationsCollection
+	}
+	if len(config.Locking.CollectionName) == 0 {
+		config.Locking.CollectionName = DefaultLockingCollection
+	}
+	if config.Locking.Timeout <= 0 {
+		config.Locking.Timeout = DefaultLockTimeout
+	}
+	if config.Locking.Interval <= 0 {
+		config.Locking.Interval = DefaultLockTimeoutInterval
+	}
+
+	mc := &Mongo{
+		client: instance,
+		db:     instance.Database(config.DatabaseName),
+		config: config,
+	}
+
+	if mc.config.Locking.Enabled {
+		if err := mc.ensureLockTable(); err != nil {
+			return nil, err
+		}
+	}
+	if err := mc.ensureVersionTable(); err != nil {
+		return nil, err
+	}
+
+	return mc, nil
+}
+
+func (m *Mongo) Open(dsn string) (database.Driver, error) {
+	// connstring is experimental package, but it used for parse connection string in mongo.Connect function
+	uri, err := connstring.Parse(dsn)
+	if err != nil {
+		return nil, err
+	}
+	if len(uri.Database) == 0 {
+		return nil, ErrNoDatabaseName
+	}
+	unknown := url.Values(uri.UnknownOptions)
+
+	migrationsCollection := unknown.Get("x-migrations-collection")
+	lockCollection := unknown.Get("x-advisory-lock-collection")
+	transactionMode, err := parseBoolean(unknown.Get("x-transaction-mode"), false)
+	if err != nil {
+		return nil, err
+	}
+	advisoryLockingFlag, err := parseBoolean(unknown.Get("x-advisory-locking"), DefaultAdvisoryLockingFlag)
+	if err != nil {
+		return nil, err
+	}
+	lockingTimout, err := parseInt(unknown.Get("x-advisory-lock-timeout"), DefaultLockTimeout)
+	if err != nil {
+		return nil, err
+	}
+
+	lockTimeoutIntervalValue := unknown.Get("x-advisory-lock-timeout-interval")
+	// The initial release had a typo for this argument but for backwards compatibility sake, we will keep supporting it
+	// and we will error out if both values are set.
+	lockTimeoutIntervalValueFromTypo := unknown.Get("x-advisory-lock-timout-interval")
+
+	lockTimeout := lockTimeoutIntervalValue
+
+	if lockTimeoutIntervalValue != "" && lockTimeoutIntervalValueFromTypo != "" {
+		return nil, ErrLockTimeoutConfigConflict
+	} else if lockTimeoutIntervalValueFromTypo != "" {
+		lockTimeout = lockTimeoutIntervalValueFromTypo
+	}
+
+	maxLockCheckInterval, err := parseInt(lockTimeout, DefaultLockTimeoutInterval)
+
+	if err != nil {
+		return nil, err
+	}
+	client, err := mongo.Connect(context.TODO(), options.Client().ApplyURI(dsn))
+	if err != nil {
+		return nil, err
+	}
+
+	if err = client.Ping(context.TODO(), nil); err != nil {
+		return nil, err
+	}
+	mc, err := WithInstance(client, &Config{
+		DatabaseName:         uri.Database,
+		MigrationsCollection: migrationsCollection,
+		TransactionMode:      transactionMode,
+		Locking: Locking{
+			CollectionName: lockCollection,
+			Timeout:        lockingTimout,
+			Enabled:        advisoryLockingFlag,
+			Interval:       maxLockCheckInterval,
+		},
+	})
+	if err != nil {
+		return nil, err
+	}
+	return mc, nil
+}
+
+// Parse the url param, convert it to boolean
+// returns error if param invalid. returns defaultValue if param not present
+func parseBoolean(urlParam string, defaultValue bool) (bool, error) {
+
+	// if parameter passed, parse it (otherwise return default value)
+	if urlParam != "" {
+		result, err := strconv.ParseBool(urlParam)
+		if err != nil {
+			return false, err
+		}
+		return result, nil
+	}
+
+	// if no url Param passed, return default value
+	return defaultValue, nil
+}
+
+// Parse the url param, convert it to int
+// returns error if param invalid. returns defaultValue if param not present
+func parseInt(urlParam string, defaultValue int) (int, error) {
+
+	// if parameter passed, parse it (otherwise return default value)
+	if urlParam != "" {
+		result, err := strconv.Atoi(urlParam)
+		if err != nil {
+			return -1, err
+		}
+		return result, nil
+	}
+
+	// if no url Param passed, return default value
+	return defaultValue, nil
+}
+func (m *Mongo) SetVersion(version int, dirty bool) error {
+	migrationsCollection := m.db.Collection(m.config.MigrationsCollection)
+	if err := migrationsCollection.Drop(context.TODO()); err != nil {
+		return &database.Error{OrigErr: err, Err: "drop migrations collection failed"}
+	}
+	_, err := migrationsCollection.InsertOne(context.TODO(), bson.M{"version": version, "dirty": dirty})
+	if err != nil {
+		return &database.Error{OrigErr: err, Err: "save version failed"}
+	}
+	return nil
+}
+
+func (m *Mongo) Version() (version int, dirty bool, err error) {
+	var versionInfo versionInfo
+	err = m.db.Collection(m.config.MigrationsCollection).FindOne(context.TODO(), bson.M{}).Decode(&versionInfo)
+	switch {
+	case err == mongo.ErrNoDocuments:
+		return database.NilVersion, false, nil
+	case err != nil:
+		return 0, false, &database.Error{OrigErr: err, Err: "failed to get migration version"}
+	default:
+		return versionInfo.Version, versionInfo.Dirty, nil
+	}
+}
+
+func (m *Mongo) Run(migration io.Reader) error {
+	migr, err := io.ReadAll(migration)
+	if err != nil {
+		return err
+	}
+	var cmds []bson.D
+	err = bson.UnmarshalExtJSON(migr, true, &cmds)
+	if err != nil {
+		return fmt.Errorf("unmarshaling json error: %s", err)
+	}
+	if m.config.TransactionMode {
+		if err := m.executeCommandsWithTransaction(context.TODO(), cmds); err != nil {
+			return err
+		}
+	} else {
+		if err := m.executeCommands(context.TODO(), cmds); err != nil {
+			return err
+		}
+	}
+	return nil
+}
+
+func (m *Mongo) executeCommandsWithTransaction(ctx context.Context, cmds []bson.D) error {
+	err := m.db.Client().UseSession(ctx, func(sessionContext mongo.SessionContext) error {
+		if err := sessionContext.StartTransaction(); err != nil {
+			return &database.Error{OrigErr: err, Err: "failed to start transaction"}
+		}
+		if err := m.executeCommands(sessionContext, cmds); err != nil {
+			// When command execution is failed, it's aborting transaction
+			// If you tried to call abortTransaction, it`s return error that transaction already aborted
+			return err
+		}
+		if err := sessionContext.CommitTransaction(sessionContext); err != nil {
+			return &database.Error{OrigErr: err, Err: "failed to commit transaction"}
+		}
+		return nil
+	})
+	if err != nil {
+		return err
+	}
+	return nil
+}
+
+func (m *Mongo) executeCommands(ctx context.Context, cmds []bson.D) error {
+	for _, cmd := range cmds {
+		err := m.db.RunCommand(ctx, cmd).Err()
+		if err != nil {
+			return &database.Error{OrigErr: err, Err: fmt.Sprintf("failed to execute command:%v", cmd)}
+		}
+	}
+	return nil
+}
+
+func (m *Mongo) Close() error {
+	return m.client.Disconnect(context.TODO())
+}
+
+func (m *Mongo) Drop() error {
+	return m.db.Drop(context.TODO())
+}
+
+func (m *Mongo) ensureLockTable() error {
+	indexes := m.db.Collection(m.config.Locking.CollectionName).Indexes()
+
+	indexOptions := options.Index().SetUnique(true).SetName(LockIndexName)
+	_, err := indexes.CreateOne(context.TODO(), mongo.IndexModel{
+		Options: indexOptions,
+		Keys:    findFilter{Key: -1},
+	})
+	if err != nil {
+		return err
+	}
+	return nil
+}
+
+// ensureVersionTable checks if versions table exists and, if not, creates it.
+// Note that this function locks the database, which deviates from the usual
+// convention of "caller locks" in the MongoDb type.
+func (m *Mongo) ensureVersionTable() (err error) {
+	if err = m.Lock(); err != nil {
+		return err
+	}
+
+	defer func() {
+		if e := m.Unlock(); e != nil {
+			err = errors.Join(err, e)
+		}
+	}()
+
+	if err != nil {
+		return err
+	}
+	if _, _, err = m.Version(); err != nil {
+		return err
+	}
+	return nil
+}
+
+// Utilizes advisory locking on the config.LockingCollection collection
+// This uses a unique index on the `locking_key` field.
+func (m *Mongo) Lock() error {
+	return database.CasRestoreOnErr(&m.isLocked, false, true, database.ErrLocked, func() error {
+		if !m.config.Locking.Enabled {
+			return nil
+		}
+
+		pid := os.Getpid()
+		hostname, err := os.Hostname()
+		if err != nil {
+			hostname = fmt.Sprintf("Could not determine hostname. Error: %s", err.Error())
+		}
+
+		newLockObj := lockObj{
+			Key:       lockKeyUniqueValue,
+			Pid:       pid,
+			Hostname:  hostname,
+			CreatedAt: time.Now(),
+		}
+		operation := func() error {
+			timeout, cancelFunc := context.WithTimeout(context.Background(), contextWaitTimeout)
+			_, err := m.db.Collection(m.config.Locking.CollectionName).InsertOne(timeout, newLockObj)
+			defer cancelFunc()
+			return err
+		}
+		exponentialBackOff := backoff.NewExponentialBackOff()
+		duration := time.Duration(m.config.Locking.Timeout) * time.Second
+		exponentialBackOff.MaxElapsedTime = duration
+		exponentialBackOff.MaxInterval = time.Duration(m.config.Locking.Interval) * time.Second
+
+		err = backoff.Retry(operation, exponentialBackOff)
+		if err != nil {
+			return database.ErrLocked
+		}
+
+		return nil
+	})
+}
+
+func (m *Mongo) Unlock() error {
+	return database.CasRestoreOnErr(&m.isLocked, true, false, database.ErrNotLocked, func() error {
+		if !m.config.Locking.Enabled {
+			return nil
+		}
+
+		filter := findFilter{
+			Key: lockKeyUniqueValue,
+		}
+
+		ctx, cancel := context.WithTimeout(context.Background(), contextWaitTimeout)
+		_, err := m.db.Collection(m.config.Locking.CollectionName).DeleteMany(ctx, filter)
+		defer cancel()
+
+		if err != nil {
+			return err
+		}
+		return nil
+	})
+}