Exporter: export permissions for databricks_notebook & databricks_directory (#1908)

Author: qili86

exporter/exporter_test.go

@@ -460,6 +460,7 @@ func TestImportingUsersGroupsSecretScopes(t *testing.T) {
 				Resource: "/api/2.0/secrets/acls/get?principal=users&scope=a",
 				Response: secrets.ACLItem{Permission: "READ", Principal: "users"},
 			},
+			emptyWorkspace,
 		}, func(ctx context.Context, client *common.DatabricksClient) {
 			tmpDir := fmt.Sprintf("/tmp/tf-%s", qa.RandomName())
 			defer os.RemoveAll(tmpDir)
@@ -527,6 +528,7 @@ func TestImportingNoResourcesError(t *testing.T) {
 					Scopes: []secrets.SecretScope{},
 				},
 			},
+			emptyWorkspace,
 		}, func(ctx context.Context, client *common.DatabricksClient) {
 			tmpDir := fmt.Sprintf("/tmp/tf-%s", qa.RandomName())
 			defer os.RemoveAll(tmpDir)
@@ -1588,6 +1590,11 @@ func TestImportingDLTPipelines(t *testing.T) {
 				Resource: "/api/2.0/permissions/pipelines/123",
 				Response: getJSONObject("test-data/get-pipeline-permissions.json"),
 			},
+			{
+				Method:   "GET",
+				Resource: "/api/2.0/permissions/notebooks/123",
+				Response: getJSONObject("test-data/get-notebook-permissions.json"),
+			},
 			{
 				Method:   "GET",
 				Resource: "/api/2.0/workspace/get-status?path=%2FUsers%2Fuser%40domain.com%2FTest%20DLT",

exporter/importables.go

@@ -749,8 +749,8 @@ var resourcesMap map[string]importable = map[string]importable{
 			{Path: "registered_model_id", Resource: "databricks_mlflow_model"},
 			{Path: "experiment_id", Resource: "databricks_mlflow_experiment"},
 			{Path: "repo_id", Resource: "databricks_repo"},
-			{Path: "directory_path", Resource: "databricks_directory"},
-			{Path: "notebook_path", Resource: "databricks_notebook"},
+			{Path: "directory_id", Resource: "databricks_directory", Match: "object_id"},
+			{Path: "notebook_id", Resource: "databricks_notebook", Match: "object_id"},
 			{Path: "access_control.user_name", Resource: "databricks_user", Match: "user_name"},
 			{Path: "access_control.group_name", Resource: "databricks_group", Match: "display_name"},
 			{Path: "access_control.service_principal_name", Resource: "databricks_service_principal", Match: "application_id"},
@@ -1163,6 +1163,28 @@ var resourcesMap map[string]importable = map[string]importable{
 			name := r.ID[1:] + ext[language] // todo: replace non-alphanum+/ with _
 			content, _ := base64.StdEncoding.DecodeString(contentB64)
 			fileName, err := ic.createFileIn("notebooks", name, []byte(content))
+			splits := strings.Split(r.Name, "_")
+			notebookId := splits[len(splits)-1]
+			directorySplits := strings.Split(r.ID, "/")
+			directorySplits = directorySplits[:len(directorySplits)-1]
+			directoryPath := strings.Join(directorySplits, "/")
+
+			if ic.meAdmin {
+				ic.Emit(&resource{
+					Resource: "databricks_permissions",
+					ID:       fmt.Sprintf("/notebooks/%s", notebookId),
+					Name:     "notebook_" + ic.Importables["databricks_notebook"].Name(ic, r.Data),
+				})
+			}
+
+			if ic.meAdmin {
+				ic.Emit(&resource{
+					Resource:  "databricks_directory",
+					Attribute: "path",
+					Value:     directoryPath,
+				})
+			}
+
 			if err != nil {
 				return err
 			}
@@ -1515,4 +1537,64 @@ var resourcesMap map[string]importable = map[string]importable{
 			{Path: "library.whl", Resource: "databricks_dbfs_file", Match: "dbfs_path"},
 		},
 	},
+	"databricks_directory": {
+		Service: "notebooks",
+		Name: func(ic *importContext, d *schema.ResourceData) string {
+			name := d.Get("path").(string)
+			if name == "" {
+				return d.Id()
+			} else {
+				name = nameNormalizationRegex.ReplaceAllString(name[1:], "_") + "_" +
+					strconv.FormatInt(int64(d.Get("object_id").(int)), 10)
+			}
+			return name
+		},
+		List: func(ic *importContext) error {
+			notebooksAPI := workspace.NewNotebooksAPI(ic.Context, ic.Client)
+			directoryList, err := notebooksAPI.ListDirectories("/", true)
+			if err != nil {
+				return err
+			}
+			for offset, directory := range directoryList {
+				if strings.HasPrefix(directory.Path, "/Repos") {
+					continue
+				}
+				ic.Emit(&resource{
+					Resource: "databricks_directory",
+					ID:       directory.Path,
+				})
+				if offset%50 == 0 {
+					log.Printf("[INFO] Scanned %d of %d directories",
+						offset+1, len(directoryList))
+				}
+			}
+			return nil
+		},
+		Import: func(ic *importContext, r *resource) error {
+
+			ic.emitUserOrServicePrincipalForPath(r.ID, "/Users")
+			splits := strings.Split(r.Name, "_")
+			directoryId := splits[len(splits)-1]
+
+			if ic.meAdmin {
+				ic.Emit(&resource{
+					Resource: "databricks_permissions",
+					ID:       fmt.Sprintf("/directories/%s", directoryId),
+					Name:     "directory_" + ic.Importables["databricks_directory"].Name(ic, r.Data),
+				})
+			}
+
+			if r.ID == "/Shared" || r.ID == "/Users" || ic.IsUserOrServicePrincipalDirectory(r.ID, "/Users") {
+				r.Mode = "data"
+			}
+
+			return nil
+
+		},
+		Body: resourceOrDataBlockBody,
+		Depends: []reference{
+			{Path: "path", Resource: "databricks_user", Match: "home", MatchType: MatchPrefix},
+			{Path: "path", Resource: "databricks_service_principal", Match: "home", MatchType: MatchPrefix},
+		},
+	},
 }

exporter/importables_test.go

@@ -783,6 +783,51 @@ func TestNotebookGeneration(t *testing.T) {
 	})
 }
 
