Skip to content

Commit f2460dc

Browse files
alexottmgyucht
alexott
and
mgyucht
authored
Allow search SQL Warehouses by name in databricks_sql_warehouse data source (#2458)
* Allow search SQL Warehouses by name in `databricks_sql_warehouse` data source Right now it's possible to search only by the warehouse ID, but it's not always convenient although it's possible by using `databricks_sql_warehouses` data source + explicit filtering. This PR adds a capability to search by either SQL warehouse name or ID. This fixes #2443 * Update docs/data-sources/sql_warehouse.md Co-authored-by: Miles Yucht <[email protected]> * Address review comments also change documentation a bit to better match the data source - it was copied from the resource as-is. * More fixes from review * code review comments --------- Co-authored-by: Miles Yucht <[email protected]>
1 parent 570911e commit f2460dc

File tree

4 files changed

+166
-11
lines changed

4 files changed

+166
-11
lines changed

‎docs/data-sources/sql_warehouse.md‎

Lines changed: 16 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -9,7 +9,7 @@ Retrieves information about a [databricks_sql_warehouse](../resources/sql_wareho
99

1010
## Example usage
1111

12-
Retrieve attributes of each SQL warehouses in a workspace
12+
* Retrieve attributes of each SQL warehouses in a workspace:
1313

1414
```hcl
1515
data "databricks_sql_warehouses" "all" {
@@ -19,32 +19,41 @@ data "databricks_sql_warehouse" "all" {
1919
for_each = data.databricks_sql.warehouses.ids
2020
id = each.value
2121
}
22+
```
23+
24+
* Search for a specific SQL Warehouse by name:
2225

26+
```hcl
27+
data "databricks_sql_warehouse" "all" {
28+
name = "Starter Warehouse"
29+
}
2330
```
2431

2532
## Argument reference
2633

27-
* `id` - (Required) The ID of the SQL warehouse
34+
* `id` - (Required, if `name` isn't specified) The ID of the SQL warehouse.
35+
* `name` - (Required, if `id` isn't specified) Name of the SQL warehouse to search (case-sensitive).
2836

2937
## Attribute reference
3038

3139
This data source exports the following attributes:
3240

33-
* `name` - Name of the SQL warehouse. Must be unique.
41+
* `id` - The ID of the SQL warehouse.
42+
* `name` - Name of the SQL warehouse.
3443
* `cluster_size` - The size of the clusters allocated to the warehouse: "2X-Small", "X-Small", "Small", "Medium", "Large", "X-Large", "2X-Large", "3X-Large", "4X-Large".
3544
* `min_num_clusters` - Minimum number of clusters available when a SQL warehouse is running.
3645
* `max_num_clusters` - Maximum number of clusters available when a SQL warehouse is running.
3746
* `auto_stop_mins` - Time in minutes until an idle SQL warehouse terminates all clusters and stops.
38-
* `tags` - Databricks tags all warehouse resources with these tags.
47+
* `tags` - tags used for SQL warehouse resources.
3948
* `spot_instance_policy` - The spot policy to use for allocating instances to clusters: `COST_OPTIMIZED` or `RELIABILITY_OPTIMIZED`.
40-
* `enable_photon` - Whether to enable [Photon](https://databricks.com/product/delta-engine).
41-
* `enable_serverless_compute` - Whether this SQL warehouse is a serverless SQL warehouse. If this value is `true`, `warehouse_type` must be `PRO`.
49+
* `enable_photon` - Whether [Photon](https://databricks.com/product/delta-engine) is enabled.
50+
* `enable_serverless_compute` - Whether this SQL warehouse is a serverless SQL warehouse.
4251

4352
- **For AWS**: If your account needs updated [terms of use](https://docs.databricks.com/sql/admin/serverless.html#accept-terms), workspace admins are prompted in the Databricks SQL UI. A workspace must meet the [requirements](https://docs.databricks.com/sql/admin/serverless.html#requirements) and might require an update its instance profile role to [add a trust relationship](https://docs.databricks.com/sql/admin/serverless.html#aws-instance-profile-setup).
4453

4554
- **For Azure**, you must [enable your workspace for serverless SQL warehouse](https://learn.microsoft.com/azure/databricks/sql/admin/serverless).
4655

47-
* `warehouse_type` - SQL warehouse type. See for [AWS](https://docs.databricks.com/sql/index.html#warehouse-types) or [Azure](https://learn.microsoft.com/azure/databricks/sql/#warehouse-types). Set to `PRO` or `CLASSIC`. If the field `enable_serverless_compute` has the value `true`, this needs to be set to `PRO`.
56+
* `warehouse_type` - SQL warehouse type. See for [AWS](https://docs.databricks.com/sql/index.html#warehouse-types) or [Azure](https://learn.microsoft.com/azure/databricks/sql/#warehouse-types).
4857
* `channel` block, consisting of following fields:
4958
* `name` - Name of the Databricks SQL release channel. Possible values are: `CHANNEL_NAME_PREVIEW` and `CHANNEL_NAME_CURRENT`. Default is `CHANNEL_NAME_CURRENT`.
5059
* `jdbc_url` - JDBC connection string.

‎sql/data_sql_warehouse.go‎

Lines changed: 33 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -10,7 +10,7 @@ import (
1010

1111
func DataSourceWarehouse() *schema.Resource {
1212
type SQLWarehouseInfo struct {
13-
ID string `json:"id"`
13+
ID string `json:"id,omitempty" tf:"computed"`
1414
Name string `json:"name,omitempty" tf:"computed"`
1515
ClusterSize string `json:"cluster_size,omitempty" tf:"computed"`
1616
AutoStopMinutes int `json:"auto_stop_mins,omitempty" tf:"computed"`
@@ -31,15 +31,44 @@ func DataSourceWarehouse() *schema.Resource {
3131

3232
return common.DataResource(SQLWarehouseInfo{}, func(ctx context.Context, e interface{}, c *common.DatabricksClient) error {
3333
data := e.(*SQLWarehouseInfo)
34-
err := c.Get(ctx, fmt.Sprintf("/sql/warehouses/%s", data.ID), nil, data)
34+
var id string
35+
if data.ID == "" && data.Name == "" {
36+
return fmt.Errorf("either 'id' or 'name' should be provided")
37+
}
38+
endpointsAPI := NewSQLEndpointsAPI(ctx, c)
39+
selected := []DataSource{}
40+
dataSources, err := endpointsAPI.listDataSources()
3541
if err != nil {
3642
return err
3743
}
38-
endpointsAPI := NewSQLEndpointsAPI(ctx, c)
39-
data.DataSourceID, err = endpointsAPI.ResolveDataSourceID(data.ID)
44+
for _, source := range dataSources {
45+
if data.Name != "" && source.Name == data.Name {
46+
selected = append(selected, source)
47+
} else if data.ID != "" && source.EndpointID == data.ID {
48+
selected = append(selected, source)
49+
break
50+
}
51+
}
52+
if len(selected) == 0 {
53+
if data.Name != "" {
54+
return fmt.Errorf("can't find SQL warehouse with the name '%s'", data.Name)
55+
} else {
56+
return fmt.Errorf("can't find SQL warehouse with the ID '%s'", data.ID)
57+
}
58+
}
59+
if len(selected) > 1 {
60+
if data.Name != "" {
61+
return fmt.Errorf("there are multiple SQL warehouses with the name '%s'", data.Name)
62+
} else {
63+
return fmt.Errorf("there are multiple SQL warehouses with the ID '%s'", data.ID)
64+
}
65+
}
66+
id = selected[0].EndpointID
67+
err = c.Get(ctx, fmt.Sprintf("/sql/warehouses/%s", id), nil, data)
4068
if err != nil {
4169
return err
4270
}
71+
data.DataSourceID = selected[0].ID
4372
return nil
4473
})
4574
}

‎sql/data_sql_warehouse_test.go‎

Lines changed: 110 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -46,3 +46,113 @@ func TestWarehouseData_Error(t *testing.T) {
4646
ID: "_",
4747
}.ExpectError(t, "I'm a teapot")
4848
}
49+
50+
func TestWarehouseDataByName_ListError(t *testing.T) {
51+
qa.ResourceFixture{
52+
Fixtures: qa.HTTPFailures,
53+
Resource: DataSourceWarehouse(),
54+
Read: true,
55+
NonWritable: true,
56+
HCL: `name = "abc"`,
57+
ID: "_",
58+
}.ExpectError(t, "I'm a teapot")
59+
}
60+
61+
func TestWarehouseDataByName_NotFoundError(t *testing.T) {
62+
qa.ResourceFixture{
63+
Fixtures: []qa.HTTPFixture{
64+
{
65+
Method: "GET",
66+
Resource: "/api/2.0/preview/sql/data_sources",
67+
Response: []DataSource{
68+
{
69+
ID: "d7c9d05c-7496-4c69-b089-48823edad401",
70+
EndpointID: "def",
71+
Name: "test",
72+
},
73+
{
74+
ID: "d7c9d05c-7496-4c69-b089-48823edad40c",
75+
EndpointID: "abc",
76+
Name: "abc2",
77+
},
78+
},
79+
},
80+
},
81+
Resource: DataSourceWarehouse(),
82+
Read: true,
83+
NonWritable: true,
84+
HCL: `name = "abc"`,
85+
ID: "_",
86+
}.ExpectError(t, "can't find SQL warehouse with the name 'abc'")
87+
}
88+
89+
func TestWarehouseDataByName_DuplicatesError(t *testing.T) {
90+
qa.ResourceFixture{
91+
Fixtures: []qa.HTTPFixture{
92+
{
93+
Method: "GET",
94+
Resource: "/api/2.0/preview/sql/data_sources",
95+
Response: []DataSource{
96+
{
97+
ID: "d7c9d05c-7496-4c69-b089-48823edad401",
98+
EndpointID: "def",
99+
Name: "abc",
100+
},
101+
{
102+
ID: "d7c9d05c-7496-4c69-b089-48823edad40c",
103+
EndpointID: "abc",
104+
Name: "abc",
105+
},
106+
},
107+
},
108+
},
109+
Resource: DataSourceWarehouse(),
110+
Read: true,
111+
NonWritable: true,
112+
HCL: `name = "abc"`,
113+
ID: "_",
114+
}.ExpectError(t, "there are multiple SQL warehouses with the name 'abc'")
115+
}
116+
117+
func TestWarehouseDataByName(t *testing.T) {
118+
d, err := qa.ResourceFixture{
119+
Fixtures: []qa.HTTPFixture{
120+
{
121+
Method: "GET",
122+
Resource: "/api/2.0/preview/sql/data_sources",
123+
Response: []DataSource{
124+
{
125+
ID: "d7c9d05c-7496-4c69-b089-48823edad401",
126+
EndpointID: "def",
127+
Name: "abc",
128+
},
129+
{
130+
ID: "d7c9d05c-7496-4c69-b089-48823edad40c",
131+
EndpointID: "abc",
132+
Name: "test",
133+
},
134+
},
135+
},
136+
{
137+
Method: "GET",
138+
Resource: "/api/2.0/sql/warehouses/abc",
139+
ReuseRequest: true,
140+
Response: SQLEndpoint{
141+
Name: "test",
142+
ClusterSize: "Small",
143+
ID: "abc",
144+
State: "RUNNING",
145+
},
146+
},
147+
},
148+
Resource: DataSourceWarehouse(),
149+
Read: true,
150+
NonWritable: true,
151+
HCL: `name = "test"`,
152+
ID: "_",
153+
}.Apply(t)
154+
require.NoError(t, err)
155+
assert.Equal(t, "abc", d.Id())
156+
assert.Equal(t, "RUNNING", d.Get("state"))
157+
assert.Equal(t, "d7c9d05c-7496-4c69-b089-48823edad40c", d.Get("data_source_id"))
158+
}

‎sql/resource_sql_endpoint.go‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -80,6 +80,7 @@ type Tag struct {
8080
type DataSource struct {
8181
ID string `json:"id"`
8282
EndpointID string `json:"endpoint_id"`
83+
Name string `json:"name"`
8384
}
8485

8586
// endpointList ...
@@ -138,6 +139,12 @@ func (a SQLEndpointsAPI) Create(se *SQLEndpoint, timeout time.Duration) error {
138139
return a.waitForRunning(se.ID, timeout)
139140
}
140141

142+
func (a SQLEndpointsAPI) listDataSources() ([]DataSource, error) {
143+
var dss []DataSource
144+
err := a.client.Get(a.context, "/preview/sql/data_sources", nil, &dss)
145+
return dss, err
146+
}
147+
141148
// ResolveDataSourceID ...
142149
func (a SQLEndpointsAPI) ResolveDataSourceID(endpointID string) (dataSourceID string, err error) {
143150
var dss []DataSource

0 commit comments

Comments
 (0)