github.com/kubeshop/testkube@v1.17.23/pkg/dbmigrator/database.go (about)

     1  package dbmigrator
     2  
     3  import (
     4  	"context"
     5  	"time"
     6  
     7  	"go.mongodb.org/mongo-driver/bson"
     8  	"go.mongodb.org/mongo-driver/mongo"
     9  	"go.mongodb.org/mongo-driver/mongo/options"
    10  )
    11  
    12  //go:generate mockgen -destination=./mock_database.go -package=dbmigrator "github.com/kubeshop/testkube/pkg/dbmigrator" Database
    13  type Database interface {
    14  	RunCommands(ctx context.Context, commands []bson.D) error
    15  	InsertMigrationState(ctx context.Context, migration *DbMigration) error
    16  	DeleteMigrationState(ctx context.Context, migration *DbMigration) error
    17  	GetAppliedMigrations(ctx context.Context) ([]DbMigration, error)
    18  }
    19  
    20  type database struct {
    21  	db             *mongo.Database
    22  	migrationsColl *mongo.Collection
    23  }
    24  
    25  // TODO: Consider locks
    26  func NewDatabase(db *mongo.Database, migrationsColl string) Database {
    27  	return &database{db: db, migrationsColl: db.Collection(migrationsColl)}
    28  }
    29  
    30  // TODO: Consider transactions, but it requires MongoDB with replicaset
    31  func (d *database) RunCommands(ctx context.Context, commands []bson.D) error {
    32  	for _, cmd := range commands {
    33  		err := d.db.RunCommand(ctx, cmd).Err()
    34  		if err != nil {
    35  			return err
    36  		}
    37  	}
    38  	return nil
    39  }
    40  
    41  func (d *database) InsertMigrationState(ctx context.Context, migration *DbMigration) error {
    42  	_, err := d.migrationsColl.InsertOne(ctx, bson.M{
    43  		"name":      migration.Name,
    44  		"up":        migration.UpScript,
    45  		"down":      migration.DownScript,
    46  		"timestamp": time.Now(),
    47  	})
    48  	return err
    49  }
    50  
    51  func (d *database) DeleteMigrationState(ctx context.Context, migration *DbMigration) error {
    52  	_, err := d.migrationsColl.DeleteOne(ctx, bson.M{"name": migration.Name})
    53  	return err
    54  }
    55  
    56  func (d *database) GetAppliedMigrations(ctx context.Context) (results []DbMigration, err error) {
    57  	cursor, err := d.migrationsColl.Find(ctx, bson.M{}, &options.FindOptions{Sort: bson.M{"name": 1}})
    58  	if err != nil {
    59  		return nil, err
    60  	}
    61  	err = cursor.All(ctx, &results)
    62  	return results, err
    63  }