Merge branch '243-custom-timeout' into 'master'

feat: allow changing requests timeout for the DLE client package (#243)

Closes #243

See merge request postgres-ai/database-lab!276

Author: agneum

cmd/cli/commands/client.go

@@ -15,13 +15,14 @@ import (
 
 // CLI configuration keys.
 const (
-	EnvironmentIDKey = "environment-id"
-	URLKey           = "url"
-	TokenKey         = "token"
-	InsecureKey      = "insecure"
-	FwServerURLKey   = "forwarding-server-url"
-	FwLocalPortKey   = "forwarding-local-port"
-	IdentityFileKey  = "identity-file"
+	EnvironmentIDKey  = "environment-id"
+	URLKey            = "url"
+	TokenKey          = "token"
+	InsecureKey       = "insecure"
+	RequestTimeoutKey = "request-timeout"
+	FwServerURLKey    = "forwarding-server-url"
+	FwLocalPortKey    = "forwarding-local-port"
+	IdentityFileKey   = "identity-file"
 )
 
 // ClientByCLIContext creates a new Database Lab API client.
@@ -39,6 +40,7 @@ func ClientByCLIContext(cliCtx *cli.Context) (*dblabapi.Client, error) {
 		Host:              remoteURL.String(),
 		VerificationToken: cliCtx.String(TokenKey),
 		Insecure:          cliCtx.Bool(InsecureKey),
+		RequestTimeout:    cliCtx.Duration(RequestTimeoutKey),
 	}
 
 	// TODO(akartasov): Init and use logger.

cmd/cli/commands/config/command_list.go

@@ -37,6 +37,10 @@ func CommandList() []*cli.Command {
 							Name:  "insecure",
 							Usage: "allow insecure server connections when using SSL",
 						},
+						&cli.DurationFlag{
+							Name:  "request-timeout",
+							Usage: "allow changing requests timeout",
+						},
 						&cli.StringFlag{
 							Name:  "forwarding-server-url",
 							Usage: "forwarding server URL of Database Lab instance",
@@ -69,6 +73,10 @@ func CommandList() []*cli.Command {
 							Name:  "insecure",
 							Usage: "allow insecure server connections when using SSL",
 						},
+						&cli.DurationFlag{
+							Name:  "request-timeout",
+							Usage: "allow changing requests timeout",
+						},
 						&cli.StringFlag{
 							Name:  "forwarding-server-url",
 							Usage: "forwarding server URL of Database Lab instance",

cmd/cli/commands/config/duration.go

@@ -0,0 +1,64 @@
+/*
+2021 © Postgres.ai
+*/
+
+package config
+
+import (
+	"bytes"
+	"fmt"
+	"time"
+)
+
+// Duration defines a custom duration type.
+type Duration time.Duration
+
+// String returns formatted duration.
+func (d Duration) String() string {
+	return time.Duration(d).String()
+}
+
+// UnmarshalJSON un-marshals json duration.
+func (d *Duration) UnmarshalJSON(b []byte) error {
+	timeDuration, err := time.ParseDuration(string(bytes.Trim(b, `"`)))
+	if err != nil {
+		return err
+	}
+
+	*d = Duration(timeDuration)
+
+	return err
+}
+
+// MarshalJSON marshals json duration.
+func (d Duration) MarshalJSON() ([]byte, error) {
+	return []byte(fmt.Sprintf("%q", d.String())), nil
+}
+
+// MarshalYAML marshals yaml duration.
+func (d Duration) MarshalYAML() (interface{}, error) {
+	return d.String(), nil
+}
+
+// UnmarshalYAML marshals yaml duration.
+func (d *Duration) UnmarshalYAML(unmarshal func(interface{}) error) error {
+	var durationString string
+	if err := unmarshal(&durationString); err != nil {
+		return err
+	}
+
+	timeDuration, err := time.ParseDuration(durationString)
+	if err != nil {
+		*d = 0
+		return err
+	}
+
+	*d = Duration(timeDuration)
+
+	return nil
+}
+
+// Duration retrieves duration from a custom type.
+func (d *Duration) Duration() time.Duration {
+	return time.Duration(*d)
+}

cmd/cli/commands/config/environment.go

@@ -19,11 +19,12 @@ type CLIConfig struct {
 
 // Environment defines a format of environment configuration.
 type Environment struct {
-	EnvironmentID string     `yaml:"-" json:"environment_id"`
-	URL           string     `yaml:"url" json:"url"`
-	Token         string     `yaml:"token" json:"token"`
-	Insecure      bool       `yaml:"insecure" json:"insecure"`
-	Forwarding    Forwarding `yaml:"forwarding" json:"forwarding"`
+	EnvironmentID  string     `yaml:"-" json:"environment_id"`
+	URL            string     `yaml:"url" json:"url"`
+	Token          string     `yaml:"token" json:"token"`
+	Insecure       bool       `yaml:"insecure" json:"insecure"`
+	RequestTimeout Duration   `yaml:"request_timeout,omitempty" json:"request_timeout,omitempty"`
+	Forwarding     Forwarding `yaml:"forwarding" json:"forwarding"`
 }
 
 // Forwarding defines configuration for port forwarding.
@@ -44,9 +45,10 @@ func AddEnvironmentToConfig(c *cli.Context, cfg *CLIConfig, environmentID string
 	}
 
 	env := Environment{
-		URL:      c.String(commands.URLKey),
-		Token:    c.String(commands.TokenKey),
-		Insecure: c.Bool(commands.InsecureKey),
+		URL:            c.String(commands.URLKey),
+		Token:          c.String(commands.TokenKey),
+		Insecure:       c.Bool(commands.InsecureKey),
+		RequestTimeout: Duration(c.Duration(commands.RequestTimeoutKey)),
 		Forwarding: Forwarding{
 			ServerURL:    c.String(commands.FwServerURLKey),
 			LocalPort:    c.String(commands.FwLocalPortKey),
@@ -93,6 +95,10 @@ func updateEnvironmentInConfig(c *cli.Context, cfg *CLIConfig, environmentID str
 		newEnvironment.Insecure = c.Bool(commands.InsecureKey)
 	}
 
+	if c.IsSet(commands.RequestTimeoutKey) {
+		newEnvironment.RequestTimeout = Duration(c.Duration(commands.RequestTimeoutKey))
+	}
+
 	if c.IsSet(commands.FwServerURLKey) {
 		newEnvironment.Forwarding.ServerURL = c.String(commands.FwServerURLKey)
 	}

cmd/cli/commands/global/command_list.go

@@ -38,6 +38,10 @@ func List() []*cli.Command {
 					Name:  "insecure",
 					Usage: "allow insecure server connections when using SSL",
 				},
+				&cli.DurationFlag{
+					Name:  "request-timeout",
+					Usage: "allow changing requests timeout",
+				},
 				&cli.StringFlag{
 					Name:  "forwarding-server-url",
 					Usage: "forwarding server URL of Database Lab instance",

cmd/cli/main.go

@@ -55,6 +55,11 @@ func main() {
 				Usage:   "allow insecure server connections when using SSL",
 				EnvVars: []string{"DBLAB_INSECURE_SKIP_VERIFY"},
 			},
+			&cli.DurationFlag{
+				Name:    "request-timeout",
+				Usage:   "allow changing requests timeout",
+				EnvVars: []string{"DBLAB_REQUEST_TIMEOUT"},
+			},
 			&cli.StringFlag{
 				Name:    "forwarding-server-url",
 				Usage:   "forwarding server URL of Database Lab instance",
@@ -121,6 +126,12 @@ func loadEnvironmentParams(c *cli.Context) error {
 			}
 		}
 
+		if !c.IsSet(commands.RequestTimeoutKey) {
+			if err := c.Set(commands.RequestTimeoutKey, env.RequestTimeout.String()); err != nil {
+				return err
+			}
+		}
+
 		if !c.IsSet(commands.FwServerURLKey) {
 			if err := c.Set(commands.FwServerURLKey, env.Forwarding.ServerURL); err != nil {
 				return err

pkg/client/dblabapi/client.go

@@ -37,6 +37,7 @@ type Client struct {
 	url               *url.URL
 	verificationToken string
 	client            *http.Client
+	requestTimeout    time.Duration
 	pollingInterval   time.Duration
 }
 
@@ -45,6 +46,7 @@ type Options struct {
 	Host              string
 	VerificationToken string
 	Insecure          bool
+	RequestTimeout    time.Duration
 }
 
 const (
@@ -65,10 +67,15 @@ func NewClient(options Options) (*Client, error) {
 		TLSClientConfig: &tls.Config{InsecureSkipVerify: options.Insecure},
 	}
 
+	if options.RequestTimeout == 0 {
+		options.RequestTimeout = defaultPollingTimeout
+	}
+
 	return &Client{
 		url:               u,
 		verificationToken: options.VerificationToken,
-		client:            &http.Client{Transport: tr, Timeout: defaultPollingTimeout},
+		client:            &http.Client{Transport: tr},
+		requestTimeout:    options.RequestTimeout,
 		pollingInterval:   defaultPollingInterval,
 	}, nil
 }

pkg/client/dblabapi/client_test.go

@@ -3,6 +3,7 @@ package dblabapi
 import (
 	"net/http"
 	"testing"
+	"time"
 
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
@@ -28,12 +29,14 @@ func TestNewClient(t *testing.T) {
 	c, err := NewClient(Options{
 		Host:              "https://example.com//",
 		VerificationToken: "testVerify",
+		RequestTimeout:    30 * time.Second,
 	})
 	require.NoError(t, err)
 
 	assert.IsType(t, &Client{}, c)
 	assert.Equal(t, "https://example.com", c.url.String())
 	assert.Equal(t, "testVerify", c.verificationToken)
+	assert.Equal(t, 30*time.Second, c.requestTimeout)
 	assert.IsType(t, &http.Client{}, c.client)
 }
 

pkg/client/dblabapi/clone.go

@@ -133,7 +133,7 @@ func (c *Client) watchCloneStatus(ctx context.Context, cloneID string, initialSt
 	var cancel context.CancelFunc
 
 	if _, ok := ctx.Deadline(); !ok {
-		ctx, cancel = context.WithTimeout(ctx, defaultPollingTimeout)
+		ctx, cancel = context.WithTimeout(ctx, c.requestTimeout)
 		defer cancel()
 	}