add cmd.CustomFileSystemFunc function

cmd/apply.go

@@ -23,8 +23,7 @@ import (
 	"context"
 	"errors"
 	"fmt"
-	"os"
-
+	"github.com/cloudspannerecosystem/wrench/internal/fs"
 	"github.com/cloudspannerecosystem/wrench/pkg/spanner"
 	"github.com/spf13/cobra"
 )
@@ -57,7 +56,7 @@ func apply(c *cobra.Command, _ []string) error {
 			return errors.New("cannot specify DDL and DML at same time")
 		}
 
-		ddl, err := os.ReadFile(ddlFile)
+		ddl, err := fs.ReadFile(ctx, ddlFile)
 		if err != nil {
 			return &Error{
 				err: err,
@@ -81,7 +80,7 @@ func apply(c *cobra.Command, _ []string) error {
 	}
 
 	// apply dml
-	dml, err := os.ReadFile(dmlFile)
+	dml, err := fs.ReadFile(ctx, dmlFile)
 	if err != nil {
 		return &Error{
 			err: err,

cmd/create.go

@@ -21,8 +21,7 @@ package cmd
 
 import (
 	"context"
-	"os"
-
+	"github.com/cloudspannerecosystem/wrench/internal/fs"
 	"github.com/spf13/cobra"
 )
 
@@ -43,7 +42,7 @@ func create(c *cobra.Command, _ []string) error {
 	defer client.Close()
 
 	filename := schemaFilePath(c)
-	ddl, err := os.ReadFile(filename)
+	ddl, err := fs.ReadFile(ctx, filename)
 	if err != nil {
 		return &Error{
 			err: err,

cmd/root.go

@@ -21,10 +21,12 @@ package cmd
 
 import (
 	"context"
+	"io/fs"
 	"os"
 	"runtime/debug"
 	"time"
 
+	wrenchfs "github.com/cloudspannerecosystem/wrench/internal/fs"
 	"github.com/spf13/cobra"
 )
 
@@ -44,8 +46,20 @@ var (
 	timeout         time.Duration
 )
 
+// CustomFileSystemFunc is a function that returns a custom fs.FS.
+// This variable allows customizing what kind of fs.FS should be use in wrench CLI execution.
+// e.g. embed.FS for use.
+var CustomFileSystemFunc func() fs.FS
+
 var rootCmd = &cobra.Command{
 	Use: "wrench",
+	PersistentPreRun: func(cmd *cobra.Command, args []string) {
+		if CustomFileSystemFunc != nil {
+			ctx := cmd.Context()
+			ctx = wrenchfs.WithContext(ctx, CustomFileSystemFunc())
+			cmd.SetContext(ctx)
+		}
+	},
 }
 
 func Execute(ctx context.Context) error {

internal/fs/fs.go

@@ -0,0 +1,31 @@
+package fs
+
+import (
+	"context"
+	"io/fs"
+	"os"
+)
+
+type contextFSKey struct{}
+
+func WithContext(ctx context.Context, fsys fs.FS) context.Context {
+	return context.WithValue(ctx, contextFSKey{}, fsys)
+}
+
+func FromContext(ctx context.Context) fs.FS {
+	fsys, ok := ctx.Value(contextFSKey{}).(fs.FS)
+	if ok && fsys != nil {
+		return fsys
+	}
+	return os.DirFS(".")
+}
+
+func ReadFile(ctx context.Context, path string) ([]byte, error) {
+	fsys := FromContext(ctx)
+	return fs.ReadFile(fsys, path)
+}
+
+func ReadDir(ctx context.Context, path string) ([]fs.DirEntry, error) {
+	fsys := FromContext(ctx)
+	return fs.ReadDir(fsys, path)
+}

pkg/spanner/client_test.go

@@ -219,7 +219,7 @@ func TestExecuteMigrations(t *testing.T) {
 		t.Fatalf("failed to apply mutation: %v", err)
 	}
 
-	migrations, err := LoadMigrations("testdata/migrations")
+	migrations, err := ReadMigrations(ctx, "testdata/migrations")
 	if err != nil {
 		t.Fatalf("failed to load migrations: %v", err)
 	}

pkg/spanner/migration.go

@@ -20,12 +20,14 @@
 package spanner
 
 import (
+	"context"
 	"errors"
 	"fmt"
-	"os"
 	"path/filepath"
 	"regexp"
 	"strconv"
+
+	"github.com/cloudspannerecosystem/wrench/internal/fs"
 )
 
 var (
@@ -76,8 +78,8 @@ func (ms Migrations) Less(i, j int) bool {
 	return ms[i].Version < ms[j].Version
 }
 
-func LoadMigrations(dir string) (Migrations, error) {
-	files, err := os.ReadDir(dir)
+func ReadMigrations(ctx context.Context, dir string) (Migrations, error) {
+	files, err := fs.ReadDir(ctx, dir)
 	if err != nil {
 		return nil, err
 	}
@@ -103,7 +105,7 @@ func LoadMigrations(dir string) (Migrations, error) {
 			continue
 		}
 
-		file, err := os.ReadFile(filepath.Join(dir, filename))
+		file, err := fs.ReadFile(ctx, filepath.Join(dir, filename))
 		if err != nil {
 			continue
 		}
@@ -138,6 +140,11 @@ func LoadMigrations(dir string) (Migrations, error) {
 	return migrations, nil
 }
 
+// Deprecated: use ReadMigrations instead.
+func LoadMigrations(dir string) (Migrations, error) {
+	return ReadMigrations(context.Background(), dir)
+}
+
 func ddlToStatements(filename string, data []byte) ([]string, error) {
 	return toStatements(filename, data)
 }
@@ -174,4 +181,3 @@ func inspectStatementsKind(statements []string) (statementKind, error) {
 		return "", errors.New("DDL, DML (INSERT), and partitioned DML (UPDATE or DELETE) must not be combined in the same migration file")
 	}
 }
-

pkg/spanner/migration_test.go

@@ -20,14 +20,17 @@
 package spanner_test
 
 import (
+	"context"
 	"path/filepath"
 	"testing"
 
 	"github.com/cloudspannerecosystem/wrench/pkg/spanner"
 )
 
 func TestLoadMigrations(t *testing.T) {
-	ms, err := spanner.LoadMigrations(filepath.Join("testdata", "migrations"))
+	ctx := context.Background()
+
+	ms, err := spanner.ReadMigrations(ctx, filepath.Join("testdata", "migrations"))
 	if err != nil {
 		t.Fatal(err)
 	}