Skip to content

Commit

Permalink
[Feature] Add databricks_mlflow_models data source (#3874)
Browse files Browse the repository at this point in the history
## 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 <alexott@gmail.com>
  • Loading branch information
840 and alexott authored Oct 2, 2024
1 parent f757db0 commit 1da7d93
Show file tree
Hide file tree
Showing 5 changed files with 145 additions and 0 deletions.
42 changes: 42 additions & 0 deletions docs/data-sources/mlflow_models.md
Original file line number Diff line number Diff line change
@@ -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)
44 changes: 44 additions & 0 deletions internal/acceptance/data_mlflow_models_test.go
Original file line number Diff line number Diff line change
@@ -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
},
})
}
1 change: 1 addition & 0 deletions internal/providers/sdkv2/sdkv2.go
Original file line number Diff line number Diff line change
Expand Up @@ -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(),
Expand Down
26 changes: 26 additions & 0 deletions mlflow/data_mlflow_models.go
Original file line number Diff line number Diff line change
@@ -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
})
}
32 changes: 32 additions & 0 deletions mlflow/data_mlflow_models_test.go
Original file line number Diff line number Diff line change
@@ -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"},
})
}

0 comments on commit 1da7d93

Please sign in to comment.