Skip to content

Commit

Permalink
[Fix] Skip calling get-status for notebooks in Read after Create/Up…
Browse files Browse the repository at this point in the history 
…date operations

It was found that `import` API returns `object_id` as result of its
execution so we don't really need to call `get-status` to fill all
attributes.

This PR adds internal attribute `skip_get_status` that is set by
Create/Update operations, and then checked in `Read` so we can avoid
performing another API call. This should help in cases when we're
importing a huge number of notebooks, i.e., as result of applying
exported resources.
  • Loading branch information
@alexott
alexott committed Oct 30, 2024
1 parent 92357dc commit 7ecd15c
Show file tree
Hide file tree
Showing 2 changed files with 78 additions and 93 deletions.
59 changes: 45 additions & 14 deletions workspace/resource_notebook.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -60,6 +60,10 @@ type ObjectStatus struct {
Size int64 `json:"size,omitempty"`
}

type ImportResponce struct {
ObjectID int64 `json:"object_id,omitempty"`
}

// ExportPath contains the base64 content of the notebook
type ExportPath struct {
Content string `json:"content,omitempty"`
Expand Down Expand Up @@ -98,12 +102,14 @@ type NotebooksAPI struct {
var mtx = &sync.Mutex{}

// Create creates a notebook given the content and path
func (a NotebooksAPI) Create(r ImportPath) error {
func (a NotebooksAPI) Create(r ImportPath) (ImportResponce, error) {
if r.Format == "DBC" {
mtx.Lock()
defer mtx.Unlock()
}
return a.client.Post(a.context, "/workspace/import", r, nil)
var responce ImportResponce
err := a.client.Post(a.context, "/workspace/import", r, &responce)
return responce, err
}

// Read returns the notebook metadata and not the contents
Expand Down Expand Up @@ -203,6 +209,11 @@ func (a NotebooksAPI) Delete(path string, recursive bool) error {
}, nil)
}

func setComputedProperties(d *schema.ResourceData, c *common.DatabricksClient) {
d.Set("url", c.FormatURL("#workspace", d.Id()))
d.Set("workspace_path", "/Workspace"+d.Id())
}

// ResourceNotebook manages notebooks
func ResourceNotebook() common.Resource {
s := FileContentSchema(map[string]*schema.Schema{
Expand Down Expand Up @@ -253,6 +264,10 @@ func ResourceNotebook() common.Resource {
Type: schema.TypeString,
Computed: true,
},
"skip_get_status": {
Type: schema.TypeBool,
Computed: true,
},
})
s["content_base64"].RequiredWith = []string{"language"}
return common.Resource{
Expand Down Expand Up @@ -282,7 +297,7 @@ func ResourceNotebook() common.Resource {
// by default it's SOURCE, but for DBC we have to change it
d.Set("format", createNotebook.Format)
}
err = notebooksAPI.Create(createNotebook)
resp, err := notebooksAPI.Create(createNotebook)
if err != nil {
if isParentDoesntExistError(err) {
parent := filepath.ToSlash(filepath.Dir(path))
Expand All @@ -291,16 +306,24 @@ func ResourceNotebook() common.Resource {
if err != nil {
return err
}
err = notebooksAPI.Create(createNotebook)
resp, err = notebooksAPI.Create(createNotebook)
}
if err != nil {
return err
}
}
d.Set("object_id", resp.ObjectID)
d.Set("skip_get_status", true)
d.SetId(path)
setComputedProperties(d, c)
return nil
},
Read: func(ctx context.Context, d *schema.ResourceData, c *common.DatabricksClient) error {
if d.Get("skip_get_status").(bool) {
log.Printf("[DEBUG] Skipping get-status for '%s'", d.Id())
d.Set("skip_get_status", false)
return nil
}
w, err := c.WorkspaceClient()
if err != nil {
return err
Expand All @@ -311,8 +334,8 @@ func ResourceNotebook() common.Resource {
if err != nil {
return err
}
d.Set("url", c.FormatURL("#workspace", d.Id()))
d.Set("workspace_path", "/Workspace"+objectStatus.Path)
d.Set("skip_get_status", false)
setComputedProperties(d, c)
return common.StructToData(objectStatus, s, d)
},
Update: func(ctx context.Context, d *schema.ResourceData, c *common.DatabricksClient) error {
Expand All @@ -322,25 +345,33 @@ func ResourceNotebook() common.Resource {
return err
}
format := d.Get("format").(string)
var resp ImportResponce
if format == "DBC" {
// Overwrite cannot be used for source format when importing a folder
err = notebooksAPI.Delete(d.Id(), true)
if err != nil {
return err
}
return notebooksAPI.Create(ImportPath{
resp, err = notebooksAPI.Create(ImportPath{
Content: base64.StdEncoding.EncodeToString(content),
Format: format,
Path: d.Id(),
})
} else {
resp, err = notebooksAPI.Create(ImportPath{
Content: base64.StdEncoding.EncodeToString(content),
Language: d.Get("language").(string),
Format: format,
Overwrite: true,
Path: d.Id(),
})
}
return notebooksAPI.Create(ImportPath{
Content: base64.StdEncoding.EncodeToString(content),
Language: d.Get("language").(string),
Format: format,
Overwrite: true,
Path: d.Id(),
})
if err == nil {
d.Set("object_id", resp.ObjectID)
d.Set("skip_get_status", true)
setComputedProperties(d, c)
}
return err
},
Delete: func(ctx context.Context, d *schema.ResourceData, c *common.DatabricksClient) error {
objType := d.Get("object_type")
Expand Down
112 changes: 33 additions & 79 deletions workspace/resource_notebook_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -119,22 +119,8 @@ func TestResourceNotebookCreate_DirectoryExist(t *testing.T) {
Overwrite: true,
Format: "SOURCE",
},
},
{
Method: http.MethodGet,
Resource: "/api/2.0/workspace/export?format=SOURCE&path=%2Ffoo%2Fpath.py",
Response: ExportPath{
Content: "YWJjCg==",
},
},
{
Method: http.MethodGet,
Resource: "/api/2.0/workspace/get-status?path=%2Ffoo%2Fpath.py",
Response: ObjectStatus{
ObjectID: 4567,
ObjectType: "NOTEBOOK",
Path: "/foo/path.py",
Language: "PYTHON",
Response: ImportResponce{
ObjectID: 12345,
},
},
},
Expand All @@ -146,8 +132,9 @@ func TestResourceNotebookCreate_DirectoryExist(t *testing.T) {
},
Create: true,
}.ApplyAndExpectData(t, map[string]any{
"path": "/foo/path.py",
"id": "/foo/path.py",
"path": "/foo/path.py",
"id": "/foo/path.py",
"object_id": 12345,
})
}

Expand Down Expand Up @@ -187,22 +174,8 @@ func TestResourceNotebookCreate_DirectoryDoesntExist(t *testing.T) {
Overwrite: true,
Format: "SOURCE",
},
},
{
Method: http.MethodGet,
Resource: "/api/2.0/workspace/export?format=SOURCE&path=%2Ffoo%2Fpath.py",
Response: ExportPath{
Content: "YWJjCg==",
},
},
{
Method: http.MethodGet,
Resource: "/api/2.0/workspace/get-status?path=%2Ffoo%2Fpath.py",
Response: ObjectStatus{
ObjectID: 4567,
ObjectType: "NOTEBOOK",
Path: "/foo/path.py",
Language: "PYTHON",
Response: ImportResponce{
ObjectID: 12345,
},
},
},
Expand All @@ -214,8 +187,9 @@ func TestResourceNotebookCreate_DirectoryDoesntExist(t *testing.T) {
},
Create: true,
}.ApplyAndExpectData(t, map[string]any{
"path": "/foo/path.py",
"id": "/foo/path.py",
"path": "/foo/path.py",
"id": "/foo/path.py",
"object_id": 12345,
})
}

Expand Down Expand Up @@ -294,14 +268,8 @@ func TestResourceNotebookCreateSource_Jupyter(t *testing.T) {
Overwrite: true,
Format: "JUPYTER",
},
},
{
Method: http.MethodGet,
Resource: "/api/2.0/workspace/get-status?path=%2FMars",
Response: ObjectStatus{
ObjectID: 4567,
ObjectType: "NOTEBOOK",
Path: "/Mars",
Response: ImportResponce{
ObjectID: 12345,
},
},
},
Expand All @@ -312,7 +280,8 @@ func TestResourceNotebookCreateSource_Jupyter(t *testing.T) {
},
Create: true,
}.ApplyAndExpectData(t, map[string]any{
"id": "/Mars",
"id": "/Mars",
"object_id": 12345,
})
}

