Skip to content

Commit

Permalink
[Feature] Add databricks_query resource instead of `databricks_sql_…
Browse files Browse the repository at this point in the history
…query`

This PR is built on top of #4051 that should be merged first.

The new resource uses the new [Queries API](https://docs.databricks.com/api/workspace/queries/create) instead of the legacy one that will be deprecated. Since the new resource has a slightly different set of parameters, it was decided to create a new resource and deprecate the old one.

This resource uses old TF SDK to be compatible with TF exporter (until #4050 is implemented).

TODOs:

- Add documentation
- Need to discuss how to handle permissions - `sql_query` permissions look like working, but not sure if we should continue to use that API
- Support in the exporter will be in a separate PR
  • Loading branch information
alexott committed Oct 14, 2024
1 parent 9b4fc63 commit e79a28e
Show file tree
Hide file tree
Showing 7 changed files with 386 additions and 14 deletions.
48 changes: 48 additions & 0 deletions internal/acceptance/query_test.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,48 @@
package acceptance

import (
"testing"
)

func TestAccQuery(t *testing.T) {
WorkspaceLevel(t, Step{
Template: `
resource "databricks_query" "this" {
warehouse_id = "{env.TEST_DEFAULT_WAREHOUSE_ID}"
display_name = "tf-{var.RANDOM}"
query_text = "SELECT 1 AS p1, 2 as p2"
}
resource "databricks_permissions" "query_usage" {
sql_query_id = databricks_query.this.id
access_control {
group_name = "users"
permission_level = "CAN_RUN"
}
}
`,
}, Step{
Template: `
resource "databricks_sql_query" "this" {
warehouse_id = "{env.TEST_DEFAULT_WAREHOUSE_ID}"
display_name = "tf-{var.RANDOM}"
query_text = "SELECT 1 AS p1, 2 as p2"
parameter {
name = "foo"
text_value {
value = "bar"
}
title = "foo"
}
}
resource "databricks_permissions" "query_usage" {
sql_query_id = databricks_query.this.id
access_control {
group_name = "users"
permission_level = "CAN_RUN"
}
}
`,
})
}
2 changes: 1 addition & 1 deletion internal/acceptance/sql_query_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,7 @@ import (
"testing"
)

func TestAccQuery(t *testing.T) {
func TestAccSqlQuery(t *testing.T) {
WorkspaceLevel(t, Step{
Template: `
resource "databricks_sql_query" "q1" {
Expand Down
1 change: 1 addition & 0 deletions internal/providers/sdkv2/sdkv2.go
Original file line number Diff line number Diff line change
Expand Up @@ -189,6 +189,7 @@ func DatabricksProvider() *schema.Provider {
"databricks_pipeline": pipelines.ResourcePipeline().ToResource(),
"databricks_provider": sharing.ResourceProvider().ToResource(),
"databricks_quality_monitor": catalog.ResourceQualityMonitor().ToResource(),
"databricks_query": sql.ResourceQuery().ToResource(),
"databricks_recipient": sharing.ResourceRecipient().ToResource(),
"databricks_registered_model": catalog.ResourceRegisteredModel().ToResource(),
"databricks_repo": repos.ResourceRepo().ToResource(),
Expand Down
169 changes: 169 additions & 0 deletions sql/resource_query.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,169 @@
package sql

import (
"context"
"log"
"strings"

"github.com/databricks/databricks-sdk-go/service/sql"
"github.com/databricks/terraform-provider-databricks/common"
"github.com/hashicorp/terraform-plugin-sdk/v2/helper/schema"
"github.com/hashicorp/terraform-plugin-sdk/v2/helper/validation"
)

// Need a struct for Query because there are aliases we need and it'll be needed in the create method.
type queryStruct struct {
sql.Query
}

var queryAliasMap = map[string]string{
"parameters": "parameter",
}

func (queryStruct) Aliases() map[string]map[string]string {
return map[string]map[string]string{
"sql.queryStruct": queryAliasMap,
}
}

func (queryStruct) CustomizeSchema(m *common.CustomizableSchema) *common.CustomizableSchema {
m.SchemaPath("display_name").SetRequired().SetValidateFunc(validation.StringIsNotWhiteSpace)
m.SchemaPath("query_text").SetRequired()
m.SchemaPath("warehouse_id").SetRequired().SetValidateFunc(validation.StringIsNotWhiteSpace)
m.SchemaPath("parent_path").SetCustomSuppressDiff(common.WorkspaceOrEmptyPathPrefixDiffSuppress).SetForceNew()
m.SchemaPath("owner_user_name").SetSuppressDiff()
m.SchemaPath("run_as_mode").SetSuppressDiff()
//m.SchemaPath("").SetSuppressDiff()
//m.SchemaPath("").SetSuppressDiff()
m.SchemaPath("id").SetReadOnly()
m.SchemaPath("create_time").SetReadOnly()
m.SchemaPath("lifecycle_state").SetReadOnly()
m.SchemaPath("last_modifier_user_name").SetReadOnly()
m.SchemaPath("update_time").SetReadOnly()

// customize parameters
m.SchemaPath("parameter", "name").SetRequired().SetValidateFunc(validation.StringIsNotWhiteSpace)
m.SchemaPath("parameter", "date_range_value", "precision").SetSuppressDiff()
m.SchemaPath("parameter", "date_value", "precision").SetSuppressDiff()
m.SchemaPath("parameter", "query_backed_value", "query_id").SetRequired()
m.SchemaPath("parameter", "text_value", "value").SetRequired()
m.SchemaPath("parameter", "numeric_value", "value").SetRequired()
// TODO: fix setting of AtLeastOneOf
// valuesAlof := []string{
// "parameter.0.date_range_value",
// "parameter.0.date_value",
// "parameter.0.query_backed_value",
// "parameter.0.text_value",
// "parameter.0.numeric_value",
// "parameter.0.enum_value",
// }
// for _, f := range valuesAlof {
// m.SchemaPath("parameter", strings.TrimPrefix(f, "parameter.0.")).SetAtLeastOneOf(valuesAlof)
// }
return m
}

type queryCreateStruct struct {
sql.CreateQueryRequestQuery
}

func (queryCreateStruct) Aliases() map[string]map[string]string {
return map[string]map[string]string{
"sql.queryCreateStruct": queryAliasMap,
}
}

func (queryCreateStruct) CustomizeSchema(s *common.CustomizableSchema) *common.CustomizableSchema {
return s
}

type queryUpdateStruct struct {
sql.UpdateQueryRequestQuery
}

func (queryUpdateStruct) Aliases() map[string]map[string]string {
return map[string]map[string]string{
"sql.queryUpdateStruct": queryAliasMap,
}
}

func (queryUpdateStruct) CustomizeSchema(s *common.CustomizableSchema) *common.CustomizableSchema {
return s
}

func ResourceQuery() common.Resource {
s := common.StructToSchema(queryStruct{}, nil)
return common.Resource{
Create: func(ctx context.Context, d *schema.ResourceData, c *common.DatabricksClient) error {
w, err := c.WorkspaceClient()
if err != nil {
return err
}
var q queryCreateStruct
common.DataToStructPointer(d, s, &q)
apiQuery, err := w.Queries.Create(ctx, sql.CreateQueryRequest{
Query: &q.CreateQueryRequestQuery,
})
if err != nil {
return err
}
d.SetId(apiQuery.Id)
owner := d.Get("owner_user_name").(string)
if owner != "" {
_, err = w.Queries.Update(ctx, sql.UpdateQueryRequest{
Query: &sql.UpdateQueryRequestQuery{
OwnerUserName: owner,
},
Id: apiQuery.Id,
UpdateMask: "owner_user_name",
})
}
return err
},
Read: func(ctx context.Context, d *schema.ResourceData, c *common.DatabricksClient) error {
w, err := c.WorkspaceClient()
if err != nil {
return err
}
apiQuery, err := w.Queries.GetById(ctx, d.Id())
if err != nil {
log.Printf("[WARN] error getting query by ID: %v", err)
return err
}
parentPath := d.Get("parent_path").(string)
if parentPath != "" && strings.HasPrefix(apiQuery.ParentPath, "/Workspace") && !strings.HasPrefix(parentPath, "/Workspace") {
apiQuery.ParentPath = strings.TrimPrefix(parentPath, "/Workspace")
}
return common.StructToData(queryStruct{Query: *apiQuery}, s, d)
},
Update: func(ctx context.Context, d *schema.ResourceData, c *common.DatabricksClient) error {
w, err := c.WorkspaceClient()
if err != nil {
return err
}
var q queryUpdateStruct
common.DataToStructPointer(d, s, &q)
updateMask := "display_name,query_text,warehouse_id,parameters"
for _, f := range []string{"run_as_mode", "owner_user_name", "description", "tags",
"apply_auto_limit", "catalog", "schema"} {
if d.HasChange(f) {
updateMask += "," + f
}
}
_, err = w.Queries.Update(ctx, sql.UpdateQueryRequest{
Query: &q.UpdateQueryRequestQuery,
Id: d.Id(),
UpdateMask: updateMask,
})
return err
},
Delete: func(ctx context.Context, d *schema.ResourceData, c *common.DatabricksClient) error {
w, err := c.WorkspaceClient()
if err != nil {
return err
}
return w.Queries.DeleteById(ctx, d.Id())
},
Schema: s,
}
}
153 changes: 153 additions & 0 deletions sql/resource_query_test.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,153 @@
package sql

import (
"net/http"
"testing"

"github.com/databricks/databricks-sdk-go/apierr"
"github.com/databricks/databricks-sdk-go/experimental/mocks"
"github.com/databricks/databricks-sdk-go/service/sql"
"github.com/databricks/terraform-provider-databricks/qa"
"github.com/stretchr/testify/mock"
)

var (
queryResponse = sql.Query{
Id: "7890",
WarehouseId: "123456",
DisplayName: "TF new query",
OwnerUserName: "user@domain.com",
ParentPath: "/Workspace/Shared/Querys",
QueryText: "select 42 as value",
}
createQueryHcl = `warehouse_id = "123456"
query_text = "select 42 as value"
display_name = "TF new query"
parent_path = "/Shared/Querys"
owner_user_name = "user@domain.com"
`
createQueryRequest = sql.CreateQueryRequest{
Query: &sql.CreateQueryRequestQuery{
WarehouseId: "123456",
QueryText: "select 42 as value",
DisplayName: "TF new query",
ParentPath: "/Shared/Querys",
}}
)

func TestQueryCreate(t *testing.T) {
qa.ResourceFixture{
MockWorkspaceClientFunc: func(w *mocks.MockWorkspaceClient) {
e := w.GetMockQueriesAPI().EXPECT()
e.Create(mock.Anything, createQueryRequest).Return(&queryResponse, nil)
e.Update(mock.Anything, sql.UpdateQueryRequest{
Id: "7890",
UpdateMask: "owner_user_name",
Query: &sql.UpdateQueryRequestQuery{
OwnerUserName: "user@domain.com",
},
}).Return(&queryResponse, nil)
e.GetById(mock.Anything, "7890").Return(&queryResponse, nil)
},
Resource: ResourceQuery(),
Create: true,
HCL: createQueryHcl,
}.ApplyAndExpectData(t, map[string]any{
"id": "7890",
"warehouse_id": "123456",
"display_name": "TF new query",
"owner_user_name": "user@domain.com",
})
}

func TestQueryCreate_Error(t *testing.T) {
qa.ResourceFixture{
MockWorkspaceClientFunc: func(w *mocks.MockWorkspaceClient) {
e := w.GetMockQueriesAPI().EXPECT()
e.Create(mock.Anything, createQueryRequest).Return(nil, &apierr.APIError{
StatusCode: http.StatusBadRequest,
Message: "bad payload",
})
},
Resource: ResourceQuery(),
Create: true,
HCL: createQueryHcl,
}.ExpectError(t, "bad payload")
}

func TestQueryRead_Import(t *testing.T) {
qa.ResourceFixture{
MockWorkspaceClientFunc: func(w *mocks.MockWorkspaceClient) {
w.GetMockQueriesAPI().EXPECT().GetById(mock.Anything, "7890").Return(&queryResponse, nil)
},
Resource: ResourceQuery(),
Read: true,
ID: "7890",
New: true,
}.ApplyAndExpectData(t, map[string]any{
"id": "7890",
"warehouse_id": "123456",
"query_text": "select 42 as value",
"display_name": "TF new query",
"owner_user_name": "user@domain.com",
})
}

func TestQueryRead_Error(t *testing.T) {
qa.ResourceFixture{
MockWorkspaceClientFunc: func(w *mocks.MockWorkspaceClient) {
w.GetMockQueriesAPI().EXPECT().GetById(mock.Anything, "7890").Return(nil, &apierr.APIError{
StatusCode: http.StatusBadRequest,
Message: "bad payload",
})
},
Resource: ResourceQuery(),
Read: true,
ID: "7890",
New: true,
}.ExpectError(t, "bad payload")
}

func TestQueryDelete(t *testing.T) {
qa.ResourceFixture{
MockWorkspaceClientFunc: func(w *mocks.MockWorkspaceClient) {
w.GetMockQueriesAPI().EXPECT().DeleteById(mock.Anything, "7890").Return(nil)
},
Resource: ResourceQuery(),
Delete: true,
ID: "7890",
New: true,
}.ApplyNoError(t)
}

func TestQueryUpdate(t *testing.T) {
qa.ResourceFixture{
MockWorkspaceClientFunc: func(w *mocks.MockWorkspaceClient) {
e := w.GetMockQueriesAPI().EXPECT()
e.Update(mock.Anything, sql.UpdateQueryRequest{
Id: "7890",
UpdateMask: "display_name,query_text,warehouse_id,parameters,owner_user_name",
Query: &sql.UpdateQueryRequestQuery{
WarehouseId: "123456",
DisplayName: "TF new query",
OwnerUserName: "user@domain.com",
QueryText: "select 42 as value",
}}).Return(&queryResponse, nil)
e.GetById(mock.Anything, "7890").Return(&queryResponse, nil)
},
Resource: ResourceQuery(),
Update: true,
ID: "7890",
HCL: `warehouse_id = "123456"
query_text = "select 42 as value"
display_name = "TF new query"
owner_user_name = "user@domain.com"
`,
}.ApplyAndExpectData(t, map[string]any{
"id": "7890",
"warehouse_id": "123456",
"query_text": "select 42 as value",
"display_name": "TF new query",
"owner_user_name": "user@domain.com",
})
}
3 changes: 2 additions & 1 deletion sql/resource_sql_query.go
Original file line number Diff line number Diff line change
Expand Up @@ -587,6 +587,7 @@ func ResourceSqlQuery() common.Resource {
Delete: func(ctx context.Context, data *schema.ResourceData, c *common.DatabricksClient) error {
return NewQueryAPI(ctx, c).Delete(data.Id())
},
Schema: s,
Schema: s,
DeprecationMessage: "This resource is deprecated and will be removed in the future. Please use the `databricks_query` resource instead.",
}
}
Loading

0 comments on commit e79a28e

Please sign in to comment.