+func TestDirectoryGeneration(t *testing.T) {
+	testGenerate(t, []qa.HTTPFixture{
+		{
+			Method:   "GET",
+			Resource: "/api/2.0/workspace/list?path=%2F",
+			Response: workspace.ObjectList{
+				Objects: []workspace.ObjectStatus{
+					{
+						ObjectID:   1234,
+						Path:       "/first",
+						ObjectType: "DIRECTORY",
+					},
+				},
+			},
+		},
+		{
+			Method:   "GET",
+			Resource: "/api/2.0/workspace/list?path=%2Ffirst",
+			Response: workspace.ObjectList{
+				Objects: []workspace.ObjectStatus{
+					{},
+				},
+			},
+		},
+		{
+			Method:   "GET",
+			Resource: "/api/2.0/workspace/get-status?path=%2Ffirst",
+			Response: workspace.ObjectStatus{
+				ObjectID:   1234,
+				ObjectType: "DIRECTORY",
+				Path:       "/first",
+			},
+		},
+	}, "notebooks", func(ic *importContext) {
+		err := resourcesMap["databricks_directory"].List(ic)
+		assert.NoError(t, err)
+
+		ic.generateHclForResources(nil)
+		assert.Equal(t, commands.TrimLeadingWhitespace(`
+		resource "databricks_directory" "first_1234" {
+		  path = "/first"
+		}`), string(ic.Files["notebooks"].Bytes()))
+	})
+}
+
 func TestGlobalInitScriptGen(t *testing.T) {
 	testGenerate(t, []qa.HTTPFixture{
 		{

exporter/test-data/get-notebook-permissions.json

@@ -0,0 +1,16 @@
+{
+  "object_id": "/notebooks/123",
+  "object_type": "notebook",
+  "access_control_list": [
+    {
+      "user_name": "[email protected]",
+      "all_permissions": [
+        {
+          "permission_level": "CAN_RUN",
+          "inherited": true,
+          "inherited_from_object": "/directories/"
+        }
+      ]
+    }
+  ]
+}

exporter/util.go

@@ -106,6 +106,24 @@ func (ic *importContext) emitUserOrServicePrincipalForPath(path, prefix string)
 	}
 }
 
+func (ic *importContext) IsUserOrServicePrincipalDirectory(path, prefix string) bool {
+	if strings.HasPrefix(path, prefix) {
+		parts := strings.SplitN(path, "/", 4)
+		if len(parts) == 3 {
+			userOrSPName := parts[2]
+			if strings.Contains(userOrSPName, "@") || uuidRegex.MatchString(userOrSPName) {
+				return true
+			} else {
+				return false
+			}
+		} else {
+			return false
+		}
+	} else {
+		return false
+	}
+}
+
 func (ic *importContext) emitNotebookOrRepo(path string) {
 	if strings.HasPrefix(path, "/Repos") {
 		ic.Emit(&resource{

exporter/util_test.go

@@ -84,3 +84,30 @@ func TestEmitNotebookOrRepo(t *testing.T) {
 	assert.True(t, len(ic.testEmits) == 1)
 	assert.True(t, ic.testEmits["databricks_repo[<unknown>] (path: /Repos/[email protected]/repo)"])
 }
+
+func TestIsUserOrServicePrincipalDirectory(t *testing.T) {
+
+	ic := importContextForTest()
+	result_false_partslength_more_than_3 := ic.IsUserOrServicePrincipalDirectory("/Users/[email protected]/abc", "/Users")
+	assert.False(t, result_false_partslength_more_than_3)
+
+	ic = importContextForTest()
+	result_false_partslength_less_than_3 := ic.IsUserOrServicePrincipalDirectory("/Users", "/Users")
+	assert.False(t, result_false_partslength_less_than_3)
+
+	ic = importContextForTest()
+	result_false_part2_empty := ic.IsUserOrServicePrincipalDirectory("/Users/", "/Users")
+	assert.False(t, result_false_part2_empty)
+
+	ic = importContextForTest()
+	result_false_notprefix_with_user := ic.IsUserOrServicePrincipalDirectory("/Shared", "/Users")
+	assert.False(t, result_false_notprefix_with_user)
+
+	ic = importContextForTest()
+	result_true_user_directory := ic.IsUserOrServicePrincipalDirectory("/Users/[email protected]", "/Users")
+	assert.True(t, result_true_user_directory)
+
+	ic = importContextForTest()
+	result_true_sp_directory := ic.IsUserOrServicePrincipalDirectory("/Users/0e561119-c5a0-4f29-b246-5a953adb9575", "/Users")
+	assert.True(t, result_true_sp_directory)
+}

permissions/resource_permissions_test.go

@@ -1527,3 +1527,86 @@ func TestResourcePermissionsUpdate_Sql_Queries(t *testing.T) {
 	assert.Equal(t, TestingUser, firstElem["user_name"])
 	assert.Equal(t, "CAN_RUN", firstElem["permission_level"])
 }
+
+func TestResourcePermissionsCreate_DirectoryPath(t *testing.T) {
+	d, err := qa.ResourceFixture{
+		Fixtures: []qa.HTTPFixture{
+			me,
+			{
+				Method:   http.MethodGet,
+				Resource: "/api/2.0/workspace/get-status?path=%2FFirst",
+				Response: workspace.ObjectStatus{
+					ObjectID:   123456,
+					ObjectType: "directory",
+				},
+			},
+			{
+				Method:   http.MethodPut,
+				Resource: "/api/2.0/permissions/directories/123456",
+				ExpectedRequest: AccessControlChangeList{
+					AccessControlList: []AccessControlChange{
+						{
+							UserName:        TestingUser,
+							PermissionLevel: "CAN_READ",
+						},
+					},
+				},
+			},
+			{
+				Method:   http.MethodGet,
+				Resource: "/api/2.0/permissions/directories/123456",
+				Response: ObjectACL{
+					ObjectID:   "/directories/123456",
+					ObjectType: "directory",
+					AccessControlList: []AccessControl{
+						{
+							UserName: TestingUser,
+							AllPermissions: []Permission{
+								{
+									PermissionLevel: "CAN_READ",
+									Inherited:       false,
+								},
+							},
+						},
+						{
+							UserName: TestingAdminUser,
+							AllPermissions: []Permission{
+								{
+									PermissionLevel: "CAN_RUN",
+									Inherited:       false,
+								},
+							},
+						},
+						{
+							UserName: TestingAdminUser,
+							AllPermissions: []Permission{
+								{
+									PermissionLevel: "CAN_MANAGE",
+									Inherited:       false,
+								},
+							},
+						},
+					},
+				},
+			},
+		},
+		Resource: ResourcePermissions(),
+		State: map[string]any{
+			"directory_path": "/First",
+			"access_control": []any{
+				map[string]any{
+					"user_name":        TestingUser,
+					"permission_level": "CAN_READ",
+				},
+			},
+		},
+		Create: true,
+	}.Apply(t)
+
+	assert.NoError(t, err, err)
+	ac := d.Get("access_control").(*schema.Set)
+	require.Equal(t, 1, len(ac.List()))
+	firstElem := ac.List()[0].(map[string]any)
+	assert.Equal(t, TestingUser, firstElem["user_name"])
+	assert.Equal(t, "CAN_READ", firstElem["permission_level"])
+}

workspace/resource_notebook.go

@@ -161,6 +161,35 @@ func (a NotebooksAPI) recursiveAddPaths(path string, pathList *[]ObjectStatus) e
 	return err
 }
 
+func (a NotebooksAPI) ListDirectories(path string, recursive bool) ([]ObjectStatus, error) {
+	if recursive {
+		var paths []ObjectStatus
+		err := a.recursiveAddDirectoryPaths(path, &paths)
+		if err != nil {
+			return nil, err
+		}
+		return paths, err
+	}
+	return a.list(path)
+}
+
+func (a NotebooksAPI) recursiveAddDirectoryPaths(path string, pathList *[]ObjectStatus) error {
+	directoryInfoList, err := a.list(path)
+	if err != nil {
+		return err
+	}
+	for _, v := range directoryInfoList {
+		if v.ObjectType == Directory {
+			*pathList = append(*pathList, v)
+			err := a.recursiveAddDirectoryPaths(v.Path, pathList)
+			if err != nil {
+				return err
+			}
+		}
+	}
+	return err
+}
+
 type ObjectList struct {
 	Objects []ObjectStatus `json:"objects,omitempty"`
 }