Exporter: add support for databricks_artifact_allowlist

Also refactored a bit to avoid duplicate calls to metastore summary API
databricks · Jan 7, 2024 · 4e6c9d5 · 4e6c9d5
1 parent 2fe11af
commit 4e6c9d5
Show file tree

Hide file tree

Showing 5 changed files with 84 additions and 23 deletions.
diff --git a/docs/guides/experimental-exporter.md b/docs/guides/experimental-exporter.md
@@ -77,6 +77,7 @@ Services are just logical groups of resources used for filtering and organizatio
 * `sql-endpoints` - **listing** [databricks_sql_endpoint](../resources/sql_endpoint.md) along with [databricks_sql_global_config](../resources/sql_global_config.md).
 * `sql-queries` - **listing** [databricks_sql_query](../resources/sql_query.md).
 * `storage` - only [databricks_dbfs_file](../resources/dbfs_file.md) referenced in other resources (libraries, init scripts, ...) will be downloaded locally and properly arranged into terraform state.
+* `uc-artifact-allowlist` - exports [databricks_artifact_allowlist](../resources/artifact_allowlist.md) resources for Unity Catalog Allow Lists attached to the current metastore.
 * `uc-system-schemas` - exports [databricks_system_schema](../resources/system_schema.md) resources for the UC metastore of the current workspace.
 * `users` - [databricks_user](../resources/user.md) and [databricks_service_principal](../resources/service_principal.md) are written to their own file, simply because of their amount. If you use SCIM provisioning, migrating workspaces is the only use case for importing `users` service.
 * `workspace` - [databricks_workspace_conf](../resources/workspace_conf.md) and [databricks_global_init_script](../resources/global_init_script.md)
@@ -101,6 +102,7 @@ Exporter aims to generate HCL code for most of the resources within the Databric
 | Resource | Generated code | Incremental |
 | --- | --- | --- |
 | [databricks_access_control_rule_set](../resources/access_control_rule_set.md) | Yes | No |
+| [databricks_artifact_allowlist](../resources/artifact_allowlist.md) | Yes | No |
 | [databricks_cluster](../resources/cluster.md) | Yes | No |
 | [databricks_cluster_policy](../resources/cluster_policy.md) | Yes | No |
 | [databricks_dbfs_file](../resources/dbfs_file.md) | Yes | No |

