Exporter: generate references that are matching only to a prefix

There are cases when reference isn't a full match of a given value but a prefix of it. For
example, user's notebook has a path of `/Users/user@domain/notebook`.  To make it working
correctly we need to emit user or service principal, and then make a correct reference
with dependency on emitted user or notebook, like,
`${databricks_user.user_domain.home}/notebook`.

This PR adds a new field to dependency specification: `MatchType` that may have following values:

- `prefix` - value of attribute of resource is a prefix of a given attribute
- `exact` (or empty string) - must have exact matching (previous behaviour)

this fixes #1777

exporter/context.go

@@ -300,7 +300,7 @@ func (ic *importContext) MatchesName(n string) bool {
  return strings.Contains(strings.ToLower(n), strings.ToLower(ic.match))
 }
 
-func (ic *importContext) Find(r *resource, pick string) hcl.Traversal {
+func (ic *importContext) Find(r *resource, pick, matchType string) (string, hcl.Traversal) {
  for _, sr := range ic.State.Resources {
  if sr.Type != r.Resource {
  continue
@@ -312,24 +312,28 @@ func (ic *importContext) Find(r *resource, pick string) hcl.Traversal {
  r.Attribute, r.Resource, r.Name, r.ID)
  continue
  }
- if v.(string) == r.Value {
- if sr.Mode == "data" {
- return hcl.Traversal{
- hcl.TraverseRoot{Name: "data"},
- hcl.TraverseAttr{Name: sr.Type},
- hcl.TraverseAttr{Name: sr.Name},
- hcl.TraverseAttr{Name: pick},
- }
- }
- return hcl.Traversal{
- hcl.TraverseRoot{Name: sr.Type},
+ strValue := v.(string)
+ if ((matchType == "" || matchType == "exact") && strValue != r.Value) ||
+ (matchType == "prefix" && !strings.HasPrefix(r.Value, strValue)) {
+ continue
+ }
+ if sr.Mode == "data" {
+ return strValue, hcl.Traversal{
+ hcl.TraverseRoot{Name: "data"},
+ hcl.TraverseAttr{Name: sr.Type},
  hcl.TraverseAttr{Name: sr.Name},
  hcl.TraverseAttr{Name: pick},
  }
  }
+ return strValue, hcl.Traversal{
+ hcl.TraverseRoot{Name: sr.Type},
+ hcl.TraverseAttr{Name: sr.Name},
+ hcl.TraverseAttr{Name: pick},
+ }
+
  }
  }
- return nil
+ return "", nil
 }
 
 func (ic *importContext) Has(r *resource) bool {
@@ -509,13 +513,22 @@ func (ic *importContext) reference(i importable, path []string, value string) hc
  if d.Match != "" {
  attr = d.Match
  }
- traversal := ic.Find(&resource{
+ attrValue, traversal := ic.Find(&resource{
  Resource: d.Resource,
  Attribute: attr,
  Value: value,
- }, attr)
- //at least one invocation of ic.Find will assign Nil to traversal if resource with value is not found
+ }, attr, d.MatchType)
+ // at least one invocation of ic.Find will assign Nil to traversal if resource with value is not found
  if traversal != nil {
+ if d.MatchType == "prefix" { // we're replacing the found prefix with the HCL reference using the "${reference}rest" syntax
+ rest := value[len(attrValue):]
+ tokens := hclwrite.Tokens{&hclwrite.Token{Type: hclsyntax.TokenOQuote, Bytes: []byte{'"', '$', '{'}}}
+ tokens = append(tokens, hclwrite.TokensForTraversal(traversal)...)
+ tokens = append(tokens, &hclwrite.Token{Type: hclsyntax.TokenCQuote, Bytes: []byte{'}'}})
+ tokens = append(tokens, &hclwrite.Token{Type: hclsyntax.TokenQuotedLit, Bytes: []byte(rest)})
+ tokens = append(tokens, &hclwrite.Token{Type: hclsyntax.TokenCQuote, Bytes: []byte{'"'}})
+ return tokens
+ }
  return hclwrite.TokensForTraversal(traversal)
  }
  }

exporter/context_test.go

@@ -17,7 +17,7 @@ func TestMatchesName(t *testing.T) {
 }
 
 func TestImportContextFindSkips(t *testing.T) {
- assert.Nil(t, (&importContext{
+ _, traversal := (&importContext{
  State: stateApproximation{
  Resources: []resourceApproximation{
  {
@@ -36,7 +36,8 @@ func TestImportContextFindSkips(t *testing.T) {
  Resource: "a",
  Attribute: "b",
  Name: "c",
- }, "x"))
+ }, "x", "")
+ assert.Nil(t, traversal)
 }
 
 func TestImportContextHas(t *testing.T) {

exporter/importables.go

@@ -41,6 +41,7 @@ var (
  nameNormalizationRegex = regexp.MustCompile(`\W+`)
  jobClustersRegex = regexp.MustCompile(`^((job_cluster|task)\.[0-9]+\.new_cluster\.[0-9]+\.)`)
  dltClusterRegex = regexp.MustCompile(`^(cluster\.[0-9]+\.)`)
+ uuidRegex = regexp.MustCompile(`^[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}$`)
  predefinedClusterPolicies = []string{"Personal Compute", "Job Compute", "Power User Compute", "Shared Compute"}
 )
 
@@ -286,6 +287,7 @@ var resourcesMap map[string]importable = map[string]importable{
  {Path: "spark_python_task.parameters", Resource: "databricks_dbfs_file", Match: "dbfs_path"},
  {Path: "spark_jar_task.jar_uri", Resource: "databricks_dbfs_file", Match: "dbfs_path"},
  {Path: "notebook_task.notebook_path", Resource: "databricks_notebook"},
+ {Path: "notebook_task.notebook_path", Resource: "databricks_repo", Match: "path", MatchType: "prefix"},
  {Path: "pipeline_task.pipeline_id", Resource: "databricks_pipeline"},
  {Path: "task.library.jar", Resource: "databricks_dbfs_file", Match: "dbfs_path"},
  {Path: "task.library.whl", Resource: "databricks_dbfs_file", Match: "dbfs_path"},
@@ -294,6 +296,7 @@ var resourcesMap map[string]importable = map[string]importable{
  {Path: "task.spark_python_task.parameters", Resource: "databricks_dbfs_file", Match: "dbfs_path"},
  {Path: "task.spark_jar_task.jar_uri", Resource: "databricks_dbfs_file", Match: "dbfs_path"},
  {Path: "task.notebook_task.notebook_path", Resource: "databricks_notebook"},
+ {Path: "task.notebook_task.notebook_path", Resource: "databricks_repo", Match: "path", MatchType: "prefix"},
  {Path: "task.pipeline_task.pipeline_id", Resource: "databricks_pipeline"},
  {Path: "task.sql_task.query.query_id", Resource: "databricks_sql_query"},
  {Path: "task.sql_task.dashboard.dashboard_id", Resource: "databricks_sql_dashboard"},
@@ -358,10 +361,7 @@ var resourcesMap map[string]importable = map[string]importable{
  }
  }
  if job.NotebookTask != nil {
- ic.Emit(&resource{
- Resource: "databricks_notebook",
- ID: job.NotebookTask.NotebookPath,
- })
+ ic.emitNotebookOrRepo(job.NotebookTask.NotebookPath)
  }
  if job.PipelineTask != nil {
  ic.Emit(&resource{
@@ -372,10 +372,7 @@ var resourcesMap map[string]importable = map[string]importable{
  // Support for multitask jobs
  for _, task := range job.Tasks {
  if task.NotebookTask != nil {
- ic.Emit(&resource{
- Resource: "databricks_notebook",
- ID: task.NotebookTask.NotebookPath,
- })
+ ic.emitNotebookOrRepo(task.NotebookTask.NotebookPath)
  }
  if task.PipelineTask != nil {
  ic.Emit(&resource{
@@ -957,7 +954,21 @@ var resourcesMap map[string]importable = map[string]importable{
  if name == "" {
  return d.Id()
  }
- return strings.TrimPrefix(name, "/")
+ return nameNormalizationRegex.ReplaceAllString(name[7:], "_")
+ },
+ Search: func(ic *importContext, r *resource) error {
+ reposAPI := repos.NewReposAPI(ic.Context, ic.Client)
+ notebooksAPI := workspace.NewNotebooksAPI(ic.Context, ic.Client)
+ repoDir, err := notebooksAPI.Read(r.Value)
+ if err != nil {
+ return err
+ }
+ repo, err := reposAPI.Read(fmt.Sprintf("%d", repoDir.ObjectID))
+ if err != nil {
+ return err
+ }
+ r.ID = fmt.Sprintf("%d", repo.ID)
+ return nil
  },
  List: func(ic *importContext) error {
  repoList, err := repos.NewReposAPI(ic.Context, ic.Client).ListAll()
@@ -1093,6 +1104,7 @@ var resourcesMap map[string]importable = map[string]importable{
  return nil
  },
  Import: func(ic *importContext, r *resource) error {
+ ic.emitUserOrServicePrincipalForPath(r.ID, "/Users")
  notebooksAPI := workspace.NewNotebooksAPI(ic.Context, ic.Client)
  contentB64, err := notebooksAPI.Export(r.ID, "SOURCE")
  if err != nil {
@@ -1117,6 +1129,8 @@ var resourcesMap map[string]importable = map[string]importable{
  },
  Depends: []reference{
  {Path: "source", File: true},
+ {Path: "path", Resource: "databricks_user", Match: "home", MatchType: "prefix"},
+ {Path: "path", Resource: "databricks_service_principal", Match: "home", MatchType: "prefix"},
  },
  },
  "databricks_sql_query": {
@@ -1390,10 +1404,7 @@ var resourcesMap map[string]importable = map[string]importable{
  common.DataToStructPointer(r.Data, s, &pipeline)
  for _, lib := range pipeline.Libraries {
  if lib.Notebook != nil {
- ic.Emit(&resource{
- Resource: "databricks_notebook",
- ID: lib.Notebook.Path,
- })
+ ic.emitNotebookOrRepo(lib.Notebook.Path)
  }
  ic.emitIfDbfsFile(lib.Jar)
  ic.emitIfDbfsFile(lib.Whl)
@@ -1449,6 +1460,7 @@ var resourcesMap map[string]importable = map[string]importable{
  {Path: "cluster.instance_pool_id", Resource: "databricks_instance_pool"},
  {Path: "cluster.driver_instance_pool_id", Resource: "databricks_instance_pool"},
  {Path: "library.notebook.path", Resource: "databricks_notebook"},
+ {Path: "library.notebook.path", Resource: "databricks_repo", Match: "path", MatchType: "prefix"},
  {Path: "library.jar", Resource: "databricks_dbfs_file", Match: "dbfs_path"},
  {Path: "library.whl", Resource: "databricks_dbfs_file", Match: "dbfs_path"},
  },

exporter/model.go

@@ -57,11 +57,12 @@ type importable struct {
 }
 
 type reference struct {
- Path string
- Resource string
- Match string
- Variable bool
- File bool
+ Path string
+ Resource string
+ Match string
+ MatchType string // type of match, `prefix` - reference is embedded into string, `` (or `exact`) - full match
+ Variable bool
+ File bool
 }
 
 type resource struct {

exporter/util.go

@@ -74,7 +74,7 @@ func (ic *importContext) emitUserOrServicePrincipal(userOrSPName string) {
  Attribute: "user_name",
  Value: userOrSPName,
  })
- } else {
+ } else if uuidRegex.MatchString(userOrSPName) {
  ic.Emit(&resource{
  Resource: "databricks_service_principal",
  Attribute: "application_id",
@@ -83,6 +83,30 @@ func (ic *importContext) emitUserOrServicePrincipal(userOrSPName string) {
  }
 }
 
+func (ic *importContext) emitUserOrServicePrincipalForPath(path, prefix string) {
+ if strings.HasPrefix(path, prefix) {
+ parts := strings.SplitN(path, "/", 4)
+ if len(parts) >= 3 {
+ ic.emitUserOrServicePrincipal(parts[2])
+ }
+ }
+}
+
+func (ic *importContext) emitNotebookOrRepo(path string) {
+ if strings.HasPrefix(path, "/Repos") {
+ ic.Emit(&resource{
+ Resource: "databricks_repo",
+ Attribute: "path",
+ Value: strings.Join(strings.SplitN(path, "/", 5)[:4], "/"),
+ })
+ } else {
+ ic.Emit(&resource{
+ Resource: "databricks_notebook",
+ ID: path,
+ })
+ }
+}
+
 func (ic *importContext) emitGroups(u scim.User, principal string) {
  for _, g := range u.Groups {
  if g.Type != "direct" {

exporter/util_test.go

@@ -45,4 +45,42 @@ func TestEmitUserOrServicePrincipal(t *testing.T) {
  ic.emitUserOrServicePrincipal("21aab5a7-ee70-4385-34d4-a77278be5cb6")
  assert.True(t, len(ic.testEmits) == 1)
  assert.True(t, ic.testEmits["databricks_service_principal[<unknown>] (application_id: 21aab5a7-ee70-4385-34d4-a77278be5cb6)"])
+
+ // unsuccessfull test
+ ic = importContextForTest()
+ ic.emitUserOrServicePrincipal("abc")
+ assert.True(t, len(ic.testEmits) == 0)
+}
+
+func TestEmitUserOrServicePrincipalForPath(t *testing.T) {
+ ic := importContextForTest()
+
+ ic.emitUserOrServicePrincipalForPath("/Users/user@domain.com/abc", "/Users")
+ assert.True(t, len(ic.testEmits) == 1)
+ assert.True(t, ic.testEmits["databricks_user[<unknown>] (user_name: user@domain.com)"])
+
+ // Negative cases
+ ic = importContextForTest()
+ ic.emitUserOrServicePrincipalForPath("/Shared/abc", "/Users")
+ assert.True(t, len(ic.testEmits) == 0)
+
+ ic = importContextForTest()
+ ic.emitUserOrServicePrincipalForPath("/Users/abc", "/Users")
+ assert.True(t, len(ic.testEmits) == 0)
+ ic = importContextForTest()
+ ic.emitUserOrServicePrincipalForPath("/Users/", "/Users")
+ assert.True(t, len(ic.testEmits) == 0)
+}
+
+func TestEmitNotebookOrRepo(t *testing.T) {
+ ic := importContextForTest()
+ ic.emitNotebookOrRepo("/Users/user@domain.com/abc")
+ assert.True(t, len(ic.testEmits) == 1)
+ assert.True(t, ic.testEmits["databricks_notebook[<unknown>] (id: /Users/user@domain.com/abc)"])
+
+ // test for repository
+ ic = importContextForTest()
+ ic.emitNotebookOrRepo("/Repos/user@domain.com/repo/abc")
+ assert.True(t, len(ic.testEmits) == 1)
+ assert.True(t, ic.testEmits["databricks_repo[<unknown>] (path: /Repos/user@domain.com/repo)"])
 }