From 1da7d9342ba758c865f6506b48ff9713ee0a2f78 Mon Sep 17 00:00:00 2001 From: Cedric <23346008+840@users.noreply.github.com> Date: Wed, 2 Oct 2024 15:29:53 +0200 Subject: [PATCH] [Feature] Add `databricks_mlflow_models` data source (#3874) ## Changes Add `databricks_mlflow_models` data source ## Tests Add integration and acceptance test - [x] `make test` run locally - [x] relevant change in `docs/` folder - [x] covered with integration tests in `internal/acceptance` - [x] relevant acceptance tests are passing - [x] using Go SDK Resolves #3791 --------- Co-authored-by: Alex Ott --- docs/data-sources/mlflow_models.md | 42 ++++++++++++++++++ .../acceptance/data_mlflow_models_test.go | 44 +++++++++++++++++++ internal/providers/sdkv2/sdkv2.go | 1 + mlflow/data_mlflow_models.go | 26 +++++++++++ mlflow/data_mlflow_models_test.go | 32 ++++++++++++++ 5 files changed, 145 insertions(+) create mode 100644 docs/data-sources/mlflow_models.md create mode 100644 internal/acceptance/data_mlflow_models_test.go create mode 100644 mlflow/data_mlflow_models.go create mode 100644 mlflow/data_mlflow_models_test.go diff --git a/docs/data-sources/mlflow_models.md b/docs/data-sources/mlflow_models.md new file mode 100644 index 0000000000..b8b67c9096 --- /dev/null +++ b/docs/data-sources/mlflow_models.md @@ -0,0 +1,42 @@ +--- +subcategory: "MLflow" +--- +# databricks_mlflow_models Data Source + +-> **Note** This data source could be only used with workspace-level provider! + +Retrieves a list of [databricks_mlflow_model](../resources/mlflow_model.md) objects, that were created by Terraform or manually, so that special handling could be applied. + +## Example Usage + +```hcl +data "databricks_mlflow_models" "this" {} + +output "model" { + value = data.databricks_mlflow_models.this +} +``` + +```hcl +data "databricks_mlflow_models" "this" {} + +check "model_list_not_empty" { + assert { + condition = length(data.databricks_mlflow_models.this.names) != 0 + error_message = "Model list is empty." + } +} + +check "model_list_contains_model" { + assert { + condition = contains(data.databricks_mlflow_models.this.names, "model_1") + error_message = "model_1 is missing in model list." + } +} +``` + +## Attribute Reference + +This data source exports the following attributes: + +* `names` - List of names of [databricks_mlflow_model](./mlflow_model.md) \ No newline at end of file diff --git a/internal/acceptance/data_mlflow_models_test.go b/internal/acceptance/data_mlflow_models_test.go new file mode 100644 index 0000000000..7cd13a2a02 --- /dev/null +++ b/internal/acceptance/data_mlflow_models_test.go @@ -0,0 +1,44 @@ +package acceptance + +import ( + "fmt" + "testing" + + "github.com/hashicorp/terraform-plugin-testing/terraform" +) + +func TestAccDataMlflowModels(t *testing.T) { + WorkspaceLevel(t, + Step{ + Template: ` + resource "databricks_mlflow_model" "this" { + name = "model-{var.RANDOM}" + + description = "My MLflow model description" + + tags { + key = "key1" + value = "value1" + } + tags { + key = "key2" + value = "value2" + } + } + + data "databricks_mlflow_models" "this" { + depends_on = [databricks_mlflow_model.this] + }`, + Check: func(s *terraform.State) error { + r, ok := s.RootModule().Resources["data.databricks_mlflow_models.this"] + if !ok { + return fmt.Errorf("data not found in state") + } + names := r.Primary.Attributes["names.#"] + if names == "" { + return fmt.Errorf("names are empty: %v", r.Primary.Attributes) + } + return nil + }, + }) +} diff --git a/internal/providers/sdkv2/sdkv2.go b/internal/providers/sdkv2/sdkv2.go index 6d1b712cb0..65ab1f4973 100644 --- a/internal/providers/sdkv2/sdkv2.go +++ b/internal/providers/sdkv2/sdkv2.go @@ -100,6 +100,7 @@ func DatabricksProvider() *schema.Provider { "databricks_metastores": catalog.DataSourceMetastores().ToResource(), "databricks_mlflow_experiment": mlflow.DataSourceExperiment().ToResource(), "databricks_mlflow_model": mlflow.DataSourceModel().ToResource(), + "databricks_mlflow_models": mlflow.DataSourceModels().ToResource(), "databricks_mws_credentials": mws.DataSourceMwsCredentials().ToResource(), "databricks_mws_workspaces": mws.DataSourceMwsWorkspaces().ToResource(), "databricks_node_type": clusters.DataSourceNodeType().ToResource(), diff --git a/mlflow/data_mlflow_models.go b/mlflow/data_mlflow_models.go new file mode 100644 index 0000000000..96ce6bed05 --- /dev/null +++ b/mlflow/data_mlflow_models.go @@ -0,0 +1,26 @@ +package mlflow + +import ( + "context" + "github.com/databricks/databricks-sdk-go/service/ml" + + "github.com/databricks/databricks-sdk-go" + "github.com/databricks/terraform-provider-databricks/common" +) + +type modelsData struct { + Names []string `json:"names,omitempty" tf:"computed"` +} + +func DataSourceModels() common.Resource { + return common.WorkspaceData(func(ctx context.Context, data *modelsData, w *databricks.WorkspaceClient) error { + list, err := w.ModelRegistry.ListModelsAll(ctx, ml.ListModelsRequest{}) + if err != nil { + return err + } + for _, m := range list { + data.Names = append(data.Names, m.Name) + } + return nil + }) +} diff --git a/mlflow/data_mlflow_models_test.go b/mlflow/data_mlflow_models_test.go new file mode 100644 index 0000000000..b115ce9e23 --- /dev/null +++ b/mlflow/data_mlflow_models_test.go @@ -0,0 +1,32 @@ +package mlflow + +import ( + "github.com/databricks/databricks-sdk-go/experimental/mocks" + "github.com/stretchr/testify/mock" + "testing" + + "github.com/databricks/databricks-sdk-go/service/ml" + "github.com/databricks/terraform-provider-databricks/qa" +) + +func TestDataSourceModels(t *testing.T) { + qa.ResourceFixture{ + MockWorkspaceClientFunc: func(w *mocks.MockWorkspaceClient) { + api := w.GetMockModelRegistryAPI() + api.EXPECT().ListModelsAll(mock.Anything, ml.ListModelsRequest{}).Return([]ml.Model{ + { + Name: "model-01", + }, + { + Name: "model-02", + }, + }, nil) + }, + Read: true, + NonWritable: true, + Resource: DataSourceModels(), + ID: ".", + }.ApplyAndExpectData(t, map[string]interface{}{ + "names": []interface{}{"model-01", "model-02"}, + }) +}