refactor databricks_clusters data source to Go SDK

databricks · Jun 20, 2024 · c792976 · c792976
1 parent c91b65f
commit c792976
Show file tree

Hide file tree

Showing 2 changed files with 53 additions and 73 deletions.
diff --git a/clusters/data_clusters.go b/clusters/data_clusters.go
@@ -4,42 +4,32 @@ import (
 	"context"
 	"strings"
 
+	"github.com/databricks/databricks-sdk-go"
+	"github.com/databricks/databricks-sdk-go/service/compute"
 	"github.com/databricks/terraform-provider-databricks/common"
-	"github.com/hashicorp/terraform-plugin-sdk/v2/helper/schema"
 )
 
 func DataSourceClusters() common.Resource {
-	return common.Resource{
-		Read: func(ctx context.Context, d *schema.ResourceData, i *common.DatabricksClient) error {
-			clusters, err := NewClustersAPI(ctx, i).List()
-			if err != nil {
-				return err
+	return common.WorkspaceData(func(ctx context.Context, data *struct {
+		Id                  string   `json:"id,omitempty" tf:"computed"`
+		Ids                 []string `json:"ids,omitempty" tf:"computed,slice_set"`
+		ClusterNameContains string   `json:"cluster_name_contains"`
+	}, w *databricks.WorkspaceClient) error {
+		clusters, err := w.Clusters.ListAll(ctx, compute.ListClustersRequest{})
+		if err != nil {
+			return err
+		}
+		ids := make([]string, 0, len(clusters))
+		name_contains := strings.ToLower(data.ClusterNameContains)
+		for _, v := range clusters {
+			match_name := strings.Contains(strings.ToLower(v.ClusterName), name_contains)
+			if name_contains != "" && !match_name {
+				continue
 			}
-			ids := schema.NewSet(schema.HashString, []any{})
-			name_contains := strings.ToLower(d.Get("cluster_name_contains").(string))
-			for _, v := range clusters {
-				match_name := strings.Contains(strings.ToLower(v.ClusterName), name_contains)
-				if name_contains != "" && !match_name {
-					continue
-				}
-				ids.Add(v.ClusterID)
-			}
-			d.Set("ids", ids)
-			d.SetId("_")
-			return nil
-		},
-		Schema: map[string]*schema.Schema{
-			"ids": {
-				Computed: true,
-				Type:     schema.TypeSet,
-				Elem: &schema.Schema{
-					Type: schema.TypeString,
-				},
-			},
-			"cluster_name_contains": {
-				Optional: true,
-				Type:     schema.TypeString,
-			},
-		},
-	}
+			ids = append(ids, v.ClusterId)
+		}
+		data.Ids = ids
+		data.Id = "_"
+		return nil
+	})
 }
diff --git a/clusters/data_clusters_test.go b/clusters/data_clusters_test.go
@@ -6,69 +6,59 @@ import (
 
 	"github.com/databricks/databricks-sdk-go/client"
 	"github.com/databricks/databricks-sdk-go/config"
+	"github.com/databricks/databricks-sdk-go/experimental/mocks"
+	"github.com/databricks/databricks-sdk-go/service/compute"
 	"github.com/databricks/terraform-provider-databricks/common"
 	"github.com/databricks/terraform-provider-databricks/qa"
-	"github.com/hashicorp/terraform-plugin-sdk/v2/helper/schema"
 	"github.com/stretchr/testify/assert"
-	"github.com/stretchr/testify/require"
+	"github.com/stretchr/testify/mock"
 )
 
 func TestClustersDataSource(t *testing.T) {
 	qa.ResourceFixture{
-		Fixtures: []qa.HTTPFixture{
-			{
-				Method:   "GET",
-				Resource: "/api/2.0/clusters/list",
-
-				Response: ClusterList{
-					Clusters: []ClusterInfo{
-						{
-							ClusterID: "b",
-						},
-						{
-							ClusterID: "a",
-						},
-					},
+		MockWorkspaceClientFunc: func(m *mocks.MockWorkspaceClient) {
+			e := m.GetMockClustersAPI().EXPECT()
+			e.ListAll(mock.Anything, compute.ListClustersRequest{}).Return([]compute.ClusterDetails{
+				{
+					ClusterId: "b",
+				},
+				{
+					ClusterId: "a",
 				},
-			},
+			}, nil)
 		},
 		Resource:    DataSourceClusters(),
 		NonWritable: true,
 		Read:        true,
 		ID:          "_",
-	}.ApplyNoError(t)
+	}.ApplyAndExpectData(t, map[string]any{
+		"ids": []string{"a", "b"},
+	})
 }
 
 func TestClustersDataSourceContainsName(t *testing.T) {
-	d, err := qa.ResourceFixture{
-		Fixtures: []qa.HTTPFixture{
-			{
-				Method:   "GET",
-				Resource: "/api/2.0/clusters/list",
-				Response: ClusterList{
-					Clusters: []ClusterInfo{
-						{
-							ClusterID:   "b",
-							ClusterName: "THIS NAME",
-						},
-						{
-							ClusterID:   "a",
-							ClusterName: "that name",
-						},
-					},
+	qa.ResourceFixture{
+		MockWorkspaceClientFunc: func(m *mocks.MockWorkspaceClient) {
+			e := m.GetMockClustersAPI().EXPECT()
+			e.ListAll(mock.Anything, compute.ListClustersRequest{}).Return([]compute.ClusterDetails{
+				{
+					ClusterId:   "b",
+					ClusterName: "THIS NAME",
+				},
+				{
+					ClusterId:   "a",
+					ClusterName: "that name",
 				},
-			},
+			}, nil)
 		},
 		Resource:    DataSourceClusters(),
 		NonWritable: true,
 		Read:        true,
 		ID:          "_",
 		HCL:         `cluster_name_contains = "this"`,
-	}.Apply(t)
-	require.NoError(t, err)
-	ids := d.Get("ids").(*schema.Set)
-	assert.True(t, ids.Contains("b"))
-	assert.Equal(t, 1, ids.Len())
+	}.ApplyAndExpectData(t, map[string]any{
+		"ids": []string{"b"},
+	})
 }
 
 func TestClustersDataSourceErrorsOut(t *testing.T) {