Skip to content

Commit a30518e

Browse files
mgyuchtmgyucht
authored
Robust retries for workspace get-status (#3550)
* Robust retries for workspace get-status * debug log * refactor * fix test
1 parent 774cc97 commit a30518e

File tree

8 files changed

+112
-16
lines changed

8 files changed

+112
-16
lines changed

‎common/retry.go‎

Lines changed: 25 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,25 @@
1+
package common
2+
3+
import (
4+
"context"
5+
"log"
6+
"regexp"
7+
8+
"github.com/databricks/databricks-sdk-go/retries"
9+
)
10+
11+
var timeoutRegex = regexp.MustCompile(`request timed out after .* of inactivity`)
12+
13+
func RetryOnTimeout[T any](ctx context.Context, f func(context.Context) (*T, error)) (*T, error) {
14+
r := retries.New[T](retries.WithRetryFunc(func(err error) bool {
15+
msg := err.Error()
16+
isTimeout := timeoutRegex.MatchString(msg)
17+
if isTimeout {
18+
log.Printf("[DEBUG] Retrying due to timeout: %s", msg)
19+
}
20+
return isTimeout
21+
}))
22+
return r.Run(ctx, func(ctx context.Context) (*T, error) {
23+
return f(ctx)
24+
})
25+
}

‎common/retry_test.go‎

Lines changed: 49 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,49 @@
1+
package common
2+
3+
import (
4+
"context"
5+
"errors"
6+
"testing"
7+
8+
"github.com/databricks/databricks-sdk-go/experimental/mocks"
9+
"github.com/databricks/databricks-sdk-go/service/workspace"
10+
"github.com/stretchr/testify/assert"
11+
"github.com/stretchr/testify/mock"
12+
)
13+
14+
func TestRetryOnTimeout_NoError(t *testing.T) {
15+
w := mocks.NewMockWorkspaceClient(t)
16+
expected := &workspace.ObjectInfo{}
17+
api := w.GetMockWorkspaceAPI().EXPECT()
18+
api.GetStatusByPath(mock.Anything, mock.Anything).Return(expected, nil)
19+
res, err := RetryOnTimeout(context.Background(), func(ctx context.Context) (*workspace.ObjectInfo, error) {
20+
return w.WorkspaceClient.Workspace.GetStatusByPath(ctx, "path")
21+
})
22+
assert.NoError(t, err)
23+
assert.Equal(t, expected, res)
24+
}
25+
26+
func TestRetryOnTimeout_OneError(t *testing.T) {
27+
w := mocks.NewMockWorkspaceClient(t)
28+
expected := &workspace.ObjectInfo{}
29+
api := w.GetMockWorkspaceAPI().EXPECT()
30+
call1 := api.GetStatusByPath(mock.Anything, mock.Anything).Return(nil, errors.New("request failed: request timed out after 1m0s of inactivity"))
31+
call1.Repeatability = 1
32+
api.GetStatusByPath(mock.Anything, mock.Anything).Return(expected, nil)
33+
res, err := RetryOnTimeout(context.Background(), func(ctx context.Context) (*workspace.ObjectInfo, error) {
34+
return w.WorkspaceClient.Workspace.GetStatusByPath(ctx, "path")
35+
})
36+
assert.NoError(t, err)
37+
assert.Equal(t, expected, res)
38+
}
39+
40+
func TestRetryOnTimeout_NonRetriableError(t *testing.T) {
41+
w := mocks.NewMockWorkspaceClient(t)
42+
expected := errors.New("request failed: non-retriable error")
43+
api := w.GetMockWorkspaceAPI().EXPECT()
44+
api.GetStatusByPath(mock.Anything, mock.Anything).Return(nil, expected)
45+
_, err := RetryOnTimeout(context.Background(), func(ctx context.Context) (*workspace.ObjectInfo, error) {
46+
return w.WorkspaceClient.Workspace.GetStatusByPath(ctx, "path")
47+
})
48+
assert.ErrorIs(t, err, expected)
49+
}

‎exporter/exporter_test.go‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2007,7 +2007,7 @@ func TestImportingDLTPipelines(t *testing.T) {
20072007
},
20082008
{
20092009
Method: "GET",
2010-
Resource: "/api/2.0/workspace/get-status?path=%2FUsers%2Fuser%40domain.com%2FTest%20DLT",
2010+
Resource: "/api/2.0/workspace/get-status?path=%2FUsers%2Fuser%40domain.com%2FTest+DLT",
20112011
Response: workspace.ObjectStatus{
20122012
Language: workspace.Python,
20132013
ObjectID: 123,

‎workspace/data_directory.go‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -18,7 +18,9 @@ func DataSourceDirectory() common.Resource {
1818
WorkspacePath string `json:"workspace_path,omitempty" tf:"computed"`
1919
}
2020
return common.WorkspaceData(func(ctx context.Context, d *Directory, client *databricks.WorkspaceClient) error {
21-
data, err := client.Workspace.GetStatusByPath(ctx, d.Path)
21+
data, err := common.RetryOnTimeout(ctx, func(ctx context.Context) (*workspace.ObjectInfo, error) {
22+
return client.Workspace.GetStatusByPath(ctx, d.Path)
23+
})
2224
if err != nil {
2325
return err
2426
}

‎workspace/data_notebook.go‎

Lines changed: 8 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@ package workspace
33
import (
44
"context"
55

6+
"github.com/databricks/databricks-sdk-go/service/workspace"
67
"github.com/databricks/terraform-provider-databricks/common"
78
"github.com/hashicorp/terraform-plugin-sdk/v2/helper/schema"
89
"github.com/hashicorp/terraform-plugin-sdk/v2/helper/validation"
@@ -49,6 +50,10 @@ func DataSourceNotebook() common.Resource {
4950
return common.Resource{
5051
Schema: s,
5152
Read: func(ctx context.Context, d *schema.ResourceData, m *common.DatabricksClient) error {
53+
w, err := m.WorkspaceClient()
54+
if err != nil {
55+
return err
56+
}
5257
notebooksAPI := NewNotebooksAPI(ctx, m)
5358
path := d.Get("path").(string)
5459
format := d.Get("format").(string)
@@ -59,7 +64,9 @@ func DataSourceNotebook() common.Resource {
5964
d.SetId(path)
6065
// nolint
6166
d.Set("content", notebookContent)
62-
objectStatus, err := notebooksAPI.Read(d.Id())
67+
objectStatus, err := common.RetryOnTimeout(ctx, func(ctx context.Context) (*workspace.ObjectInfo, error) {
68+
return w.Workspace.GetStatusByPath(ctx, d.Id())
69+
})
6370
if err != nil {
6471
return err
6572
}

‎workspace/resource_directory.go‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -46,7 +46,9 @@ func ResourceDirectory() common.Resource {
4646
if err != nil {
4747
return err
4848
}
49-
objectStatus, err := client.Workspace.GetStatusByPath(ctx, d.Id())
49+
objectStatus, err := common.RetryOnTimeout(ctx, func(ctx context.Context) (*workspace.ObjectInfo, error) {
50+
return client.Workspace.GetStatusByPath(ctx, d.Id())
51+
})
5052
if err != nil {
5153
return err
5254
}

‎workspace/resource_notebook.go‎

Lines changed: 14 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@ import (
88
"strings"
99
"sync"
1010

11+
"github.com/databricks/databricks-sdk-go/service/workspace"
1112
"github.com/databricks/terraform-provider-databricks/common"
1213

1314
"github.com/hashicorp/terraform-plugin-sdk/v2/helper/schema"
@@ -108,9 +109,12 @@ func (a NotebooksAPI) Create(r ImportPath) error {
108109
// Read returns the notebook metadata and not the contents
109110
func (a NotebooksAPI) Read(path string) (ObjectStatus, error) {
110111
var notebookInfo ObjectStatus
111-
err := a.client.Get(a.context, "/workspace/get-status", map[string]string{
112-
"path": path,
113-
}, &notebookInfo)
112+
_, err := common.RetryOnTimeout(a.context, func(ctx context.Context) (*ObjectStatus, error) {
113+
err := a.client.Get(a.context, "/workspace/get-status", map[string]string{
114+
"path": path,
115+
}, &notebookInfo)
116+
return nil, err
117+
})
114118
return notebookInfo, err
115119
}
116120

@@ -293,8 +297,13 @@ func ResourceNotebook() common.Resource {
293297
return nil
294298
},
295299
Read: func(ctx context.Context, d *schema.ResourceData, c *common.DatabricksClient) error {
296-
notebooksAPI := NewNotebooksAPI(ctx, c)
297-
objectStatus, err := notebooksAPI.Read(d.Id())
300+
w, err := c.WorkspaceClient()
301+
if err != nil {
302+
return err
303+
}
304+
objectStatus, err := common.RetryOnTimeout(ctx, func(ctx context.Context) (*workspace.ObjectInfo, error) {
305+
return w.Workspace.GetStatusByPath(ctx, d.Id())
306+
})
298307
if err != nil {
299308
return err
300309
}

‎workspace/resource_workspace_file.go‎

Lines changed: 9 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,7 @@ import (
66
"log"
77
"path/filepath"
88

9-
ws_api "github.com/databricks/databricks-sdk-go/service/workspace"
9+
"github.com/databricks/databricks-sdk-go/service/workspace"
1010
"github.com/databricks/terraform-provider-databricks/common"
1111

1212
"github.com/hashicorp/terraform-plugin-sdk/v2/helper/schema"
@@ -42,9 +42,9 @@ func ResourceWorkspaceFile() common.Resource {
4242
return err
4343
}
4444
path := d.Get("path").(string)
45-
importReq := ws_api.Import{
45+
importReq := workspace.Import{
4646
Content: base64.StdEncoding.EncodeToString(content),
47-
Format: ws_api.ImportFormatAuto,
47+
Format: workspace.ImportFormatAuto,
4848
Path: path,
4949
Overwrite: true,
5050
ForceSendFields: []string{"Content"},
@@ -72,7 +72,9 @@ func ResourceWorkspaceFile() common.Resource {
7272
if err != nil {
7373
return err
7474
}
75-
objectStatus, err := client.Workspace.GetStatusByPath(ctx, d.Id())
75+
objectStatus, err := common.RetryOnTimeout(ctx, func(ctx context.Context) (*workspace.ObjectInfo, error) {
76+
return client.Workspace.GetStatusByPath(ctx, d.Id())
77+
})
7678
if err != nil {
7779
return err
7880
}
@@ -89,9 +91,9 @@ func ResourceWorkspaceFile() common.Resource {
8991
if err != nil {
9092
return err
9193
}
92-
return client.Workspace.Import(ctx, ws_api.Import{
94+
return client.Workspace.Import(ctx, workspace.Import{
9395
Content: base64.StdEncoding.EncodeToString(content),
94-
Format: ws_api.ImportFormatAuto,
96+
Format: workspace.ImportFormatAuto,
9597
Overwrite: true,
9698
Path: d.Id(),
9799
ForceSendFields: []string{"Content"},
@@ -102,7 +104,7 @@ func ResourceWorkspaceFile() common.Resource {
102104
if err != nil {
103105
return err
104106
}
105-
return client.Workspace.Delete(ctx, ws_api.Delete{Path: d.Id(), Recursive: false})
107+
return client.Workspace.Delete(ctx, workspace.Delete{Path: d.Id(), Recursive: false})
106108
},
107109
}
108110
}

0 commit comments

Comments
 (0)