Skip to content

Commit b2318d2

Browse files
alexottalexott
authored
Improve performance of databricks_notebook and databricks_workspace_file resources (#2902)
* Improve performance of `databricks_notebook` and `databricks_workspace_file` resources There were problems with performance of `databricks_notebook` resource: * There is an explicit Mutex there to coordinate creation & deletion of directories - use of it lead to the situation when notebooks were created sequentially even with high parallelism. * Both `databricks_notebook` and `databricks_workspace_file` always called `mkdirs` API to make sure that parent directory exist. This PR includes following changes: * In both `databricks_notebook` and `databricks_workspace_file`, the `mkdirs` API is called only when the error about "Parent folder doesn't exist". This should decrease the number of API calls. * Decrease use of locks in `databricks_notebook` * Address review comment
1 parent ddff8ea commit b2318d2

File tree

4 files changed

+268
-27
lines changed

4 files changed

+268
-27
lines changed

‎workspace/resource_notebook.go‎

Lines changed: 28 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -101,8 +101,10 @@ var mtx = &sync.Mutex{}
101101

102102
// Create creates a notebook given the content and path
103103
func (a NotebooksAPI) Create(r ImportPath) error {
104-
mtx.Lock()
105-
defer mtx.Unlock()
104+
if r.Format == "DBC" {
105+
mtx.Lock()
106+
defer mtx.Unlock()
107+
}
106108
return a.client.Post(a.context, "/workspace/import", r, nil)
107109
}
108110

@@ -297,8 +299,11 @@ func (a NotebooksAPI) list(path string) ([]ObjectStatus, error) {
297299

298300
// Delete will delete folders given a path and recursive flag
299301
func (a NotebooksAPI) Delete(path string, recursive bool) error {
300-
mtx.Lock()
301-
defer mtx.Unlock()
302+
if recursive {
303+
log.Printf("[DEBUG] Doing recursive delete of path '%s'", path)
304+
mtx.Lock()
305+
defer mtx.Unlock()
306+
}
302307
return a.client.Post(a.context, "/workspace/delete", DeletePath{
303308
Path: path,
304309
Recursive: recursive,
@@ -363,13 +368,6 @@ func ResourceNotebook() *schema.Resource {
363368
}
364369
notebooksAPI := NewNotebooksAPI(ctx, c)
365370
path := d.Get("path").(string)
366-
parent := filepath.ToSlash(filepath.Dir(path))
367-
if parent != "/" {
368-
err = notebooksAPI.Mkdirs(parent)
369-
if err != nil {
370-
return err
371-
}
372-
}
373371
createNotebook := ImportPath{
374372
Content: base64.StdEncoding.EncodeToString(content),
375373
Language: d.Get("language").(string),
@@ -389,7 +387,18 @@ func ResourceNotebook() *schema.Resource {
389387
}
390388
err = notebooksAPI.Create(createNotebook)
391389
if err != nil {
392-
return err
390+
if isParentDoesntExistError(err) {
391+
parent := filepath.ToSlash(filepath.Dir(path))
392+
log.Printf("[DEBUG] Parent folder '%s' doesn't exist, creating...", parent)
393+
err = notebooksAPI.Mkdirs(parent)
394+
if err != nil {
395+
return err
396+
}
397+
err = notebooksAPI.Create(createNotebook)
398+
}
399+
if err != nil {
400+
return err
401+
}
393402
}
394403
d.SetId(path)
395404
return nil
@@ -431,7 +440,13 @@ func ResourceNotebook() *schema.Resource {
431440
})
432441
},
433442
Delete: func(ctx context.Context, d *schema.ResourceData, c *common.DatabricksClient) error {
434-
return NewNotebooksAPI(ctx, c).Delete(d.Id(), true)
443+
objType := d.Get("object_type")
444+
return NewNotebooksAPI(ctx, c).Delete(d.Id(), !(objType == Notebook || objType == File))
435445
},
436446
}.ToResource()
437447
}
448+
449+
func isParentDoesntExistError(err error) bool {
450+
errStr := err.Error()
451+
return strings.HasPrefix(errStr, "The parent folder ") && strings.HasSuffix(errStr, " does not exist.")
452+
}

‎workspace/resource_notebook_test.go‎

Lines changed: 111 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -101,7 +101,7 @@ func TestResourceNotebookRead_Error(t *testing.T) {
101101
assert.Equal(t, "/test/path", d.Id(), "Id should not be empty for error reads")
102102
}
103103

104-
func TestResourceNotebookCreate(t *testing.T) {
104+
func TestResourceNotebookCreate_DirectoryExist(t *testing.T) {
105105
d, err := qa.ResourceFixture{
106106
Fixtures: []qa.HTTPFixture{
107107
{
@@ -152,6 +152,116 @@ func TestResourceNotebookCreate(t *testing.T) {
152152
assert.Equal(t, "/foo/path.py", d.Id())
153153
}
154154

155+
func TestResourceNotebookCreate_DirectoryDoesntExist(t *testing.T) {
156+
d, err := qa.ResourceFixture{
157+
Fixtures: []qa.HTTPFixture{
158+
{
159+
Method: "POST",
160+
Resource: "/api/2.0/workspace/mkdirs",
161+
ExpectedRequest: map[string]string{
162+
"path": "/foo",
163+
},
164+
},
165+
{
166+
Method: http.MethodPost,
167+
Resource: "/api/2.0/workspace/import",
168+
ExpectedRequest: ImportPath{
169+
Content: "YWJjCg==",
170+
Path: "/foo/path.py",
171+
Language: "PYTHON",
172+
Overwrite: true,
173+
Format: "SOURCE",
174+
},
175+
Response: map[string]string{
176+
"error_code": "RESOURCE_DOES_NOT_EXIST",
177+
"message": "The parent folder (/foo) does not exist.",
178+
},
179+
Status: 404,
180+
},
181+
{
182+
Method: http.MethodPost,
183+
Resource: "/api/2.0/workspace/import",
184+
ExpectedRequest: ImportPath{
185+
Content: "YWJjCg==",
186+
Path: "/foo/path.py",
187+
Language: "PYTHON",
188+
Overwrite: true,
189+
Format: "SOURCE",
190+
},
191+
},
192+
{
193+
Method: http.MethodGet,
194+
Resource: "/api/2.0/workspace/export?format=SOURCE&path=%2Ffoo%2Fpath.py",
195+
Response: ExportPath{
196+
Content: "YWJjCg==",
197+
},
198+
},
199+
{
200+
Method: http.MethodGet,
201+
Resource: "/api/2.0/workspace/get-status?path=%2Ffoo%2Fpath.py",
202+
Response: ObjectStatus{
203+
ObjectID: 4567,
204+
ObjectType: "NOTEBOOK",
205+
Path: "/foo/path.py",
206+
Language: "PYTHON",
207+
},
208+
},
209+
},
210+
Resource: ResourceNotebook(),
211+
State: map[string]any{
212+
"content_base64": "YWJjCg==",
213+
"language": "PYTHON",
214+
"path": "/foo/path.py",
215+
},
216+
Create: true,
217+
}.Apply(t)
218+
assert.NoError(t, err)
219+
assert.Equal(t, "/foo/path.py", d.Id())
220+
}
221+
222+
func TestResourceNotebookCreate_DirectoryCreateError(t *testing.T) {
223+
_, err := qa.ResourceFixture{
224+
Fixtures: []qa.HTTPFixture{
225+
{
226+
Method: "POST",
227+
Resource: "/api/2.0/workspace/mkdirs",
228+
ExpectedRequest: map[string]string{
229+
"path": "/foo",
230+
},
231+
Response: apierr.APIErrorBody{
232+
ErrorCode: "INVALID_REQUEST",
233+
Message: "Internal error happened",
234+
},
235+
Status: 400,
236+
},
237+
{
238+
Method: http.MethodPost,
239+
Resource: "/api/2.0/workspace/import",
240+
ExpectedRequest: ImportPath{
241+
Content: "YWJjCg==",
242+
Path: "/foo/path.py",
243+
Language: "PYTHON",
244+
Overwrite: true,
245+
Format: "SOURCE",
246+
},
247+
Response: map[string]string{
248+
"error_code": "RESOURCE_DOES_NOT_EXIST",
249+
"message": "The parent folder (/foo) does not exist.",
250+
},
251+
Status: 404,
252+
},
253+
},
254+
Resource: ResourceNotebook(),
255+
State: map[string]any{
256+
"content_base64": "YWJjCg==",
257+
"language": "PYTHON",
258+
"path": "/foo/path.py",
259+
},
260+
Create: true,
261+
}.Apply(t)
262+
assert.Error(t, err, "Internal error happened")
263+
}
264+
155265
func TestResourceNotebookCreateSource_Jupyter(t *testing.T) {
156266
d, err := qa.ResourceFixture{
157267
Fixtures: []qa.HTTPFixture{

‎workspace/resource_workspace_file.go‎

Lines changed: 17 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@ package workspace
33
import (
44
"context"
55
"encoding/base64"
6+
"log"
67
"path/filepath"
78

89
ws_api "github.com/databricks/databricks-sdk-go/service/workspace"
@@ -37,21 +38,26 @@ func ResourceWorkspaceFile() *schema.Resource {
3738
return err
3839
}
3940
path := d.Get("path").(string)
40-
parent := filepath.ToSlash(filepath.Dir(path))
41-
if parent != "/" {
42-
err = client.Workspace.MkdirsByPath(ctx, parent)
43-
if err != nil {
44-
return err
45-
}
46-
}
47-
err = client.Workspace.Import(ctx, ws_api.Import{
41+
importReq := ws_api.Import{
4842
Content: base64.StdEncoding.EncodeToString(content),
4943
Format: ws_api.ImportFormatAuto,
5044
Path: path,
5145
Overwrite: true,
52-
})
46+
}
47+
err = client.Workspace.Import(ctx, importReq)
5348
if err != nil {
54-
return err
49+
if isParentDoesntExistError(err) {
50+
parent := filepath.ToSlash(filepath.Dir(path))
51+
log.Printf("[DEBUG] Parent folder '%s' doesn't exist, creating...", parent)
52+
err = client.Workspace.MkdirsByPath(ctx, parent)
53+
if err != nil {
54+
return err
55+
}
56+
err = client.Workspace.Import(ctx, importReq)
57+
}
58+
if err != nil {
59+
return err
60+
}
5561
}
5662
d.SetId(path)
5763
return nil
@@ -89,7 +95,7 @@ func ResourceWorkspaceFile() *schema.Resource {
8995
if err != nil {
9096
return err
9197
}
92-
return client.Workspace.Delete(ctx, ws_api.Delete{Path: d.Id(), Recursive: true})
98+
return client.Workspace.Delete(ctx, ws_api.Delete{Path: d.Id(), Recursive: false})
9399
},
94100
}.ToResource()
95101
}

‎workspace/resource_workspace_file_test.go‎

Lines changed: 112 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -45,7 +45,7 @@ func TestResourceWorkspaceFileDelete(t *testing.T) {
4545
Method: http.MethodPost,
4646
Resource: "/api/2.0/workspace/delete",
4747
Status: http.StatusOK,
48-
ExpectedRequest: DeletePath{Path: path, Recursive: true},
48+
ExpectedRequest: DeletePath{Path: path, Recursive: false},
4949
},
5050
},
5151
Resource: ResourceWorkspaceFile(),
@@ -97,7 +97,7 @@ func TestResourceWorkspaceFileRead_Error(t *testing.T) {
9797
assert.Equal(t, "/test/path", d.Id(), "Id should not be empty for error reads")
9898
}
9999

100-
func TestResourceWorkspaceFileCreate(t *testing.T) {
100+
func TestResourceWorkspaceFileCreate_DirectoryExist(t *testing.T) {
101101
d, err := qa.ResourceFixture{
102102
Fixtures: []qa.HTTPFixture{
103103
{
@@ -145,6 +145,110 @@ func TestResourceWorkspaceFileCreate(t *testing.T) {
145145
assert.Equal(t, "/foo/path.py", d.Id())
146146
}
147147

148+
func TestResourceWorkspaceFileCreate_DirectoryDoesntExist(t *testing.T) {
149+
d, err := qa.ResourceFixture{
150+
Fixtures: []qa.HTTPFixture{
151+
{
152+
Method: "POST",
153+
Resource: "/api/2.0/workspace/mkdirs",
154+
ExpectedRequest: map[string]string{
155+
"path": "/foo",
156+
},
157+
},
158+
{
159+
Method: http.MethodPost,
160+
Resource: "/api/2.0/workspace/import",
161+
ExpectedRequest: ws_api.Import{
162+
Content: "YWJjCg==",
163+
Path: "/foo/path.py",
164+
Overwrite: true,
165+
Format: "AUTO",
166+
},
167+
Response: map[string]string{
168+
"error_code": "RESOURCE_DOES_NOT_EXIST",
169+
"message": "The parent folder (/foo) does not exist.",
170+
},
171+
Status: 404,
172+
},
173+
{
174+
Method: http.MethodPost,
175+
Resource: "/api/2.0/workspace/import",
176+
ExpectedRequest: ws_api.Import{
177+
Content: "YWJjCg==",
178+
Path: "/foo/path.py",
179+
Overwrite: true,
180+
Format: "AUTO",
181+
},
182+
},
183+
{
184+
Method: http.MethodGet,
185+
Resource: "/api/2.0/workspace/export?format=SOURCE&path=%2Ffoo%2Fpath.py",
186+
Response: ExportPath{
187+
Content: "YWJjCg==",
188+
},
189+
},
190+
{
191+
Method: http.MethodGet,
192+
Resource: "/api/2.0/workspace/get-status?path=%2Ffoo%2Fpath.py",
193+
Response: ObjectStatus{
194+
ObjectID: 4567,
195+
ObjectType: File,
196+
Path: "/foo/path.py",
197+
},
198+
},
199+
},
200+
Resource: ResourceWorkspaceFile(),
201+
State: map[string]any{
202+
"content_base64": "YWJjCg==",
203+
"path": "/foo/path.py",
204+
},
205+
Create: true,
206+
}.Apply(t)
207+
assert.NoError(t, err)
208+
assert.Equal(t, "/foo/path.py", d.Id())
209+
}
210+
211+
func TestResourceWorkspaceFileCreate_DirectoryCreateError(t *testing.T) {
212+
_, err := qa.ResourceFixture{
213+
Fixtures: []qa.HTTPFixture{
214+
{
215+
Method: "POST",
216+
Resource: "/api/2.0/workspace/mkdirs",
217+
ExpectedRequest: map[string]string{
218+
"path": "/foo",
219+
},
220+
Response: apierr.APIErrorBody{
221+
ErrorCode: "INVALID_REQUEST",
222+
Message: "Internal error happened",
223+
},
224+
Status: 400,
225+
},
226+
{
227+
Method: http.MethodPost,
228+
Resource: "/api/2.0/workspace/import",
229+
ExpectedRequest: ws_api.Import{
230+
Content: "YWJjCg==",
231+
Path: "/foo/path.py",
232+
Overwrite: true,
233+
Format: "AUTO",
234+
},
235+
Response: map[string]string{
236+
"error_code": "RESOURCE_DOES_NOT_EXIST",
237+
"message": "The parent folder (/foo) does not exist.",
238+
},
239+
Status: 404,
240+
},
241+
},
242+
Resource: ResourceWorkspaceFile(),
243+
State: map[string]any{
244+
"content_base64": "YWJjCg==",
245+
"path": "/foo/path.py",
246+
},
247+
Create: true,
248+
}.Apply(t)
249+
assert.Error(t, err, "Internal error happened")
250+
}
251+
148252
func TestResourceWorkspaceFileCreateSource(t *testing.T) {
149253
d, err := qa.ResourceFixture{
150254
Fixtures: []qa.HTTPFixture{
@@ -187,6 +291,12 @@ func TestResourceWorkspaceFileCreate_Error(t *testing.T) {
187291
{
188292
Method: http.MethodPost,
189293
Resource: "/api/2.0/workspace/import",
294+
ExpectedRequest: map[string]interface{}{
295+
"content": "YWJjCg==",
296+
"format": "AUTO",
297+
"overwrite": true,
298+
"path": "/path.py",
299+
},
190300
Response: apierr.APIErrorBody{
191301
ErrorCode: "INVALID_REQUEST",
192302
Message: "Internal error happened",

0 commit comments

Comments
 (0)