diff --git a/exporter/context.go b/exporter/context.go
@@ -19,6 +19,7 @@ import (
 	"time"
 
 	"github.com/databricks/databricks-sdk-go"
+	"github.com/databricks/databricks-sdk-go/service/catalog"
 	"github.com/databricks/databricks-sdk-go/service/compute"
 
 	"github.com/databricks/terraform-provider-databricks/commands"
@@ -140,6 +141,9 @@ type importContext struct {
 
 	builtInPolicies      map[string]compute.PolicyFamily
 	builtInPoliciesMutex sync.Mutex
+
+	// Workspace-level UC Metastore information
+	currentMetastore *catalog.GetMetastoreSummaryResponse
 }
 
 type mount struct {
@@ -321,6 +325,12 @@ func (ic *importContext) Run() error {
 				break
 			}
 		}
+		currentMetastore, err := ic.workspaceClient.Metastores.Summary(ic.Context)
+		if err == nil {
+			ic.currentMetastore = currentMetastore
+		} else {
+			log.Printf("[WARN] can't get current UC metastore: %v", err)
+		}
 	}
 	// Concurrent execution part
 	if ic.waitGroup == nil {

diff --git a/exporter/exporter_test.go b/exporter/exporter_test.go
@@ -12,6 +12,7 @@ import (
 	"time"
 
 	"github.com/databricks/databricks-sdk-go/apierr"
+	"github.com/databricks/databricks-sdk-go/service/catalog"
 	"github.com/databricks/databricks-sdk-go/service/compute"
 	"github.com/databricks/databricks-sdk-go/service/iam"
 	sdk_jobs "github.com/databricks/databricks-sdk-go/service/jobs"
@@ -382,6 +383,18 @@ var noCurrentMetastoreAttached = qa.HTTPFixture{
 	ReuseRequest: true,
 }
 
+var currentMetastoreResponse = &catalog.GetMetastoreSummaryResponse{
+	MetastoreId: "12345678-1234",
+	Name:        "test",
+}
+
+var currentMetastoreSuccess = qa.HTTPFixture{
+	Method:       "GET",
+	Resource:     "/api/2.1/unity-catalog/metastore_summary",
+	Response:     currentMetastoreResponse,
+	ReuseRequest: true,
+}
+
 func TestImportingUsersGroupsSecretScopes(t *testing.T) {
 	listSpFixtures := qa.ListServicePrincipalsFixtures([]iam.ServicePrincipal{
 		{
@@ -723,6 +736,7 @@ func TestImportingClusters(t *testing.T) {
 	qa.HTTPFixturesApply(t,
 		[]qa.HTTPFixture{
 			meAdminFixture,
+			noCurrentMetastoreAttached,
 			emptyRepos,
 			{
 				Method:   "GET",
@@ -1402,6 +1416,7 @@ func TestImportingSecrets(t *testing.T) {
 	qa.HTTPFixturesApply(t,
 		[]qa.HTTPFixture{
 			meAdminFixture,
+			noCurrentMetastoreAttached,
 			emptyRepos,
 			{
 				Method:   "GET",
@@ -1485,6 +1500,7 @@ func TestImportingGlobalInitScripts(t *testing.T) {
 	qa.HTTPFixturesApply(t,
 		[]qa.HTTPFixture{
 			meAdminFixture,
+			noCurrentMetastoreAttached,
 			emptyRepos,
 			emptyWorkspaceConf,
 			dummyWorkspaceConf,
@@ -1587,6 +1603,7 @@ func TestImportingRepos(t *testing.T) {
 	qa.HTTPFixturesApply(t,
 		[]qa.HTTPFixture{
 			meAdminFixture,
+			noCurrentMetastoreAttached,
 			userListIdUsernameFixture,
 			userListIdUsernameFixture2,
 			userListFixture,
@@ -1641,6 +1658,7 @@ func TestImportingIPAccessLists(t *testing.T) {
 	qa.HTTPFixturesApply(t,
 		[]qa.HTTPFixture{
 			meAdminFixture,
+			noCurrentMetastoreAttached,
 			emptyRepos,
 			emptyWorkspaceConf,
 			dummyWorkspaceConf,
@@ -1700,6 +1718,7 @@ func TestImportingSqlObjects(t *testing.T) {
 	qa.HTTPFixturesApply(t,
 		[]qa.HTTPFixture{
 			meAdminFixture,
+			noCurrentMetastoreAttached,
 			emptyRepos,
 			emptyIpAccessLIst,
 			emptyGlobalSQLConfig,
@@ -1838,6 +1857,7 @@ func TestImportingDLTPipelines(t *testing.T) {
 	qa.HTTPFixturesApply(t,
 		[]qa.HTTPFixture{
 			meAdminFixture,
+			noCurrentMetastoreAttached,
 			emptyRepos,
 			emptyWorkspace,
 			emptyIpAccessLIst,
@@ -2013,6 +2033,7 @@ func TestImportingDLTPipelinesMatchingOnly(t *testing.T) {
 	qa.HTTPFixturesApply(t,
 		[]qa.HTTPFixture{
 			meAdminFixture,
+			noCurrentMetastoreAttached,
 			emptyRepos,
 			emptyIpAccessLIst,
 			userListIdUsernameFixture,
@@ -2070,6 +2091,7 @@ func TestImportingGlobalSqlConfig(t *testing.T) {
 	qa.HTTPFixturesApply(t,
 		[]qa.HTTPFixture{
 			meAdminFixture,
+			noCurrentMetastoreAttached,
 			{
 				Method:   "GET",
 				Resource: "/api/2.0/sql/warehouses",
@@ -2113,6 +2135,7 @@ func TestImportingNotebooksWorkspaceFiles(t *testing.T) {
 	qa.HTTPFixturesApply(t,
 		[]qa.HTTPFixture{
 			meAdminFixture,
+			noCurrentMetastoreAttached,
 			emptyRepos,
 			emptyIpAccessLIst,
 			{
@@ -2165,6 +2188,7 @@ func TestImportingModelServing(t *testing.T) {
 	qa.HTTPFixturesApply(t,
 		[]qa.HTTPFixture{
 			meAdminFixture,
+			noCurrentMetastoreAttached,
 			emptyRepos,
 			emptyIpAccessLIst,
 			emptyWorkspace,
@@ -2215,6 +2239,7 @@ func TestImportingMlfloweWebhooks(t *testing.T) {
 	qa.HTTPFixturesApply(t,
 		[]qa.HTTPFixture{
 			meAdminFixture,
+			noCurrentMetastoreAttached,
 			emptyRepos,
 			emptyIpAccessLIst,
 			emptyWorkspace,
@@ -2308,6 +2333,7 @@ func TestIncrementalDLTAndMLflowWebhooks(t *testing.T) {
 	qa.HTTPFixturesApply(t,
 		[]qa.HTTPFixture{
 			meAdminFixture,
+			noCurrentMetastoreAttached,
 			emptyRepos,
 			emptyIpAccessLIst,
 			emptyWorkspace,
@@ -2418,6 +2444,7 @@ func TestImportingRunJobTask(t *testing.T) {
 	qa.HTTPFixturesApply(t,
 		[]qa.HTTPFixture{
 			meAdminFixture,
+			noCurrentMetastoreAttached,
 			emptyRepos,
 			emptyIpAccessLIst,
 			emptyWorkspace,

diff --git a/exporter/importables.go b/exporter/importables.go
@@ -2165,11 +2165,10 @@ var resourcesMap map[string]importable = map[string]importable{
 		WorkspaceLevel: true,
 		Service:        "uc-system-schemas",
 		List: func(ic *importContext) error {
-			summary, err := ic.workspaceClient.Metastores.Summary(ic.Context)
-			if err != nil {
-				return err
+			if ic.currentMetastore == nil {
+				return fmt.Errorf("there is no UC metastore information")
 			}
-			currentMetastore := summary.MetastoreId
+			currentMetastore := ic.currentMetastore.MetastoreId
 			systemSchemas, err := ic.workspaceClient.SystemSchemas.ListAll(ic.Context,
 				catalog.ListSystemSchemasRequest{MetastoreId: currentMetastore})
 			if err != nil {
@@ -2199,4 +2198,25 @@ var resourcesMap map[string]importable = map[string]importable{
 			return nil
 		},
 	},
+	"databricks_artifact_allowlist": {
+		WorkspaceLevel: true,
+		Service:        "uc-artifact-allowlist",
+		List: func(ic *importContext) error {
+			if ic.currentMetastore == nil {
+				return fmt.Errorf("there is no UC metastore information")
+			}
+			artifactTypes := []string{"INIT_SCRIPT", "LIBRARY_JAR", "LIBRARY_MAVEN"}
+			for _, v := range artifactTypes {
+				id := fmt.Sprintf("%s|%s", ic.currentMetastore.MetastoreId, v)
+				name := fmt.Sprintf("%s_%s_%s", v, ic.currentMetastore.Name, ic.currentMetastore.MetastoreId[:8])
+				ic.Emit(&resource{
+					Resource: "databricks_artifact_allowlist",
+					ID:       id,
+					Name:     nameNormalizationRegex.ReplaceAllString(name, "_"),
+				})
+			}
+			return nil
+		},
+		// TODO: add Depends & Import to emit corresponding UC Volumes when support for them is added
+	},
 }
diff --git a/exporter/importables_test.go b/exporter/importables_test.go
@@ -1297,21 +1297,12 @@ func TestIncrementalListDLT(t *testing.T) {
 	})
 }
 
-var currentMetastoreSuccess = qa.HTTPFixture{
-	Method:   "GET",
-	Resource: "/api/2.1/unity-catalog/metastore_summary",
-	Response: catalog.GetMetastoreSummaryResponse{
-		MetastoreId: "1234",
-	},
-	ReuseRequest: true,
-}
-
 func TestListSystemSchemasSuccess(t *testing.T) {
 	qa.HTTPFixturesApply(t, []qa.HTTPFixture{
 		currentMetastoreSuccess,
 		{
 			Method:   "GET",
-			Resource: "/api/2.1/unity-catalog/metastores/1234/systemschemas?",
+			Resource: fmt.Sprintf("/api/2.1/unity-catalog/metastores/%s/systemschemas?", currentMetastoreResponse.MetastoreId),
 			Response: catalog.ListSystemSchemasResponse{
 				Schemas: []catalog.SystemSchemaInfo{
 					{
@@ -1327,34 +1318,45 @@ func TestListSystemSchemasSuccess(t *testing.T) {
 		},
 	}, func(ctx context.Context, client *common.DatabricksClient) {
 		ic := importContextForTestWithClient(ctx, client)
+		ic.currentMetastore = currentMetastoreResponse
 		err := resourcesMap["databricks_system_schema"].List(ic)
 		assert.NoError(t, err)
 		assert.Equal(t, len(ic.testEmits), 1)
 	})
 }
 
 func TestListSystemSchemasErrorGetMetastore(t *testing.T) {
-	qa.HTTPFixturesApply(t, []qa.HTTPFixture{
-		noCurrentMetastoreAttached,
-	}, func(ctx context.Context, client *common.DatabricksClient) {
-		ic := importContextForTestWithClient(ctx, client)
-		err := resourcesMap["databricks_system_schema"].List(ic)
-		assert.EqualError(t, err, "nope")
-	})
+	ic := importContextForTest()
+	err := resourcesMap["databricks_system_schema"].List(ic)
+	assert.EqualError(t, err, "there is no UC metastore information")
 }
 
 func TestListSystemSchemasErrorListing(t *testing.T) {
 	qa.HTTPFixturesApply(t, []qa.HTTPFixture{
-		currentMetastoreSuccess,
 		{
 			Method:   "GET",
-			Resource: "/api/2.1/unity-catalog/metastores/1234/systemschemas?",
+			Resource: fmt.Sprintf("/api/2.1/unity-catalog/metastores/%s/systemschemas?", currentMetastoreResponse.MetastoreId),
 			Status:   404,
 			Response: apierr.NotFound("nope"),
 		},
 	}, func(ctx context.Context, client *common.DatabricksClient) {
 		ic := importContextForTestWithClient(ctx, client)
+		ic.currentMetastore = currentMetastoreResponse
 		err := resourcesMap["databricks_system_schema"].List(ic)
 		assert.EqualError(t, err, "nope")
 	})
 }
+
+func TestListUcAllowListError(t *testing.T) {
+	ic := importContextForTest()
+	err := resourcesMap["databricks_artifact_allowlist"].List(ic)
+	assert.EqualError(t, err, "there is no UC metastore information")
+}
+
+func TestListUcAllowListSuccess(t *testing.T) {
+	ic := importContextForTest()
+	ic.currentMetastore = currentMetastoreResponse
+	err := resourcesMap["databricks_artifact_allowlist"].List(ic)
+	assert.NoError(t, err)
+	assert.Equal(t, len(ic.testEmits), 3)
+}