configurable timeout and CLI flag

canonical · May 13, 2024 · 556cbef · 556cbef
1 parent 034248c
commit 556cbef
Showing 1 changed file with 7 additions and 6 deletions.
diff --git a/cmd/dqlite/dqlite.go b/cmd/dqlite/dqlite.go
@@ -11,20 +11,20 @@ import (
 	"strings"
 	"time"
 
+	"github.com/peterh/liner"
+	"github.com/spf13/cobra"
+
 	"github.com/canonical/go-dqlite/app"
 	"github.com/canonical/go-dqlite/client"
 	"github.com/canonical/go-dqlite/internal/shell"
-	"github.com/peterh/liner"
-	"github.com/spf13/cobra"
 )
 
-const queryTimeout = time.Second * 2
-
 func main() {
 	var crt string
 	var key string
 	var servers *[]string
 	var format string
+	var timeoutMsec uint
 
 	cmd := &cobra.Command{
 		Use:   "dqlite -s <servers> <database> [command]",
@@ -94,7 +94,7 @@ func main() {
 
 			if len(args) > 1 {
 				for _, input := range strings.Split(args[1], ";") {
-					ctx, cancel := context.WithTimeout(context.Background(), queryTimeout)
+					ctx, cancel := context.WithTimeout(context.Background(), time.Duration(timeoutMsec)*time.Millisecond)
 					defer cancel()
 					result, err := sh.Process(ctx, input)
 					if err != nil {
@@ -118,7 +118,7 @@ func main() {
 					return err
 				}
 
-				ctx, cancel := context.WithTimeout(context.Background(), queryTimeout)
+				ctx, cancel := context.WithTimeout(context.Background(), time.Duration(timeoutMsec)*time.Millisecond)
 				defer cancel()
 				result, err := sh.Process(ctx, input)
 				if err != nil {
@@ -140,6 +140,7 @@ func main() {
 	flags.StringVarP(&crt, "cert", "c", "", "public TLS cert")
 	flags.StringVarP(&key, "key", "k", "", "private TLS key")
 	flags.StringVarP(&format, "format", "f", "tabular", "output format (tabular, json)")
+	flags.UintVar(&timeoutMsec, "timeout", 2000, "timeout of each request (msec)")
 
 	cmd.MarkFlagRequired("servers")