contextを引数に取るようにした refs #9

gcpug · May 23, 2019 · a77ebad · a77ebad
1 parent 7259354
commit a77ebad
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 7 deletions.
diff --git a/spanner/spanner_service.go b/spanner/spanner_service.go
@@ -57,8 +57,8 @@ func (s *SpannerEntityService) QueryToWrite(ctx context.Context, sql string, hea
 }
 
 type UpdateMutationer interface {
-	GetKey(row *spanner.Row) spanner.Key
-	GetMutation(row *spanner.Row) *spanner.Mutation
+	GetKey(ctx context.Context, row *spanner.Row) spanner.Key
+	GetMutation(ctx context.Context, row *spanner.Row) *spanner.Mutation
 }
 
 func (s *SpannerEntityService) UpdateExperiment(ctx context.Context, table string, columns []string, sql string, mutationer UpdateMutationer) (int, error) {
@@ -105,7 +105,7 @@ func (s *SpannerEntityService) UpdateExperiment(ctx context.Context, table strin
 func (s *SpannerEntityService) update(ctx context.Context, table string, columns []string, rows []*spanner.Row, mutationer UpdateMutationer) error {
 	var keys spanner.KeySet
 	for _, v := range rows {
-		key := mutationer.GetKey(v)
+		key := mutationer.GetKey(ctx, v)
 		keys = append(key)
 	}
 
@@ -122,7 +122,7 @@ func (s *SpannerEntityService) update(ctx context.Context, table string, columns
 			if err != nil {
 				return errors.WithStack(err)
 			}
-			m := mutationer.GetMutation(row)
+			m := mutationer.GetMutation(ctx, row)
 			ml = append(ml, m)
 		}
 

diff --git a/spanner/spanner_service_test.go b/spanner/spanner_service_test.go
@@ -3,12 +3,12 @@ package spanner_test
 import (
 	"context"
 	"fmt"
-	"github.com/google/uuid"
 	"testing"
 	"time"
 
 	"cloud.google.com/go/spanner"
 	shovels "github.com/gcpug/spshovel/spanner"
+	"github.com/google/uuid"
 )
 
 type SampleMutationer struct{}
@@ -18,15 +18,15 @@ var _ shovels.UpdateMutationer = &SampleMutationer{}
 const TableName = "TweetHashKey"
 const PrimaryKeyName = "Id"
 
-func (m *SampleMutationer) GetKey(row *spanner.Row) spanner.Key {
+func (m *SampleMutationer) GetKey(ctx context.Context, row *spanner.Row) spanner.Key {
 	var id string
 	if err := row.ColumnByName(PrimaryKeyName, &id); err != nil {
 		panic(err)
 	}
 	return spanner.Key{id}
 }
 
-func (m *SampleMutationer) GetMutation(row *spanner.Row) *spanner.Mutation {
+func (m *SampleMutationer) GetMutation(ctx context.Context, row *spanner.Row) *spanner.Mutation {
 	var id string
 	if err := row.ColumnByName(PrimaryKeyName, &id); err != nil {
 		panic(err)