feat: use *bun.DB in MigratorDialect

uptrace · Oct 27, 2024 · a8788bf · a8788bf
1 parent 27a5a2b
commit a8788bf
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 8 deletions.
diff --git a/dialect/pgdialect/alter_table.go b/dialect/pgdialect/alter_table.go
@@ -2,23 +2,23 @@ package pgdialect
 
 import (
 	"context"
-	"database/sql"
 	"fmt"
 
+	"github.com/uptrace/bun"
 	"github.com/uptrace/bun/migrate/sqlschema"
 )
 
-func (d *Dialect) Migrator(sqldb *sql.DB) sqlschema.Migrator {
-	return &Migrator{sqldb: sqldb}
+func (d *Dialect) Migrator(db *bun.DB) sqlschema.Migrator {
+	return &Migrator{db: db}
 }
 
 type Migrator struct {
-	sqldb *sql.DB
+	db *bun.DB
 }
 
 func (m *Migrator) RenameTable(ctx context.Context, oldName, newName string) error {
 	query := fmt.Sprintf("ALTER TABLE %s RENAME TO %s", oldName, newName)
-	_, err := m.sqldb.ExecContext(ctx, query)
+	_, err := m.db.ExecContext(ctx, query)
 	if err != nil {
 		return err
 	}

diff --git a/migrate/sqlschema/migrator.go b/migrate/sqlschema/migrator.go
@@ -2,7 +2,6 @@ package sqlschema
 
 import (
 	"context"
-	"database/sql"
 	"fmt"
 
 	"github.com/uptrace/bun"
@@ -11,7 +10,7 @@ import (
 
 type MigratorDialect interface {
 	schema.Dialect
-	Migrator(*sql.DB) Migrator
+	Migrator(*bun.DB) Migrator
 }
 
 type Migrator interface {
@@ -28,6 +27,6 @@ func NewMigrator(db *bun.DB) (Migrator, error) {
 		return nil, fmt.Errorf("%q dialect does not implement sqlschema.Migrator", db.Dialect().Name())
 	}
 	return &migrator{
-		Migrator: md.Migrator(db.DB),
+		Migrator: md.Migrator(db),
 	}, nil
 }