Skip to content

Commit

Permalink
Skip calling Read after Create/Update operations for notebooks
Browse files Browse the repository at this point in the history 
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.

This also changes `format` and `language` attributes to `optional,computed` to avoid
having issues with defaults, and suppress diff.
  • Loading branch information
@alexott
alexott committed Nov 5, 2024
1 parent 1e067f7 commit e80cda6
Show file tree
Hide file tree
Showing 3 changed files with 150 additions and 104 deletions.
10 changes: 9 additions & 1 deletion exporter/importables.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -1614,7 +1614,15 @@ 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 {
switch pathString {
case "language":
return d.Get("language") == ""
case "format":
return d.Get("format") == "SOURCE"
}
return shouldOmitMd5Field(ic, pathString, as, d)
},
Depends: []reference{
{Path: "source", File: true},
{Path: "path", Resource: "databricks_directory", MatchType: MatchLongestPrefix,
Expand Down
80 changes: 56 additions & 24 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 ImportResponse 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) (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
Expand Down Expand Up @@ -203,31 +209,29 @@ 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,
Optional: true,
Default: "SOURCE",
Computed: true,
ValidateFunc: validation.StringInSlice([]string{
"SOURCE",
"DBC",
Expand Down Expand Up @@ -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 {
Expand All @@ -272,6 +279,10 @@ func ResourceNotebook() common.Resource {
Path: path,
Overwrite: true,
}
if createNotebook.Format == "" && createNotebook.Language != "" {
createNotebook.Format = "SOURCE"
d.Set("format", createNotebook.Format)
}
if createNotebook.Language == "" {
// TODO: check what happens with empty source
ext := strings.ToLower(filepath.Ext(d.Get("source").(string)))
Expand All @@ -281,8 +292,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))
Expand All @@ -291,16 +303,25 @@ 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 {
oldFormat := d.Get("format").(string)
if oldFormat == "" {
oldFormat = "SOURCE"
}
w, err := c.WorkspaceClient()
if err != nil {
return err
Expand All @@ -311,9 +332,12 @@ func ResourceNotebook() common.Resource {
if err != nil {
return err
}
d.Set("url", c.FormatURL("#workspace", d.Id()))
d.Set("workspace_path", "/Workspace"+objectStatus.Path)
return common.StructToData(objectStatus, s, d)
setComputedProperties(d, c)
err = common.StructToData(objectStatus, s, d)
if err != nil {
return err
}
return d.Set("format", oldFormat)
},
Update: func(ctx context.Context, d *schema.ResourceData, c *common.DatabricksClient) error {
notebooksAPI := NewNotebooksAPI(ctx, c)
Expand All @@ -322,25 +346,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")
Expand Down
Loading

0 comments on commit e80cda6

Please sign in to comment.