Add a flag to set context timeout (#64)

As you can see in the following code, gRPC client defines 30 seconds. https://github.com/googleapis/google-cloud-go/blob/spanner/v1.41.0/spanner/apiv1/spanner_client.go#L551-L555 It's too short in some cases.
cloudspannerecosystem · Dec 9, 2022 · 58c2c06 · 58c2c06
1 parent bd910ea
commit 58c2c06
Show file tree

Hide file tree

Showing 9 changed files with 35 additions and 13 deletions.
diff --git a/cmd/apply.go b/cmd/apply.go
@@ -23,6 +23,7 @@ import (
 	"errors"
 	"fmt"
 	"io/ioutil"
+	"context"
 
 	"github.com/cloudspannerecosystem/wrench/pkg/spanner"
 	"github.com/spf13/cobra"
@@ -42,7 +43,8 @@ var applyCmd = &cobra.Command{
 }
 
 func apply(c *cobra.Command, _ []string) error {
-	ctx := c.Context()
+	ctx, cancel := context.WithTimeout(c.Context(), timeout)
+	defer cancel()
 
 	client, err := newSpannerClient(ctx, c)
 	if err != nil {

diff --git a/cmd/cmd.go b/cmd/cmd.go
@@ -40,6 +40,7 @@ const (
 	flagPartitioned       = "partitioned"
 	flagPriority          = "priority"
 	flagNode              = "node"
+	flagTimeout           = "timeout"
 	defaultSchemaFileName = "schema.sql"
 )
 

diff --git a/cmd/create.go b/cmd/create.go
@@ -20,6 +20,7 @@
 package cmd
 
 import (
+	"context"
 	"io/ioutil"
 
 	"github.com/spf13/cobra"
@@ -32,7 +33,8 @@ var createCmd = &cobra.Command{
 }
 
 func create(c *cobra.Command, _ []string) error {
-	ctx := c.Context()
+	ctx, cancel := context.WithTimeout(c.Context(), timeout)
+	defer cancel()
 
 	client, err := newSpannerClient(ctx, c)
 	if err != nil {

diff --git a/cmd/drop.go b/cmd/drop.go
@@ -20,6 +20,8 @@
 package cmd
 
 import (
+	"context"
+
 	"github.com/spf13/cobra"
 )
 
@@ -30,7 +32,8 @@ var dropCmd = &cobra.Command{
 }
 
 func drop(c *cobra.Command, _ []string) error {
-	ctx := c.Context()
+	ctx, cancel := context.WithTimeout(c.Context(), timeout)
+	defer cancel()
 
 	client, err := newSpannerClient(ctx, c)
 	if err != nil {

diff --git a/cmd/instance.go b/cmd/instance.go
@@ -19,7 +19,8 @@ var instanceCreateCmd = &cobra.Command{
 }
 
 func instanceCreate(c *cobra.Command, _ []string) error {
-	ctx := c.Context()
+	ctx, cancel := context.WithTimeout(c.Context(), timeout)
+	defer cancel()
 
 	client, err := newSpannerAdminClient(ctx, c)
 	if err != nil {
@@ -46,7 +47,8 @@ var instanceDeleteCmd = &cobra.Command{
 }
 
 func instanceDelete(c *cobra.Command, _ []string) error {
-	ctx := context.Background()
+	ctx, cancel := context.WithTimeout(c.Context(), timeout)
+	defer cancel()
 
 	client, err := newSpannerAdminClient(ctx, c)
 	if err != nil {

diff --git a/cmd/load.go b/cmd/load.go
@@ -20,6 +20,7 @@
 package cmd
 
 import (
+	"context"
 	"io/ioutil"
 
 	"github.com/spf13/cobra"
@@ -31,8 +32,9 @@ var loadCmd = &cobra.Command{
 	RunE:  load,
 }
 
-func load(c *cobra.Command, args []string) error {
-	ctx := c.Context()
+func load(c *cobra.Command, _ []string) error {
+	ctx, cancel := context.WithTimeout(c.Context(), timeout)
+	defer cancel()
 
 	client, err := newSpannerClient(ctx, c)
 	if err != nil {
@@ -48,7 +50,7 @@ func load(c *cobra.Command, args []string) error {
 		}
 	}
 
-	err = ioutil.WriteFile(schemaFilePath(c), ddl, 0664)
+	err = ioutil.WriteFile(schemaFilePath(c), ddl, 0o664)
 	if err != nil {
 		return &Error{
 			err: err,

diff --git a/cmd/migrate.go b/cmd/migrate.go
@@ -20,6 +20,7 @@
 package cmd
 
 import (
+	"context"
 	"errors"
 	"fmt"
 	"os"
@@ -106,7 +107,8 @@ func migrateCreate(c *cobra.Command, args []string) error {
 }
 
 func migrateUp(c *cobra.Command, args []string) error {
-	ctx := c.Context()
+	ctx, cancel := context.WithTimeout(c.Context(), timeout)
+	defer cancel()
 
 	limit := -1
 	if len(args) > 0 {
@@ -145,8 +147,9 @@ func migrateUp(c *cobra.Command, args []string) error {
 	return client.ExecuteMigrations(ctx, migrations, limit, migrationTableName)
 }
 
-func migrateVersion(c *cobra.Command, args []string) error {
-	ctx := c.Context()
+func migrateVersion(c *cobra.Command, _ []string) error {
+	ctx, cancel := context.WithTimeout(c.Context(), timeout)
+	defer cancel()
 
 	client, err := newSpannerClient(ctx, c)
 	if err != nil {
@@ -180,7 +183,8 @@ func migrateVersion(c *cobra.Command, args []string) error {
 }
 
 func migrateSet(c *cobra.Command, args []string) error {
-	ctx := c.Context()
+	ctx, cancel := context.WithTimeout(c.Context(), timeout)
+	defer cancel()
 
 	if len(args) == 0 {
 		return &Error{

diff --git a/cmd/root.go b/cmd/root.go
@@ -22,6 +22,7 @@ package cmd
 import (
 	"context"
 	"os"
+	"time"
 
 	"github.com/spf13/cobra"
 )
@@ -39,6 +40,7 @@ var (
 	directory       string
 	schemaFile      string
 	credentialsFile string
+	timeout         time.Duration
 )
 
 var rootCmd = &cobra.Command{
@@ -70,6 +72,7 @@ func init() {
 	rootCmd.PersistentFlags().StringVar(&directory, flagNameDirectory, "", "Directory that schema file placed (required)")
 	rootCmd.PersistentFlags().StringVar(&schemaFile, flagNameSchemaFile, "", "Name of schema file (optional. if not set, will use default 'schema.sql' file name)")
 	rootCmd.PersistentFlags().StringVar(&credentialsFile, flagCredentialsFile, "", "Specify Credentials File")
+	rootCmd.PersistentFlags().DurationVar(&timeout, flagTimeout, time.Hour, "Context timeout")
 
 	rootCmd.Version = Version
 	rootCmd.SetVersionTemplate(versionTemplate)

diff --git a/cmd/truncate.go b/cmd/truncate.go
@@ -20,6 +20,8 @@
 package cmd
 
 import (
+	"context"
+
 	"github.com/spf13/cobra"
 )
 
@@ -30,7 +32,8 @@ var truncateCmd = &cobra.Command{
 }
 
 func truncate(c *cobra.Command, _ []string) error {
-	ctx := c.Context()
+	ctx, cancel := context.WithTimeout(c.Context(), timeout)
+	defer cancel()
 
 	client, err := newSpannerClient(ctx, c)
 	if err != nil {