Merge branch 'uptrace:master' into jeffreydwalter-fix-migration-sort

Author: jeffreydwalter

driver/pgdriver/listener.go

@@ -237,21 +237,30 @@ type Notification struct {
 
 type ChannelOption func(c *channel)
 
+type channelOverflowHandler func(n Notification)
+
 func WithChannelSize(size int) ChannelOption {
 	return func(c *channel) {
 		c.size = size
 	}
 }
 
+func WithChannelOverflowHandler(handler channelOverflowHandler) ChannelOption {
+	return func(c *channel) {
+		c.overflowHandler = handler
+	}
+}
+
 type channel struct {
 	ctx context.Context
 	ln  *Listener
 
 	size        int
 	pingTimeout time.Duration
 
-	ch     chan Notification
-	pingCh chan struct{}
+	ch              chan Notification
+	pingCh          chan struct{}
+	overflowHandler channelOverflowHandler
 }
 
 func newChannel(ln *Listener, opts []ChannelOption) *channel {
@@ -310,6 +319,9 @@ func (c *channel) startReceive() {
 			case c.ch <- Notification{channel, payload}:
 			default:
 				Logger.Printf(c.ctx, "pgdriver: Listener buffer is full (message is dropped)")
+				if c.overflowHandler != nil {
+					c.overflowHandler(Notification{channel, payload})
+				}
 			}
 		}
 	}

driver/pgdriver/listener_test.go

@@ -0,0 +1,28 @@
+package pgdriver
+
+import (
+	"reflect"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestWithChannelOverflowHandler(t *testing.T) {
+	// Create a test handler
+	testHandler := func(n Notification) {}
+
+	// Create a channel instance
+	c := &channel{}
+
+	// Apply the option
+	opt := WithChannelOverflowHandler(testHandler)
+	opt(c)
+
+	// Verify the handler was set correctly
+	assert.NotNil(t, c.overflowHandler, "overflow handler should be set")
+	assert.Equal(t,
+		reflect.ValueOf(testHandler).Pointer(),
+		reflect.ValueOf(c.overflowHandler).Pointer(),
+		"overflow handler should match the provided handler",
+	)
+}

example/migrate/main.go

@@ -30,11 +30,14 @@ func main() {
 		bundebug.FromEnv(),
 	))
 
