Skip to content

Commit cd40526

Browse files
agneumagneum
committed
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
2 parents 988d91f + dd81417 commit cd40526

File tree

9 files changed

+122
-17
lines changed

9 files changed

+122
-17
lines changed

‎cmd/cli/commands/client.go‎

Lines changed: 9 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -15,13 +15,14 @@ import (
1515

1616
// CLI configuration keys.
1717
const (
18-
EnvironmentIDKey = "environment-id"
19-
URLKey = "url"
20-
TokenKey = "token"
21-
InsecureKey = "insecure"
22-
FwServerURLKey = "forwarding-server-url"
23-
FwLocalPortKey = "forwarding-local-port"
24-
IdentityFileKey = "identity-file"
18+
EnvironmentIDKey = "environment-id"
19+
URLKey = "url"
20+
TokenKey = "token"
21+
InsecureKey = "insecure"
22+
RequestTimeoutKey = "request-timeout"
23+
FwServerURLKey = "forwarding-server-url"
24+
FwLocalPortKey = "forwarding-local-port"
25+
IdentityFileKey = "identity-file"
2526
)
2627

2728
// ClientByCLIContext creates a new Database Lab API client.
@@ -39,6 +40,7 @@ func ClientByCLIContext(cliCtx *cli.Context) (*dblabapi.Client, error) {
3940
Host: remoteURL.String(),
4041
VerificationToken: cliCtx.String(TokenKey),
4142
Insecure: cliCtx.Bool(InsecureKey),
43+
RequestTimeout: cliCtx.Duration(RequestTimeoutKey),
4244
}
4345

4446
// TODO(akartasov): Init and use logger.

‎cmd/cli/commands/config/command_list.go‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -37,6 +37,10 @@ func CommandList() []*cli.Command {
3737
Name: "insecure",
3838
Usage: "allow insecure server connections when using SSL",
3939
},
40+
&cli.DurationFlag{
41+
Name: "request-timeout",
42+
Usage: "allow changing requests timeout",
43+
},
4044
&cli.StringFlag{
4145
Name: "forwarding-server-url",
4246
Usage: "forwarding server URL of Database Lab instance",
@@ -69,6 +73,10 @@ func CommandList() []*cli.Command {
6973
Name: "insecure",
7074
Usage: "allow insecure server connections when using SSL",
7175
},
76+
&cli.DurationFlag{
77+
Name: "request-timeout",
78+
Usage: "allow changing requests timeout",
79+
},
7280
&cli.StringFlag{
7381
Name: "forwarding-server-url",
7482
Usage: "forwarding server URL of Database Lab instance",

‎cmd/cli/commands/config/duration.go‎

Lines changed: 64 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,64 @@
1+
/*
2+
2021 © Postgres.ai
3+
*/
4+
5+
package config
6+
7+
import (
8+
"bytes"
9+
"fmt"
10+
"time"
11+
)
12+
13+
// Duration defines a custom duration type.
14+
type Duration time.Duration
15+
16+
// String returns formatted duration.
17+
func (d Duration) String() string {
18+
return time.Duration(d).String()
19+
}
20+
21+
// UnmarshalJSON un-marshals json duration.
22+
func (d *Duration) UnmarshalJSON(b []byte) error {
23+
timeDuration, err := time.ParseDuration(string(bytes.Trim(b, `"`)))
24+
if err != nil {
25+
return err
26+
}
27+
28+
*d = Duration(timeDuration)
29+
30+
return err
31+
}
32+
33+
// MarshalJSON marshals json duration.
34+
func (d Duration) MarshalJSON() ([]byte, error) {
35+
return []byte(fmt.Sprintf("%q", d.String())), nil
36+
}
37+
38+
// MarshalYAML marshals yaml duration.
39+
func (d Duration) MarshalYAML() (interface{}, error) {
40+
return d.String(), nil
41+
}
42+
43+
// UnmarshalYAML marshals yaml duration.
44+
func (d *Duration) UnmarshalYAML(unmarshal func(interface{}) error) error {
45+
var durationString string
46+
if err := unmarshal(&durationString); err != nil {
47+
return err
48+
}
49+
50+
timeDuration, err := time.ParseDuration(durationString)
51+
if err != nil {
52+
*d = 0
53+
return err
54+
}
55+
56+
*d = Duration(timeDuration)
57+
58+
return nil
59+
}
60+
61+
// Duration retrieves duration from a custom type.
62+
func (d *Duration) Duration() time.Duration {
63+
return time.Duration(*d)
64+
}

‎cmd/cli/commands/config/environment.go‎

Lines changed: 14 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -19,11 +19,12 @@ type CLIConfig struct {
1919

2020
// Environment defines a format of environment configuration.
2121
type Environment struct {
22-
EnvironmentID string `yaml:"-" json:"environment_id"`
23-
URL string `yaml:"url" json:"url"`
24-
Token string `yaml:"token" json:"token"`
25-
Insecure bool `yaml:"insecure" json:"insecure"`
26-
Forwarding Forwarding `yaml:"forwarding" json:"forwarding"`
22+
EnvironmentID string `yaml:"-" json:"environment_id"`
23+
URL string `yaml:"url" json:"url"`
24+
Token string `yaml:"token" json:"token"`
25+
Insecure bool `yaml:"insecure" json:"insecure"`
26+
RequestTimeout Duration `yaml:"request_timeout,omitempty" json:"request_timeout,omitempty"`
27+
Forwarding Forwarding `yaml:"forwarding" json:"forwarding"`
2728
}
2829

2930
// Forwarding defines configuration for port forwarding.
@@ -44,9 +45,10 @@ func AddEnvironmentToConfig(c *cli.Context, cfg *CLIConfig, environmentID string
4445
}
4546

4647
env := Environment{
47-
URL: c.String(commands.URLKey),
48-
Token: c.String(commands.TokenKey),
49-
Insecure: c.Bool(commands.InsecureKey),
48+
URL: c.String(commands.URLKey),
49+
Token: c.String(commands.TokenKey),
50+
Insecure: c.Bool(commands.InsecureKey),
51+
RequestTimeout: Duration(c.Duration(commands.RequestTimeoutKey)),
5052
Forwarding: Forwarding{
5153
ServerURL: c.String(commands.FwServerURLKey),
5254
LocalPort: c.String(commands.FwLocalPortKey),
@@ -93,6 +95,10 @@ func updateEnvironmentInConfig(c *cli.Context, cfg *CLIConfig, environmentID str
9395
newEnvironment.Insecure = c.Bool(commands.InsecureKey)
9496
}
9597

98+
if c.IsSet(commands.RequestTimeoutKey) {
99+
newEnvironment.RequestTimeout = Duration(c.Duration(commands.RequestTimeoutKey))
100+
}
101+
96102
if c.IsSet(commands.FwServerURLKey) {
97103
newEnvironment.Forwarding.ServerURL = c.String(commands.FwServerURLKey)
98104
}

‎cmd/cli/commands/global/command_list.go‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -38,6 +38,10 @@ func List() []*cli.Command {
3838
Name: "insecure",
3939
Usage: "allow insecure server connections when using SSL",
4040
},
41+
&cli.DurationFlag{
42+
Name: "request-timeout",
43+
Usage: "allow changing requests timeout",
44+
},
4145
&cli.StringFlag{
4246
Name: "forwarding-server-url",
4347
Usage: "forwarding server URL of Database Lab instance",

‎cmd/cli/main.go‎

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -55,6 +55,11 @@ func main() {
5555
Usage: "allow insecure server connections when using SSL",
5656
EnvVars: []string{"DBLAB_INSECURE_SKIP_VERIFY"},
5757
},
58+
&cli.DurationFlag{
59+
Name: "request-timeout",
60+
Usage: "allow changing requests timeout",
61+
EnvVars: []string{"DBLAB_REQUEST_TIMEOUT"},
62+
},
5863
&cli.StringFlag{
5964
Name: "forwarding-server-url",
6065
Usage: "forwarding server URL of Database Lab instance",
@@ -121,6 +126,12 @@ func loadEnvironmentParams(c *cli.Context) error {
121126
}
122127
}
123128

129+
if !c.IsSet(commands.RequestTimeoutKey) {
130+
if err := c.Set(commands.RequestTimeoutKey, env.RequestTimeout.String()); err != nil {
131+
return err
132+
}
133+
}
134+
124135
if !c.IsSet(commands.FwServerURLKey) {
125136
if err := c.Set(commands.FwServerURLKey, env.Forwarding.ServerURL); err != nil {
126137
return err

‎pkg/client/dblabapi/client.go‎

Lines changed: 8 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -37,6 +37,7 @@ type Client struct {
3737
url *url.URL
3838
verificationToken string
3939
client *http.Client
40+
requestTimeout time.Duration
4041
pollingInterval time.Duration
4142
}
4243

@@ -45,6 +46,7 @@ type Options struct {
4546
Host string
4647
VerificationToken string
4748
Insecure bool
49+
RequestTimeout time.Duration
4850
}
4951

5052
const (
@@ -65,10 +67,15 @@ func NewClient(options Options) (*Client, error) {
6567
TLSClientConfig: &tls.Config{InsecureSkipVerify: options.Insecure},
6668
}
6769

70+
if options.RequestTimeout == 0 {
71+
options.RequestTimeout = defaultPollingTimeout
72+
}
73+
6874
return &Client{
6975
url: u,
7076
verificationToken: options.VerificationToken,
71-
client: &http.Client{Transport: tr, Timeout: defaultPollingTimeout},
77+
client: &http.Client{Transport: tr},
78+
requestTimeout: options.RequestTimeout,
7279
pollingInterval: defaultPollingInterval,
7380
}, nil
7481
}

‎pkg/client/dblabapi/client_test.go‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@ package dblabapi
33
import (
44
"net/http"
55
"testing"
6+
"time"
67

78
"github.com/stretchr/testify/assert"
89
"github.com/stretchr/testify/require"
@@ -28,12 +29,14 @@ func TestNewClient(t *testing.T) {
2829
c, err := NewClient(Options{
2930
Host: "https://example.com//",
3031
VerificationToken: "testVerify",
32+
RequestTimeout: 30 * time.Second,
3133
})
3234
require.NoError(t, err)
3335

3436
assert.IsType(t, &Client{}, c)
3537
assert.Equal(t, "https://example.com", c.url.String())
3638
assert.Equal(t, "testVerify", c.verificationToken)
39+
assert.Equal(t, 30*time.Second, c.requestTimeout)
3740
assert.IsType(t, &http.Client{}, c.client)
3841
}
3942

‎pkg/client/dblabapi/clone.go‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -133,7 +133,7 @@ func (c *Client) watchCloneStatus(ctx context.Context, cloneID string, initialSt
133133
var cancel context.CancelFunc
134134

135135
if _, ok := ctx.Deadline(); !ok {
136-
ctx, cancel = context.WithTimeout(ctx, defaultPollingTimeout)
136+
ctx, cancel = context.WithTimeout(ctx, c.requestTimeout)
137137
defer cancel()
138138
}
139139

0 commit comments

Comments
 (0)