Skip to content

Commit

Permalink
refactor: remove AlterTableQuery bloat and fix test errors
Browse files Browse the repository at this point in the history 
- Do not implement AppendQuery on Operation level. This is a leaky abstraction as queries are dialect-specific and migrate package should not be concerned with how they are constructed.
- AlterTableQuery also is an unnecessary abstraction. Now pgdialect will just build a simple string-query for each Operation.
- Moved operations.go to migrate/ package and deleted alt/ package.
- Minor clean-ups and documentation.

testChangeColumnType is commented out because the implementation is missing.
  • Loading branch information
@bevzzz
bevzzz committed Aug 3, 2024
1 parent a8e5b3d commit 8708552
Show file tree
Hide file tree
Showing 7 changed files with 187 additions and 306 deletions.
324 changes: 114 additions & 210 deletions dialect/pgdialect/alter_table.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2,271 +2,175 @@ package pgdialect

import (
"context"
"errors"
"fmt"

"github.com/uptrace/bun"
"github.com/uptrace/bun/migrate/alt"
"github.com/uptrace/bun/internal"
"github.com/uptrace/bun/migrate"
"github.com/uptrace/bun/migrate/sqlschema"
"github.com/uptrace/bun/schema"
)

func (d *Dialect) Migrator(db *bun.DB) sqlschema.Migrator {
return &Migrator{db: db, BaseMigrator: sqlschema.NewBaseMigrator(db)}
return &migrator{db: db, BaseMigrator: sqlschema.NewBaseMigrator(db)}
}

