Get batch arg, when looking for EndTransaction, just look at the last entry

Author: veteranlu

roachpb/batch.go

@@ -60,10 +60,20 @@ func (ba *BatchRequest) IsRange() bool {
 	return (ba.flags() & isRange) != 0
 }
 
-// GetArg returns the first request of the given type, if possible.
+// GetArg returns the request of the given type.
+// If the given type is EndTransaction, just look at the last entry.
+// or else, returns the first request which match the given type
 func (ba *BatchRequest) GetArg(method Method) (Request, bool) {
-	// TODO(tschottdorf): when looking for EndTransaction, just look at the
-	// last entry.
+	// when looking for EndTransaction, just look at the last entry.
+	if method == EndTransaction {
+		if length := len(ba.Requests); length > 0 {
+			if req := ba.Requests[length-1].GetInner(); req.Method() == EndTransaction {
+				return req, true
+			}
+		}
+		return nil, false
+	}
+
 	for _, arg := range ba.Requests {
 		if req := arg.GetInner(); req.Method() == method {
 			return req, true

roachpb/batch_test.go

@@ -0,0 +1,41 @@
+// Copyright 2014 The Cockroach Authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+// implied. See the License for the specific language governing
+// permissions and limitations under the License. See the AUTHORS file
+// for names of contributors.
+//
+// Author: Veteran Lu (23907238@qq.com)
+
+package roachpb
+
+import "testing"
+
+// TestBatchRequestGetArg tests that the method for BatchRequest.GetArg
+func TestBatchRequestGetArg(t *testing.T) {
+	testCases := []struct {
+		bu  []RequestUnion
+		exp bool
+	}{
+		{[]RequestUnion{}, false},
+		{[]RequestUnion{{Get: &GetRequest{}}}, false},
+		{[]RequestUnion{{EndTransaction: &EndTransactionRequest{}}, {Get: &GetRequest{}}}, false},
+		{[]RequestUnion{{EndTransaction: &EndTransactionRequest{}}}, true},
+		{[]RequestUnion{{Get: &GetRequest{}}, {EndTransaction: &EndTransactionRequest{}}}, true},
+	}
+
+	for i, c := range testCases {
+		br := BatchRequest{Requests: c.bu}
+		if _, r := br.GetArg(EndTransaction); r != c.exp {
+			t.Errorf("%d: unexpected next bytes for %v: %v", i, c.bu, c.exp)
+		}
+	}
+}