Skip to content

Commit 918ead8

Browse files
shreyas-goenkashreyas-goenka
authored
Add support for registered models for direct deployment (#3670)
## Changes This PR adds support for registered models, matching the same behaviour as TF. This PR also adds test server handles for registered models and UC catalogs, as was necessary for the integration test. ## Why To complete coverage for direct deployment. ## Tests New integration test.
1 parent 5a78875 commit 918ead8

File tree

13 files changed

+560
-6
lines changed

13 files changed

+560
-6
lines changed

‎acceptance/bundle/deploy/registered_models/basic/databricks.yml.tmpl‎

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,10 @@
1+
bundle:
2+
name: deploy-registered-models-basic-$UNIQUE_NAME
3+
4+
resources:
5+
registered_models:
6+
my_registered_model:
7+
name: $NAME
8+
comment: $COMMENT
9+
catalog_name: $CATALOG_NAME
10+
schema_name: $SCHEMA_NAME

‎acceptance/bundle/deploy/registered_models/basic/out.test.toml‎

Lines changed: 6 additions & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎acceptance/bundle/deploy/registered_models/basic/output.txt‎

Lines changed: 132 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,132 @@
1+
2+
>>> export NAME=my-registered-model-[UNIQUE_NAME]
3+
4+
>>> export COMMENT=original comment
5+
6+
>>> export CATALOG_NAME=main
7+
8+
>>> export SCHEMA_NAME=default
9+
10+
=== create catalog and schema to test diff functionality
11+
>>> [CLI] catalogs create mycatalog-[UNIQUE_NAME]
12+
{
13+
"full_name": "mycatalog-[UNIQUE_NAME]"
14+
}
15+
16+
>>> [CLI] schemas create myschema-[UNIQUE_NAME] mycatalog-[UNIQUE_NAME]
17+
{
18+
"full_name": "mycatalog-[UNIQUE_NAME].myschema-[UNIQUE_NAME]"
19+
}
20+
21+
=== create the registered model
22+
>>> [CLI] bundle plan
23+
create registered_models.my_registered_model
24+
25+
Plan: 1 to add, 0 to change, 0 to delete, 0 unchanged
26+
27+
>>> [CLI] bundle deploy
28+
Uploading bundle files to /Workspace/Users/[USERNAME]/.bundle/deploy-registered-models-basic-[UNIQUE_NAME]/default/files...
29+
Deploying resources...
30+
Updating deployment state...
31+
Deployment complete!
32+
33+
>>> [CLI] registered-models get main.default.my-registered-model-[UNIQUE_NAME]
34+
{
35+
"name": "my-registered-model-[UNIQUE_NAME]",
36+
"comment": "original comment",
37+
"catalog_name": "main",
38+
"schema_name": "default"
39+
}
40+
41+
=== update the comment, this should not recreate
42+
>>> [CLI] bundle plan
43+
update registered_models.my_registered_model
44+
45+
Plan: 0 to add, 1 to change, 0 to delete, 0 unchanged
46+
47+
>>> [CLI] bundle deploy
48+
Uploading bundle files to /Workspace/Users/[USERNAME]/.bundle/deploy-registered-models-basic-[UNIQUE_NAME]/default/files...
49+
Deploying resources...
50+
Updating deployment state...
51+
Deployment complete!
52+
53+
>>> [CLI] registered-models get main.default.my-registered-model-[UNIQUE_NAME]
54+
{
55+
"name": "my-registered-model-[UNIQUE_NAME]",
56+
"comment": "updated comment",
57+
"catalog_name": "main",
58+
"schema_name": "default"
59+
}
60+
61+
=== update the name, this should recreate
62+
>>> [CLI] bundle plan
63+
recreate registered_models.my_registered_model
64+
65+
Plan: 1 to add, 0 to change, 1 to delete, 0 unchanged
66+
67+
>>> [CLI] bundle deploy
68+
Uploading bundle files to /Workspace/Users/[USERNAME]/.bundle/deploy-registered-models-basic-[UNIQUE_NAME]/default/files...
69+
Deploying resources...
70+
Updating deployment state...
71+
Deployment complete!
72+
73+
>>> [CLI] registered-models get main.default.my-registered-model-updated-[UNIQUE_NAME]
74+
{
75+
"name": "my-registered-model-updated-[UNIQUE_NAME]",
76+
"comment": "updated comment",
77+
"catalog_name": "main",
78+
"schema_name": "default"
79+
}
80+
81+
=== update the catalog name, this should recreate
82+
>>> [CLI] bundle plan
83+
recreate registered_models.my_registered_model
84+
85+
Plan: 1 to add, 0 to change, 1 to delete, 0 unchanged
86+
87+
>>> [CLI] bundle deploy
88+
Uploading bundle files to /Workspace/Users/[USERNAME]/.bundle/deploy-registered-models-basic-[UNIQUE_NAME]/default/files...
89+
Deploying resources...
90+
Updating deployment state...
91+
Deployment complete!
92+
93+
>>> [CLI] registered-models get mycatalog-[UNIQUE_NAME].default.my-registered-model-updated-[UNIQUE_NAME]
94+
{
95+
"name": "my-registered-model-updated-[UNIQUE_NAME]",
96+
"comment": "updated comment",
97+
"catalog_name": "mycatalog-[UNIQUE_NAME]",
98+
"schema_name": "default"
99+
}
100+
101+
=== update the schema name, this should recreate
102+
>>> [CLI] bundle plan
103+
recreate registered_models.my_registered_model
104+
105+
Plan: 1 to add, 0 to change, 1 to delete, 0 unchanged
106+
107+
>>> [CLI] bundle deploy
108+
Uploading bundle files to /Workspace/Users/[USERNAME]/.bundle/deploy-registered-models-basic-[UNIQUE_NAME]/default/files...
109+
Deploying resources...
110+
Updating deployment state...
111+
Deployment complete!
112+
113+
>>> [CLI] registered-models get mycatalog-[UNIQUE_NAME].myschema-[UNIQUE_NAME].my-registered-model-updated-[UNIQUE_NAME]
114+
{
115+
"name": "my-registered-model-updated-[UNIQUE_NAME]",
116+
"comment": "updated comment",
117+
"catalog_name": "mycatalog-[UNIQUE_NAME]",
118+
"schema_name": "myschema-[UNIQUE_NAME]"
119+
}
120+
121+
>>> [CLI] bundle destroy --auto-approve
122+
The following resources will be deleted:
123+
delete registered_model my_registered_model
124+
125+
All files and directories at the following location will be deleted: /Workspace/Users/[USERNAME]/.bundle/deploy-registered-models-basic-[UNIQUE_NAME]/default
126+
127+
Deleting files...
128+
Destroy complete!
129+
130+
>>> [CLI] schemas delete mycatalog-[UNIQUE_NAME].myschema-[UNIQUE_NAME] --force
131+
132+
>>> [CLI] catalogs delete mycatalog-[UNIQUE_NAME] --force

‎acceptance/bundle/deploy/registered_models/basic/script‎

Lines changed: 50 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,50 @@
1+
trace export NAME="my-registered-model-$UNIQUE_NAME"
2+
trace export COMMENT="original comment"
3+
trace export CATALOG_NAME="main"
4+
trace export SCHEMA_NAME="default"
5+
envsubst < databricks.yml.tmpl > databricks.yml
6+
7+
title "create catalog and schema to test diff functionality"
8+
catalog_name="mycatalog-${UNIQUE_NAME}"
9+
schema_name="myschema-${UNIQUE_NAME}"
10+
trace $CLI catalogs create ${catalog_name} | jq '{full_name}'
11+
trace $CLI schemas create ${schema_name} ${catalog_name} | jq '{full_name}'
12+
13+
cleanup() {
14+
trace $CLI bundle destroy --auto-approve
15+
trace $CLI schemas delete ${catalog_name}.${schema_name} --force
16+
trace $CLI catalogs delete ${catalog_name} --force
17+
}
18+
trap cleanup EXIT
19+
20+
deploy_registered_model() {
21+
trace $CLI bundle plan
22+
trace $CLI bundle deploy
23+
registered_model_id=$($CLI bundle summary --output json | jq -r '.resources.registered_models.my_registered_model.id')
24+
trace $CLI registered-models get "${registered_model_id}" | jq '{name, comment, catalog_name, schema_name}'
25+
}
26+
27+
title "create the registered model"
28+
deploy_registered_model
29+
30+
export COMMENT="updated comment"
31+
envsubst < databricks.yml.tmpl > databricks.yml
32+
33+
title "update the comment, this should not recreate"
34+
deploy_registered_model
35+
36+
export NAME="my-registered-model-updated-$UNIQUE_NAME"
37+
envsubst < databricks.yml.tmpl > databricks.yml
38+
39+
title "update the name, this should recreate"
40+
deploy_registered_model
41+
42+
title "update the catalog name, this should recreate"
43+
export CATALOG_NAME="${catalog_name}"
44+
envsubst < databricks.yml.tmpl > databricks.yml
45+
deploy_registered_model
46+
47+
title "update the schema name, this should recreate"
48+
export SCHEMA_NAME="${schema_name}"
49+
envsubst < databricks.yml.tmpl > databricks.yml
50+
deploy_registered_model

‎acceptance/bundle/deploy/registered_models/basic/test.toml‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,3 @@
1+
Cloud = true
2+
Local = true
3+
RequiresUnityCatalog = true

‎acceptance/bundle/refschema/out.fields.txt‎

Lines changed: 27 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2550,6 +2550,33 @@ resources.pipelines.*.trigger.cron.quartz_cron_schedule string INPUT STATE
25502550
resources.pipelines.*.trigger.cron.timezone_id string INPUT STATE
25512551
resources.pipelines.*.trigger.manual *pipelines.ManualTrigger INPUT STATE
25522552
resources.pipelines.*.url string INPUT
2553+
resources.registered_models.*.aliases []catalog.RegisteredModelAlias REMOTE
2554+
resources.registered_models.*.aliases[*] catalog.RegisteredModelAlias REMOTE
2555+
resources.registered_models.*.aliases[*].alias_name string REMOTE
2556+
resources.registered_models.*.aliases[*].version_num int REMOTE
2557+
resources.registered_models.*.browse_only bool REMOTE
2558+
resources.registered_models.*.catalog_name string ALL
2559+
resources.registered_models.*.comment string ALL
2560+
resources.registered_models.*.created_at int64 REMOTE
2561+
resources.registered_models.*.created_by string REMOTE
2562+
resources.registered_models.*.full_name string REMOTE
2563+
resources.registered_models.*.grants []resources.Grant INPUT
2564+
resources.registered_models.*.grants[*] resources.Grant INPUT
2565+
resources.registered_models.*.grants[*].principal string INPUT
2566+
resources.registered_models.*.grants[*].privileges []string INPUT
2567+
resources.registered_models.*.grants[*].privileges[*] string INPUT
2568+
resources.registered_models.*.id string INPUT
2569+
resources.registered_models.*.lifecycle resources.Lifecycle INPUT
2570+
resources.registered_models.*.lifecycle.prevent_destroy bool INPUT
2571+
resources.registered_models.*.metastore_id string REMOTE
2572+
resources.registered_models.*.modified_status string INPUT
2573+
resources.registered_models.*.name string ALL
2574+
resources.registered_models.*.owner string REMOTE
2575+
resources.registered_models.*.schema_name string ALL
2576+
resources.registered_models.*.storage_location string ALL
2577+
resources.registered_models.*.updated_at int64 REMOTE
2578+
resources.registered_models.*.updated_by string REMOTE
2579+
resources.registered_models.*.url string INPUT
25532580
resources.schemas.*.browse_only bool REMOTE
25542581
resources.schemas.*.catalog_name string ALL
25552582
resources.schemas.*.catalog_type catalog.CatalogType REMOTE

‎bundle/direct/dresources/all.go‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@ var SupportedResources = map[string]any{
1818
"database_catalogs": (*ResourceDatabaseCatalog)(nil),
1919
"synced_database_tables": (*ResourceSyncedDatabaseTable)(nil),
2020
"alerts": (*ResourceAlert)(nil),
21+
"registered_models": (*ResourceRegisteredModel)(nil),
2122
}
2223

2324
func InitAll(client *databricks.WorkspaceClient) (map[string]*Adapter, error) {

‎bundle/direct/dresources/all_test.go‎

Lines changed: 21 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -61,6 +61,17 @@ var testConfig map[string]any = map[string]any{
6161
Name: "main.myschema.my_synced_table",
6262
},
6363
},
64+
65+
"registered_models": &resources.RegisteredModel{
66+
CreateRegisteredModelRequest: catalog.CreateRegisteredModelRequest{
67+
Name: "my_registered_model",
68+
Comment: "Test registered model",
69+
CatalogName: "main",
70+
SchemaName: "default",
71+
StorageLocation: "s3://my-bucket/my-path",
72+
},
73+
},
74+
6475
"models": &resources.MlflowModel{
6576
CreateModelRequest: ml.CreateModelRequest{
6677
Name: "my_mlflow_model",
@@ -150,22 +161,26 @@ func testCRUD(t *testing.T, group string, adapter *Adapter, client *databricks.W
150161
require.Equal(t, remote, remoteStateFromWaitCreate)
151162
}
152163

164+
remappedState, err := adapter.RemapState(remote)
165+
require.NoError(t, err)
166+
require.NotNil(t, remappedState)
167+
153168
remoteStateFromUpdate, err := adapter.DoUpdate(ctx, createdID, newState)
154169
require.NoError(t, err, "DoUpdate failed")
155170
if remoteStateFromUpdate != nil {
156-
require.Equal(t, remote, remoteStateFromUpdate)
171+
remappedStateFromUpdate, err := adapter.RemapState(remoteStateFromUpdate)
172+
require.NoError(t, err)
173+
require.Equal(t, remappedState, remappedStateFromUpdate)
157174
}
158175

159176
remoteStateFromWaitUpdate, err := adapter.WaitAfterUpdate(ctx, newState)
160177
require.NoError(t, err)
161178
if remoteStateFromWaitUpdate != nil {
162-
require.Equal(t, remote, remoteStateFromWaitUpdate)
179+
remappedStateFromWaitUpdate, err := adapter.RemapState(remoteStateFromWaitUpdate)
180+
require.NoError(t, err)
181+
require.Equal(t, remappedState, remappedStateFromWaitUpdate)
163182
}
164183

165-
remappedState, err := adapter.RemapState(remote)
166-
require.NoError(t, err)
167-
require.NotNil(t, remappedState)
168-
169184
require.NoError(t, structwalk.Walk(newState, func(path *structpath.PathNode, val any, field *reflect.StructField) {
170185
remoteValue, err := structaccess.Get(remappedState, path)
171186
if err != nil {

‎bundle/direct/dresources/registered_model.go‎

Lines changed: 93 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,93 @@
1+
package dresources
2+
3+
import (
4+
"context"
5+
6+
"github.com/databricks/cli/bundle/config/resources"
7+
"github.com/databricks/cli/bundle/deployplan"
8+
"github.com/databricks/databricks-sdk-go"
9+
"github.com/databricks/databricks-sdk-go/service/catalog"
10+
)
11+
12+
type ResourceRegisteredModel struct {
13+
client *databricks.WorkspaceClient
14+
}
15+
16+
func (*ResourceRegisteredModel) New(client *databricks.WorkspaceClient) *ResourceRegisteredModel {
17+
return &ResourceRegisteredModel{
18+
client: client,
19+
}
20+
}
21+
22+
func (*ResourceRegisteredModel) PrepareState(input *resources.RegisteredModel) *catalog.CreateRegisteredModelRequest {
23+
return &input.CreateRegisteredModelRequest
24+
}
25+
26+
func (*ResourceRegisteredModel) RemapState(model *catalog.RegisteredModelInfo) *catalog.CreateRegisteredModelRequest {
27+
return &catalog.CreateRegisteredModelRequest{
28+
CatalogName: model.CatalogName,
29+
Comment: model.Comment,
30+
Name: model.Name,
31+
SchemaName: model.SchemaName,
32+
StorageLocation: model.StorageLocation,
33+
ForceSendFields: filterFields[catalog.CreateRegisteredModelRequest](model.ForceSendFields),
34+
}
35+
}
36+
37+
func (r *ResourceRegisteredModel) DoRefresh(ctx context.Context, id string) (*catalog.RegisteredModelInfo, error) {
38+
return r.client.RegisteredModels.Get(ctx, catalog.GetRegisteredModelRequest{
39+
FullName: id,
40+
IncludeAliases: false,
41+
IncludeBrowse: false,
42+
ForceSendFields: nil,
43+
})
44+
}
45+
46+
func (r *ResourceRegisteredModel) DoCreate(ctx context.Context, config *catalog.CreateRegisteredModelRequest) (string, *catalog.RegisteredModelInfo, error) {
47+
response, err := r.client.RegisteredModels.Create(ctx, *config)
48+
if err != nil {
49+
return "", nil, err
50+
}
51+
52+
return response.FullName, response, nil
53+
}
54+
55+
func (r *ResourceRegisteredModel) DoUpdate(ctx context.Context, id string, config *catalog.CreateRegisteredModelRequest) (*catalog.RegisteredModelInfo, error) {
56+
updateRequest := catalog.UpdateRegisteredModelRequest{
57+
FullName: id,
58+
Comment: config.Comment,
59+
ForceSendFields: filterFields[catalog.UpdateRegisteredModelRequest](config.ForceSendFields, "Owner", "NewName"),
60+
61+
// Owner is not part of the configuration tree
62+
Owner: "",
63+
64+
// Name updates are not supported yet without recreating. Can be added as a follow-up.
65+
// Note: TF also does not support changing name without a recreate so the current behavior matches TF.
66+
NewName: "",
67+
}
68+
69+
response, err := r.client.RegisteredModels.Update(ctx, updateRequest)
70+
if err != nil {
71+
return nil, err
72+
}
73+
74+
return response, nil
75+
}
76+
77+
func (r *ResourceRegisteredModel) DoDelete(ctx context.Context, id string) error {
78+
return r.client.RegisteredModels.Delete(ctx, catalog.DeleteRegisteredModelRequest{
79+
FullName: id,
80+
})
81+
}
82+
83+
func (*ResourceRegisteredModel) FieldTriggers() map[string]deployplan.ActionType {
84+
return map[string]deployplan.ActionType{
85+
// The name can technically be updated without recreated. We recreate for now though
86+
// to match TF implementation.
87+
"name": deployplan.ActionTypeRecreate,
88+
89+
"catalog_name": deployplan.ActionTypeRecreate,
90+
"schema_name": deployplan.ActionTypeRecreate,
91+
"storage_location": deployplan.ActionTypeRecreate,
92+
}
93+
}

0 commit comments

Comments
 (0)