Fix code after latest changes in the framework

databricks · Oct 12, 2024 · 3eb939a · 3eb939a
1 parent e19c8c8
commit 3eb939a
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 4 deletions.
diff --git a/docs/data-sources/registered_model.md b/docs/data-sources/registered_model.md
@@ -21,6 +21,7 @@ The following arguments are supported:
 
 * `full_name` - (Required, String) The fully-qualified name of the registered model (`catalog_name.schema_name.name`).
 * `include_aliases` - (Optional, Boolean) flag to specify if list of aliases should be included into output.
+* `include_browse` - (Optional, Boolean) flag to specify if include registered models in the response for which the principal can only access selective metadata for.
 
 ## Attribute Reference
 

diff --git a/internal/providers/pluginfw/resources/registered_model/data_registered_model.go b/internal/providers/pluginfw/resources/registered_model/data_registered_model.go
@@ -27,9 +27,10 @@ type RegisteredModelDataSource struct {
 }
 
 type RegisteredModelData struct {
-	FullName       types.String                    `tfsdk:"full_name"`
-	IncludeAliases types.Bool                      `tfsdk:"include_aliases" tf:"optional"`
-	ModelInfo      *catalog_tf.RegisteredModelInfo `tfsdk:"model_info" tf:"optional,computed"`
+	FullName       types.String                     `tfsdk:"full_name"`
+	IncludeAliases types.Bool                       `tfsdk:"include_aliases" tf:"optional"`
+	IncludeBrowse  types.Bool                       `tfsdk:"include_browse" tf:"optional"`
+	ModelInfo      []catalog_tf.RegisteredModelInfo `tfsdk:"model_info" tf:"optional,computed"`
 }
 
 func (d *RegisteredModelDataSource) Metadata(ctx context.Context, req datasource.MetadataRequest, resp *datasource.MetadataResponse) {
@@ -67,6 +68,7 @@ func (d *RegisteredModelDataSource) Read(ctx context.Context, req datasource.Rea
 	modelInfoSdk, err := w.RegisteredModels.Get(ctx, catalog.GetRegisteredModelRequest{
 		FullName:       modelFullName,
 		IncludeAliases: registeredModel.IncludeAliases.ValueBool(),
+		IncludeBrowse:  registeredModel.IncludeBrowse.ValueBool(),
 	})
 	if err != nil {
 		if apierr.IsMissing(err) {
@@ -83,6 +85,6 @@ func (d *RegisteredModelDataSource) Read(ctx context.Context, req datasource.Rea
 	if modelInfo.Aliases == nil {
 		modelInfo.Aliases = []catalog_tf.RegisteredModelAlias{}
 	}
-	registeredModel.ModelInfo = &modelInfo
+	registeredModel.ModelInfo = append(registeredModel.ModelInfo, modelInfo)
 	resp.Diagnostics.Append(resp.State.Set(ctx, registeredModel)...)
 }