+	templateData := map[string]string{
+		"Prefix": "example_",
+	}
 	app := &cli.App{
 		Name: "bun",
 
 		Commands: []*cli.Command{
-			newDBCommand(migrate.NewMigrator(db, migrations.Migrations)),
+			newDBCommand(migrate.NewMigrator(db, migrations.Migrations, migrate.WithTemplateData(templateData))),
 		},
 	}
 	if err := app.Run(os.Args); err != nil {

example/migrate/migrations/20210505110847_bar.down.sql

@@ -1 +1 @@
-DROP TABLE IF EXISTS test;
+DROP TABLE IF EXISTS {{.Prefix}}test;

example/migrate/migrations/20210505110847_bar.up.sql

@@ -1,7 +1,7 @@
-CREATE TABLE test (
+CREATE TABLE {{.Prefix}}test (
   id bigint PRIMARY KEY
 );
 
 --bun:split
 
-ALTER TABLE test ADD COLUMN name varchar(100);
+ALTER TABLE {{.Prefix}}test ADD COLUMN name varchar(100);

internal/dbtest/listener_test.go

@@ -2,6 +2,8 @@ package dbtest_test
 
 import (
 	"context"
+	"fmt"
+	"sync/atomic"
 	"testing"
 	"time"
 
@@ -84,3 +86,45 @@ func TestListenerChannel(t *testing.T) {
 		return !ok
 	}, 3*time.Second, 100*time.Millisecond)
 }
+
+func TestListenerChannelOverflowHandler(t *testing.T) {
+	ctx := context.Background()
+	channelSize := 1
+	overflowMessagesCount := channelSize * 3
+
+	db := pg(t)
+	defer db.Close()
+
+	ln := pgdriver.NewListener(db)
+	defer ln.Close()
+
+	var overflowCount atomic.Int32
+
+	// Create channel with small buffer and overflow handler
+	ch := ln.Channel(
+		pgdriver.WithChannelSize(channelSize),
+		pgdriver.WithChannelOverflowHandler(func(n pgdriver.Notification) {
+			overflowCount.Add(1)
+		}),
+	)
+
+	err := ln.Listen(ctx, "test_channel")
+	require.NoError(t, err)
+
+	// Fill the channel buffer
+	_, err = db.ExecContext(ctx, "NOTIFY test_channel, ?", "msg1")
+	require.NoError(t, err)
+
+	// Wait for the first message to be received
+	<-ch
+
+	// Send more messages to trigger overflow
+	for i := 0; i < overflowMessagesCount; i++ {
+		_, err = db.ExecContext(ctx, "NOTIFY test_channel, ?", fmt.Sprintf("msg%d", i+2))
+		require.NoError(t, err)
+	}
+
+	require.Eventually(t, func() bool {
+		return overflowCount.Load() > 0
+	}, time.Second, 10*time.Millisecond, "overflow handler should have been called")
+}

internal/dbtest/migrate_test.go

@@ -75,22 +75,22 @@ func testMigrateUpAndDown(t *testing.T, db *bun.DB) {
 	migrations := migrate.NewMigrations()
 	migrations.Add(migrate.Migration{
 		Name: "20060102150405",
-		Up: func(ctx context.Context, db *bun.DB) error {
+		Up: func(ctx context.Context, db *bun.DB, templateData any) error {
 			history = append(history, "up1")
 			return nil
 		},
-		Down: func(ctx context.Context, db *bun.DB) error {
+		Down: func(ctx context.Context, db *bun.DB, templateData any) error {
 			history = append(history, "down1")
 			return nil
 		},
 	})
 	migrations.Add(migrate.Migration{
 		Name: "20060102160405",
-		Up: func(ctx context.Context, db *bun.DB) error {
+		Up: func(ctx context.Context, db *bun.DB, templateData any) error {
 			history = append(history, "up2")
 			return nil
 		},
-		Down: func(ctx context.Context, db *bun.DB) error {
+		Down: func(ctx context.Context, db *bun.DB, templateData any) error {
 			history = append(history, "down2")
 			return nil
 		},
@@ -125,33 +125,33 @@ func testMigrateUpError(t *testing.T, db *bun.DB) {
 	migrations := migrate.NewMigrations()
 	migrations.Add(migrate.Migration{
 		Name: "20060102150405",
-		Up: func(ctx context.Context, db *bun.DB) error {
+		Up: func(ctx context.Context, db *bun.DB, templateData any) error {
 			history = append(history, "up1")
 			return nil
 		},
-		Down: func(ctx context.Context, db *bun.DB) error {
+		Down: func(ctx context.Context, db *bun.DB, templateData any) error {
 			history = append(history, "down1")
 			return nil
 		},
 	})
 	migrations.Add(migrate.Migration{
 		Name: "20060102160405",
-		Up: func(ctx context.Context, db *bun.DB) error {
+		Up: func(ctx context.Context, db *bun.DB, templateData any) error {
 			history = append(history, "up2")
 			return errors.New("failed")
 		},
-		Down: func(ctx context.Context, db *bun.DB) error {
+		Down: func(ctx context.Context, db *bun.DB, templateData any) error {
 			history = append(history, "down2")
 			return nil
 		},
 	})
 	migrations.Add(migrate.Migration{
 		Name: "20060102170405",
-		Up: func(ctx context.Context, db *bun.DB) error {
+		Up: func(ctx context.Context, db *bun.DB, templateData any) error {
 			history = append(history, "up3")
 			return errors.New("failed")
 		},
-		Down: func(ctx context.Context, db *bun.DB) error {
+		Down: func(ctx context.Context, db *bun.DB, templateData any) error {
 			history = append(history, "down3")
 			return nil
 		},

migrate/auto.go

@@ -268,8 +268,8 @@ func (am *AutoMigrator) createSQLMigrations(ctx context.Context, transactional b
 	migrations := NewMigrations(am.migrationsOpts...)
 	migrations.Add(Migration{
 		Name:    name,
-		Up:      changes.Up(am.dbMigrator),
-		Down:    changes.Down(am.dbMigrator),
+		Up:      wrapMigrationFunc(changes.Up(am.dbMigrator)),
+		Down:    wrapMigrationFunc(changes.Down(am.dbMigrator)),
 		Comment: "Changes detected by bun.AutoMigrator",
 	})
 

migrate/migration.go

@@ -10,6 +10,7 @@ import (
 	"sort"
 	"strconv"
 	"strings"
+	"text/template"
 	"time"
 
 	"github.com/uptrace/bun"
@@ -24,8 +25,8 @@ type Migration struct {
 	GroupID    int64
 	MigratedAt time.Time `bun:",notnull,nullzero,default:current_timestamp"`
 
-	Up   MigrationFunc `bun:"-"`
-	Down MigrationFunc `bun:"-"`
+	Up   internalMigrationFunc `bun:"-"`
+	Down internalMigrationFunc `bun:"-"`
 }
 
 func (m Migration) String() string {
@@ -36,23 +37,57 @@ func (m Migration) IsApplied() bool {
 	return m.ID > 0
 }
 
+type internalMigrationFunc func(ctx context.Context, db *bun.DB, templateData any) error
+
 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 {
+func NewSQLMigrationFunc(fsys fs.FS, name string) internalMigrationFunc {
+	return func(ctx context.Context, db *bun.DB, templateData any) 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)
+		return Exec(ctx, db, f, templateData, isTx)
+	}
+}
+
+func wrapMigrationFunc(fn MigrationFunc) internalMigrationFunc {
+	return func(ctx context.Context, db *bun.DB, templateData any) error {
+		return fn(ctx, db)
+	}
+}
+
+func renderTemplate(contents []byte, templateData any) (*bytes.Buffer, error) {
+	tmpl, err := template.New("migration").Parse(string(contents))
+	if err != nil {
+		return nil, fmt.Errorf("failed to parse template: %w", err)
+	}
+
+	var rendered bytes.Buffer
+	if err := tmpl.Execute(&rendered, templateData); err != nil {
+		return nil, fmt.Errorf("failed to execute template: %w", err)
 	}
+
+	return &rendered, nil
 }
 
 // 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)
+func Exec(ctx context.Context, db *bun.DB, f io.Reader, templateData any, isTx bool) error {
+	contents, err := io.ReadAll(f)
+	if err != nil {
+		return err
+	}
+	var reader io.Reader = bytes.NewReader(contents)
+	if templateData != nil {
+		buf, err := renderTemplate(contents, templateData)
+		if err != nil {
+			return err
+		}
+		reader = buf
+	}
+	scanner := bufio.NewScanner(reader)
 	var queries []string
 
 	var query []byte

migrate/migrations.go

@@ -58,8 +58,8 @@ func (m *Migrations) Register(up, down MigrationFunc) error {
 	m.Add(Migration{
 		Name:    name,
 		Comment: comment,
-		Up:      up,
-		Down:    down,
+		Up:      wrapMigrationFunc(up),
+		Down:    wrapMigrationFunc(down),
 	})
 
 	return nil