gotosocial/vendor/github.com/uptrace/bun/migrate/migration.go

298 lines
5.5 KiB
Go
Raw Normal View History

2021-08-31 17:27:02 +00:00
package migrate
import (
"bufio"
"bytes"
"context"
"fmt"
"io"
2021-08-31 17:27:02 +00:00
"io/fs"
"sort"
"strings"
"time"
"github.com/uptrace/bun"
)
type Migration struct {
bun.BaseModel
ID int64 `bun:",pk,autoincrement"`
2021-08-31 17:27:02 +00:00
Name string
Comment string `bun:"-"`
2021-08-31 17:27:02 +00:00
GroupID int64
MigratedAt time.Time `bun:",notnull,nullzero,default:current_timestamp"`
Up MigrationFunc `bun:"-"`
Down MigrationFunc `bun:"-"`
}
func (m Migration) String() string {
return fmt.Sprintf("%s_%s", m.Name, m.Comment)
2021-08-31 17:27:02 +00:00
}
func (m Migration) IsApplied() bool {
2021-08-31 17:27:02 +00:00
return m.ID > 0
}
type MigrationFunc func(ctx context.Context, db *bun.DB) error
func NewSQLMigrationFunc(fsys fs.FS, name string) MigrationFunc {
return func(ctx context.Context, db *bun.DB) error {
f, err := fsys.Open(name)
if err != nil {
return err
}
isTx := strings.HasSuffix(name, ".tx.up.sql") || strings.HasSuffix(name, ".tx.down.sql")
return Exec(ctx, db, f, isTx)
}
}
2021-08-31 17:27:02 +00:00
// Exec reads and executes the SQL migration in the f.
func Exec(ctx context.Context, db *bun.DB, f io.Reader, isTx bool) error {
scanner := bufio.NewScanner(f)
var queries []string
var query []byte
for scanner.Scan() {
b := scanner.Bytes()
const prefix = "--bun:"
if bytes.HasPrefix(b, []byte(prefix)) {
b = b[len(prefix):]
if bytes.Equal(b, []byte("split")) {
queries = append(queries, string(query))
query = query[:0]
continue
}
return fmt.Errorf("bun: unknown directive: %q", b)
2021-08-31 17:27:02 +00:00
}
query = append(query, b...)
query = append(query, '\n')
}
if len(query) > 0 {
queries = append(queries, string(query))
}
if err := scanner.Err(); err != nil {
return err
}
var idb bun.IConn
if isTx {
tx, err := db.BeginTx(ctx, nil)
if err != nil {
return err
2021-08-31 17:27:02 +00:00
}
idb = tx
} else {
conn, err := db.Conn(ctx)
if err != nil {
2021-08-31 17:27:02 +00:00
return err
}
idb = conn
}
2021-08-31 17:27:02 +00:00
var retErr error
var execErr error
2021-08-31 17:27:02 +00:00
defer func() {
if tx, ok := idb.(bun.Tx); ok {
if execErr != nil {
retErr = tx.Rollback()
} else {
retErr = tx.Commit()
}
return
2021-08-31 17:27:02 +00:00
}
if conn, ok := idb.(bun.Conn); ok {
retErr = conn.Close()
return
}
panic("not reached")
}()
for _, q := range queries {
_, execErr = idb.ExecContext(ctx, q)
if execErr != nil {
return execErr
2021-08-31 17:27:02 +00:00
}
}
return retErr
2021-08-31 17:27:02 +00:00
}
const goTemplate = `package %s
import (
"context"
"fmt"
"github.com/uptrace/bun"
)
func init() {
Migrations.MustRegister(func(ctx context.Context, db *bun.DB) error {
fmt.Print(" [up migration] ")
return nil
}, func(ctx context.Context, db *bun.DB) error {
fmt.Print(" [down migration] ")
return nil
})
}
`
2021-10-24 11:14:37 +00:00
const sqlTemplate = `SET statement_timeout = 0;
--bun:split
SELECT 1
2021-08-31 17:27:02 +00:00
--bun:split
SELECT 2
`
//------------------------------------------------------------------------------
type MigrationSlice []Migration
func (ms MigrationSlice) String() string {
if len(ms) == 0 {
return "empty"
}
if len(ms) > 5 {
return fmt.Sprintf("%d migrations (%s ... %s)", len(ms), ms[0].Name, ms[len(ms)-1].Name)
}
var sb strings.Builder
for i := range ms {
if i > 0 {
sb.WriteString(", ")
}
sb.WriteString(ms[i].String())
2021-08-31 17:27:02 +00:00
}
return sb.String()
}
// Applied returns applied migrations in descending order
// (the order is important and is used in Rollback).
func (ms MigrationSlice) Applied() MigrationSlice {
var applied MigrationSlice
for i := range ms {
if ms[i].IsApplied() {
applied = append(applied, ms[i])
}
}
sortDesc(applied)
return applied
}
// Unapplied returns unapplied migrations in ascending order
// (the order is important and is used in Migrate).
func (ms MigrationSlice) Unapplied() MigrationSlice {
var unapplied MigrationSlice
for i := range ms {
if !ms[i].IsApplied() {
unapplied = append(unapplied, ms[i])
}
}
sortAsc(unapplied)
return unapplied
}
// LastGroupID returns the last applied migration group id.
// The id is 0 when there are no migration groups.
func (ms MigrationSlice) LastGroupID() int64 {
var lastGroupID int64
for i := range ms {
groupID := ms[i].GroupID
2021-09-23 09:13:28 +00:00
if groupID > lastGroupID {
2021-08-31 17:27:02 +00:00
lastGroupID = groupID
}
}
return lastGroupID
}
// LastGroup returns the last applied migration group.
func (ms MigrationSlice) LastGroup() *MigrationGroup {
group := &MigrationGroup{
ID: ms.LastGroupID(),
}
if group.ID == 0 {
return group
}
for i := range ms {
if ms[i].GroupID == group.ID {
group.Migrations = append(group.Migrations, ms[i])
}
}
return group
}
type MigrationGroup struct {
ID int64
Migrations MigrationSlice
}
func (g MigrationGroup) IsZero() bool {
2021-08-31 17:27:02 +00:00
return g.ID == 0 && len(g.Migrations) == 0
}
func (g MigrationGroup) String() string {
2021-08-31 17:27:02 +00:00
if g.IsZero() {
return "nil"
}
return fmt.Sprintf("group #%d (%s)", g.ID, g.Migrations)
}
type MigrationFile struct {
Name string
Path string
Content string
}
//------------------------------------------------------------------------------
type migrationConfig struct {
nop bool
}
func newMigrationConfig(opts []MigrationOption) *migrationConfig {
cfg := new(migrationConfig)
for _, opt := range opts {
opt(cfg)
}
return cfg
}
type MigrationOption func(cfg *migrationConfig)
func WithNopMigration() MigrationOption {
return func(cfg *migrationConfig) {
cfg.nop = true
}
}
//------------------------------------------------------------------------------
func sortAsc(ms MigrationSlice) {
sort.Slice(ms, func(i, j int) bool {
return ms[i].Name < ms[j].Name
})
}
func sortDesc(ms MigrationSlice) {
sort.Slice(ms, func(i, j int) bool {
return ms[i].Name > ms[j].Name
})
}