feat(source): add embed source support

leozeli · leozeli · commit 59bd95e803b0 · 2025-06-10T17:48:38.000+08:00
diff --git a/source/embed/README.md b/source/embed/README.md
@@ -0,0 +1,21 @@
+# embed
+
+```golang
+//go:embed *.sql
+var MigrationFiles embed.FS 
+
+```
+
+```golang
+embed, err := migrations.NewEmbed(migrations.MigrationFiles, ".")
+if err != nil {
+    klog.Error(fmt.Sprintf("newConnectionEngine migrations.NewEmbed error:%v", err))
+    return
+}
+m, err := migrate.NewWithInstance(
+    "embed",
+    embed,
+    "mysql",
+    dbdriver,
+)
+```
diff --git a/source/embed/embed.go b/source/embed/embed.go
@@ -0,0 +1,162 @@
+package embed
+
+import (
+	"embed"
+	"errors"
+	"fmt"
+	"io"
+	"io/fs"
+	"path"
+	"strconv"
+
+	"github.com/golang-migrate/migrate/v4/source"
+)
+
+type Embed struct {
+	FS         embed.FS
+	migrations *source.Migrations
+	path       string
+}
+
+// NewEmbed returns a new Driver using the embed.FS and a relative path.
+func NewEmbed(fsys embed.FS, path string) (source.Driver, error) {
+	var e Embed
+	if err := e.Init(fsys, path); err != nil {
+		return nil, fmt.Errorf("failed to init embed driver with path %s: %w", path, err)
+	}
+	return &e, nil
+}
+
+// Open is part of source.Driver interface implementation.
+// Open cannot be called on the embed driver directly as it's designed to use embed.FS.
+func (e *Embed) Open(url string) (source.Driver, error) {
+	return nil, errors.New("Open() cannot be called on the embed driver")
+}
+
+// Init prepares Embed instance to read migrations from embed.FS and a relative path.
+func (e *Embed) Init(fsys embed.FS, path string) error {
+	entries, err := fs.ReadDir(fsys, path)
+	if err != nil {
+		return err
+	}
+
+	ms := source.NewMigrations()
+	for _, e := range entries {
+		if e.IsDir() {
+			continue
+		}
+		m, err := source.DefaultParse(e.Name())
+		if err != nil {
+			continue
+		}
+		file, err := e.Info()
+		if err != nil {
+			return err
+		}
+		if !ms.Append(m) {
+			return source.ErrDuplicateMigration{
+				Migration: *m,
+				FileInfo:  file,
+			}
+		}
+	}
+
+	e.FS = fsys
+	e.path = path
+	e.migrations = ms
+	return nil
+}
+
+// Close is part of source.Driver interface implementation.
+func (e *Embed) Close() error {
+	// Since embed.FS doesn't support Close(), this method is a no-op
+	return nil
+}
+
+// First is part of source.Driver interface implementation.
+func (e *Embed) First() (version uint, err error) {
+	if version, ok := e.migrations.First(); ok {
+		return version, nil
+	}
+	return 0, &fs.PathError{
+		Op:   "first",
+		Path: e.path,
+		Err:  fs.ErrNotExist,
+	}
+}
+
+// Prev is part of source.Driver interface implementation.
+func (e *Embed) Prev(version uint) (prevVersion uint, err error) {
+	if version, ok := e.migrations.Prev(version); ok {
+		return version, nil
+	}
+	return 0, &fs.PathError{
+		Op:   "prev for version " + strconv.FormatUint(uint64(version), 10),
+		Path: e.path,
+		Err:  fs.ErrNotExist,
+	}
+}
+
+// Next is part of source.Driver interface implementation.
+func (e *Embed) Next(version uint) (nextVersion uint, err error) {
+	if version, ok := e.migrations.Next(version); ok {
+		return version, nil
+	}
+	return 0, &fs.PathError{
+		Op:   "next for version " + strconv.FormatUint(uint64(version), 10),
+		Path: e.path,
+		Err:  fs.ErrNotExist,
+	}
+}
+
+// ReadUp is part of source.Driver interface implementation.
+func (e *Embed) ReadUp(version uint) (r io.ReadCloser, identifier string, err error) {
+	if m, ok := e.migrations.Up(version); ok {
+		body, err := e.FS.ReadFile(path.Join(e.path, m.Raw))
+		if err != nil {
+			return nil, "", err
+		}
+		return io.NopCloser(&fileReader{data: body}), m.Identifier, nil
+	}
+	return nil, "", &fs.PathError{
+		Op:   "read up for version " + strconv.FormatUint(uint64(version), 10),
+		Path: e.path,
+		Err:  fs.ErrNotExist,
+	}
+}
+
+// ReadDown is part of source.Driver interface implementation.
+func (e *Embed) ReadDown(version uint) (r io.ReadCloser, identifier string, err error) {
+	if m, ok := e.migrations.Down(version); ok {
+		body, err := e.FS.ReadFile(path.Join(e.path, m.Raw))
+		if err != nil {
+			return nil, "", err
+		}
+		return io.NopCloser(&fileReader{data: body}), m.Identifier, nil
+	}
+	return nil, "", &fs.PathError{
+		Op:   "read down for version " + strconv.FormatUint(uint64(version), 10),
+		Path: e.path,
+		Err:  fs.ErrNotExist,
+	}
+}
+
+// fileReader []byte to io.ReadCloser
+type fileReader struct {
+	data []byte
+	pos  int
+}
+
+func (fr *fileReader) Read(p []byte) (n int, err error) {
+	if fr.pos >= len(fr.data) {
+		return 0, io.EOF
+	}
+	n = copy(p, fr.data[fr.pos:])
+	fr.pos += n
+	return n, nil
+}
+
+func (fr *fileReader) Close() error {
+	// do nothing, as embed.FS does not require closing
+	return nil
+}
diff --git a/source/embed/embed_test.go b/source/embed/embed_test.go
@@ -0,0 +1,198 @@
+package embed
+
+import (
+	"embed"
+	"errors"
+	"io"
+	"io/fs"
+	"testing"
+
+	"github.com/golang-migrate/migrate/v4/source"
+	st "github.com/golang-migrate/migrate/v4/source/testing"
+)
+
+//go:embed testmigrations/*.sql
+var testFS embed.FS
+
+const testPath = "testmigrations"
+
+func Test(t *testing.T) {
+	driver, err := NewEmbed(testFS, testPath)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	st.Test(t, driver)
+}
+
+func TestNewEmbed_Success(t *testing.T) {
+	driver, err := NewEmbed(testFS, testPath)
+	if err != nil {
+		t.Fatalf("expected no error, got %v", err)
+	}
+	if driver == nil {
+		t.Fatal("expected driver, got nil")
+	}
+}
+
+func TestNewEmbed_InvalidPath(t *testing.T) {
+	_, err := NewEmbed(testFS, "doesnotexist")
+	if err == nil {
+		t.Fatal("expected error for invalid path, got nil")
+	}
+}
+
+func TestEmbed_Open(t *testing.T) {
+	driver, _ := NewEmbed(testFS, "testmigrations")
+	_, err := driver.(*Embed).Open("someurl")
+	if err == nil || err.Error() != "Open() cannot be called on the embed driver" {
+		t.Fatalf("expected Open() error, got %v", err)
+	}
+}
+
+func TestEmbed_First(t *testing.T) {
+	driver, _ := NewEmbed(testFS, "testmigrations")
+	version, err := driver.First()
+	if err != nil {
+		t.Fatalf("expected no error, got %v", err)
+	}
+	if version == 0 {
+		t.Fatal("expected non-zero version")
+	}
+}
+
+func TestEmbed_First_Empty(t *testing.T) {
+	emptyFS := embed.FS{}
+	e := &Embed{}
+	e.FS = emptyFS
+	e.path = "empty"
+	e.migrations = source.NewMigrations()
+	_, err := e.First()
+	if err == nil {
+		t.Fatal("expected error for empty migrations")
+	}
+}
+
+func TestEmbed_PrevNext(t *testing.T) {
+	driver, _ := NewEmbed(testFS, "testmigrations")
+	first, _ := driver.First()
+	_, err := driver.Prev(first)
+	if err == nil {
+		t.Fatal("expected error for prev of first migration")
+	}
+	next, err := driver.Next(first)
+	if err != nil {
+		t.Fatalf("expected no error for next, got %v", err)
+	}
+	if next == 0 {
+		t.Fatal("expected next version to be non-zero")
+	}
+}
+
+func TestEmbed_ReadUpDown(t *testing.T) {
+	driver, _ := NewEmbed(testFS, "testmigrations")
+	first, _ := driver.First()
+	r, id, err := driver.ReadUp(first)
+	if err != nil {
+		t.Fatalf("expected no error, got %v", err)
+	}
+	if r == nil || id == "" {
+		t.Fatal("expected valid reader and identifier")
+	}
+	b, err := io.ReadAll(r)
+	if err != nil {
+		t.Fatalf("failed to read: %v", err)
+	}
+	if len(b) == 0 {
+		t.Fatal("expected file content")
+	}
+	r.Close()
+
+	// Down migration may not exist for first, so test with next if available
+	next, _ := driver.Next(first)
+	rd, idd, err := driver.ReadDown(next)
+	if err == nil {
+		if rd == nil || idd == "" {
+			t.Fatal("expected valid reader and identifier for down")
+		}
+		rd.Close()
+	}
+}
+
+func TestEmbed_ReadUp_NotExist(t *testing.T) {
+	driver, _ := NewEmbed(testFS, "testmigrations")
+	_, _, err := driver.ReadUp(999999)
+	if err == nil {
+		t.Fatal("expected error for non-existent migration")
+	}
+	var pathErr *fs.PathError
+	if !errors.As(err, &pathErr) {
+		t.Fatalf("expected fs.PathError, got %T", err)
+	}
+}
+
+func TestEmbed_Close(t *testing.T) {
+	driver, _ := NewEmbed(testFS, "testmigrations")
+	if err := driver.Close(); err != nil {
+		t.Fatalf("expected no error, got %v", err)
+	}
+}
+
+func TestFileReader_ReadClose(t *testing.T) {
+	data := []byte("hello world")
+	fr := &fileReader{data: data}
+	buf := make([]byte, 5)
+	n, err := fr.Read(buf)
+	if n != 5 || err != nil {
+		t.Fatalf("expected to read 5 bytes, got %d, err %v", n, err)
+	}
+	n, err = fr.Read(buf)
+	if n != 5 || err != nil {
+		t.Fatalf("expected to read next 5 bytes, got %d, err %v", n, err)
+	}
+	n, err = fr.Read(buf)
+	if n != 1 || err != nil {
+		t.Fatalf("expected to read last byte, got %d, err %v", n, err)
+	}
+	n, err = fr.Read(buf)
+	if n != 0 || err != io.EOF {
+		t.Fatalf("expected EOF, got %d, err %v", n, err)
+	}
+	if err := fr.Close(); err != nil {
+		t.Fatalf("expected no error on close, got %v", err)
+	}
+}
+
+// createBenchmarkEmbed creates an Embed driver with test migrations
+// This is a helper function for benchmarks
+func createBenchmarkEmbed(b *testing.B) *Embed {
+	driver, err := NewEmbed(testFS, testPath)
+	if err != nil {
+		b.Fatal(err)
+	}
+	return driver.(*Embed)
+}
+
+func BenchmarkFirst(b *testing.B) {
+	e := createBenchmarkEmbed(b)
+	b.ResetTimer()
+	for n := 0; n < b.N; n++ {
+		_, err := e.First()
+		if err != nil {
+			b.Error(err)
+		}
+	}
+	b.StopTimer()
+}
+
+func BenchmarkNext(b *testing.B) {
+	e := createBenchmarkEmbed(b)
+	b.ResetTimer()
+	v, err := e.First()
+	for n := 0; n < b.N; n++ {
+		for !errors.Is(err, fs.ErrNotExist) {
+			v, err = e.Next(v)
+		}
+	}
+	b.StopTimer()
+}
