feat(bigtable/cmd/cbt): Add a timeout option (#4276)

cbt.go

@@ -49,6 +49,7 @@ var (
 
 	config              *cbtconfig.Config
 	client              *bigtable.Client
+	table               tableLike
 	adminClient         *bigtable.AdminClient
 	instanceAdminClient *bigtable.InstanceAdminClient
 
@@ -58,6 +59,10 @@ var (
 	cliUserAgent = "cbt-cli-go/unknown"
 )
 
+type tableLike interface {
+	ReadRows(ctx context.Context, arg bigtable.RowSet, f func(bigtable.Row) bool, opts ...bigtable.ReadOption) (err error)
+}
+
 func getCredentialOpts(opts []option.ClientOption) []option.ClientOption {
 	if ts := config.TokenSource; ts != nil {
 		opts = append(opts, option.WithTokenSource(ts))
@@ -85,6 +90,14 @@ func getClient(clientConf bigtable.ClientConfig) *bigtable.Client {
 	return client
 }
 
+func getTable(clientConf bigtable.ClientConfig, tableName string) tableLike {
+	if table != nil {
+		return table
+	}
+	table = getClient(clientConf).Open(tableName)
+	return table
+}
+
 func getAdminClient() *bigtable.AdminClient {
 	if adminClient == nil {
 		var opts []option.ClientOption
@@ -146,25 +159,37 @@ func main() {
 		os.Stdout = f
 	}
 
+	doMain(config, flag.Args())
+}
+
+func doMain(config *cbtconfig.Config, args []string) {
 	if config.UserAgent != "" {
 		cliUserAgent = config.UserAgent
 	}
 
-	ctx := context.Background()
+	var ctx context.Context
+	if config.Timeout > 0 {
+		var cancel context.CancelFunc
+		ctx, cancel = context.WithTimeout(context.Background(), config.Timeout)
+		defer cancel()
+	} else {
+		ctx = context.Background()
+	}
+
 	if config.AuthToken != "" {
 		ctx = metadata.AppendToOutgoingContext(ctx, "x-goog-iam-authorization-token", config.AuthToken)
 	}
 
 	for _, cmd := range commands {
-		if cmd.Name == flag.Arg(0) {
+		if cmd.Name == args[0] {
 			if err := config.CheckFlags(cmd.Required); err != nil {
 				log.Fatal(err)
 			}
-			cmd.do(ctx, flag.Args()[1:]...)
+			cmd.do(ctx, args[1:]...)
 			return
 		}
 	}
-	log.Fatalf("Unknown command %q", flag.Arg(0))
+	log.Fatalf("Unknown command %q", args[0])
 }
 
 func usage(w io.Writer) {
@@ -212,6 +237,7 @@ options to your ~/.cbtrc file in the following format:
     admin-endpoint = hostname:port
     data-endpoint = hostname:port
     auth-token = AJAvW039NO1nDcijk_J6_rFXG_...
+    timeout = 30s
 
 All values are optional and can be overridden at the command prompt.
 `
@@ -581,7 +607,7 @@ func doCount(ctx context.Context, args ...string) {
 	if len(args) != 1 {
 		log.Fatal("usage: cbt count <table>")
 	}
-	tbl := getClient(bigtable.ClientConfig{}).Open(args[0])
+	tbl := getTable(bigtable.ClientConfig{}, args[0])
 
 	n := 0
 	err := tbl.ReadRows(ctx, bigtable.InfiniteRange(""), func(_ bigtable.Row) bool {
@@ -841,7 +867,7 @@ func doMDDoc(ctx context.Context, args ...string) { doMDDocFn(ctx, args...) }
 func docFlags() []*flag.Flag {
 	// Only include specific flags, in a specific order.
 	var flags []*flag.Flag
-	for _, name := range []string{"project", "instance", "creds"} {
+	for _, name := range []string{"project", "instance", "creds", "timeout"} {
 		f := flag.Lookup(name)
 		if f == nil {
 			log.Fatalf("Flag not linked: -%s", name)

cbtdoc.go

@@ -71,6 +71,8 @@ The options are:
         Cloud Bigtable instance
     -creds string
         Path to the credentials file. If set, uses the application credentials in this file
+    -timeout string
+        Timeout (e.g. 10s, 100ms, 5m )
 
 Example:  cbt -instance=my-instance ls
 
@@ -100,6 +102,7 @@ options to your ~/.cbtrc file in the following format:
     admin-endpoint = hostname:port
     data-endpoint = hostname:port
     auth-token = AJAvW039NO1nDcijk_J6_rFXG_...
+    timeout = 30s
 
 All values are optional and can be overridden at the command prompt.
 

testing.go

@@ -0,0 +1,50 @@
+/*
+Copyright 2021 Google LLC
+
+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.
+*/
+
+package main
+
+import (
+	"bytes"
+	"io"
+	"os"
+)
+
+func captureStdout(f func()) string {
+	/*
+	   Capture standard output to facilitate testing code that prints
+
+	   or useless print output in running tests.
+	*/
+	saved := os.Stdout
+	r, w, _ := os.Pipe()
+	os.Stdout = w
+	defer func() { os.Stdout = saved }()
+
+	outC := make(chan string)
+	// https://stackoverflow.com/questions/10473800/in-go-how-do-i-capture-stdout-of-a-function-into-a-string
+	// copy the output in a separate goroutine so printing can't block indefinitely
+	go func() {
+		var buf bytes.Buffer
+		io.Copy(&buf, r)
+		outC <- buf.String()
+	}()
+
+	f()
+
+	// back to normal state
+	w.Close()
+	return <-outC
+}

timeout_test.go

@@ -0,0 +1,67 @@
+/*
+Copyright 2021 Google LLC
+
+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.
+*/
+
+package main
+
+import (
+	"context"
+	"testing"
+	"time"
+
+	"cloud.google.com/go/bigtable"
+	"cloud.google.com/go/bigtable/internal/cbtconfig"
+)
+
+type ctxTable struct {
+	ctx context.Context
+}
+
+func (ct *ctxTable) ReadRows(
+	ctx context.Context,
+	arg bigtable.RowSet,
+	f func(bigtable.Row) bool,
+	opts ...bigtable.ReadOption,
+) (err error) {
+	ct.ctx = ctx
+	return nil
+}
+
+func TestTimeout(t *testing.T) {
+	ctxt := ctxTable{}
+	table = &ctxt
+	defer func() { table = nil }()
+
+	config := cbtconfig.Config{Creds: "c", Project: "p", Instance: "i"}
+	captureStdout(func() { doMain(&config, []string{"count", "mytable"}) })
+
+	_, deadlineSet := ctxt.ctx.Deadline()
+	if deadlineSet {
+		t.Errorf("Deadline set with no timeout in config")
+	}
+
+	config.Timeout = time.Duration(42e9)
+	now := time.Now()
+	captureStdout(func() { doMain(&config, []string{"count", "mytable"}) })
+
+	deadline, deadlineSet := ctxt.ctx.Deadline()
+	if !deadlineSet {
+		t.Errorf("No deadline set, even though the config set one")
+	}
+	timeout := deadline.Sub(now).Nanoseconds()
+	if !(timeout > 42e9 && timeout < 43e9) {
+		t.Errorf("Bad actual timeout nanoseconds %d", timeout)
+	}
+}