From a937717bef1afda14eadb172c8715a35b15ebb14 Mon Sep 17 00:00:00 2001 From: Arpit Jasapara Date: Sun, 8 Oct 2023 17:10:04 -0700 Subject: [PATCH 1/5] Support UC Registered Models in Bundles Signed-off-by: Arpit Jasapara --- Makefile | 4 +- bundle/config/mutator/process_target_mode.go | 6 ++ .../mutator/process_target_mode_test.go | 17 +++++- bundle/config/resources.go | 17 ++++++ bundle/config/resources/grant.go | 9 +++ bundle/config/resources/registered_model.go | 25 +++++++++ bundle/deploy/terraform/convert.go | 40 +++++++++++++ bundle/deploy/terraform/convert_test.go | 56 +++++++++++++++++++ bundle/deploy/terraform/interpolate.go | 3 + .../tf/schema/resource_registered_model.go | 11 ++++ bundle/internal/tf/schema/resources.go | 2 + bundle/schema/docs/bundle_descriptions.json | 37 ++++++++++++ bundle/schema/openapi.go | 18 ++++++ bundle/tests/registered_model/databricks.yml | 32 +++++++++++ bundle/tests/registered_model_test.go | 47 ++++++++++++++++ 15 files changed, 321 insertions(+), 3 deletions(-) create mode 100644 bundle/config/resources/grant.go create mode 100644 bundle/config/resources/registered_model.go create mode 100644 bundle/internal/tf/schema/resource_registered_model.go create mode 100644 bundle/tests/registered_model/databricks.yml create mode 100644 bundle/tests/registered_model_test.go diff --git a/Makefile b/Makefile index 6067d45b996..477190fe941 100644 --- a/Makefile +++ b/Makefile @@ -8,7 +8,7 @@ fmt: lint: vendor @echo "✓ Linting source code with https://staticcheck.io/ ..." - @staticcheck ./... + @go run honnef.co/go/tools/cmd/staticcheck@v0.4.0 ./... test: lint @echo "✓ Running tests ..." @@ -30,4 +30,4 @@ vendor: @echo "✓ Filling vendor folder with library code ..." @go mod vendor -.PHONY: build vendor coverage test lint fmt \ No newline at end of file +.PHONY: build vendor coverage test lint fmt diff --git a/bundle/config/mutator/process_target_mode.go b/bundle/config/mutator/process_target_mode.go index 2f80fe3b300..c11bd1c5a6e 100644 --- a/bundle/config/mutator/process_target_mode.go +++ b/bundle/config/mutator/process_target_mode.go @@ -87,6 +87,12 @@ func transformDevelopmentMode(b *bundle.Bundle) error { // (model serving doesn't yet support tags) } + for i := range r.RegisteredModels { + prefix = "dev_" + b.Config.Workspace.CurrentUser.ShortName + "_" + r.RegisteredModels[i].Name = prefix + r.RegisteredModels[i].Name + // (registered models in Unity Catalog don't yet support tags) + } + return nil } diff --git a/bundle/config/mutator/process_target_mode_test.go b/bundle/config/mutator/process_target_mode_test.go index a0b2bac8537..fd786992953 100644 --- a/bundle/config/mutator/process_target_mode_test.go +++ b/bundle/config/mutator/process_target_mode_test.go @@ -11,6 +11,7 @@ import ( "github.com/databricks/cli/bundle/config/resources" "github.com/databricks/cli/libs/tags" sdkconfig "github.com/databricks/databricks-sdk-go/config" + "github.com/databricks/databricks-sdk-go/service/catalog" "github.com/databricks/databricks-sdk-go/service/iam" "github.com/databricks/databricks-sdk-go/service/jobs" "github.com/databricks/databricks-sdk-go/service/ml" @@ -59,6 +60,9 @@ func mockBundle(mode config.Mode) *bundle.Bundle { ModelServingEndpoints: map[string]*resources.ModelServingEndpoint{ "servingendpoint1": {CreateServingEndpoint: &serving.CreateServingEndpoint{Name: "servingendpoint1"}}, }, + RegisteredModels: map[string]*resources.RegisteredModel{ + "registeredmodel1": {CreateRegisteredModelRequest: &catalog.CreateRegisteredModelRequest{Name: "registeredmodel1", Comment: "comment", CatalogName: "catalog", SchemaName: "schema"}}, + }, }, }, // Use AWS implementation for testing. @@ -86,6 +90,7 @@ func TestProcessTargetModeDevelopment(t *testing.T) { // Experiment 1 assert.Equal(t, "/Users/lennart.kats@databricks.com/[dev lennart] experiment1", bundle.Config.Resources.Experiments["experiment1"].Name) assert.Contains(t, bundle.Config.Resources.Experiments["experiment1"].Experiment.Tags, ml.ExperimentTag{Key: "dev", Value: "lennart"}) + assert.Equal(t, "dev", bundle.Config.Resources.Experiments["experiment1"].Experiment.Tags[0].Key) // Experiment 2 assert.Equal(t, "[dev lennart] experiment2", bundle.Config.Resources.Experiments["experiment2"].Name) @@ -96,7 +101,9 @@ func TestProcessTargetModeDevelopment(t *testing.T) { // Model serving endpoint 1 assert.Equal(t, "dev_lennart_servingendpoint1", bundle.Config.Resources.ModelServingEndpoints["servingendpoint1"].Name) - assert.Equal(t, "dev", bundle.Config.Resources.Experiments["experiment1"].Experiment.Tags[0].Key) + + // Registered Model 1 + assert.Equal(t, "dev_lennart_registeredmodel1", bundle.Config.Resources.RegisteredModels["registeredmodel1"].Name) } func TestProcessTargetModeDevelopmentTagNormalizationForAws(t *testing.T) { @@ -151,6 +158,7 @@ func TestProcessTargetModeDefault(t *testing.T) { assert.Equal(t, "pipeline1", bundle.Config.Resources.Pipelines["pipeline1"].Name) assert.False(t, bundle.Config.Resources.Pipelines["pipeline1"].PipelineSpec.Development) assert.Equal(t, "servingendpoint1", bundle.Config.Resources.ModelServingEndpoints["servingendpoint1"].Name) + assert.Equal(t, "registeredmodel1", bundle.Config.Resources.RegisteredModels["registeredmodel1"].Name) } func TestProcessTargetModeProduction(t *testing.T) { @@ -179,6 +187,12 @@ func TestProcessTargetModeProduction(t *testing.T) { bundle.Config.Resources.Experiments["experiment2"].Permissions = permissions bundle.Config.Resources.Models["model1"].Permissions = permissions bundle.Config.Resources.ModelServingEndpoints["servingendpoint1"].Permissions = permissions + bundle.Config.Resources.RegisteredModels["registeredmodel1"].Grants = []resources.Grant{ + { + Privileges: []string{"ALL_PRIVILEGES"}, + Principal: "user@company.com", + }, + } err = validateProductionMode(context.Background(), bundle, false) require.NoError(t, err) @@ -187,6 +201,7 @@ func TestProcessTargetModeProduction(t *testing.T) { assert.Equal(t, "pipeline1", bundle.Config.Resources.Pipelines["pipeline1"].Name) assert.False(t, bundle.Config.Resources.Pipelines["pipeline1"].PipelineSpec.Development) assert.Equal(t, "servingendpoint1", bundle.Config.Resources.ModelServingEndpoints["servingendpoint1"].Name) + assert.Equal(t, "registeredmodel1", bundle.Config.Resources.RegisteredModels["registeredmodel1"].Name) } func TestProcessTargetModeProductionOkForPrincipal(t *testing.T) { diff --git a/bundle/config/resources.go b/bundle/config/resources.go index ad1d6e9a3c9..2b453c666e5 100644 --- a/bundle/config/resources.go +++ b/bundle/config/resources.go @@ -14,6 +14,7 @@ type Resources struct { Models map[string]*resources.MlflowModel `json:"models,omitempty"` Experiments map[string]*resources.MlflowExperiment `json:"experiments,omitempty"` ModelServingEndpoints map[string]*resources.ModelServingEndpoint `json:"model_serving_endpoints,omitempty"` + RegisteredModels map[string]*resources.RegisteredModel `json:"registered_models,omitempty"` } type UniqueResourceIdTracker struct { @@ -107,6 +108,19 @@ func (r *Resources) VerifyUniqueResourceIdentifiers() (*UniqueResourceIdTracker, tracker.Type[k] = "model_serving_endpoint" tracker.ConfigPath[k] = r.ModelServingEndpoints[k].ConfigFilePath } + for k := range r.RegisteredModels { + if _, ok := tracker.Type[k]; ok { + return tracker, fmt.Errorf("multiple resources named %s (%s at %s, %s at %s)", + k, + tracker.Type[k], + tracker.ConfigPath[k], + "registered_model", + r.RegisteredModels[k].ConfigFilePath, + ) + } + tracker.Type[k] = "registered_model" + tracker.ConfigPath[k] = r.RegisteredModels[k].ConfigFilePath + } return tracker, nil } @@ -129,6 +143,9 @@ func (r *Resources) SetConfigFilePath(path string) { for _, e := range r.ModelServingEndpoints { e.ConfigFilePath = path } + for _, e := range r.RegisteredModels { + e.ConfigFilePath = path + } } // Merge iterates over all resources and merges chunks of the diff --git a/bundle/config/resources/grant.go b/bundle/config/resources/grant.go new file mode 100644 index 00000000000..5522033d163 --- /dev/null +++ b/bundle/config/resources/grant.go @@ -0,0 +1,9 @@ +package resources + +// Grant holds the grant level settings for a single principal in Unity Catalog. +// Multiple of these can be defined on any resource. +type Grant struct { + Privileges []string `json:"privileges"` + + Principal string `json:"principal"` +} diff --git a/bundle/config/resources/registered_model.go b/bundle/config/resources/registered_model.go new file mode 100644 index 00000000000..46cb169de4b --- /dev/null +++ b/bundle/config/resources/registered_model.go @@ -0,0 +1,25 @@ +package resources + +import ( + "github.com/databricks/cli/bundle/config/paths" + "github.com/databricks/databricks-sdk-go/service/catalog" +) + +type RegisteredModel struct { + // This is a resource agnostic implementation of grants for ACLs. + // Implementation could be different based on the resource type. + Grants []Grant `json:"grants,omitempty"` + + // This represents the id which is the full name of the model + // (catalog_name.schema_name.model_name) that can be used + // as a reference in other resources. This value is returned by terraform. + ID string + + // Local path where the bundle is defined. All bundle resources include + // this for interpolation purposes. + paths.Paths + + // This represents the input args for terraform, and will get converted + // to a HCL representation for CRUD + *catalog.CreateRegisteredModelRequest +} diff --git a/bundle/deploy/terraform/convert.go b/bundle/deploy/terraform/convert.go index 7d95e719d5d..d844f4b4af8 100644 --- a/bundle/deploy/terraform/convert.go +++ b/bundle/deploy/terraform/convert.go @@ -44,6 +44,26 @@ func convPermission(ac resources.Permission) schema.ResourcePermissionsAccessCon return dst } +func convGrants(acl []resources.Grant) *schema.ResourceGrants { + if len(acl) == 0 { + return nil + } + + resource := schema.ResourceGrants{} + for _, ac := range acl { + resource.Grant = append(resource.Grant, convGrant(ac)) + } + + return &resource +} + +func convGrant(ac resources.Grant) schema.ResourceGrantsGrant { + return schema.ResourceGrantsGrant{ + Privileges: ac.Privileges, + Principal: ac.Principal, + } +} + // BundleToTerraform converts resources in a bundle configuration // to the equivalent Terraform JSON representation. // @@ -174,6 +194,19 @@ func BundleToTerraform(config *config.Root) *schema.Root { } } + for k, src := range config.Resources.RegisteredModels { + noResources = false + var dst schema.ResourceRegisteredModel + conv(src, &dst) + tfroot.Resource.RegisteredModel[k] = &dst + + // Configure permissions for this resource. + if rp := convGrants(src.Grants); rp != nil { + rp.Function = fmt.Sprintf("${databricks_registered_model.%s.id}", k) + tfroot.Resource.Grants["registered_model_"+k] = rp + } + } + // We explicitly set "resource" to nil to omit it from a JSON encoding. // This is required because the terraform CLI requires >= 1 resources defined // if the "resource" property is used in a .tf.json file. @@ -221,7 +254,14 @@ func TerraformToBundle(state *tfjson.State, config *config.Root) error { cur := config.Resources.ModelServingEndpoints[resource.Name] conv(tmp, &cur) config.Resources.ModelServingEndpoints[resource.Name] = cur + case "databricks_registered_model": + var tmp schema.ResourceRegisteredModel + conv(resource.AttributeValues, &tmp) + cur := config.Resources.RegisteredModels[resource.Name] + conv(tmp, &cur) + config.Resources.RegisteredModels[resource.Name] = cur case "databricks_permissions": + case "databricks_grants": // Ignore; no need to pull these back into the configuration. default: return fmt.Errorf("missing mapping for %s", resource.Type) diff --git a/bundle/deploy/terraform/convert_test.go b/bundle/deploy/terraform/convert_test.go index b6b29f35ab0..bb5a63ecce5 100644 --- a/bundle/deploy/terraform/convert_test.go +++ b/bundle/deploy/terraform/convert_test.go @@ -5,6 +5,7 @@ import ( "github.com/databricks/cli/bundle/config" "github.com/databricks/cli/bundle/config/resources" + "github.com/databricks/databricks-sdk-go/service/catalog" "github.com/databricks/databricks-sdk-go/service/compute" "github.com/databricks/databricks-sdk-go/service/jobs" "github.com/databricks/databricks-sdk-go/service/ml" @@ -366,3 +367,58 @@ func TestConvertModelServingPermissions(t *testing.T) { assert.Equal(t, "CAN_VIEW", p.PermissionLevel) } + +func TestConvertRegisteredModel(t *testing.T) { + var src = resources.RegisteredModel{ + CreateRegisteredModelRequest: &catalog.CreateRegisteredModelRequest{ + Name: "name", + CatalogName: "catalog", + SchemaName: "schema", + Comment: "comment", + }, + } + + var config = config.Root{ + Resources: config.Resources{ + RegisteredModels: map[string]*resources.RegisteredModel{ + "my_registered_model": &src, + }, + }, + } + + out := BundleToTerraform(&config) + resource := out.Resource.RegisteredModel["my_registered_model"] + assert.Equal(t, "name", resource.Name) + assert.Equal(t, "catalog", resource.CatalogName) + assert.Equal(t, "schema", resource.SchemaName) + assert.Equal(t, "comment", resource.Comment) + assert.Nil(t, out.Data) +} + +func TestConvertRegisteredModelGrants(t *testing.T) { + var src = resources.RegisteredModel{ + Grants: []resources.Grant{ + { + Privileges: []string{"EXECUTE"}, + Principal: "jane@doe.com", + }, + }, + } + + var config = config.Root{ + Resources: config.Resources{ + RegisteredModels: map[string]*resources.RegisteredModel{ + "my_registered_model": &src, + }, + }, + } + + out := BundleToTerraform(&config) + assert.NotEmpty(t, out.Resource.Grants["registered_model_my_registered_model"].Function) + assert.Len(t, out.Resource.Grants["registered_model_my_registered_model"].Grant, 1) + + p := out.Resource.Grants["registered_model_my_registered_model"].Grant[0] + assert.Equal(t, "jane@doe.com", p.Principal) + assert.Equal(t, "EXECUTE", p.Privileges[0]) + +} diff --git a/bundle/deploy/terraform/interpolate.go b/bundle/deploy/terraform/interpolate.go index ea3c99aa10f..4f00c27ebb6 100644 --- a/bundle/deploy/terraform/interpolate.go +++ b/bundle/deploy/terraform/interpolate.go @@ -28,6 +28,9 @@ func interpolateTerraformResourceIdentifiers(path string, lookup map[string]stri case "model_serving_endpoints": path = strings.Join(append([]string{"databricks_model_serving"}, parts[2:]...), interpolation.Delimiter) return fmt.Sprintf("${%s}", path), nil + case "registered_models": + path = strings.Join(append([]string{"databricks_registered_model"}, parts[2:]...), interpolation.Delimiter) + return fmt.Sprintf("${%s}", path), nil default: panic("TODO: " + parts[1]) } diff --git a/bundle/internal/tf/schema/resource_registered_model.go b/bundle/internal/tf/schema/resource_registered_model.go new file mode 100644 index 00000000000..0d2c65cd9c9 --- /dev/null +++ b/bundle/internal/tf/schema/resource_registered_model.go @@ -0,0 +1,11 @@ +// Generated from Databricks Terraform provider schema. DO NOT EDIT. + +package schema + +type ResourceRegisteredModel struct { + Comment string `json:"comment,omitempty"` + Id string `json:"id,omitempty"` + Name string `json:"name"` + CatalogName string `json:"catalog_name"` + SchemaName string `json:"schema_name"` +} diff --git a/bundle/internal/tf/schema/resources.go b/bundle/internal/tf/schema/resources.go index c2361254a5a..26bfc05813f 100644 --- a/bundle/internal/tf/schema/resources.go +++ b/bundle/internal/tf/schema/resources.go @@ -52,6 +52,7 @@ type Resources struct { Pipeline map[string]*ResourcePipeline `json:"databricks_pipeline,omitempty"` Provider map[string]*ResourceProvider `json:"databricks_provider,omitempty"` Recipient map[string]*ResourceRecipient `json:"databricks_recipient,omitempty"` + RegisteredModel map[string]*ResourceRegisteredModel `json:"databricks_registered_model,omitempty"` Repo map[string]*ResourceRepo `json:"databricks_repo,omitempty"` Schema map[string]*ResourceSchema `json:"databricks_schema,omitempty"` Secret map[string]*ResourceSecret `json:"databricks_secret,omitempty"` @@ -132,6 +133,7 @@ func NewResources() *Resources { Pipeline: make(map[string]*ResourcePipeline), Provider: make(map[string]*ResourceProvider), Recipient: make(map[string]*ResourceRecipient), + RegisteredModel: make(map[string]*ResourceRegisteredModel), Repo: make(map[string]*ResourceRepo), Schema: make(map[string]*ResourceSchema), Secret: make(map[string]*ResourceSecret), diff --git a/bundle/schema/docs/bundle_descriptions.json b/bundle/schema/docs/bundle_descriptions.json index 98f3cf8d0a0..b2bcd627feb 100644 --- a/bundle/schema/docs/bundle_descriptions.json +++ b/bundle/schema/docs/bundle_descriptions.json @@ -1521,6 +1521,43 @@ } } }, + "registered_models": { + "description": "List of Registered Models", + "additionalproperties": { + "description": "", + "properties": { + "name": { + "description": "The name of the registered model." + }, + "catalog_name": { + "description": "The name of the catalog where the schema and the registered model reside." + }, + "schema_name": { + "description": "The name of the schema where the registered model resides." + }, + "comment": { + "description": "The comment attached to the registered model." + }, + "grants": { + "description": "The grants that can be applied to the registered model.", + "items": { + "description": "", + "properties": { + "privileges": { + "description": "The privileges that should be granted to the principal.", + "items": { + "description": "Must be one of `ALL_PRIVILEGES`, `APPLY_TAG`, or `EXECUTE`." + } + }, + "principal": { + "description": "User name, group name or service principal application ID." + } + } + } + } + } + } + }, "pipelines": { "description": "List of DLT pipelines", "additionalproperties": { diff --git a/bundle/schema/openapi.go b/bundle/schema/openapi.go index 1a8b76ed9ea..0b64c43e3c7 100644 --- a/bundle/schema/openapi.go +++ b/bundle/schema/openapi.go @@ -223,6 +223,19 @@ func (reader *OpenapiReader) modelServingEndpointsDocs() (*Docs, error) { return modelServingEndpointsAllDocs, nil } +func (reader *OpenapiReader) registeredModelDocs() (*Docs, error) { + registeredModelsSpecSchema, err := reader.readResolvedSchema(SchemaPathPrefix + "catalog.CreateRegisteredModelRequest") + if err != nil { + return nil, err + } + registeredModelsDocs := schemaToDocs(registeredModelsSpecSchema) + registeredModelsAllDocs := &Docs{ + Description: "List of Registered Models", + AdditionalProperties: registeredModelsDocs, + } + return registeredModelsAllDocs, nil +} + func (reader *OpenapiReader) ResourcesDocs() (*Docs, error) { jobsDocs, err := reader.jobsDocs() if err != nil { @@ -244,6 +257,10 @@ func (reader *OpenapiReader) ResourcesDocs() (*Docs, error) { if err != nil { return nil, err } + registeredModelsDocs, err := reader.registeredModelDocs() + if err != nil { + return nil, err + } return &Docs{ Description: "Collection of Databricks resources to deploy.", @@ -253,6 +270,7 @@ func (reader *OpenapiReader) ResourcesDocs() (*Docs, error) { "experiments": experimentsDocs, "models": modelsDocs, "model_serving_endpoints": modelServingEndpointsDocs, + "registered_models": registeredModelsDocs, }, }, nil } diff --git a/bundle/tests/registered_model/databricks.yml b/bundle/tests/registered_model/databricks.yml new file mode 100644 index 00000000000..b7b8ea5d246 --- /dev/null +++ b/bundle/tests/registered_model/databricks.yml @@ -0,0 +1,32 @@ +resources: + registered_models: + my_registered_model: + name: "my-model" + comment: "comment" + catalog_name: "main" + schema_name: "default" + grants: + - privileges: + - EXECUTE + principal: "account users" + +targets: + development: + mode: development + resources: + registered_models: + my_registered_model: + name: "my-dev-model" + + staging: + resources: + registered_models: + my_registered_model: + name: "my-staging-model" + + production: + mode: production + resources: + registered_models: + my_registered_model: + name: "my-prod-model" diff --git a/bundle/tests/registered_model_test.go b/bundle/tests/registered_model_test.go new file mode 100644 index 00000000000..920a2ac78e6 --- /dev/null +++ b/bundle/tests/registered_model_test.go @@ -0,0 +1,47 @@ +package config_tests + +import ( + "path/filepath" + "testing" + + "github.com/databricks/cli/bundle/config" + "github.com/databricks/cli/bundle/config/resources" + "github.com/stretchr/testify/assert" +) + +func assertExpectedModel(t *testing.T, p *resources.RegisteredModel) { + assert.Equal(t, "registered_model/databricks.yml", filepath.ToSlash(p.ConfigFilePath)) + assert.Equal(t, "main", p.CatalogName) + assert.Equal(t, "default", p.SchemaName) + assert.Equal(t, "comment", p.Comment) + assert.Equal(t, "account users", p.Grants[0].Principal) + assert.Equal(t, "EXECUTE", p.Grants[0].Privileges[0]) +} + +func TestRegisteredModelDevelopment(t *testing.T) { + b := loadTarget(t, "./registered_model", "development") + assert.Len(t, b.Config.Resources.RegisteredModels, 1) + assert.Equal(t, b.Config.Bundle.Mode, config.Development) + + p := b.Config.Resources.RegisteredModels["my_registered_model"] + assert.Equal(t, "my-dev-model", p.Name) + assertExpectedModel(t, p) +} + +func TestRegisteredModelStaging(t *testing.T) { + b := loadTarget(t, "./registered_model", "staging") + assert.Len(t, b.Config.Resources.RegisteredModels, 1) + + p := b.Config.Resources.RegisteredModels["my_registered_model"] + assert.Equal(t, "my-staging-model", p.Name) + assertExpectedModel(t, p) +} + +func TestRegisteredModelProduction(t *testing.T) { + b := loadTarget(t, "./registered_model", "production") + assert.Len(t, b.Config.Resources.RegisteredModels, 1) + + p := b.Config.Resources.RegisteredModels["my_registered_model"] + assert.Equal(t, "my-prod-model", p.Name) + assertExpectedModel(t, p) +} From b5b758a8ce963c4639301d0c04c78636a7574eea Mon Sep 17 00:00:00 2001 From: Arpit Jasapara Date: Tue, 10 Oct 2023 12:16:35 -0700 Subject: [PATCH 2/5] Apply comments Signed-off-by: Arpit Jasapara --- .../mutator/process_target_mode_test.go | 6 ---- .../resources/model_serving_endpoint.go | 4 +-- bundle/config/resources/registered_model.go | 4 +-- bundle/config/resources_test.go | 30 +++++++++++++++++++ bundle/deploy/terraform/convert.go | 12 +++----- 5 files changed, 38 insertions(+), 18 deletions(-) diff --git a/bundle/config/mutator/process_target_mode_test.go b/bundle/config/mutator/process_target_mode_test.go index fd786992953..b8d157e7bdd 100644 --- a/bundle/config/mutator/process_target_mode_test.go +++ b/bundle/config/mutator/process_target_mode_test.go @@ -187,12 +187,6 @@ func TestProcessTargetModeProduction(t *testing.T) { bundle.Config.Resources.Experiments["experiment2"].Permissions = permissions bundle.Config.Resources.Models["model1"].Permissions = permissions bundle.Config.Resources.ModelServingEndpoints["servingendpoint1"].Permissions = permissions - bundle.Config.Resources.RegisteredModels["registeredmodel1"].Grants = []resources.Grant{ - { - Privileges: []string{"ALL_PRIVILEGES"}, - Principal: "user@company.com", - }, - } err = validateProductionMode(context.Background(), bundle, false) require.NoError(t, err) diff --git a/bundle/config/resources/model_serving_endpoint.go b/bundle/config/resources/model_serving_endpoint.go index dccecaa6f86..4d774873eb9 100644 --- a/bundle/config/resources/model_serving_endpoint.go +++ b/bundle/config/resources/model_serving_endpoint.go @@ -14,8 +14,8 @@ type ModelServingEndpoint struct { // as a reference in other resources. This value is returned by terraform. ID string - // Local path where the bundle is defined. All bundle resources include - // this for interpolation purposes. + // Path to config file where the resource is defined. All bundle resources + // include this for interpolation purposes. paths.Paths // This is a resource agnostic implementation of permissions for ACLs. diff --git a/bundle/config/resources/registered_model.go b/bundle/config/resources/registered_model.go index 46cb169de4b..8633da29d0c 100644 --- a/bundle/config/resources/registered_model.go +++ b/bundle/config/resources/registered_model.go @@ -15,8 +15,8 @@ type RegisteredModel struct { // as a reference in other resources. This value is returned by terraform. ID string - // Local path where the bundle is defined. All bundle resources include - // this for interpolation purposes. + // Path to config file where the resource is defined. All bundle resources + // include this for interpolation purposes. paths.Paths // This represents the input args for terraform, and will get converted diff --git a/bundle/config/resources_test.go b/bundle/config/resources_test.go index 82cb9f4549e..9c4104e4d48 100644 --- a/bundle/config/resources_test.go +++ b/bundle/config/resources_test.go @@ -95,3 +95,33 @@ func TestVerifySafeMergeForSameResourceType(t *testing.T) { err := r.VerifySafeMerge(&other) assert.ErrorContains(t, err, "multiple resources named foo (job at foo.yml, job at foo2.yml)") } + +func TestVerifySafeMergeForRegisteredModels(t *testing.T) { + r := Resources{ + Jobs: map[string]*resources.Job{ + "foo": { + Paths: paths.Paths{ + ConfigFilePath: "foo.yml", + }, + }, + }, + RegisteredModels: map[string]*resources.RegisteredModel{ + "bar": { + Paths: paths.Paths{ + ConfigFilePath: "bar.yml", + }, + }, + }, + } + other := Resources{ + RegisteredModels: map[string]*resources.RegisteredModel{ + "bar": { + Paths: paths.Paths{ + ConfigFilePath: "bar2.yml", + }, + }, + }, + } + err := r.VerifySafeMerge(&other) + assert.ErrorContains(t, err, "multiple resources named bar (registered_model at bar.yml, registered_model at bar2.yml)") +} diff --git a/bundle/deploy/terraform/convert.go b/bundle/deploy/terraform/convert.go index d844f4b4af8..3bfc8b83ba8 100644 --- a/bundle/deploy/terraform/convert.go +++ b/bundle/deploy/terraform/convert.go @@ -51,19 +51,15 @@ func convGrants(acl []resources.Grant) *schema.ResourceGrants { resource := schema.ResourceGrants{} for _, ac := range acl { - resource.Grant = append(resource.Grant, convGrant(ac)) + resource.Grant = append(resource.Grant, schema.ResourceGrantsGrant{ + Privileges: ac.Privileges, + Principal: ac.Principal, + }) } return &resource } -func convGrant(ac resources.Grant) schema.ResourceGrantsGrant { - return schema.ResourceGrantsGrant{ - Privileges: ac.Privileges, - Principal: ac.Principal, - } -} - // BundleToTerraform converts resources in a bundle configuration // to the equivalent Terraform JSON representation. // From 7cad3761d5ea4e0ea6edb7271ac0049e365dc275 Mon Sep 17 00:00:00 2001 From: Arpit Jasapara Date: Fri, 13 Oct 2023 08:58:42 -0700 Subject: [PATCH 3/5] Revert certain files and apply comments --- Makefile | 2 +- bundle/config/mutator/process_target_mode.go | 6 --- .../mutator/process_target_mode_test.go | 9 ----- bundle/config/resources/grant.go | 2 +- bundle/config/resources/registered_model.go | 2 +- .../tf/schema/resource_registered_model.go | 11 ------ bundle/internal/tf/schema/resources.go | 2 - bundle/schema/docs/bundle_descriptions.json | 37 ------------------- 8 files changed, 3 insertions(+), 68 deletions(-) delete mode 100644 bundle/internal/tf/schema/resource_registered_model.go diff --git a/Makefile b/Makefile index 477190fe941..81aac2b4b3e 100644 --- a/Makefile +++ b/Makefile @@ -8,7 +8,7 @@ fmt: lint: vendor @echo "✓ Linting source code with https://staticcheck.io/ ..." - @go run honnef.co/go/tools/cmd/staticcheck@v0.4.0 ./... + @staticcheck ./... test: lint @echo "✓ Running tests ..." diff --git a/bundle/config/mutator/process_target_mode.go b/bundle/config/mutator/process_target_mode.go index c11bd1c5a6e..2f80fe3b300 100644 --- a/bundle/config/mutator/process_target_mode.go +++ b/bundle/config/mutator/process_target_mode.go @@ -87,12 +87,6 @@ func transformDevelopmentMode(b *bundle.Bundle) error { // (model serving doesn't yet support tags) } - for i := range r.RegisteredModels { - prefix = "dev_" + b.Config.Workspace.CurrentUser.ShortName + "_" - r.RegisteredModels[i].Name = prefix + r.RegisteredModels[i].Name - // (registered models in Unity Catalog don't yet support tags) - } - return nil } diff --git a/bundle/config/mutator/process_target_mode_test.go b/bundle/config/mutator/process_target_mode_test.go index b8d157e7bdd..eb671aedf7b 100644 --- a/bundle/config/mutator/process_target_mode_test.go +++ b/bundle/config/mutator/process_target_mode_test.go @@ -11,7 +11,6 @@ import ( "github.com/databricks/cli/bundle/config/resources" "github.com/databricks/cli/libs/tags" sdkconfig "github.com/databricks/databricks-sdk-go/config" - "github.com/databricks/databricks-sdk-go/service/catalog" "github.com/databricks/databricks-sdk-go/service/iam" "github.com/databricks/databricks-sdk-go/service/jobs" "github.com/databricks/databricks-sdk-go/service/ml" @@ -60,9 +59,6 @@ func mockBundle(mode config.Mode) *bundle.Bundle { ModelServingEndpoints: map[string]*resources.ModelServingEndpoint{ "servingendpoint1": {CreateServingEndpoint: &serving.CreateServingEndpoint{Name: "servingendpoint1"}}, }, - RegisteredModels: map[string]*resources.RegisteredModel{ - "registeredmodel1": {CreateRegisteredModelRequest: &catalog.CreateRegisteredModelRequest{Name: "registeredmodel1", Comment: "comment", CatalogName: "catalog", SchemaName: "schema"}}, - }, }, }, // Use AWS implementation for testing. @@ -101,9 +97,6 @@ func TestProcessTargetModeDevelopment(t *testing.T) { // Model serving endpoint 1 assert.Equal(t, "dev_lennart_servingendpoint1", bundle.Config.Resources.ModelServingEndpoints["servingendpoint1"].Name) - - // Registered Model 1 - assert.Equal(t, "dev_lennart_registeredmodel1", bundle.Config.Resources.RegisteredModels["registeredmodel1"].Name) } func TestProcessTargetModeDevelopmentTagNormalizationForAws(t *testing.T) { @@ -158,7 +151,6 @@ func TestProcessTargetModeDefault(t *testing.T) { assert.Equal(t, "pipeline1", bundle.Config.Resources.Pipelines["pipeline1"].Name) assert.False(t, bundle.Config.Resources.Pipelines["pipeline1"].PipelineSpec.Development) assert.Equal(t, "servingendpoint1", bundle.Config.Resources.ModelServingEndpoints["servingendpoint1"].Name) - assert.Equal(t, "registeredmodel1", bundle.Config.Resources.RegisteredModels["registeredmodel1"].Name) } func TestProcessTargetModeProduction(t *testing.T) { @@ -195,7 +187,6 @@ func TestProcessTargetModeProduction(t *testing.T) { assert.Equal(t, "pipeline1", bundle.Config.Resources.Pipelines["pipeline1"].Name) assert.False(t, bundle.Config.Resources.Pipelines["pipeline1"].PipelineSpec.Development) assert.Equal(t, "servingendpoint1", bundle.Config.Resources.ModelServingEndpoints["servingendpoint1"].Name) - assert.Equal(t, "registeredmodel1", bundle.Config.Resources.RegisteredModels["registeredmodel1"].Name) } func TestProcessTargetModeProductionOkForPrincipal(t *testing.T) { diff --git a/bundle/config/resources/grant.go b/bundle/config/resources/grant.go index 5522033d163..f0ecd876740 100644 --- a/bundle/config/resources/grant.go +++ b/bundle/config/resources/grant.go @@ -1,7 +1,7 @@ package resources // Grant holds the grant level settings for a single principal in Unity Catalog. -// Multiple of these can be defined on any resource. +// Multiple of these can be defined on any Unity Catalog resource. type Grant struct { Privileges []string `json:"privileges"` diff --git a/bundle/config/resources/registered_model.go b/bundle/config/resources/registered_model.go index 8633da29d0c..d37a982e34c 100644 --- a/bundle/config/resources/registered_model.go +++ b/bundle/config/resources/registered_model.go @@ -6,7 +6,7 @@ import ( ) type RegisteredModel struct { - // This is a resource agnostic implementation of grants for ACLs. + // This is a resource agnostic implementation of grants. // Implementation could be different based on the resource type. Grants []Grant `json:"grants,omitempty"` diff --git a/bundle/internal/tf/schema/resource_registered_model.go b/bundle/internal/tf/schema/resource_registered_model.go deleted file mode 100644 index 0d2c65cd9c9..00000000000 --- a/bundle/internal/tf/schema/resource_registered_model.go +++ /dev/null @@ -1,11 +0,0 @@ -// Generated from Databricks Terraform provider schema. DO NOT EDIT. - -package schema - -type ResourceRegisteredModel struct { - Comment string `json:"comment,omitempty"` - Id string `json:"id,omitempty"` - Name string `json:"name"` - CatalogName string `json:"catalog_name"` - SchemaName string `json:"schema_name"` -} diff --git a/bundle/internal/tf/schema/resources.go b/bundle/internal/tf/schema/resources.go index 26bfc05813f..c2361254a5a 100644 --- a/bundle/internal/tf/schema/resources.go +++ b/bundle/internal/tf/schema/resources.go @@ -52,7 +52,6 @@ type Resources struct { Pipeline map[string]*ResourcePipeline `json:"databricks_pipeline,omitempty"` Provider map[string]*ResourceProvider `json:"databricks_provider,omitempty"` Recipient map[string]*ResourceRecipient `json:"databricks_recipient,omitempty"` - RegisteredModel map[string]*ResourceRegisteredModel `json:"databricks_registered_model,omitempty"` Repo map[string]*ResourceRepo `json:"databricks_repo,omitempty"` Schema map[string]*ResourceSchema `json:"databricks_schema,omitempty"` Secret map[string]*ResourceSecret `json:"databricks_secret,omitempty"` @@ -133,7 +132,6 @@ func NewResources() *Resources { Pipeline: make(map[string]*ResourcePipeline), Provider: make(map[string]*ResourceProvider), Recipient: make(map[string]*ResourceRecipient), - RegisteredModel: make(map[string]*ResourceRegisteredModel), Repo: make(map[string]*ResourceRepo), Schema: make(map[string]*ResourceSchema), Secret: make(map[string]*ResourceSecret), diff --git a/bundle/schema/docs/bundle_descriptions.json b/bundle/schema/docs/bundle_descriptions.json index b2bcd627feb..98f3cf8d0a0 100644 --- a/bundle/schema/docs/bundle_descriptions.json +++ b/bundle/schema/docs/bundle_descriptions.json @@ -1521,43 +1521,6 @@ } } }, - "registered_models": { - "description": "List of Registered Models", - "additionalproperties": { - "description": "", - "properties": { - "name": { - "description": "The name of the registered model." - }, - "catalog_name": { - "description": "The name of the catalog where the schema and the registered model reside." - }, - "schema_name": { - "description": "The name of the schema where the registered model resides." - }, - "comment": { - "description": "The comment attached to the registered model." - }, - "grants": { - "description": "The grants that can be applied to the registered model.", - "items": { - "description": "", - "properties": { - "privileges": { - "description": "The privileges that should be granted to the principal.", - "items": { - "description": "Must be one of `ALL_PRIVILEGES`, `APPLY_TAG`, or `EXECUTE`." - } - }, - "principal": { - "description": "User name, group name or service principal application ID." - } - } - } - } - } - } - }, "pipelines": { "description": "List of DLT pipelines", "additionalproperties": { From 601a57331d080ab90961d52fa90cb995229ad379 Mon Sep 17 00:00:00 2001 From: Arpit Jasapara <87999496+arpitjasa-db@users.noreply.github.com> Date: Fri, 13 Oct 2023 09:04:13 -0700 Subject: [PATCH 4/5] Update Makefile Co-authored-by: Andrew Nester --- Makefile | 1 + 1 file changed, 1 insertion(+) diff --git a/Makefile b/Makefile index 81aac2b4b3e..3c55b8cf87e 100644 --- a/Makefile +++ b/Makefile @@ -31,3 +31,4 @@ vendor: @go mod vendor .PHONY: build vendor coverage test lint fmt + From 7e6c36586883d6278d022dfef7925a88015d9dfa Mon Sep 17 00:00:00 2001 From: Pieter Noordhuis Date: Mon, 16 Oct 2023 16:56:46 +0200 Subject: [PATCH 5/5] Fix tests --- bundle/config/mutator/process_target_mode.go | 6 ++++++ bundle/config/mutator/process_target_mode_test.go | 9 +++++++++ bundle/config/resources/registered_model.go | 9 +++++++++ 3 files changed, 24 insertions(+) diff --git a/bundle/config/mutator/process_target_mode.go b/bundle/config/mutator/process_target_mode.go index 2f80fe3b300..c11bd1c5a6e 100644 --- a/bundle/config/mutator/process_target_mode.go +++ b/bundle/config/mutator/process_target_mode.go @@ -87,6 +87,12 @@ func transformDevelopmentMode(b *bundle.Bundle) error { // (model serving doesn't yet support tags) } + for i := range r.RegisteredModels { + prefix = "dev_" + b.Config.Workspace.CurrentUser.ShortName + "_" + r.RegisteredModels[i].Name = prefix + r.RegisteredModels[i].Name + // (registered models in Unity Catalog don't yet support tags) + } + return nil } diff --git a/bundle/config/mutator/process_target_mode_test.go b/bundle/config/mutator/process_target_mode_test.go index eb671aedf7b..a9da0b0f315 100644 --- a/bundle/config/mutator/process_target_mode_test.go +++ b/bundle/config/mutator/process_target_mode_test.go @@ -11,6 +11,7 @@ import ( "github.com/databricks/cli/bundle/config/resources" "github.com/databricks/cli/libs/tags" sdkconfig "github.com/databricks/databricks-sdk-go/config" + "github.com/databricks/databricks-sdk-go/service/catalog" "github.com/databricks/databricks-sdk-go/service/iam" "github.com/databricks/databricks-sdk-go/service/jobs" "github.com/databricks/databricks-sdk-go/service/ml" @@ -59,6 +60,9 @@ func mockBundle(mode config.Mode) *bundle.Bundle { ModelServingEndpoints: map[string]*resources.ModelServingEndpoint{ "servingendpoint1": {CreateServingEndpoint: &serving.CreateServingEndpoint{Name: "servingendpoint1"}}, }, + RegisteredModels: map[string]*resources.RegisteredModel{ + "registeredmodel1": {CreateRegisteredModelRequest: &catalog.CreateRegisteredModelRequest{Name: "registeredmodel1"}}, + }, }, }, // Use AWS implementation for testing. @@ -97,6 +101,9 @@ func TestProcessTargetModeDevelopment(t *testing.T) { // Model serving endpoint 1 assert.Equal(t, "dev_lennart_servingendpoint1", bundle.Config.Resources.ModelServingEndpoints["servingendpoint1"].Name) + + // Registered model 1 + assert.Equal(t, "dev_lennart_registeredmodel1", bundle.Config.Resources.RegisteredModels["registeredmodel1"].Name) } func TestProcessTargetModeDevelopmentTagNormalizationForAws(t *testing.T) { @@ -151,6 +158,7 @@ func TestProcessTargetModeDefault(t *testing.T) { assert.Equal(t, "pipeline1", bundle.Config.Resources.Pipelines["pipeline1"].Name) assert.False(t, bundle.Config.Resources.Pipelines["pipeline1"].PipelineSpec.Development) assert.Equal(t, "servingendpoint1", bundle.Config.Resources.ModelServingEndpoints["servingendpoint1"].Name) + assert.Equal(t, "registeredmodel1", bundle.Config.Resources.RegisteredModels["registeredmodel1"].Name) } func TestProcessTargetModeProduction(t *testing.T) { @@ -187,6 +195,7 @@ func TestProcessTargetModeProduction(t *testing.T) { assert.Equal(t, "pipeline1", bundle.Config.Resources.Pipelines["pipeline1"].Name) assert.False(t, bundle.Config.Resources.Pipelines["pipeline1"].PipelineSpec.Development) assert.Equal(t, "servingendpoint1", bundle.Config.Resources.ModelServingEndpoints["servingendpoint1"].Name) + assert.Equal(t, "registeredmodel1", bundle.Config.Resources.RegisteredModels["registeredmodel1"].Name) } func TestProcessTargetModeProductionOkForPrincipal(t *testing.T) { diff --git a/bundle/config/resources/registered_model.go b/bundle/config/resources/registered_model.go index d37a982e34c..32a451a2ab3 100644 --- a/bundle/config/resources/registered_model.go +++ b/bundle/config/resources/registered_model.go @@ -2,6 +2,7 @@ package resources import ( "github.com/databricks/cli/bundle/config/paths" + "github.com/databricks/databricks-sdk-go/marshal" "github.com/databricks/databricks-sdk-go/service/catalog" ) @@ -23,3 +24,11 @@ type RegisteredModel struct { // to a HCL representation for CRUD *catalog.CreateRegisteredModelRequest } + +func (s *RegisteredModel) UnmarshalJSON(b []byte) error { + return marshal.Unmarshal(b, s) +} + +func (s RegisteredModel) MarshalJSON() ([]byte, error) { + return marshal.Marshal(s) +}