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

bigtable/cmd/cbt/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)

bigtable/cmd/cbt/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.
 

bigtable/cmd/cbt/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
+}

bigtable/cmd/cbt/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)
+	}
+}

bigtable/internal/cbtconfig/cbtconfig.go

@@ -47,6 +47,7 @@ type Config struct {
 	CertFile          string                           // optional
 	UserAgent         string                           // optional
 	AuthToken         string                           // optional
+	Timeout           time.Duration                    // optional
 	TokenSource       oauth2.TokenSource               // derived
 	TLSCreds          credentials.TransportCredentials // derived
 }
@@ -76,6 +77,8 @@ func (c *Config) RegisterFlags() {
 	flag.StringVar(&c.CertFile, "cert-file", c.CertFile, "Override the TLS certificates file")
 	flag.StringVar(&c.UserAgent, "user-agent", c.UserAgent, "Override the user agent string")
 	flag.StringVar(&c.AuthToken, "auth-token", c.AuthToken, "if set, use IAM Auth Token for requests")
+	flag.DurationVar(&c.Timeout, "timeout", c.Timeout,
+		"Timeout (e.g. 10s, 100ms, 5m )")
 }
 
 // CheckFlags checks that the required config values are set.
@@ -163,6 +166,12 @@ func readConfig(s *bufio.Scanner, filename string) (*Config, error) {
 			c.UserAgent = val
 		case "auth-token":
 			c.AuthToken = val
+		case "timeout":
+			timeout, err := time.ParseDuration(val)
+			if err != nil {
+				return nil, err
+			}
+			c.Timeout = timeout
 		}
 
 	}

bigtable/internal/cbtconfig/cbtconfig_test.go

@@ -21,6 +21,7 @@ import (
 	"fmt"
 	"strings"
 	"testing"
+	"time"
 )
 
 func TestReadConfig(t *testing.T) {
@@ -32,11 +33,13 @@ func TestReadConfig(t *testing.T) {
 	certificateFile := "test-certificate-file"
 	userAgent := "test-user-agent"
 	authToken := "test-auth-token="
+	timeout := time.Duration(42e9)
 	// Read configuration from string containing spaces, tabs and empty lines.
 	validConfig := fmt.Sprintf(`
         project=%s
         instance=%s
         creds=%s
+        timeout=42s
 
         admin-endpoint =%s
         data-endpoint= %s
@@ -72,6 +75,9 @@ func TestReadConfig(t *testing.T) {
 	if g, w := c.AuthToken, authToken; g != w {
 		t.Errorf("AuthToken mismatch\nGot: %s\nWant: %s", g, w)
 	}
+	if g, w := c.Timeout, timeout; g != w {
+		t.Errorf("AuthToken mismatch\nGot: %s\nWant: %s", g, w)
+	}
 
 	// Try to read an invalid config file and verify that it fails.
 	unknownKey := fmt.Sprintf("%s\nunknown-key=some-value", validConfig)