Print mutation count in verbose mode (#107)

* Support commit stats

* Update go.mod

* Refactor

* Fix tests for mutation count

* Fix wording

* Narrow variable scope

cli.go

@@ -398,12 +398,18 @@ func resultLine(result *Result, verbose bool) string {
 			// See https://cloud.google.com/spanner/docs/reference/rpc/google.spanner.v1?hl=en#resultsetstats
 			affectedRowsPrefix = "at least "
 		}
-		if verbose && timestamp != "" {
-			return fmt.Sprintf("Query OK, %s%d rows affected (%s)\ntimestamp: %s\n",
-				affectedRowsPrefix, result.AffectedRows, result.Stats.ElapsedTime, timestamp)
+
+		var detail string
+		if verbose {
+			if timestamp != "" {
+				detail += fmt.Sprintf("timestamp:      %s\n", timestamp)
+			}
+			if result.CommitStats != nil {
+				detail += fmt.Sprintf("mutation_count: %d\n", result.CommitStats.GetMutationCount())
+			}
 		}
-		return fmt.Sprintf("Query OK, %s%d rows affected (%s)\n",
-			affectedRowsPrefix, result.AffectedRows, result.Stats.ElapsedTime)
+		return fmt.Sprintf("Query OK, %s%d rows affected (%s)\n%s",
+			affectedRowsPrefix, result.AffectedRows, result.Stats.ElapsedTime, detail)
 	}
 
 	var set string

cli_test.go

@@ -27,6 +27,7 @@ import (
 
 	"github.com/chzyer/readline"
 	"github.com/google/go-cmp/cmp"
+	sppb "google.golang.org/genproto/googleapis/spanner/v1"
 )
 
 type nopCloser struct {
@@ -259,7 +260,21 @@ func TestResultLine(t *testing.T) {
 				Timestamp: ts,
 			},
 			verbose: true,
-			want:    fmt.Sprintf("Query OK, 3 rows affected (10 msec)\ntimestamp: %s\n", timestamp),
+			want:    fmt.Sprintf("Query OK, 3 rows affected (10 msec)\ntimestamp:      %s\n", timestamp),
+		},
+		{
+			desc: "mutation in verbose mode (both of timestamp and mutation count exist)",
+			result: &Result{
+				AffectedRows: 3,
+				IsMutation:   true,
+				Stats: QueryStats{
+					ElapsedTime: "10 msec",
+				},
+				CommitStats: &sppb.CommitResponse_CommitStats{MutationCount: 6},
+				Timestamp:   ts,
+			},
+			verbose: true,
+			want:    fmt.Sprintf("Query OK, 3 rows affected (10 msec)\ntimestamp:      %s\nmutation_count: 6\n", timestamp),
 		},
 		{
 			desc: "mutation in verbose mode (timestamp not exist)",

go.mod

@@ -3,16 +3,16 @@ module github.com/cloudspannerecosystem/spanner-cli
 go 1.13
 
 require (
-	cloud.google.com/go v0.65.0
-	cloud.google.com/go/spanner v1.10.0
+	cloud.google.com/go v0.75.0
+	cloud.google.com/go/spanner v1.14.0
 	github.com/chzyer/readline v0.0.0-20180603132655-2972be24d48e
-	github.com/google/go-cmp v0.5.2
+	github.com/google/go-cmp v0.5.4
 	github.com/jessevdk/go-flags v1.4.0
 	github.com/mattn/go-runewidth v0.0.8 // indirect
 	github.com/olekukonko/tablewriter v0.0.4
 	github.com/xlab/treeprint v1.0.1-0.20200715141336-10e0bc383e01
-	google.golang.org/api v0.31.0
-	google.golang.org/genproto v0.0.0-20200904004341-0bd0a958aa1d
-	google.golang.org/grpc v1.31.1
+	google.golang.org/api v0.39.0
+	google.golang.org/genproto v0.0.0-20210207032614-bba0dbe2a9ea
+	google.golang.org/grpc v1.35.0
 	google.golang.org/protobuf v1.25.0
 )

integration_test.go

@@ -143,6 +143,8 @@ func compareResult(t *testing.T, got *Result, expected *Result) {
 	opts := []cmp.Option{
 		cmpopts.IgnoreFields(Result{}, "Stats"),
 		cmpopts.IgnoreFields(Result{}, "Timestamp"),
+		// Commit Stats is only provided by real instances
+		cmpopts.IgnoreFields(Result{}, "CommitStats"),
 	}
 	if !cmp.Equal(got, expected, opts...) {
 		t.Errorf("diff: %s", cmp.Diff(got, expected, opts...))

session.go

@@ -106,7 +106,8 @@ func (s *Session) BeginReadWriteTransaction() error {
 		return errors.New("read-write transaction is already running")
 	}
 
-	txn, err := spanner.NewReadWriteStmtBasedTransaction(s.ctx, s.client)
+	txn, err := spanner.NewReadWriteStmtBasedTransactionWithOptions(s.ctx, s.client,
+		spanner.TransactionOptions{CommitOptions: spanner.CommitOptions{ReturnCommitStats: true}})
 	if err != nil {
 		return err
 	}
@@ -115,18 +116,18 @@ func (s *Session) BeginReadWriteTransaction() error {
 }
 
 // CommitReadWriteTransaction commits read-write transaction and returns commit timestamp if successful.
-func (s *Session) CommitReadWriteTransaction() (time.Time, error) {
+func (s *Session) CommitReadWriteTransaction() (spanner.CommitResponse, error) {
 	if s.rwTxn == nil {
-		return time.Time{}, errors.New("read-write transaction is not running")
+		return spanner.CommitResponse{}, errors.New("read-write transaction is not running")
 	}
 
 	s.rwTxnMutex.Lock()
 	defer s.rwTxnMutex.Unlock()
 
-	ts, err := s.rwTxn.Commit(s.ctx)
+	resp, err := s.rwTxn.CommitWithReturnResp(s.ctx)
 	s.rwTxn = nil
 	s.sendHeartbeat = false
-	return ts, err
+	return resp, err
 }
 
 // RollbackReadWriteTransaction rollbacks read-write transaction.

statement.go

@@ -60,6 +60,7 @@ type Result struct {
 	IsMutation       bool
 	Timestamp        time.Time
 	ForceVerbose     bool
+	CommitStats      *pb.CommitResponse_CommitStats
 }
 
 type Row struct {
@@ -754,6 +755,7 @@ func (s *DmlStatement) Execute(session *Session) (*Result, error) {
 			return nil, err
 		}
 		result.Timestamp = txnResult.Timestamp
+		result.CommitStats = txnResult.CommitStats
 	}
 
 	result.AffectedRows = int(numRows)
@@ -920,12 +922,13 @@ func (s *CommitStatement) Execute(session *Session) (*Result, error) {
 		return result, nil
 	}
 
-	ts, err := session.CommitReadWriteTransaction()
+	resp, err := session.CommitReadWriteTransaction()
 	if err != nil {
 		return nil, err
 	}
 
-	result.Timestamp = ts
+	result.Timestamp = resp.CommitTs
+	result.CommitStats = resp.CommitStats
 	return result, nil
 }