From 7ecd15cf274714cbab766b9613b8dc20022e813b Mon Sep 17 00:00:00 2001 From: Alex Ott Date: Wed, 30 Oct 2024 13:16:21 +0100 Subject: [PATCH] [Fix] Skip calling `get-status` for notebooks in Read after Create/Update 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. --- workspace/resource_notebook.go | 59 +++++++++++---- workspace/resource_notebook_test.go | 112 ++++++++-------------------- 2 files changed, 78 insertions(+), 93 deletions(-) diff --git a/workspace/resource_notebook.go b/workspace/resource_notebook.go index deed2caa73..333f40a7d8 100644 --- a/workspace/resource_notebook.go +++ b/workspace/resource_notebook.go @@ -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"` @@ -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 @@ -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{ @@ -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{ @@ -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)) @@ -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 @@ -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 { @@ -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") diff --git a/workspace/resource_notebook_test.go b/workspace/resource_notebook_test.go index 304ac55d8d..dd75b90e77 100644 --- a/workspace/resource_notebook_test.go +++ b/workspace/resource_notebook_test.go @@ -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, }, }, }, @@ -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, }) } @@ -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, }, }, }, @@ -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, }) } @@ -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, }, }, }, @@ -312,7 +280,8 @@ func TestResourceNotebookCreateSource_Jupyter(t *testing.T) { }, Create: true, }.ApplyAndExpectData(t, map[string]any{ - "id": "/Mars", + "id": "/Mars", + "object_id": 12345, }) } @@ -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, }, }, }, @@ -350,7 +312,8 @@ func TestResourceNotebookCreateSource(t *testing.T) { }, Create: true, }.ApplyAndExpectData(t, map[string]any{ - "id": "/Dashboard", + "id": "/Dashboard", + "object_id": 12345, }) } @@ -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, }, }, }, @@ -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", }, }, { @@ -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, }, }, }, @@ -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) {