From 2ddb757055b816b049acfae299e897c67ea54a72 Mon Sep 17 00:00:00 2001 From: Alex Ott Date: Wed, 30 Oct 2024 18:44:40 +0100 Subject: [PATCH] [Feature] Allow to skip calling Read after Create/Update operations This PR adds the ability for a resource to specify that it may not need to call `Read` after `Create` and `Update` operations so we can avoid performing another API call(s). The resource may implement `CanSkipReadAfterCreateAndUpdate` function that can decide if the Read operation should be skipped. It was found that the import API returns `object_id` as a result of its execution, so we don't really need to call get-status to fill all attributes. This should help when we import a large number of notebooks, i.e., when applying exported resources. --- common/resource.go | 29 ++++--- common/resource_test.go | 2 + exporter/importables.go | 7 +- workspace/resource_directory.go | 1 + workspace/resource_notebook.go | 64 +++++++++----- workspace/resource_notebook_test.go | 124 +++++++++++----------------- 6 files changed, 115 insertions(+), 112 deletions(-) diff --git a/common/resource.go b/common/resource.go index 4e357305db..fb8a09e5bb 100644 --- a/common/resource.go +++ b/common/resource.go @@ -16,17 +16,18 @@ import ( // Resource aims to simplify things like error & deleted entities handling type Resource struct { - Create func(ctx context.Context, d *schema.ResourceData, c *DatabricksClient) error - Read func(ctx context.Context, d *schema.ResourceData, c *DatabricksClient) error - Update func(ctx context.Context, d *schema.ResourceData, c *DatabricksClient) error - Delete func(ctx context.Context, d *schema.ResourceData, c *DatabricksClient) error - CustomizeDiff func(ctx context.Context, d *schema.ResourceDiff) error - StateUpgraders []schema.StateUpgrader - Schema map[string]*schema.Schema - SchemaVersion int - Timeouts *schema.ResourceTimeout - DeprecationMessage string - Importer *schema.ResourceImporter + Create func(ctx context.Context, d *schema.ResourceData, c *DatabricksClient) error + Read func(ctx context.Context, d *schema.ResourceData, c *DatabricksClient) error + Update func(ctx context.Context, d *schema.ResourceData, c *DatabricksClient) error + Delete func(ctx context.Context, d *schema.ResourceData, c *DatabricksClient) error + CustomizeDiff func(ctx context.Context, d *schema.ResourceDiff) error + StateUpgraders []schema.StateUpgrader + Schema map[string]*schema.Schema + SchemaVersion int + Timeouts *schema.ResourceTimeout + DeprecationMessage string + Importer *schema.ResourceImporter + CanSkipReadAfterCreateAndUpdate func(d *schema.ResourceData) bool } func nicerError(ctx context.Context, err error, action string) error { @@ -94,6 +95,9 @@ func (r Resource) ToResource() *schema.Resource { err = nicerError(ctx, err, "update") return diag.FromErr(err) } + if r.CanSkipReadAfterCreateAndUpdate != nil && r.CanSkipReadAfterCreateAndUpdate(d) { + return nil + } if err := recoverable(r.Read)(ctx, d, c); err != nil { err = nicerError(ctx, err, "read") return diag.FromErr(err) @@ -162,6 +166,9 @@ func (r Resource) ToResource() *schema.Resource { err = nicerError(ctx, err, "create") return diag.FromErr(err) } + if r.CanSkipReadAfterCreateAndUpdate != nil && r.CanSkipReadAfterCreateAndUpdate(d) { + return nil + } if err = recoverable(r.Read)(ctx, d, c); err != nil { err = nicerError(ctx, err, "read") return diag.FromErr(err) diff --git a/common/resource_test.go b/common/resource_test.go index f01f373ff5..37bba6fccb 100644 --- a/common/resource_test.go +++ b/common/resource_test.go @@ -38,6 +38,8 @@ func TestImportingCallsRead(t *testing.T) { assert.Equal(t, 1, d.Get("foo")) } +// TODO: add test for CanSkipRead and without it + func TestHTTP404TriggersResourceRemovalForReadAndDelete(t *testing.T) { nope := func(ctx context.Context, d *schema.ResourceData, diff --git a/exporter/importables.go b/exporter/importables.go index 04833df814..c5d6976996 100644 --- a/exporter/importables.go +++ b/exporter/importables.go @@ -1614,7 +1614,12 @@ var resourcesMap map[string]importable = map[string]importable{ ic.emitWorkspaceObjectParentDirectory(r) return r.Data.Set("source", fileName) }, - ShouldOmitField: shouldOmitMd5Field, + ShouldOmitField: func(ic *importContext, pathString string, as *schema.Schema, d *schema.ResourceData) bool { + if pathString == "language" { + return d.Get("language") == "" + } + return shouldOmitMd5Field(ic, pathString, as, d) + }, Depends: []reference{ {Path: "source", File: true}, {Path: "path", Resource: "databricks_directory", MatchType: MatchLongestPrefix, diff --git a/workspace/resource_directory.go b/workspace/resource_directory.go index 3c44669526..6de99eaa7d 100644 --- a/workspace/resource_directory.go +++ b/workspace/resource_directory.go @@ -74,6 +74,7 @@ func ResourceDirectory() common.Resource { path := d.Get("path").(string) err = client.Workspace.MkdirsByPath(ctx, path) if err != nil { + // TODO: ignore RESOURCE_ALREADY_EXISTS error? make it configurable? // TODO: handle RESOURCE_ALREADY_EXISTS return err } diff --git a/workspace/resource_notebook.go b/workspace/resource_notebook.go index deed2caa73..46f96a5aa2 100644 --- a/workspace/resource_notebook.go +++ b/workspace/resource_notebook.go @@ -60,6 +60,10 @@ type ObjectStatus struct { Size int64 `json:"size,omitempty"` } +type ImportResponse 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) (ImportResponse, error) { if r.Format == "DBC" { mtx.Lock() defer mtx.Unlock() } - return a.client.Post(a.context, "/workspace/import", r, nil) + var responce ImportResponse + err := a.client.Post(a.context, "/workspace/import", r, &responce) + return responce, err } // Read returns the notebook metadata and not the contents @@ -203,26 +209,24 @@ 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{ "language": { Type: schema.TypeString, Optional: true, + Computed: true, // we need it because it will be filled by the provider or backend ValidateFunc: validation.StringInSlice([]string{ Scala, Python, R, SQL, }, false), - DiffSuppressFunc: func(k, old, new string, d *schema.ResourceData) bool { - source := d.Get("source").(string) - if source == "" { - return false - } - ext := strings.ToLower(filepath.Ext(source)) - return old == extMap[ext].Language - }, }, "format": { Type: schema.TypeString, @@ -258,6 +262,9 @@ func ResourceNotebook() common.Resource { return common.Resource{ Schema: s, SchemaVersion: 1, + CanSkipReadAfterCreateAndUpdate: func(d *schema.ResourceData) bool { + return d.Get("format").(string) == "SOURCE" + }, Create: func(ctx context.Context, d *schema.ResourceData, c *common.DatabricksClient) error { content, err := ReadContent(d) if err != nil { @@ -281,8 +288,9 @@ func ResourceNotebook() common.Resource { createNotebook.Overwrite = extMap[ext].Overwrite // by default it's SOURCE, but for DBC we have to change it d.Set("format", createNotebook.Format) + d.Set("language", createNotebook.Language) } - err = notebooksAPI.Create(createNotebook) + resp, err := notebooksAPI.Create(createNotebook) if err != nil { if isParentDoesntExistError(err) { parent := filepath.ToSlash(filepath.Dir(path)) @@ -291,13 +299,18 @@ func ResourceNotebook() common.Resource { if err != nil { return err } - err = notebooksAPI.Create(createNotebook) + resp, err = notebooksAPI.Create(createNotebook) } if err != nil { return err } } + if d.Get("object_type").(string) == "" { + d.Set("object_type", Notebook) + } + d.Set("object_id", resp.ObjectID) d.SetId(path) + setComputedProperties(d, c) return nil }, Read: func(ctx context.Context, d *schema.ResourceData, c *common.DatabricksClient) error { @@ -311,8 +324,7 @@ func ResourceNotebook() common.Resource { if err != nil { return err } - d.Set("url", c.FormatURL("#workspace", d.Id())) - d.Set("workspace_path", "/Workspace"+objectStatus.Path) + setComputedProperties(d, c) return common.StructToData(objectStatus, s, d) }, Update: func(ctx context.Context, d *schema.ResourceData, c *common.DatabricksClient) error { @@ -322,25 +334,33 @@ func ResourceNotebook() common.Resource { return err } format := d.Get("format").(string) + var resp ImportResponse 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 { + return err + } + d.Set("object_id", resp.ObjectID) + setComputedProperties(d, c) + return nil }, 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..30d546f328 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: ImportResponse{ + 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: ImportResponse{ + 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,13 +268,17 @@ func TestResourceNotebookCreateSource_Jupyter(t *testing.T) { Overwrite: true, Format: "JUPYTER", }, + Response: ImportResponse{ + ObjectID: 12345, + }, }, { Method: http.MethodGet, Resource: "/api/2.0/workspace/get-status?path=%2FMars", Response: ObjectStatus{ - ObjectID: 4567, + ObjectID: 12345, ObjectType: "NOTEBOOK", + Language: "PYTHON", Path: "/Mars", }, }, @@ -312,7 +290,9 @@ func TestResourceNotebookCreateSource_Jupyter(t *testing.T) { }, Create: true, }.ApplyAndExpectData(t, map[string]any{ - "id": "/Mars", + "id": "/Mars", + "object_id": 12345, + "language": "PYTHON", }) } @@ -331,15 +311,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: ImportResponse{ + ObjectID: 12345, }, }, }, @@ -350,7 +323,9 @@ func TestResourceNotebookCreateSource(t *testing.T) { }, Create: true, }.ApplyAndExpectData(t, map[string]any{ - "id": "/Dashboard", + "id": "/Dashboard", + "object_id": 12345, + "language": "SQL", }) } @@ -410,18 +385,11 @@ func TestResourceNotebookUpdate(t *testing.T) { Format: "SOURCE", Overwrite: true, Content: "YWJjCg==", - Path: "abc", + Path: "/path.py", 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: ImportResponse{ + ObjectID: 12345, }, }, }, @@ -431,10 +399,14 @@ func TestResourceNotebookUpdate(t *testing.T) { "language": "R", "path": "/path.py", }, - ID: "abc", + ID: "/path.py", RequiresNew: true, Update: true, - }.ApplyNoError(t) + }.ApplyAndExpectData(t, map[string]any{ + "id": "/path.py", + "object_id": 12345, + "language": "R", + }) } func TestResourceNotebookUpdate_DBC(t *testing.T) { @@ -445,7 +417,7 @@ func TestResourceNotebookUpdate_DBC(t *testing.T) { Resource: "/api/2.0/workspace/delete", ExpectedRequest: DeletePath{ Recursive: true, - Path: "abc", + Path: "/path.py", }, }, { @@ -454,16 +426,16 @@ 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", + Resource: "/api/2.0/workspace/get-status?path=%2Fpath.py", Response: ObjectStatus{ - ObjectID: 4567, + ObjectID: 12345, ObjectType: Directory, - Path: "abc", + Path: "/path.py", }, }, }, @@ -472,20 +444,16 @@ 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) -} - -func TestNotebookLanguageSuppressSourceDiff(t *testing.T) { - r := ResourceNotebook() - d := r.ToResource().TestResourceData() - d.Set("source", "this.PY") - suppress := r.Schema["language"].DiffSuppressFunc - assert.True(t, suppress("language", Python, Python, d)) + }.ApplyAndExpectData(t, map[string]any{ + "id": "/path.py", + "object_id": 12345, + }) }