Skip to content

Commit d5266ec

Browse files
arpitjasa-dbnkvuongmgyucht
authored
Added databricks_registered_model resource (#2771)
* Added databricks_registered_model resource * Applied comments * Update grants * Update tests * Update update * Apply comments Signed-off-by: Arpit Jasapara <[email protected]> * Update internal/acceptance/registered_model_test.go * update deprecation comment --------- Signed-off-by: Arpit Jasapara <[email protected]> Co-authored-by: vuong-nguyen <[email protected]> Co-authored-by: Miles Yucht <[email protected]>
1 parent e320aa1 commit d5266ec

File tree

11 files changed

+444
-5
lines changed

11 files changed

+444
-5
lines changed

‎catalog/resource_grants.go‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -105,6 +105,9 @@ func getPermissionEndpoint(securable, name string) string {
105105
if securable == "foreign_connection" {
106106
return fmt.Sprintf("/unity-catalog/permissions/connection/%s", name)
107107
}
108+
if securable == "model" {
109+
return fmt.Sprintf("/unity-catalog/permissions/function/%s", name)
110+
}
108111
return fmt.Sprintf("/unity-catalog/permissions/%s/%s", securable, name)
109112
}
110113

‎catalog/resource_registered_model.go‎

Lines changed: 72 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,72 @@
1+
package catalog
2+
3+
import (
4+
"context"
5+
6+
"github.com/databricks/databricks-sdk-go/service/catalog"
7+
"github.com/databricks/terraform-provider-databricks/common"
8+
"github.com/hashicorp/terraform-plugin-sdk/v2/helper/schema"
9+
)
10+
11+
func ResourceRegisteredModel() *schema.Resource {
12+
s := common.StructToSchema(
13+
catalog.CreateRegisteredModelRequest{},
14+
func(m map[string]*schema.Schema) map[string]*schema.Schema {
15+
m["name"].ForceNew = true
16+
m["catalog_name"].ForceNew = true
17+
m["schema_name"].ForceNew = true
18+
m["storage_location"].ForceNew = true
19+
m["storage_location"].Computed = true
20+
21+
return m
22+
})
23+
24+
return common.Resource{
25+
Create: func(ctx context.Context, d *schema.ResourceData, c *common.DatabricksClient) error {
26+
w, err := c.WorkspaceClient()
27+
if err != nil {
28+
return err
29+
}
30+
var m catalog.CreateRegisteredModelRequest
31+
common.DataToStructPointer(d, s, &m)
32+
model, err := w.RegisteredModels.Create(ctx, m)
33+
if err != nil {
34+
return err
35+
}
36+
d.SetId(model.FullName)
37+
return nil
38+
},
39+
Read: func(ctx context.Context, d *schema.ResourceData, c *common.DatabricksClient) error {
40+
w, err := c.WorkspaceClient()
41+
if err != nil {
42+
return err
43+
}
44+
model, err := w.RegisteredModels.GetByFullName(ctx, d.Id())
45+
if err != nil {
46+
return err
47+
}
48+
return common.StructToData(*model, s, d)
49+
},
50+
Update: func(ctx context.Context, d *schema.ResourceData, c *common.DatabricksClient) error {
51+
w, err := c.WorkspaceClient()
52+
if err != nil {
53+
return err
54+
}
55+
var u catalog.UpdateRegisteredModelRequest
56+
common.DataToStructPointer(d, s, &u)
57+
u.FullName = d.Id()
58+
_, err = w.RegisteredModels.Update(ctx, u)
59+
return err
60+
},
61+
Delete: func(ctx context.Context, d *schema.ResourceData, c *common.DatabricksClient) error {
62+
w, err := c.WorkspaceClient()
63+
if err != nil {
64+
return err
65+
}
66+
return w.RegisteredModels.DeleteByFullName(ctx, d.Id())
67+
},
68+
StateUpgraders: []schema.StateUpgrader{},
69+
Schema: s,
70+
SchemaVersion: 0,
71+
}.ToResource()
72+
}

‎catalog/resource_registered_model_test.go‎

Lines changed: 237 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,237 @@
1+
package catalog
2+
3+
import (
4+
"net/http"
5+
"testing"
6+
7+
"github.com/databricks/databricks-sdk-go/apierr"
8+
"github.com/databricks/databricks-sdk-go/service/catalog"
9+
"github.com/databricks/terraform-provider-databricks/qa"
10+
)
11+
12+
func TestRegisteredModelCornerCases(t *testing.T) {
13+
qa.ResourceCornerCases(t, ResourceRegisteredModel())
14+
}
15+
16+
func TestRegisteredModelCreate(t *testing.T) {
17+
qa.ResourceFixture{
18+
Fixtures: []qa.HTTPFixture{
19+
{
20+
Method: http.MethodPost,
21+
Resource: "/api/2.1/unity-catalog/models",
22+
ExpectedRequest: catalog.CreateRegisteredModelRequest{
23+
Name: "model",
24+
CatalogName: "catalog",
25+
SchemaName: "schema",
26+
Comment: "comment",
27+
},
28+
Response: catalog.RegisteredModelInfo{
29+
Name: "model",
30+
CatalogName: "catalog",
31+
SchemaName: "schema",
32+
FullName: "catalog.schema.model",
33+
Comment: "comment",
34+
},
35+
},
36+
{
37+
Method: http.MethodGet,
38+
Resource: "/api/2.1/unity-catalog/models/catalog.schema.model?",
39+
Response: catalog.RegisteredModelInfo{
40+
Name: "model",
41+
CatalogName: "catalog",
42+
SchemaName: "schema",
43+
FullName: "catalog.schema.model",
44+
Comment: "comment",
45+
},
46+
},
47+
},
48+
Resource: ResourceRegisteredModel(),
49+
HCL: `
50+
name = "model"
51+
catalog_name = "catalog"
52+
schema_name = "schema"
53+
comment = "comment"
54+
`,
55+
Create: true,
56+
}.ApplyAndExpectData(t,
57+
map[string]any{
58+
"id": "catalog.schema.model",
59+
},
60+
)
61+
}
62+
63+
func TestRegisteredModelCreate_Error(t *testing.T) {
64+
qa.ResourceFixture{
65+
Fixtures: []qa.HTTPFixture{
66+
{
67+
Method: http.MethodPost,
68+
Resource: "/api/2.1/unity-catalog/models",
69+
Response: apierr.APIErrorBody{
70+
ErrorCode: "INVALID_REQUEST",
71+
Message: "Internal error happened",
72+
},
73+
Status: 400,
74+
},
75+
},
76+
Resource: ResourceRegisteredModel(),
77+
Create: true,
78+
}.ExpectError(t, "Internal error happened")
79+
}
80+
81+
func TestRegisteredModelRead(t *testing.T) {
82+
qa.ResourceFixture{
83+
Fixtures: []qa.HTTPFixture{
84+
{
85+
Method: http.MethodGet,
86+
Resource: "/api/2.1/unity-catalog/models/catalog.schema.model?",
87+
Response: catalog.RegisteredModelInfo{
88+
Name: "model",
89+
CatalogName: "catalog",
90+
SchemaName: "schema",
91+
FullName: "catalog.schema.model",
92+
Comment: "comment",
93+
},
94+
},
95+
},
96+
Resource: ResourceRegisteredModel(),
97+
Read: true,
98+
ID: "catalog.schema.model",
99+
}.ApplyAndExpectData(t,
100+
map[string]any{
101+
"id": "catalog.schema.model",
102+
},
103+
)
104+
}
105+
106+
func TestRegisteredModelRead_Error(t *testing.T) {
107+
qa.ResourceFixture{
108+
Fixtures: []qa.HTTPFixture{
109+
{
110+
Method: http.MethodGet,
111+
Resource: "/api/2.1/unity-catalog/models/catalog.schema.model?",
112+
Response: apierr.APIErrorBody{
113+
ErrorCode: "INVALID_REQUEST",
114+
Message: "Internal error happened",
115+
},
116+
Status: 400,
117+
},
118+
},
119+
Resource: ResourceRegisteredModel(),
120+
Read: true,
121+
ID: "catalog.schema.model",
122+
}.ExpectError(t, "Internal error happened")
123+
}
124+
125+
func TestRegisteredModelUpdate(t *testing.T) {
126+
qa.ResourceFixture{
127+
Fixtures: []qa.HTTPFixture{
128+
{
129+
Method: http.MethodPatch,
130+
Resource: "/api/2.1/unity-catalog/models/catalog.schema.model",
131+
ExpectedRequest: catalog.UpdateRegisteredModelRequest{
132+
FullName: "catalog.schema.model",
133+
Comment: "new comment",
134+
Name: "model",
135+
},
136+
Response: catalog.RegisteredModelInfo{
137+
Name: "model",
138+
CatalogName: "catalog",
139+
SchemaName: "schema",
140+
FullName: "catalog.schema.model",
141+
Comment: "new comment",
142+
},
143+
},
144+
{
145+
Method: http.MethodGet,
146+
Resource: "/api/2.1/unity-catalog/models/catalog.schema.model?",
147+
Response: catalog.RegisteredModelInfo{
148+
Name: "model",
149+
CatalogName: "catalog",
150+
SchemaName: "schema",
151+
FullName: "catalog.schema.model",
152+
Comment: "new comment",
153+
},
154+
},
155+
},
156+
Resource: ResourceRegisteredModel(),
157+
Update: true,
158+
ID: "catalog.schema.model",
159+
InstanceState: map[string]string{
160+
"name": "model",
161+
"catalog_name": "catalog",
162+
"schema_name": "schema",
163+
"comment": "comment",
164+
},
165+
HCL: `
166+
name = "model"
167+
catalog_name = "catalog"
168+
schema_name = "schema"
169+
comment = "new comment"
170+
`,
171+
}.ApplyNoError(t)
172+
}
173+
174+
func TestRegisteredModelUpdate_Error(t *testing.T) {
175+
qa.ResourceFixture{
176+
Fixtures: []qa.HTTPFixture{
177+
{
178+
Method: http.MethodPatch,
179+
Resource: "/api/2.1/unity-catalog/models/catalog.schema.model",
180+
Response: apierr.APIErrorBody{
181+
ErrorCode: "INVALID_REQUEST",
182+
Message: "Internal error happened",
183+
},
184+
Status: 400,
185+
},
186+
},
187+
Resource: ResourceRegisteredModel(),
188+
Update: true,
189+
ID: "catalog.schema.model",
190+
InstanceState: map[string]string{
191+
"name": "model",
192+
"catalog_name": "catalog",
193+
"schema_name": "schema",
194+
"comment": "comment",
195+
},
196+
HCL: `
197+
name = "model"
198+
catalog_name = "catalog"
199+
schema_name = "schema"
200+
comment = "new comment"
201+
`,
202+
}.ExpectError(t, "Internal error happened")
203+
}
204+
205+
func TestRegisteredModelDelete(t *testing.T) {
206+
qa.ResourceFixture{
207+
Fixtures: []qa.HTTPFixture{
208+
{
209+
Method: http.MethodDelete,
210+
Resource: "/api/2.1/unity-catalog/models/catalog.schema.model?",
211+
Response: "",
212+
},
213+
},
214+
Resource: ResourceRegisteredModel(),
215+
Delete: true,
216+
ID: "catalog.schema.model",
217+
}.ApplyNoError(t)
218+
}
219+
220+
func TestRegisteredModelDelete_Error(t *testing.T) {
221+
qa.ResourceFixture{
222+
Fixtures: []qa.HTTPFixture{
223+
{
224+
Method: http.MethodDelete,
225+
Resource: "/api/2.1/unity-catalog/models/catalog.schema.model?",
226+
Response: apierr.APIErrorBody{
227+
ErrorCode: "INVALID_REQUEST",
228+
Message: "Internal error happened",
229+
},
230+
Status: 400,
231+
},
232+
},
233+
Resource: ResourceRegisteredModel(),
234+
Delete: true,
235+
ID: "catalog.schema.model",
236+
}.ExpectError(t, "Internal error happened")
237+
}

‎docs/index.md‎

Lines changed: 5 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -52,10 +52,12 @@ Databricks SQL
5252
* Manage [dashboards](resources/sql_dashboard.md) and their [widgets](resources/sql_widget.md).
5353
* Provide [global configuration for all SQL warehouses](docs/resources/sql_global_config.md)
5454

55-
MLFlow
55+
Machine Learning
5656

57-
* Create [MLFlow models](resources/mlflow_model.md).
58-
* Create [MLFlow experiments](resources/mlflow_experiment.md).
57+
* Create [models in Unity Catalog](resources/registered_model.md).
58+
* Create [MLflow experiments](resources/mlflow_experiment.md).
59+
* Create [models in the workspace model registry](resources/mlflow_model.md).
60+
* Create [model serving endpoints](resources/model_serving.md).
5961

6062
## Example Usage
6163

‎docs/resources/grants.md‎

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -194,6 +194,24 @@ resource "databricks_grants" "volume" {
194194
}
195195
```
196196

197+
## Registered model grants
198+
199+
You can grant `ALL_PRIVILEGES`, `APPLY_TAG`, and `EXECUTE` privileges to [_`catalog.schema.model`_](registered_model.md) specified in the `model` attribute.
200+
201+
```hcl
202+
resource "databricks_grants" "customers" {
203+
model = "main.reporting.customer_model"
204+
grant {
205+
principal = "Data Engineers"
206+
privileges = ["APPLY_TAG", "EXECUTE"]
207+
}
208+
grant {
209+
principal = "Data Analysts"
210+
privileges = ["EXECUTE"]
211+
}
212+
}
213+
```
214+
197215
## Storage credential grants
198216

199217
You can grant `ALL_PRIVILEGES`, `CREATE_EXTERNAL_LOCATION`, `CREATE_EXTERNAL_TABLE`, `READ_FILES` and `WRITE_FILES` privileges to [databricks_storage_credential](storage_credential.md) id specified in `storage_credential` attribute:

‎docs/resources/mlflow_experiment.md‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -41,9 +41,10 @@ $ terraform import databricks_mlflow_experiment.this <experiment-id>
4141

4242
The following resources are often used in the same context:
4343

44+
* [databricks_registered_model](registered_model.md) to create [Models in Unity Catalog](https://docs.databricks.com/en/mlflow/models-in-uc.html) in Databricks.
4445
* [End to end workspace management](../guides/workspace-management.md) guide.
4546
* [databricks_directory](directory.md) to manage directories in [Databricks Workpace](https://docs.databricks.com/workspace/workspace-objects.html).
46-
* [databricks_mlflow_model](mlflow_model.md) to create [MLflow models](https://docs.databricks.com/applications/mlflow/models.html) in Databricks.
47+
* [databricks_mlflow_model](mlflow_model.md) to create models in the [workspace model registry](https://docs.databricks.com/en/mlflow/model-registry.html) in Databricks.
4748
* [databricks_notebook](notebook.md) to manage [Databricks Notebooks](https://docs.databricks.com/notebooks/index.html).
4849
* [databricks_notebook](../data-sources/notebook.md) data to export a notebook from Databricks Workspace.
4950
* [databricks_repo](repo.md) to manage [Databricks Repos](https://docs.databricks.com/repos.html).

‎docs/resources/mlflow_model.md‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,8 @@ subcategory: "MLflow"
55

66
This resource allows you to create [MLflow models](https://docs.databricks.com/applications/mlflow/models.html) in Databricks.
77

8+
**Note** This documentation covers the Workspace Model Registry. Databricks recommends using [Models in Unity Catalog](registered_model.md). Models in Unity Catalog provides centralized model governance, cross-workspace access, lineage, and deployment.
9+
810
## Example Usage
911

1012
```hcl
@@ -48,6 +50,7 @@ $ terraform import databricks_mlflow_model.this <name>
4850

4951
The following resources are often used in the same context:
5052

53+
* [databricks_registered_model](registered_model.md) to create [Models in Unity Catalog](https://docs.databricks.com/en/mlflow/models-in-uc.html) in Databricks.
5154
* [End to end workspace management](../guides/workspace-management.md) guide.
5255
* [databricks_model_serving](model_serving.md) to serve this model on a Databricks serving endpoint.
5356
* [databricks_directory](directory.md) to manage directories in [Databricks Workspace](https://docs.databricks.com/workspace/workspace-objects.html).

0 commit comments

Comments
 (0)