sliceをごばっと取るようにした refs #9

gcpug · May 23, 2019 · eb62bec · eb62bec
1 parent 4d226c9
commit eb62bec
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 8 deletions.
diff --git a/spanner/spanner_service.go b/spanner/spanner_service.go
@@ -65,7 +65,8 @@ type UpdateMutationer interface {
 func (s *SpannerEntityService) UpdateExperiment(ctx context.Context, table string, columns []string, sql string, mutationer UpdateMutationer) (int, error) {
 	q := spanner.NewStatement(sql)
 
-	var rows []*spanner.Row
+	const bufferCount = 1000
+	rows := make([]*spanner.Row, 0, bufferCount)
 
 	var count int
 	tx, err := s.sc.BatchReadOnlyTransaction(ctx, spanner.ReadTimestamp(time.Now()))
@@ -87,11 +88,11 @@ func (s *SpannerEntityService) UpdateExperiment(ctx context.Context, table strin
 		count++
 		rows = append(rows, row)
 
-		if count%500 == 0 {
+		if count%bufferCount == 0 {
 			if err := s.update(ctx, table, columns, rows, mutationer); err != nil {
 				return 0, errors.WithStack(err)
 			}
-			rows = []*spanner.Row{}
+			rows = make([]*spanner.Row, 0, bufferCount)
 		}
 	}
 	if len(rows) > 0 {
@@ -104,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
+	keys := spanner.KeySets()
 	for _, v := range rows {
 		key, err := mutationer.GetKey(ctx, v)
 		if err != nil {
@@ -114,11 +115,11 @@ func (s *SpannerEntityService) update(ctx context.Context, table string, columns
 			// TODO 専用のerrorを作る sinmetal
 			return fmt.Errorf("key is required")
 		}
-		keys = append(key)
+		keys = spanner.KeySets(keys, key)
 	}
 
 	_, err := s.sc.ReadWriteTransaction(ctx, func(ctx context.Context, txn *spanner.ReadWriteTransaction) error {
-		var ml []*spanner.Mutation
+		ml := make([]*spanner.Mutation, 0, len(rows))
 
 		iter := txn.Read(ctx, table, keys, columns)
 		defer iter.Stop()

diff --git a/spanner/spanner_service_test.go b/spanner/spanner_service_test.go
@@ -39,14 +39,16 @@ func TestSpannerEntityService_UpdateExperiment(t *testing.T) {
 	ctx := context.Background()
 	client := shovels.NewClient(ctx, "projects/gcpug-public-spanner/instances/merpay-sponsored-instance/databases/sinmetal")
 	s := shovels.NewSpannerEntityService(client)
-	count, err := s.UpdateExperiment(ctx, TableName, []string{PrimaryKeyName}, "SELECT Id From TweetHashKey", &SampleMutationer{})
+	count, err := s.UpdateExperiment(ctx, TableName, []string{PrimaryKeyName}, `SELECT Id From TweetHashKey WHERE Author != "sinmetal"`, &SampleMutationer{})
 	if err != nil {
 		t.Fatal(err)
 	}
 	fmt.Printf("UPDATE Count : %d\n", count)
 }
 
 func TestDataInitialize(t *testing.T) {
+	t.SkipNow()
+
 	ctx := context.Background()
 	client := shovels.NewClient(ctx, "projects/gcpug-public-spanner/instances/merpay-sponsored-instance/databases/sinmetal")
 	var ml []*spanner.Mutation
@@ -60,7 +62,6 @@ func TestDataInitialize(t *testing.T) {
 				t.Fatal(err)
 			}
 			ml = []*spanner.Mutation{}
-			fmt.Printf("INSERT COUNT : %d\n", i)
 		}
 	}