From ae8a9a103f380917ab2c8f7ccaddbe5f40670a7a Mon Sep 17 00:00:00 2001 From: Jim Fulton Date: Tue, 10 Aug 2021 16:18:42 -0400 Subject: [PATCH] feat(bigtable/cmd/cbt): Add a timeout option (#4276) --- bigtable/cmd/cbt/cbt.go | 38 +++++++++-- bigtable/cmd/cbt/cbtdoc.go | 3 + bigtable/cmd/cbt/testing.go | 50 ++++++++++++++ bigtable/cmd/cbt/timeout_test.go | 67 +++++++++++++++++++ bigtable/internal/cbtconfig/cbtconfig.go | 9 +++ bigtable/internal/cbtconfig/cbtconfig_test.go | 6 ++ 6 files changed, 167 insertions(+), 6 deletions(-) create mode 100644 bigtable/cmd/cbt/testing.go create mode 100644 bigtable/cmd/cbt/timeout_test.go diff --git a/bigtable/cmd/cbt/cbt.go b/bigtable/cmd/cbt/cbt.go index 6af809f8c63..a68ddc25cc1 100644 --- a/bigtable/cmd/cbt/cbt.go +++ b/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 ") } - 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) diff --git a/bigtable/cmd/cbt/cbtdoc.go b/bigtable/cmd/cbt/cbtdoc.go index e95c90f9801..4dc0f9b3a99 100644 --- a/bigtable/cmd/cbt/cbtdoc.go +++ b/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. diff --git a/bigtable/cmd/cbt/testing.go b/bigtable/cmd/cbt/testing.go new file mode 100644 index 00000000000..5c596f58a63 --- /dev/null +++ b/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 +} diff --git a/bigtable/cmd/cbt/timeout_test.go b/bigtable/cmd/cbt/timeout_test.go new file mode 100644 index 00000000000..3834e67040d --- /dev/null +++ b/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) + } +} diff --git a/bigtable/internal/cbtconfig/cbtconfig.go b/bigtable/internal/cbtconfig/cbtconfig.go index 7508c0ce64e..d7b4cdeb4f8 100644 --- a/bigtable/internal/cbtconfig/cbtconfig.go +++ b/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 } } diff --git a/bigtable/internal/cbtconfig/cbtconfig_test.go b/bigtable/internal/cbtconfig/cbtconfig_test.go index 4e2f23d21df..1d570a42765 100644 --- a/bigtable/internal/cbtconfig/cbtconfig_test.go +++ b/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)