Skip to content

Commit ebfc026

Browse files
alexottalexott
authored
Exporter: add support for databricks_artifact_allowlist (#3083)
Also refactored a bit to avoid duplicate calls to metastore summary API
1 parent 9a08630 commit ebfc026

File tree

5 files changed

+84
-23
lines changed

5 files changed

+84
-23
lines changed

‎docs/guides/experimental-exporter.md‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -77,6 +77,7 @@ Services are just logical groups of resources used for filtering and organizatio
7777
* `sql-endpoints` - **listing** [databricks_sql_endpoint](../resources/sql_endpoint.md) along with [databricks_sql_global_config](../resources/sql_global_config.md).
7878
* `sql-queries` - **listing** [databricks_sql_query](../resources/sql_query.md).
7979
* `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.
80+
* `uc-artifact-allowlist` - exports [databricks_artifact_allowlist](../resources/artifact_allowlist.md) resources for Unity Catalog Allow Lists attached to the current metastore.
8081
* `uc-system-schemas` - exports [databricks_system_schema](../resources/system_schema.md) resources for the UC metastore of the current workspace.
8182
* `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.
8283
* `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
101102
| Resource | Generated code | Incremental |
102103
| --- | --- | --- |
103104
| [databricks_access_control_rule_set](../resources/access_control_rule_set.md) | Yes | No |
105+
| [databricks_artifact_allowlist](../resources/artifact_allowlist.md) | Yes | No |
104106
| [databricks_cluster](../resources/cluster.md) | Yes | No |
105107
| [databricks_cluster_policy](../resources/cluster_policy.md) | Yes | No |
106108
| [databricks_dbfs_file](../resources/dbfs_file.md) | Yes | No |

‎exporter/context.go‎

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@ import (
1919
"time"
2020

2121
"github.com/databricks/databricks-sdk-go"
22+
"github.com/databricks/databricks-sdk-go/service/catalog"
2223
"github.com/databricks/databricks-sdk-go/service/compute"
2324

2425
"github.com/databricks/terraform-provider-databricks/commands"
@@ -140,6 +141,9 @@ type importContext struct {
140141

141142
builtInPolicies map[string]compute.PolicyFamily
142143
builtInPoliciesMutex sync.Mutex
144+
145+
// Workspace-level UC Metastore information
146+
currentMetastore *catalog.GetMetastoreSummaryResponse
143147
}
144148

145149
type mount struct {
@@ -321,6 +325,12 @@ func (ic *importContext) Run() error {
321325
break
322326
}
323327
}
328+
currentMetastore, err := ic.workspaceClient.Metastores.Summary(ic.Context)
329+
if err == nil {
330+
ic.currentMetastore = currentMetastore
331+
} else {
332+
log.Printf("[WARN] can't get current UC metastore: %v", err)
333+
}
324334
}
325335
// Concurrent execution part
326336
if ic.waitGroup == nil {

‎exporter/exporter_test.go‎

Lines changed: 27 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,7 @@ import (
1212
"time"
1313

1414
"github.com/databricks/databricks-sdk-go/apierr"
15+
"github.com/databricks/databricks-sdk-go/service/catalog"
1516
"github.com/databricks/databricks-sdk-go/service/compute"
1617
"github.com/databricks/databricks-sdk-go/service/iam"
1718
sdk_jobs "github.com/databricks/databricks-sdk-go/service/jobs"
@@ -382,6 +383,18 @@ var noCurrentMetastoreAttached = qa.HTTPFixture{
382383
ReuseRequest: true,
383384
}
384385

386+
var currentMetastoreResponse = &catalog.GetMetastoreSummaryResponse{
387+
MetastoreId: "12345678-1234",
388+
Name: "test",
389+
}
390+
391+
var currentMetastoreSuccess = qa.HTTPFixture{
392+
Method: "GET",
393+
Resource: "/api/2.1/unity-catalog/metastore_summary",
394+
Response: currentMetastoreResponse,
395+
ReuseRequest: true,
396+
}
397+
385398
func TestImportingUsersGroupsSecretScopes(t *testing.T) {
386399
listSpFixtures := qa.ListServicePrincipalsFixtures([]iam.ServicePrincipal{
387400
{
@@ -723,6 +736,7 @@ func TestImportingClusters(t *testing.T) {
723736
qa.HTTPFixturesApply(t,
724737
[]qa.HTTPFixture{
725738
meAdminFixture,
739+
noCurrentMetastoreAttached,
726740
emptyRepos,
727741
{
728742
Method: "GET",
@@ -1402,6 +1416,7 @@ func TestImportingSecrets(t *testing.T) {
14021416
qa.HTTPFixturesApply(t,
14031417
[]qa.HTTPFixture{
14041418
meAdminFixture,
1419+
noCurrentMetastoreAttached,
14051420
emptyRepos,
14061421
{
14071422
Method: "GET",
@@ -1485,6 +1500,7 @@ func TestImportingGlobalInitScripts(t *testing.T) {
14851500
qa.HTTPFixturesApply(t,
14861501
[]qa.HTTPFixture{
14871502
meAdminFixture,
1503+
noCurrentMetastoreAttached,
14881504
emptyRepos,
14891505
emptyWorkspaceConf,
14901506
dummyWorkspaceConf,
@@ -1587,6 +1603,7 @@ func TestImportingRepos(t *testing.T) {
15871603
qa.HTTPFixturesApply(t,
15881604
[]qa.HTTPFixture{
15891605
meAdminFixture,
1606+
noCurrentMetastoreAttached,
15901607
userListIdUsernameFixture,
15911608
userListIdUsernameFixture2,
15921609
userListFixture,
@@ -1641,6 +1658,7 @@ func TestImportingIPAccessLists(t *testing.T) {
16411658
qa.HTTPFixturesApply(t,
16421659
[]qa.HTTPFixture{
16431660
meAdminFixture,
1661+
noCurrentMetastoreAttached,
16441662
emptyRepos,
16451663
emptyWorkspaceConf,
16461664
dummyWorkspaceConf,
@@ -1700,6 +1718,7 @@ func TestImportingSqlObjects(t *testing.T) {
17001718
qa.HTTPFixturesApply(t,
17011719
[]qa.HTTPFixture{
17021720
meAdminFixture,
1721+
noCurrentMetastoreAttached,
17031722
emptyRepos,
17041723
emptyIpAccessLIst,
17051724
emptyGlobalSQLConfig,
@@ -1838,6 +1857,7 @@ func TestImportingDLTPipelines(t *testing.T) {
18381857
qa.HTTPFixturesApply(t,
18391858
[]qa.HTTPFixture{
18401859
meAdminFixture,
1860+
noCurrentMetastoreAttached,
18411861
emptyRepos,
18421862
emptyWorkspace,
18431863
emptyIpAccessLIst,
@@ -2013,6 +2033,7 @@ func TestImportingDLTPipelinesMatchingOnly(t *testing.T) {
20132033
qa.HTTPFixturesApply(t,
20142034
[]qa.HTTPFixture{
20152035
meAdminFixture,
2036+
noCurrentMetastoreAttached,
20162037
emptyRepos,
20172038
emptyIpAccessLIst,
20182039
userListIdUsernameFixture,
@@ -2070,6 +2091,7 @@ func TestImportingGlobalSqlConfig(t *testing.T) {
20702091
qa.HTTPFixturesApply(t,
20712092
[]qa.HTTPFixture{
20722093
meAdminFixture,
2094+
noCurrentMetastoreAttached,
20732095
{
20742096
Method: "GET",
20752097
Resource: "/api/2.0/sql/warehouses",
@@ -2113,6 +2135,7 @@ func TestImportingNotebooksWorkspaceFiles(t *testing.T) {
21132135
qa.HTTPFixturesApply(t,
21142136
[]qa.HTTPFixture{
21152137
meAdminFixture,
2138+
noCurrentMetastoreAttached,
21162139
emptyRepos,
21172140
emptyIpAccessLIst,
21182141
{
@@ -2165,6 +2188,7 @@ func TestImportingModelServing(t *testing.T) {
21652188
qa.HTTPFixturesApply(t,
21662189
[]qa.HTTPFixture{
21672190
meAdminFixture,
2191+
noCurrentMetastoreAttached,
21682192
emptyRepos,
21692193
emptyIpAccessLIst,
21702194
emptyWorkspace,
@@ -2215,6 +2239,7 @@ func TestImportingMlfloweWebhooks(t *testing.T) {
22152239
qa.HTTPFixturesApply(t,
22162240
[]qa.HTTPFixture{
22172241
meAdminFixture,
2242+
noCurrentMetastoreAttached,
22182243
emptyRepos,
22192244
emptyIpAccessLIst,
22202245
emptyWorkspace,
@@ -2308,6 +2333,7 @@ func TestIncrementalDLTAndMLflowWebhooks(t *testing.T) {
23082333
qa.HTTPFixturesApply(t,
23092334
[]qa.HTTPFixture{
23102335
meAdminFixture,
2336+
noCurrentMetastoreAttached,
23112337
emptyRepos,
23122338
emptyIpAccessLIst,
23132339
emptyWorkspace,
@@ -2418,6 +2444,7 @@ func TestImportingRunJobTask(t *testing.T) {
24182444
qa.HTTPFixturesApply(t,
24192445
[]qa.HTTPFixture{
24202446
meAdminFixture,
2447+
noCurrentMetastoreAttached,
24212448
emptyRepos,
24222449
emptyIpAccessLIst,
24232450
emptyWorkspace,

‎exporter/importables.go‎

Lines changed: 24 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -2180,11 +2180,10 @@ var resourcesMap map[string]importable = map[string]importable{
21802180
WorkspaceLevel: true,
21812181
Service: "uc-system-schemas",
21822182
List: func(ic *importContext) error {
2183-
summary, err := ic.workspaceClient.Metastores.Summary(ic.Context)
2184-
if err != nil {
2185-
return err
2183+
if ic.currentMetastore == nil {
2184+
return fmt.Errorf("there is no UC metastore information")
21862185
}
2187-
currentMetastore := summary.MetastoreId
2186+
currentMetastore := ic.currentMetastore.MetastoreId
21882187
systemSchemas, err := ic.workspaceClient.SystemSchemas.ListAll(ic.Context,
21892188
catalog.ListSystemSchemasRequest{MetastoreId: currentMetastore})
21902189
if err != nil {
@@ -2214,4 +2213,25 @@ var resourcesMap map[string]importable = map[string]importable{
22142213
return nil
22152214
},
22162215
},
2216+
"databricks_artifact_allowlist": {
2217+
WorkspaceLevel: true,
2218+
Service: "uc-artifact-allowlist",
2219+
List: func(ic *importContext) error {
2220+
if ic.currentMetastore == nil {
2221+
return fmt.Errorf("there is no UC metastore information")
2222+
}
2223+
artifactTypes := []string{"INIT_SCRIPT", "LIBRARY_JAR", "LIBRARY_MAVEN"}
2224+
for _, v := range artifactTypes {
2225+
id := fmt.Sprintf("%s|%s", ic.currentMetastore.MetastoreId, v)
2226+
name := fmt.Sprintf("%s_%s_%s", v, ic.currentMetastore.Name, ic.currentMetastore.MetastoreId[:8])
2227+
ic.Emit(&resource{
2228+
Resource: "databricks_artifact_allowlist",
2229+
ID: id,
2230+
Name: nameNormalizationRegex.ReplaceAllString(name, "_"),
2231+
})
2232+
}
2233+
return nil
2234+
},
2235+
// TODO: add Depends & Import to emit corresponding UC Volumes when support for them is added
2236+
},
22172237
}

‎exporter/importables_test.go‎

Lines changed: 21 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -1297,21 +1297,12 @@ func TestIncrementalListDLT(t *testing.T) {
12971297
})
12981298
}
12991299

1300-
var currentMetastoreSuccess = qa.HTTPFixture{
1301-
Method: "GET",
1302-
Resource: "/api/2.1/unity-catalog/metastore_summary",
1303-
Response: catalog.GetMetastoreSummaryResponse{
1304-
MetastoreId: "1234",
1305-
},
1306-
ReuseRequest: true,
1307-
}
1308-
13091300
func TestListSystemSchemasSuccess(t *testing.T) {
13101301
qa.HTTPFixturesApply(t, []qa.HTTPFixture{
13111302
currentMetastoreSuccess,
13121303
{
13131304
Method: "GET",
1314-
Resource: "/api/2.1/unity-catalog/metastores/1234/systemschemas?",
1305+
Resource: fmt.Sprintf("/api/2.1/unity-catalog/metastores/%s/systemschemas?", currentMetastoreResponse.MetastoreId),
13151306
Response: catalog.ListSystemSchemasResponse{
13161307
Schemas: []catalog.SystemSchemaInfo{
13171308
{
@@ -1327,34 +1318,45 @@ func TestListSystemSchemasSuccess(t *testing.T) {
13271318
},
13281319
}, func(ctx context.Context, client *common.DatabricksClient) {
13291320
ic := importContextForTestWithClient(ctx, client)
1321+
ic.currentMetastore = currentMetastoreResponse
13301322
err := resourcesMap["databricks_system_schema"].List(ic)
13311323
assert.NoError(t, err)
13321324
assert.Equal(t, len(ic.testEmits), 1)
13331325
})
13341326
}
13351327

13361328
func TestListSystemSchemasErrorGetMetastore(t *testing.T) {
1337-
qa.HTTPFixturesApply(t, []qa.HTTPFixture{
1338-
noCurrentMetastoreAttached,
1339-
}, func(ctx context.Context, client *common.DatabricksClient) {
1340-
ic := importContextForTestWithClient(ctx, client)
1341-
err := resourcesMap["databricks_system_schema"].List(ic)
1342-
assert.EqualError(t, err, "nope")
1343-
})
1329+
ic := importContextForTest()
1330+
err := resourcesMap["databricks_system_schema"].List(ic)
1331+
assert.EqualError(t, err, "there is no UC metastore information")
13441332
}
13451333

13461334
func TestListSystemSchemasErrorListing(t *testing.T) {
13471335
qa.HTTPFixturesApply(t, []qa.HTTPFixture{
1348-
currentMetastoreSuccess,
13491336
{
13501337
Method: "GET",
1351-
Resource: "/api/2.1/unity-catalog/metastores/1234/systemschemas?",
1338+
Resource: fmt.Sprintf("/api/2.1/unity-catalog/metastores/%s/systemschemas?", currentMetastoreResponse.MetastoreId),
13521339
Status: 404,
13531340
Response: apierr.NotFound("nope"),
13541341
},
13551342
}, func(ctx context.Context, client *common.DatabricksClient) {
13561343
ic := importContextForTestWithClient(ctx, client)
1344+
ic.currentMetastore = currentMetastoreResponse
13571345
err := resourcesMap["databricks_system_schema"].List(ic)
13581346
assert.EqualError(t, err, "nope")
13591347
})
13601348
}
1349+
1350+
func TestListUcAllowListError(t *testing.T) {
1351+
ic := importContextForTest()
1352+
err := resourcesMap["databricks_artifact_allowlist"].List(ic)
1353+
assert.EqualError(t, err, "there is no UC metastore information")
1354+
}
1355+
1356+
func TestListUcAllowListSuccess(t *testing.T) {
1357+
ic := importContextForTest()
1358+
ic.currentMetastore = currentMetastoreResponse
1359+
err := resourcesMap["databricks_artifact_allowlist"].List(ic)
1360+
assert.NoError(t, err)
1361+
assert.Equal(t, len(ic.testEmits), 3)
1362+
}

0 commit comments

Comments
 (0)