From eb41ec23f965052ebe8c2bd230493b4eb5ff0a0f Mon Sep 17 00:00:00 2001
From: youpy <9128+youpy@users.noreply.github.com>
Date: Wed, 22 Jun 2022 14:19:59 +0900
Subject: [PATCH] `Do` takes a context as its first argument

---
 migrator/migrator.go | 2 +-
 schema.go            | 7 +++++--
 schema_test.go       | 3 ++-
 3 files changed, 8 insertions(+), 4 deletions(-)

diff --git a/migrator/migrator.go b/migrator/migrator.go
index 9638f111..9b58f802 100644
--- a/migrator/migrator.go
+++ b/migrator/migrator.go
@@ -160,7 +160,7 @@ func (m *Migrator) run(ctx context.Context, migrations []rel.Migration) {
 	adapter := m.repo.Adapter(ctx)
 	for _, migration := range migrations {
 		if fn, ok := migration.(rel.Do); ok {
-			check(fn(m.repo))
+			check(fn(ctx, m.repo))
 		} else {
 			check(adapter.Apply(ctx, migration))
 		}
diff --git a/schema.go b/schema.go
index fe9fc8b1..b861d36a 100644
--- a/schema.go
+++ b/schema.go
@@ -1,6 +1,9 @@
 package rel
 
-import "strings"
+import (
+	"context"
+	"strings"
+)
 
 // SchemaOp type.
 type SchemaOp uint8
@@ -138,7 +141,7 @@ func (r Raw) internalMigration()       {}
 func (r Raw) internalTableDefinition() {}
 
 // Do used internally for schema migration.
-type Do func(Repository) error
+type Do func(context.Context, Repository) error
 
 func (d Do) description() string {
 	return "run go code"
diff --git a/schema_test.go b/schema_test.go
index ca93a22d..9e0b7219 100644
--- a/schema_test.go
+++ b/schema_test.go
@@ -1,6 +1,7 @@
 package rel
 
 import (
+	"context"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
@@ -221,7 +222,7 @@ func TestDo(t *testing.T) {
 		schema Schema
 	)
 
-	schema.Do(func(repo Repository) error { return nil })
+	schema.Do(func(ctx context.Context, repo Repository) error { return nil })
 	assert.NotNil(t, schema.Migrations[0])
 }