Expand All @@ -331,15 +300,8 @@ func TestResourceNotebookCreateSource(t *testing.T) {
Overwrite: true,
Format: "SOURCE",
},
},
{
Method: http.MethodGet,
Resource: "/api/2.0/workspace/get-status?path=%2FDashboard",
Response: ObjectStatus{
ObjectID: 4567,
ObjectType: "NOTEBOOK",
Path: "/Dashboard",
Language: "SQL",
Response: ImportResponce{
ObjectID: 12345,
},
},
},
Expand All @@ -350,7 +312,8 @@ func TestResourceNotebookCreateSource(t *testing.T) {
},
Create: true,
}.ApplyAndExpectData(t, map[string]any{
"id": "/Dashboard",
"id": "/Dashboard",
"object_id": 12345,
})
}

Expand Down Expand Up @@ -413,15 +376,8 @@ func TestResourceNotebookUpdate(t *testing.T) {
Path: "abc",
Language: "R",
},
},
{
Method: http.MethodGet,
Resource: "/api/2.0/workspace/get-status?path=abc",
Response: ObjectStatus{
ObjectID: 4567,
ObjectType: "NOTEBOOK",
Path: "abc",
Language: "R",
Response: ImportResponce{
ObjectID: 12345,
},
},
},
Expand All @@ -445,7 +401,7 @@ func TestResourceNotebookUpdate_DBC(t *testing.T) {
Resource: "/api/2.0/workspace/delete",
ExpectedRequest: DeletePath{
Recursive: true,
Path: "abc",
Path: "/path.py",
},
},
{
Expand All @@ -454,16 +410,10 @@ func TestResourceNotebookUpdate_DBC(t *testing.T) {
ExpectedRequest: ImportPath{
Format: "DBC",
Content: "YWJjCg==",
Path: "abc",
Path: "/path.py",
},
},
{
Method: http.MethodGet,
Resource: "/api/2.0/workspace/get-status?path=abc",
Response: ObjectStatus{
ObjectID: 4567,
ObjectType: Directory,
Path: "abc",
Response: ImportResponce{
ObjectID: 12345,
},
},
},
Expand All @@ -472,14 +422,18 @@ func TestResourceNotebookUpdate_DBC(t *testing.T) {
"content_base64": "YWJjCg==",

// technically language is not needed, but makes the test simpler
"language": "PYTHON",
"format": "DBC",
"path": "/path.py",
"language": "PYTHON",
"format": "DBC",
"path": "/path.py",
"object_id": 45678,
},
ID: "abc",
ID: "/path.py",
RequiresNew: true,
Update: true,
}.ApplyNoError(t)
}.ApplyAndExpectData(t, map[string]any{
"id": "/path.py",
"object_id": 12345,
})
}

func TestNotebookLanguageSuppressSourceDiff(t *testing.T) {
Expand Down

0 comments on commit 7ecd15c

Please sign in to comment.