Skip to content

Commit 72753b3

Browse files
authored
[Exporter] Use Go SDK to list clusters (#4385)
## Changes <!-- Summary of your changes that are easy to understand --> This helps us with the following: - use paginated output that is more efficient - automatically filter non-interactive clusters instead of doing it ourselves P.S. We still have usage of 2.0 API coming from mounts and some other places ## Tests <!-- How is this tested? Please see the checklist below and also describe any other relevant tests --> - [ ] `make test` run locally - [x] relevant change in `docs/` folder - [ ] covered with integration tests in `internal/acceptance` - [x] using Go SDK - [ ] using TF Plugin Framework
1 parent 1f7c015 commit 72753b3

File tree

6 files changed

+103
-93
lines changed

6 files changed

+103
-93
lines changed

exporter/exporter_test.go

Lines changed: 9 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -192,7 +192,7 @@ func TestImportingMounts(t *testing.T) {
192192
{
193193
Method: "POST",
194194
ReuseRequest: true,
195-
Resource: "/api/2.0/clusters/events",
195+
Resource: "/api/2.1/clusters/events",
196196
Response: clusters.EventsResponse{
197197
Events: []clusters.ClusterEvent{},
198198
},
@@ -674,7 +674,7 @@ func TestImportingUsersGroupsSecretScopes(t *testing.T) {
674674
},
675675
{
676676
Method: "GET",
677-
Resource: "/api/2.0/clusters/list",
677+
Resource: "/api/2.1/clusters/list?filter_by.cluster_sources=UI&filter_by.cluster_sources=API&page_size=100",
678678
Response: clusters.ClusterList{},
679679
},
680680
{
@@ -809,7 +809,7 @@ func TestImportingNoResourcesError(t *testing.T) {
809809
},
810810
{
811811
Method: "GET",
812-
Resource: "/api/2.0/clusters/list",
812+
Resource: "/api/2.1/clusters/list?filter_by.cluster_sources=UI&filter_by.cluster_sources=API&page_size=100",
813813
Response: clusters.ClusterList{},
814814
},
815815
{
@@ -864,7 +864,7 @@ func TestImportingClusters(t *testing.T) {
864864
},
865865
{
866866
Method: "GET",
867-
Resource: "/api/2.0/clusters/list",
867+
Resource: "/api/2.1/clusters/list?filter_by.cluster_sources=UI&filter_by.cluster_sources=API&page_size=100",
868868
Response: getJSONObject("test-data/clusters-list-response.json"),
869869
ReuseRequest: true,
870870
},
@@ -936,9 +936,10 @@ func TestImportingClusters(t *testing.T) {
936936
Response: getJSONObject("test-data/get-cluster-awscluster-response.json"),
937937
},
938938
{
939-
Method: "GET",
940-
Resource: "/api/2.0/libraries/cluster-status?cluster_id=awscluster",
941-
Response: getJSONObject("test-data/libraries-cluster-status-test2.json"),
939+
Method: "GET",
940+
Resource: "/api/2.0/libraries/cluster-status?cluster_id=awscluster",
941+
Response: getJSONObject("test-data/libraries-cluster-status-test2.json"),
942+
ReuseRequest: true,
942943
},
943944
{
944945
Method: "GET",
@@ -1593,7 +1594,7 @@ func TestImportingSecrets(t *testing.T) {
15931594
},
15941595
{
15951596
Method: "GET",
1596-
Resource: "/api/2.0/clusters/list",
1597+
Resource: "/api/2.1/clusters/list?filter_by.cluster_sources=UI&filter_by.cluster_sources=API&page_size=100",
15971598
Response: clusters.ClusterList{},
15981599
},
15991600
{

exporter/impl_compute.go

Lines changed: 91 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,91 @@
1+
package exporter
2+
3+
import (
4+
"log"
5+
"strings"
6+
7+
sdk_compute "github.com/databricks/databricks-sdk-go/service/compute"
8+
)
9+
10+
func listClusters(ic *importContext) error {
11+
lastActiveMs := ic.getLastActiveMs()
12+
interactiveClusters := []sdk_compute.ClusterSource{sdk_compute.ClusterSourceUi, sdk_compute.ClusterSourceApi}
13+
14+
it := ic.workspaceClient.Clusters.List(ic.Context, sdk_compute.ListClustersRequest{
15+
FilterBy: &sdk_compute.ListClustersFilterBy{
16+
ClusterSources: interactiveClusters,
17+
},
18+
PageSize: 100,
19+
})
20+
i := 0
21+
for it.HasNext(ic.Context) {
22+
c, err := it.Next(ic.Context)
23+
if err != nil {
24+
return err
25+
}
26+
i++
27+
28+
if strings.HasPrefix(c.ClusterName, "terraform-") {
29+
log.Printf("[INFO] Skipping terraform-specific cluster %s", c.ClusterName)
30+
continue
31+
}
32+
if !ic.MatchesName(c.ClusterName) {
33+
log.Printf("[INFO] Skipping %s because it doesn't match %s", c.ClusterName, ic.match)
34+
continue
35+
}
36+
if c.LastRestartedTime > 0 && c.LastRestartedTime < lastActiveMs {
37+
log.Printf("[INFO] Old inactive cluster %s", c.ClusterName)
38+
continue
39+
}
40+
ic.Emit(&resource{
41+
Resource: "databricks_cluster",
42+
ID: c.ClusterId,
43+
})
44+
if i%50 == 0 {
45+
log.Printf("[INFO] Scanned %d clusters", i)
46+
}
47+
}
48+
return nil
49+
}
50+
51+
func (ic *importContext) importCluster(c *sdk_compute.ClusterSpec) {
52+
if c == nil {
53+
return
54+
}
55+
if c.AwsAttributes != nil && c.AwsAttributes.InstanceProfileArn != "" {
56+
ic.Emit(&resource{
57+
Resource: "databricks_instance_profile",
58+
ID: c.AwsAttributes.InstanceProfileArn,
59+
})
60+
}
61+
if c.InstancePoolId != "" {
62+
// set enable_elastic_disk to false, and remove aws/gcp/azure_attributes
63+
ic.Emit(&resource{
64+
Resource: "databricks_instance_pool",
65+
ID: c.InstancePoolId,
66+
})
67+
}
68+
if c.DriverInstancePoolId != "" {
69+
ic.Emit(&resource{
70+
Resource: "databricks_instance_pool",
71+
ID: c.DriverInstancePoolId,
72+
})
73+
}
74+
if c.PolicyId != "" {
75+
ic.Emit(&resource{
76+
Resource: "databricks_cluster_policy",
77+
ID: c.PolicyId,
78+
})
79+
}
80+
ic.emitInitScripts(c.InitScripts)
81+
ic.emitSecretsFromSecretsPathMap(c.SparkConf)
82+
ic.emitSecretsFromSecretsPathMap(c.SparkEnvVars)
83+
ic.emitUserOrServicePrincipal(c.SingleUserName)
84+
if c.Kind.String() != "" && c.SingleUserName != "" {
85+
ic.Emit(&resource{
86+
Resource: "databricks_group",
87+
Attribute: "display_name",
88+
Value: c.SingleUserName,
89+
})
90+
}
91+
}

exporter/importables.go

Lines changed: 1 addition & 36 deletions
Original file line numberDiff line numberDiff line change
@@ -28,7 +28,6 @@ import (
2828
"github.com/databricks/databricks-sdk-go/service/vectorsearch"
2929
sdk_workspace "github.com/databricks/databricks-sdk-go/service/workspace"
3030
tfcatalog "github.com/databricks/terraform-provider-databricks/catalog"
31-
"github.com/databricks/terraform-provider-databricks/clusters"
3231
"github.com/databricks/terraform-provider-databricks/common"
3332
"github.com/databricks/terraform-provider-databricks/jobs"
3433
"github.com/databricks/terraform-provider-databricks/mws"
@@ -330,41 +329,7 @@ var resourcesMap map[string]importable = map[string]importable{
330329
{Path: "init_scripts.workspace.destination", Resource: "databricks_repo", Match: "path",
331330
MatchType: MatchPrefix, SearchValueTransformFunc: appendEndingSlashToDirName},
332331
},
333-
List: func(ic *importContext) error {
334-
clusters, err := clusters.NewClustersAPI(ic.Context, ic.Client).List()
335-
if err != nil {
336-
return err
337-
}
338-
lastActiveMs := ic.getLastActiveMs()
339-
nonInteractiveClusters := []string{"JOB", "MODELS", "PIPELINE_MAINTENANCE", "PIPELINE", "SQL"}
340-
for offset, c := range clusters {
341-
if slices.Contains(nonInteractiveClusters, string(c.ClusterSource)) {
342-
// TODO: Should we check cluster name as well?
343-
// jobRunClusterNameRegex = regexp.MustCompile(`^job-\d+-run-\d+$`)
344-
// jobRunClusterNameRegex.MatchString(c.ClusterName)
345-
log.Printf("[INFO] Skipping non-interactive cluster %s", c.ClusterID)
346-
continue
347-
}
348-
if strings.HasPrefix(c.ClusterName, "terraform-") {
349-
log.Printf("[INFO] Skipping terraform-specific cluster %s", c.ClusterName)
350-
continue
351-
}
352-
if !ic.MatchesName(c.ClusterName) {
353-
log.Printf("[INFO] Skipping %s because it doesn't match %s", c.ClusterName, ic.match)
354-
continue
355-
}
356-
if c.LastActivityTime > 0 && c.LastActivityTime < lastActiveMs {
357-
log.Printf("[INFO] Older inactive cluster %s", c.ClusterName)
358-
continue
359-
}
360-
ic.Emit(&resource{
361-
Resource: "databricks_cluster",
362-
ID: c.ClusterID,
363-
})
364-
log.Printf("[INFO] Scanned %d of %d clusters", offset+1, len(clusters))
365-
}
366-
return nil
367-
},
332+
List: listClusters,
368333
Import: func(ic *importContext, r *resource) error {
369334
var c compute.ClusterSpec
370335
s := ic.Resources["databricks_cluster"].Schema

exporter/importables_test.go

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -403,7 +403,7 @@ func TestClusterListFails(t *testing.T) {
403403
qa.HTTPFixturesApply(t, []qa.HTTPFixture{
404404
{
405405
Method: "GET",
406-
Resource: "/api/2.0/clusters/list",
406+
Resource: "/api/2.1/clusters/list?filter_by.cluster_sources=UI&filter_by.cluster_sources=API&page_size=100",
407407
Status: 404,
408408
Response: apierr.NotFound("nope"),
409409
},
@@ -418,7 +418,7 @@ func TestClusterList_NoNameMatch(t *testing.T) {
418418
qa.HTTPFixturesApply(t, []qa.HTTPFixture{
419419
{
420420
Method: "GET",
421-
Resource: "/api/2.0/clusters/list",
421+
Resource: "/api/2.1/clusters/list?filter_by.cluster_sources=UI&filter_by.cluster_sources=API&page_size=100",
422422
Response: clusters.ClusterList{
423423
Clusters: []clusters.ClusterInfo{
424424
{

exporter/test-data/clusters-list-response.json

Lines changed: 0 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -4,11 +4,6 @@
44
"cluster_id": "23443",
55
"cluster_name": "terraform-abc"
66
},
7-
{
8-
"cluster_id": "23443",
9-
"cluster_name": "job-34234234-run-32423432",
10-
"cluster_source": "JOB"
11-
},
127
{
138
"autotermination_minutes": 120,
149
"azure_attributes": {

exporter/util_compute.go

Lines changed: 0 additions & 42 deletions
Original file line numberDiff line numberDiff line change
@@ -32,48 +32,6 @@ func (ic *importContext) emitInitScripts(initScripts []compute.InitScriptInfo) {
3232
}
3333
}
3434

35-
func (ic *importContext) importCluster(c *compute.ClusterSpec) {
36-
if c == nil {
37-
return
38-
}
39-
if c.AwsAttributes != nil && c.AwsAttributes.InstanceProfileArn != "" {
40-
ic.Emit(&resource{
41-
Resource: "databricks_instance_profile",
42-
ID: c.AwsAttributes.InstanceProfileArn,
43-
})
44-
}
45-
if c.InstancePoolId != "" {
46-
// set enable_elastic_disk to false, and remove aws/gcp/azure_attributes
47-
ic.Emit(&resource{
48-
Resource: "databricks_instance_pool",
49-
ID: c.InstancePoolId,
50-
})
51-
}
52-
if c.DriverInstancePoolId != "" {
53-
ic.Emit(&resource{
54-
Resource: "databricks_instance_pool",
55-
ID: c.DriverInstancePoolId,
56-
})
57-
}
58-
if c.PolicyId != "" {
59-
ic.Emit(&resource{
60-
Resource: "databricks_cluster_policy",
61-
ID: c.PolicyId,
62-
})
63-
}
64-
ic.emitInitScripts(c.InitScripts)
65-
ic.emitSecretsFromSecretsPathMap(c.SparkConf)
66-
ic.emitSecretsFromSecretsPathMap(c.SparkEnvVars)
67-
ic.emitUserOrServicePrincipal(c.SingleUserName)
68-
if c.Kind.String() != "" && c.SingleUserName != "" {
69-
ic.Emit(&resource{
70-
Resource: "databricks_group",
71-
Attribute: "display_name",
72-
Value: c.SingleUserName,
73-
})
74-
}
75-
}
76-
7735
func (ic *importContext) emitSecretsFromSecretPathString(v string) {
7836
if res := secretPathRegex.FindStringSubmatch(v); res != nil {
7937
ic.Emit(&resource{

0 commit comments

Comments
 (0)