type Migrator struct {
type migrator struct {
*sqlschema.BaseMigrator

db *bun.DB
}

var _ sqlschema.Migrator = (*Migrator)(nil)
var _ sqlschema.Migrator = (*migrator)(nil)

func (m *Migrator) execRaw(ctx context.Context, q *bun.RawQuery) error {
if _, err := q.Exec(ctx); err != nil {
return err
}
return nil
}

func (m *Migrator) RenameTable(ctx context.Context, oldName, newName string) error {
q := m.db.NewRaw("ALTER TABLE ? RENAME TO ?", bun.Ident(oldName), bun.Ident(newName))
return m.execRaw(ctx, q)
}

func (m *Migrator) AddContraint(ctx context.Context, fk sqlschema.FK, name string) error {
q := m.db.NewRaw(
"ALTER TABLE ?.? ADD CONSTRAINT ? FOREIGN KEY (?) REFERENCES ?.? (?)",
bun.Safe(fk.From.Schema), bun.Safe(fk.From.Table), bun.Safe(name),
bun.Safe(fk.From.Column.String()),
bun.Safe(fk.To.Schema), bun.Safe(fk.To.Table),
bun.Safe(fk.To.Column.String()),
)
return m.execRaw(ctx, q)
}

func (m *Migrator) DropContraint(ctx context.Context, schema, table, name string) error {
q := m.db.NewRaw(
"ALTER TABLE ?.? DROP CONSTRAINT ?",
bun.Ident(schema), bun.Ident(table), bun.Ident(name),
)
return m.execRaw(ctx, q)
}

func (m *Migrator) RenameConstraint(ctx context.Context, schema, table, oldName, newName string) error {
q := m.db.NewRaw(
"ALTER TABLE ?.? RENAME CONSTRAINT ? TO ?",
bun.Ident(schema), bun.Ident(table), bun.Ident(oldName), bun.Ident(newName),
)
return m.execRaw(ctx, q)
}

func (m *Migrator) RenameColumn(ctx context.Context, schema, table, oldName, newName string) error {
q := m.db.NewRaw(
"ALTER TABLE ?.? RENAME COLUMN ? TO ?",
bun.Ident(schema), bun.Ident(table), bun.Ident(oldName), bun.Ident(newName),
)
return m.execRaw(ctx, q)
}

// -------------

func (m *Migrator) Apply(ctx context.Context, changes ...sqlschema.Operation) error {
func (m *migrator) Apply(ctx context.Context, changes ...sqlschema.Operation) error {
if len(changes) == 0 {
return nil
}
var conn bun.IConn
var err error

queries, err := m.buildQueries(changes...)
if err != nil {
return fmt.Errorf("apply database schema changes: %w", err)
}

for _, query := range queries {
var b []byte
if b, err = query.AppendQuery(m.db.Formatter(), b); err != nil {
return err
}
m.execRaw(ctx, m.db.NewRaw(string(b)))
if conn, err = m.db.Conn(ctx); err != nil {
return err
}

return nil
}

// buildQueries combines schema changes to a number of ALTER TABLE queries.
func (m *Migrator) buildQueries(changes ...sqlschema.Operation) ([]*AlterTableQuery, error) {
var queries []*AlterTableQuery
fmter := m.db.Formatter()
for _, change := range changes {
var b []byte // TODO(dyma): call db.MakeQueryBytes

chain := func(change sqlschema.Operation) error {
for _, query := range queries {
if err := query.Chain(change); err != errCannotChain {
return err // either nil (successful) or non-nil (failed)
switch change := change.(type) {
case *migrate.CreateTable:
err = m.CreateTable(ctx, change.Model)
if err != nil {
return fmt.Errorf("apply changes: create table %s: %w", change.FQN(), err)
}
continue
case *migrate.DropTable:
err = m.DropTable(ctx, change.Schema, change.Name)
if err != nil {
return fmt.Errorf("apply changes: drop table %s: %w", change.FQN(), err)
}
continue
case *migrate.RenameTable:
b, err = m.renameTable(fmter, b, change)
case *migrate.RenameColumn:
b, err = m.renameColumn(fmter, b, change)
case *migrate.DropConstraint:
b, err = m.dropContraint(fmter, b, change)
case *migrate.AddForeignKey:
b, err = m.addForeignKey(fmter, b, change)
case *migrate.RenameConstraint:
b, err = m.renameConstraint(fmter, b, change)
default:
return fmt.Errorf("apply changes: unknown operation %T", change)
}

// Create a new query for this change, since it cannot be chained to any of the existing ones.
q, err := newAlterTableQuery(change)
if err != nil {
return err
return fmt.Errorf("apply changes: %w", err)
}
queries = append(queries, q.Sep())
return nil
}

for _, change := range changes {
if err := chain(change); err != nil {
return nil, err
query := internal.String(b)
// log.Println("exec query: " + query)
if _, err = conn.ExecContext(ctx, query); err != nil {
return fmt.Errorf("apply changes: %w", err)
}
}
return queries, nil
}

type AlterTableQuery struct {
FQN schema.FQN

RenameTable sqlschema.Operation
RenameColumn sqlschema.Operation
RenameConstraint sqlschema.Operation
Actions Actions

separate bool
return nil
}

type Actions []*Action

var _ schema.QueryAppender = (*Actions)(nil)

type Action struct {
AddColumn sqlschema.Operation
DropColumn sqlschema.Operation
AlterColumn sqlschema.Operation
AlterType sqlschema.Operation
SetDefault sqlschema.Operation
DropDefault sqlschema.Operation
SetNotNull sqlschema.Operation
DropNotNull sqlschema.Operation
AddGenerated sqlschema.Operation
AddConstraint sqlschema.Operation
DropConstraint sqlschema.Operation
Custom sqlschema.Operation
func (m *migrator) renameTable(fmter schema.Formatter, b []byte, rename *migrate.RenameTable) (_ []byte, err error) {
b = append(b, "ALTER TABLE "...)
fqn := rename.FQN()
if b, err = fqn.AppendQuery(fmter, b); err != nil {
return b, err
}
b = append(b, " RENAME TO "...)
if b, err = bun.Ident(rename.NewName).AppendQuery(fmter, b); err != nil {
return b, err
}
return b, nil
}

var _ schema.QueryAppender = (*Action)(nil)

func newAlterTableQuery(op sqlschema.Operation) (*AlterTableQuery, error) {
q := AlterTableQuery{
FQN: op.FQN(),
func (m *migrator) renameColumn(fmter schema.Formatter, b []byte, rename *migrate.RenameColumn) (_ []byte, err error) {
b = append(b, "ALTER TABLE "...)
fqn := rename.FQN()
if b, err = fqn.AppendQuery(fmter, b); err != nil {
return b, err
}
switch op.(type) {
case *alt.RenameTable:
q.RenameTable = op
case *alt.RenameColumn:
q.RenameColumn = op
case *alt.RenameConstraint:
q.RenameConstraint = op
default:
q.Actions = append(q.Actions, newAction(op))

b = append(b, " RENAME COLUMN "...)
if b, err = bun.Ident(rename.OldName).AppendQuery(fmter, b); err != nil {
return b, err
}
return &q, nil
}

func newAction(op sqlschema.Operation) *Action {
var a Action
return &a
b = append(b, " TO "...)
if b, err = bun.Ident(rename.NewName).AppendQuery(fmter, b); err != nil {
return b, err
}
return b, nil
}

// errCannotChain is a sentinel error. To apply the change, callers should
// create a new AlterTableQuery instead and include it there.
var errCannotChain = errors.New("cannot chain change to the current query")
func (m *migrator) renameConstraint(fmter schema.Formatter, b []byte, rename *migrate.RenameConstraint) (_ []byte, err error) {
b = append(b, "ALTER TABLE "...)
fqn := rename.FQN()
if b, err = fqn.AppendQuery(fmter, b); err != nil {
return b, err
}

func (q *AlterTableQuery) Chain(op sqlschema.Operation) error {
if op.FQN() != q.FQN {
return errCannotChain
b = append(b, " RENAME CONSTRAINT "...)
if b, err = bun.Ident(rename.OldName).AppendQuery(fmter, b); err != nil {
return b, err
}

switch op.(type) {
default:
return fmt.Errorf("unsupported operation %T", op)
b = append(b, " TO "...)
if b, err = bun.Ident(rename.NewName).AppendQuery(fmter, b); err != nil {
return b, err
}
return b, nil
}

func (q *AlterTableQuery) isEmpty() bool {
return q.RenameTable == nil && q.RenameColumn == nil && q.RenameConstraint == nil && len(q.Actions) == 0
}
func (m *migrator) dropContraint(fmter schema.Formatter, b []byte, drop *migrate.DropConstraint) (_ []byte, err error) {
b = append(b, "ALTER TABLE "...)
fqn := drop.FQN()
if b, err = fqn.AppendQuery(fmter, b); err != nil {
return b, err
}

// Sep appends a ";" separator at the end of the query.
func (q *AlterTableQuery) Sep() *AlterTableQuery {
q.separate = true
return q
b = append(b, " DROP CONSTRAINT "...)
if b, err = bun.Ident(drop.ConstraintName).AppendQuery(fmter, b); err != nil {
return b, err
}
return b, nil
}

func (q *AlterTableQuery) AppendQuery(fmter schema.Formatter, b []byte) (_ []byte, err error) {
var op schema.QueryAppender
switch true {
case q.RenameTable != nil:
op = q.RenameTable
case q.RenameColumn != nil:
op = q.RenameColumn
case q.RenameConstraint != nil:
op = q.RenameConstraint
case len(q.Actions) > 0:
op = q.Actions
default:
return b, nil
}
func (m *migrator) addForeignKey(fmter schema.Formatter, b []byte, add *migrate.AddForeignKey) (_ []byte, err error) {
b = append(b, "ALTER TABLE "...)
b, _ = q.FQN.AppendQuery(fmter, b)
b = append(b, " "...)
if b, err = op.AppendQuery(fmter, b); err != nil {
fqn := add.FQN()
if b, err = fqn.AppendQuery(fmter, b); err != nil {
return b, err
}

if q.separate {
b = append(b, ";"...)
b = append(b, " ADD CONSTRAINT "...)
if b, err = bun.Ident(add.ConstraintName).AppendQuery(fmter, b); err != nil {
return b, err
}
return b, nil
}

func (actions Actions) AppendQuery(fmter schema.Formatter, b []byte) (_ []byte, err error) {
for i, a := range actions {
if i > 0 {
b = append(b, ", "...)
}
b, err = a.AppendQuery(fmter, b)
if err != nil {
return b, err
}
b = append(b, " FOREIGN KEY ("...)
if b, err = add.FK.From.Column.Safe().AppendQuery(fmter, b); err != nil {
return b, err
}
b = append(b, ") "...)

other := schema.FQN{Schema: add.FK.To.Schema, Table: add.FK.To.Table}
b = append(b, " REFERENCES "...)
if b, err = other.AppendQuery(fmter, b); err != nil {
return b, err
}
return b, nil
}

func (a *Action) AppendQuery(fmter schema.Formatter, b []byte) ([]byte, error) {
var op schema.QueryAppender
switch true {
case a.AddColumn != nil:
op = a.AddColumn
case a.DropColumn != nil:
op = a.DropColumn
case a.AlterColumn != nil:
op = a.AlterColumn
case a.AlterType != nil:
op = a.AlterType
case a.SetDefault != nil:
op = a.SetDefault
case a.DropDefault != nil:
op = a.DropDefault
case a.SetNotNull != nil:
op = a.SetNotNull
case a.DropNotNull != nil:
op = a.DropNotNull
case a.AddGenerated != nil:
op = a.AddGenerated
case a.AddConstraint != nil:
op = a.AddConstraint
case a.DropConstraint != nil:
op = a.DropConstraint
default:
return b, nil
b = append(b, " ("...)
if b, err = add.FK.To.Column.Safe().AppendQuery(fmter, b); err != nil {
return b, err
}
return op.AppendQuery(fmter, b)
b = append(b, ")"...)

return b, nil
}
Loading

0 comments on commit 8708552

Please sign in to comment.