Compare commits

..

3 Commits

Author SHA1 Message Date
shreyas-goenka abc2f3c825
Fix `TestAccBundleInitOnMlopsStacks` (#1924)
## Changes
The ML production team modified mlops-stack to use `mode: development`
for their development target here:
https://github.com/databricks/mlops-stacks/pull/174

This PR makes the integration test assertion agnostic of the prefix to
make it pass again.

## Tests
The test passes now
2024-11-21 10:46:24 +00:00
shreyas-goenka c2e2abcc35
Extend "notebook not found" error to warn about missing extension (#1920)
## Changes
The full workspace path for a notebook does not contain the notebook's
extension. If a user converts that file path to a relative path (like
`/Workspace/bundle_root/bar/nb` -> `./bar/nb`), they can be confused as
to why the new file path does not work.

The changes in this PR nudge them to add the appropriate file extension
(e.g., `./bar/nb.py` or `./bar/nb.ipynb`).

One common way users can end up in this scenario is by using the view
job as YAML functionality in the Databricks UI.

## Tests
Unit test and manually.

```
(.venv) ➜  bundle-playground git:(master) ✗ cli bundle validate 
Error: notebook ./foo not found. Local notebook references are expected
to contain one of the following file extensions: [.py, .r, .scala, .sql, .ipynb]
```
2024-11-21 16:21:21 +05:30
Pieter Noordhuis 14fe03dcb9
Breakout variable lookup into separate files and tests (#1921)
## Changes

While looking into adding variable lookups for notification destinations
([API][API]), I found the codegen approach for different classes of
variable lookups a bit complex. The template had a custom field override
(for service principals), the package had an override for the cluster
lookup, and it didn't produce tests.

The notification destinations API uses a default page size of 20 for
listing. I want to use a larger page size to limit the number of API
calls, so that would imply another customization on the template or a
manual override.

This code being rather mechanical, I used copilot to produce all
instances of the resolvers and their tests (after writing one of them
manually).

[api]: https://docs.databricks.com/api/workspace/notificationdestinations

## Tests

* Unit tests pass
* Manual confirmation that lookups of warehouses still work
2024-11-21 11:28:50 +01:00
30 changed files with 980 additions and 441 deletions

View File

@ -5,8 +5,7 @@
},
"batch": {
".codegen/cmds-workspace.go.tmpl": "cmd/workspace/cmd.go",
".codegen/cmds-account.go.tmpl": "cmd/account/cmd.go",
".codegen/lookup.go.tmpl": "bundle/config/variable/lookup.go"
".codegen/cmds-account.go.tmpl": "cmd/account/cmd.go"
},
"toolchain": {
"required": ["go"],

View File

@ -1,134 +0,0 @@
// Code generated from OpenAPI specs by Databricks SDK Generator. DO NOT EDIT.
package variable
{{ $allowlist :=
list
"alerts"
"clusters"
"cluster-policies"
"clusters"
"dashboards"
"instance-pools"
"jobs"
"metastores"
"pipelines"
"service-principals"
"queries"
"warehouses"
}}
{{ $customField :=
dict
"service-principals" "ApplicationId"
}}
import (
"context"
"fmt"
"github.com/databricks/databricks-sdk-go"
)
type Lookup struct {
{{range .Services -}}
{{- if in $allowlist .KebabName -}}
{{.Singular.PascalName}} string `json:"{{.Singular.SnakeName}},omitempty"`
{{end}}
{{- end}}
}
func LookupFromMap(m map[string]any) *Lookup {
l := &Lookup{}
{{range .Services -}}
{{- if in $allowlist .KebabName -}}
if v, ok := m["{{.Singular.SnakeName}}"]; ok {
l.{{.Singular.PascalName}} = v.(string)
}
{{end -}}
{{- end}}
return l
}
func (l *Lookup) Resolve(ctx context.Context, w *databricks.WorkspaceClient) (string, error) {
if err := l.validate(); err != nil {
return "", err
}
r := allResolvers()
{{range .Services -}}
{{- if in $allowlist .KebabName -}}
if l.{{.Singular.PascalName}} != "" {
return r.{{.Singular.PascalName}}(ctx, w, l.{{.Singular.PascalName}})
}
{{end -}}
{{- end}}
return "", fmt.Errorf("no valid lookup fields provided")
}
func (l *Lookup) String() string {
{{range .Services -}}
{{- if in $allowlist .KebabName -}}
if l.{{.Singular.PascalName}} != "" {
return fmt.Sprintf("{{.Singular.KebabName}}: %s", l.{{.Singular.PascalName}})
}
{{end -}}
{{- end}}
return ""
}
func (l *Lookup) validate() error {
// Validate that only one field is set
count := 0
{{range .Services -}}
{{- if in $allowlist .KebabName -}}
if l.{{.Singular.PascalName}} != "" {
count++
}
{{end -}}
{{- end}}
if count != 1 {
return fmt.Errorf("exactly one lookup field must be provided")
}
if strings.Contains(l.String(), "${var") {
return fmt.Errorf("lookup fields cannot contain variable references")
}
return nil
}
type resolverFunc func(ctx context.Context, w *databricks.WorkspaceClient, name string) (string, error)
type resolvers struct {
{{range .Services -}}
{{- if in $allowlist .KebabName -}}
{{.Singular.PascalName}} resolverFunc
{{end -}}
{{- end}}
}
func allResolvers() *resolvers {
r := &resolvers{}
{{range .Services -}}
{{- if in $allowlist .KebabName -}}
r.{{.Singular.PascalName}} = func(ctx context.Context, w *databricks.WorkspaceClient, name string) (string, error) {
fn, ok := lookupOverrides["{{.Singular.PascalName}}"]
if ok {
return fn(ctx, w, name)
}
entity, err := w.{{.PascalName}}.GetBy{{range .NamedIdMap.NamePath}}{{.PascalName}}{{end}}(ctx, name)
if err != nil {
return "", err
}
return fmt.Sprint(entity.{{ getOrDefault $customField .KebabName ((index .NamedIdMap.IdPath 0).PascalName) }}), nil
}
{{end -}}
{{- end}}
return r
}

1
.gitattributes vendored
View File

@ -1,4 +1,3 @@
bundle/config/variable/lookup.go linguist-generated=true
cmd/account/access-control/access-control.go linguist-generated=true
cmd/account/billable-usage/billable-usage.go linguist-generated=true
cmd/account/budgets/budgets.go linguist-generated=true

View File

@ -126,7 +126,33 @@ func (t *translateContext) rewritePath(
func (t *translateContext) translateNotebookPath(literal, localFullPath, localRelPath, remotePath string) (string, error) {
nb, _, err := notebook.DetectWithFS(t.b.SyncRoot, filepath.ToSlash(localRelPath))
if errors.Is(err, fs.ErrNotExist) {
return "", fmt.Errorf("notebook %s not found", literal)
if filepath.Ext(localFullPath) != notebook.ExtensionNone {
return "", fmt.Errorf("notebook %s not found", literal)
}
extensions := []string{
notebook.ExtensionPython,
notebook.ExtensionR,
notebook.ExtensionScala,
notebook.ExtensionSql,
notebook.ExtensionJupyter,
}
// Check whether a file with a notebook extension already exists. This
// way we can provide a more targeted error message.
for _, ext := range extensions {
literalWithExt := literal + ext
localRelPathWithExt := filepath.ToSlash(localRelPath + ext)
if _, err := fs.Stat(t.b.SyncRoot, localRelPathWithExt); err == nil {
return "", fmt.Errorf(`notebook %s not found. Did you mean %s?
Local notebook references are expected to contain one of the following
file extensions: [%s]`, literal, literalWithExt, strings.Join(extensions, ", "))
}
}
// Return a generic error message if no matching possible file is found.
return "", fmt.Errorf(`notebook %s not found. Local notebook references are expected
to contain one of the following file extensions: [%s]`, literal, strings.Join(extensions, ", "))
}
if err != nil {
return "", fmt.Errorf("unable to determine if %s is a notebook: %w", localFullPath, err)

View File

@ -2,6 +2,7 @@ package mutator_test
import (
"context"
"fmt"
"os"
"path/filepath"
"runtime"
@ -508,6 +509,59 @@ func TestPipelineNotebookDoesNotExistError(t *testing.T) {
assert.EqualError(t, diags.Error(), "notebook ./doesnt_exist.py not found")
}
func TestPipelineNotebookDoesNotExistErrorWithoutExtension(t *testing.T) {
for _, ext := range []string{
".py",
".r",
".scala",
".sql",
".ipynb",
"",
} {
t.Run("case_"+ext, func(t *testing.T) {
dir := t.TempDir()
if ext != "" {
touchEmptyFile(t, filepath.Join(dir, "foo"+ext))
}
b := &bundle.Bundle{
SyncRootPath: dir,
SyncRoot: vfs.MustNew(dir),
Config: config.Root{
Resources: config.Resources{
Pipelines: map[string]*resources.Pipeline{
"pipeline": {
PipelineSpec: &pipelines.PipelineSpec{
Libraries: []pipelines.PipelineLibrary{
{
Notebook: &pipelines.NotebookLibrary{
Path: "./foo",
},
},
},
},
},
},
},
},
}
bundletest.SetLocation(b, ".", []dyn.Location{{File: filepath.Join(dir, "fake.yml")}})
diags := bundle.Apply(context.Background(), b, mutator.TranslatePaths())
if ext == "" {
assert.EqualError(t, diags.Error(), `notebook ./foo not found. Local notebook references are expected
to contain one of the following file extensions: [.py, .r, .scala, .sql, .ipynb]`)
} else {
assert.EqualError(t, diags.Error(), fmt.Sprintf(`notebook ./foo not found. Did you mean ./foo%s?
Local notebook references are expected to contain one of the following
file extensions: [.py, .r, .scala, .sql, .ipynb]`, ext))
}
})
}
}
func TestPipelineFileDoesNotExistError(t *testing.T) {
dir := t.TempDir()

View File

@ -1,11 +1,8 @@
// Code generated from OpenAPI specs by Databricks SDK Generator. DO NOT EDIT.
package variable
import (
"context"
"fmt"
"strings"
"github.com/databricks/databricks-sdk-go"
)
@ -34,323 +31,75 @@ type Lookup struct {
Warehouse string `json:"warehouse,omitempty"`
}
func LookupFromMap(m map[string]any) *Lookup {
l := &Lookup{}
if v, ok := m["alert"]; ok {
l.Alert = v.(string)
type resolver interface {
// Resolve resolves the underlying entity's ID.
Resolve(ctx context.Context, w *databricks.WorkspaceClient) (string, error)
// String returns a human-readable representation of the resolver.
String() string
}
func (l *Lookup) constructResolver() (resolver, error) {
var resolvers []resolver
if l.Alert != "" {
resolvers = append(resolvers, resolveAlert{name: l.Alert})
}
if v, ok := m["cluster_policy"]; ok {
l.ClusterPolicy = v.(string)
if l.ClusterPolicy != "" {
resolvers = append(resolvers, resolveClusterPolicy{name: l.ClusterPolicy})
}
if v, ok := m["cluster"]; ok {
l.Cluster = v.(string)
if l.Cluster != "" {
resolvers = append(resolvers, resolveCluster{name: l.Cluster})
}
if v, ok := m["dashboard"]; ok {
l.Dashboard = v.(string)
if l.Dashboard != "" {
resolvers = append(resolvers, resolveDashboard{name: l.Dashboard})
}
if v, ok := m["instance_pool"]; ok {
l.InstancePool = v.(string)
if l.InstancePool != "" {
resolvers = append(resolvers, resolveInstancePool{name: l.InstancePool})
}
if v, ok := m["job"]; ok {
l.Job = v.(string)
if l.Job != "" {
resolvers = append(resolvers, resolveJob{name: l.Job})
}
if v, ok := m["metastore"]; ok {
l.Metastore = v.(string)
if l.Metastore != "" {
resolvers = append(resolvers, resolveMetastore{name: l.Metastore})
}
if v, ok := m["pipeline"]; ok {
l.Pipeline = v.(string)
if l.Pipeline != "" {
resolvers = append(resolvers, resolvePipeline{name: l.Pipeline})
}
if v, ok := m["query"]; ok {
l.Query = v.(string)
if l.Query != "" {
resolvers = append(resolvers, resolveQuery{name: l.Query})
}
if v, ok := m["service_principal"]; ok {
l.ServicePrincipal = v.(string)
if l.ServicePrincipal != "" {
resolvers = append(resolvers, resolveServicePrincipal{name: l.ServicePrincipal})
}
if v, ok := m["warehouse"]; ok {
l.Warehouse = v.(string)
if l.Warehouse != "" {
resolvers = append(resolvers, resolveWarehouse{name: l.Warehouse})
}
return l
switch len(resolvers) {
case 0:
return nil, fmt.Errorf("no valid lookup fields provided")
case 1:
return resolvers[0], nil
default:
return nil, fmt.Errorf("exactly one lookup field must be provided")
}
}
func (l *Lookup) Resolve(ctx context.Context, w *databricks.WorkspaceClient) (string, error) {
if err := l.validate(); err != nil {
r, err := l.constructResolver()
if err != nil {
return "", err
}
r := allResolvers()
if l.Alert != "" {
return r.Alert(ctx, w, l.Alert)
}
if l.ClusterPolicy != "" {
return r.ClusterPolicy(ctx, w, l.ClusterPolicy)
}
if l.Cluster != "" {
return r.Cluster(ctx, w, l.Cluster)
}
if l.Dashboard != "" {
return r.Dashboard(ctx, w, l.Dashboard)
}
if l.InstancePool != "" {
return r.InstancePool(ctx, w, l.InstancePool)
}
if l.Job != "" {
return r.Job(ctx, w, l.Job)
}
if l.Metastore != "" {
return r.Metastore(ctx, w, l.Metastore)
}
if l.Pipeline != "" {
return r.Pipeline(ctx, w, l.Pipeline)
}
if l.Query != "" {
return r.Query(ctx, w, l.Query)
}
if l.ServicePrincipal != "" {
return r.ServicePrincipal(ctx, w, l.ServicePrincipal)
}
if l.Warehouse != "" {
return r.Warehouse(ctx, w, l.Warehouse)
}
return "", fmt.Errorf("no valid lookup fields provided")
return r.Resolve(ctx, w)
}
func (l *Lookup) String() string {
if l.Alert != "" {
return fmt.Sprintf("alert: %s", l.Alert)
}
if l.ClusterPolicy != "" {
return fmt.Sprintf("cluster-policy: %s", l.ClusterPolicy)
}
if l.Cluster != "" {
return fmt.Sprintf("cluster: %s", l.Cluster)
}
if l.Dashboard != "" {
return fmt.Sprintf("dashboard: %s", l.Dashboard)
}
if l.InstancePool != "" {
return fmt.Sprintf("instance-pool: %s", l.InstancePool)
}
if l.Job != "" {
return fmt.Sprintf("job: %s", l.Job)
}
if l.Metastore != "" {
return fmt.Sprintf("metastore: %s", l.Metastore)
}
if l.Pipeline != "" {
return fmt.Sprintf("pipeline: %s", l.Pipeline)
}
if l.Query != "" {
return fmt.Sprintf("query: %s", l.Query)
}
if l.ServicePrincipal != "" {
return fmt.Sprintf("service-principal: %s", l.ServicePrincipal)
}
if l.Warehouse != "" {
return fmt.Sprintf("warehouse: %s", l.Warehouse)
r, _ := l.constructResolver()
if r == nil {
return ""
}
return ""
}
func (l *Lookup) validate() error {
// Validate that only one field is set
count := 0
if l.Alert != "" {
count++
}
if l.ClusterPolicy != "" {
count++
}
if l.Cluster != "" {
count++
}
if l.Dashboard != "" {
count++
}
if l.InstancePool != "" {
count++
}
if l.Job != "" {
count++
}
if l.Metastore != "" {
count++
}
if l.Pipeline != "" {
count++
}
if l.Query != "" {
count++
}
if l.ServicePrincipal != "" {
count++
}
if l.Warehouse != "" {
count++
}
if count != 1 {
return fmt.Errorf("exactly one lookup field must be provided")
}
if strings.Contains(l.String(), "${var") {
return fmt.Errorf("lookup fields cannot contain variable references")
}
return nil
}
type resolverFunc func(ctx context.Context, w *databricks.WorkspaceClient, name string) (string, error)
type resolvers struct {
Alert resolverFunc
ClusterPolicy resolverFunc
Cluster resolverFunc
Dashboard resolverFunc
InstancePool resolverFunc
Job resolverFunc
Metastore resolverFunc
Pipeline resolverFunc
Query resolverFunc
ServicePrincipal resolverFunc
Warehouse resolverFunc
}
func allResolvers() *resolvers {
r := &resolvers{}
r.Alert = func(ctx context.Context, w *databricks.WorkspaceClient, name string) (string, error) {
fn, ok := lookupOverrides["Alert"]
if ok {
return fn(ctx, w, name)
}
entity, err := w.Alerts.GetByDisplayName(ctx, name)
if err != nil {
return "", err
}
return fmt.Sprint(entity.Id), nil
}
r.ClusterPolicy = func(ctx context.Context, w *databricks.WorkspaceClient, name string) (string, error) {
fn, ok := lookupOverrides["ClusterPolicy"]
if ok {
return fn(ctx, w, name)
}
entity, err := w.ClusterPolicies.GetByName(ctx, name)
if err != nil {
return "", err
}
return fmt.Sprint(entity.PolicyId), nil
}
r.Cluster = func(ctx context.Context, w *databricks.WorkspaceClient, name string) (string, error) {
fn, ok := lookupOverrides["Cluster"]
if ok {
return fn(ctx, w, name)
}
entity, err := w.Clusters.GetByClusterName(ctx, name)
if err != nil {
return "", err
}
return fmt.Sprint(entity.ClusterId), nil
}
r.Dashboard = func(ctx context.Context, w *databricks.WorkspaceClient, name string) (string, error) {
fn, ok := lookupOverrides["Dashboard"]
if ok {
return fn(ctx, w, name)
}
entity, err := w.Dashboards.GetByName(ctx, name)
if err != nil {
return "", err
}
return fmt.Sprint(entity.Id), nil
}
r.InstancePool = func(ctx context.Context, w *databricks.WorkspaceClient, name string) (string, error) {
fn, ok := lookupOverrides["InstancePool"]
if ok {
return fn(ctx, w, name)
}
entity, err := w.InstancePools.GetByInstancePoolName(ctx, name)
if err != nil {
return "", err
}
return fmt.Sprint(entity.InstancePoolId), nil
}
r.Job = func(ctx context.Context, w *databricks.WorkspaceClient, name string) (string, error) {
fn, ok := lookupOverrides["Job"]
if ok {
return fn(ctx, w, name)
}
entity, err := w.Jobs.GetBySettingsName(ctx, name)
if err != nil {
return "", err
}
return fmt.Sprint(entity.JobId), nil
}
r.Metastore = func(ctx context.Context, w *databricks.WorkspaceClient, name string) (string, error) {
fn, ok := lookupOverrides["Metastore"]
if ok {
return fn(ctx, w, name)
}
entity, err := w.Metastores.GetByName(ctx, name)
if err != nil {
return "", err
}
return fmt.Sprint(entity.MetastoreId), nil
}
r.Pipeline = func(ctx context.Context, w *databricks.WorkspaceClient, name string) (string, error) {
fn, ok := lookupOverrides["Pipeline"]
if ok {
return fn(ctx, w, name)
}
entity, err := w.Pipelines.GetByName(ctx, name)
if err != nil {
return "", err
}
return fmt.Sprint(entity.PipelineId), nil
}
r.Query = func(ctx context.Context, w *databricks.WorkspaceClient, name string) (string, error) {
fn, ok := lookupOverrides["Query"]
if ok {
return fn(ctx, w, name)
}
entity, err := w.Queries.GetByDisplayName(ctx, name)
if err != nil {
return "", err
}
return fmt.Sprint(entity.Id), nil
}
r.ServicePrincipal = func(ctx context.Context, w *databricks.WorkspaceClient, name string) (string, error) {
fn, ok := lookupOverrides["ServicePrincipal"]
if ok {
return fn(ctx, w, name)
}
entity, err := w.ServicePrincipals.GetByDisplayName(ctx, name)
if err != nil {
return "", err
}
return fmt.Sprint(entity.ApplicationId), nil
}
r.Warehouse = func(ctx context.Context, w *databricks.WorkspaceClient, name string) (string, error) {
fn, ok := lookupOverrides["Warehouse"]
if ok {
return fn(ctx, w, name)
}
entity, err := w.Warehouses.GetByName(ctx, name)
if err != nil {
return "", err
}
return fmt.Sprint(entity.Id), nil
}
return r
return r.String()
}

View File

@ -0,0 +1,60 @@
package variable
import (
"context"
"reflect"
"testing"
"github.com/stretchr/testify/assert"
)
func TestLookup_Coverage(t *testing.T) {
var lookup Lookup
val := reflect.ValueOf(lookup)
typ := val.Type()
for i := 0; i < val.NumField(); i++ {
field := val.Field(i)
if field.Kind() != reflect.String {
t.Fatalf("Field %s is not a string", typ.Field(i).Name)
}
fieldType := typ.Field(i)
t.Run(fieldType.Name, func(t *testing.T) {
// Use a fresh instance of the struct in each test
var lookup Lookup
// Set the field to a non-empty string
reflect.ValueOf(&lookup).Elem().Field(i).SetString("value")
// Test the [String] function
assert.NotEmpty(t, lookup.String())
})
}
}
func TestLookup_Empty(t *testing.T) {
var lookup Lookup
// Resolve returns an error when no fields are provided
_, err := lookup.Resolve(context.Background(), nil)
assert.ErrorContains(t, err, "no valid lookup fields provided")
// No string representation for an invalid lookup
assert.Empty(t, lookup.String())
}
func TestLookup_Multiple(t *testing.T) {
lookup := Lookup{
Alert: "alert",
Query: "query",
}
// Resolve returns an error when multiple fields are provided
_, err := lookup.Resolve(context.Background(), nil)
assert.ErrorContains(t, err, "exactly one lookup field must be provided")
// No string representation for an invalid lookup
assert.Empty(t, lookup.String())
}

View File

@ -0,0 +1,24 @@
package variable
import (
"context"
"fmt"
"github.com/databricks/databricks-sdk-go"
)
type resolveAlert struct {
name string
}
func (l resolveAlert) Resolve(ctx context.Context, w *databricks.WorkspaceClient) (string, error) {
entity, err := w.Alerts.GetByDisplayName(ctx, l.name)
if err != nil {
return "", err
}
return fmt.Sprint(entity.Id), nil
}
func (l resolveAlert) String() string {
return fmt.Sprintf("alert: %s", l.name)
}

View File

@ -0,0 +1,49 @@
package variable
import (
"context"
"testing"
"github.com/databricks/databricks-sdk-go/apierr"
"github.com/databricks/databricks-sdk-go/experimental/mocks"
"github.com/databricks/databricks-sdk-go/service/sql"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/mock"
"github.com/stretchr/testify/require"
)
func TestResolveAlert_ResolveSuccess(t *testing.T) {
m := mocks.NewMockWorkspaceClient(t)
api := m.GetMockAlertsAPI()
api.EXPECT().
GetByDisplayName(mock.Anything, "alert").
Return(&sql.ListAlertsResponseAlert{
Id: "1234",
}, nil)
ctx := context.Background()
l := resolveAlert{name: "alert"}
result, err := l.Resolve(ctx, m.WorkspaceClient)
require.NoError(t, err)
assert.Equal(t, "1234", result)
}
func TestResolveAlert_ResolveNotFound(t *testing.T) {
m := mocks.NewMockWorkspaceClient(t)
api := m.GetMockAlertsAPI()
api.EXPECT().
GetByDisplayName(mock.Anything, "alert").
Return(nil, &apierr.APIError{StatusCode: 404})
ctx := context.Background()
l := resolveAlert{name: "alert"}
_, err := l.Resolve(ctx, m.WorkspaceClient)
require.ErrorIs(t, err, apierr.ErrNotFound)
}
func TestResolveAlert_String(t *testing.T) {
l := resolveAlert{name: "name"}
assert.Equal(t, "alert: name", l.String())
}

View File

@ -8,13 +8,13 @@ import (
"github.com/databricks/databricks-sdk-go/service/compute"
)
var lookupOverrides = map[string]resolverFunc{
"Cluster": resolveCluster,
type resolveCluster struct {
name string
}
// We added a custom resolver for the cluster to add filtering for the cluster source when we list all clusters.
// Without the filtering listing could take a very long time (5-10 mins) which leads to lookup timeouts.
func resolveCluster(ctx context.Context, w *databricks.WorkspaceClient, name string) (string, error) {
func (l resolveCluster) Resolve(ctx context.Context, w *databricks.WorkspaceClient) (string, error) {
result, err := w.Clusters.ListAll(ctx, compute.ListClustersRequest{
FilterBy: &compute.ListClustersFilterBy{
ClusterSources: []compute.ClusterSource{compute.ClusterSourceApi, compute.ClusterSourceUi},
@ -30,6 +30,8 @@ func resolveCluster(ctx context.Context, w *databricks.WorkspaceClient, name str
key := v.ClusterName
tmp[key] = append(tmp[key], v)
}
name := l.name
alternatives, ok := tmp[name]
if !ok || len(alternatives) == 0 {
return "", fmt.Errorf("cluster named '%s' does not exist", name)
@ -39,3 +41,7 @@ func resolveCluster(ctx context.Context, w *databricks.WorkspaceClient, name str
}
return alternatives[0].ClusterId, nil
}
func (l resolveCluster) String() string {
return fmt.Sprintf("cluster: %s", l.name)
}

View File

@ -0,0 +1,24 @@
package variable
import (
"context"
"fmt"
"github.com/databricks/databricks-sdk-go"
)
type resolveClusterPolicy struct {
name string
}
func (l resolveClusterPolicy) Resolve(ctx context.Context, w *databricks.WorkspaceClient) (string, error) {
entity, err := w.ClusterPolicies.GetByName(ctx, l.name)
if err != nil {
return "", err
}
return fmt.Sprint(entity.PolicyId), nil
}
func (l resolveClusterPolicy) String() string {
return fmt.Sprintf("cluster-policy: %s", l.name)
}

View File

@ -0,0 +1,49 @@
package variable
import (
"context"
"testing"
"github.com/databricks/databricks-sdk-go/apierr"
"github.com/databricks/databricks-sdk-go/experimental/mocks"
"github.com/databricks/databricks-sdk-go/service/compute"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/mock"
"github.com/stretchr/testify/require"
)
func TestResolveClusterPolicy_ResolveSuccess(t *testing.T) {
m := mocks.NewMockWorkspaceClient(t)
api := m.GetMockClusterPoliciesAPI()
api.EXPECT().
GetByName(mock.Anything, "policy").
Return(&compute.Policy{
PolicyId: "1234",
}, nil)
ctx := context.Background()
l := resolveClusterPolicy{name: "policy"}
result, err := l.Resolve(ctx, m.WorkspaceClient)
require.NoError(t, err)
assert.Equal(t, "1234", result)
}
func TestResolveClusterPolicy_ResolveNotFound(t *testing.T) {
m := mocks.NewMockWorkspaceClient(t)
api := m.GetMockClusterPoliciesAPI()
api.EXPECT().
GetByName(mock.Anything, "policy").
Return(nil, &apierr.APIError{StatusCode: 404})
ctx := context.Background()
l := resolveClusterPolicy{name: "policy"}
_, err := l.Resolve(ctx, m.WorkspaceClient)
require.ErrorIs(t, err, apierr.ErrNotFound)
}
func TestResolveClusterPolicy_String(t *testing.T) {
l := resolveClusterPolicy{name: "name"}
assert.Equal(t, "cluster-policy: name", l.String())
}

View File

@ -0,0 +1,50 @@
package variable
import (
"context"
"testing"
"github.com/databricks/databricks-sdk-go/experimental/mocks"
"github.com/databricks/databricks-sdk-go/service/compute"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/mock"
"github.com/stretchr/testify/require"
)
func TestResolveCluster_ResolveSuccess(t *testing.T) {
m := mocks.NewMockWorkspaceClient(t)
api := m.GetMockClustersAPI()
api.EXPECT().
ListAll(mock.Anything, mock.Anything).
Return([]compute.ClusterDetails{
{ClusterId: "1234", ClusterName: "cluster1"},
{ClusterId: "2345", ClusterName: "cluster2"},
}, nil)
ctx := context.Background()
l := resolveCluster{name: "cluster2"}
result, err := l.Resolve(ctx, m.WorkspaceClient)
require.NoError(t, err)
assert.Equal(t, "2345", result)
}
func TestResolveCluster_ResolveNotFound(t *testing.T) {
m := mocks.NewMockWorkspaceClient(t)
api := m.GetMockClustersAPI()
api.EXPECT().
ListAll(mock.Anything, mock.Anything).
Return([]compute.ClusterDetails{}, nil)
ctx := context.Background()
l := resolveCluster{name: "cluster"}
_, err := l.Resolve(ctx, m.WorkspaceClient)
require.Error(t, err)
assert.Contains(t, err.Error(), "cluster named 'cluster' does not exist")
}
func TestResolveCluster_String(t *testing.T) {
l := resolveCluster{name: "name"}
assert.Equal(t, "cluster: name", l.String())
}

View File

@ -0,0 +1,24 @@
package variable
import (
"context"
"fmt"
"github.com/databricks/databricks-sdk-go"
)
type resolveDashboard struct {
name string
}
func (l resolveDashboard) Resolve(ctx context.Context, w *databricks.WorkspaceClient) (string, error) {
entity, err := w.Dashboards.GetByName(ctx, l.name)
if err != nil {
return "", err
}
return fmt.Sprint(entity.Id), nil
}
func (l resolveDashboard) String() string {
return fmt.Sprintf("dashboard: %s", l.name)
}

View File

@ -0,0 +1,49 @@
package variable
import (
"context"
"testing"
"github.com/databricks/databricks-sdk-go/apierr"
"github.com/databricks/databricks-sdk-go/experimental/mocks"
"github.com/databricks/databricks-sdk-go/service/sql"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/mock"
"github.com/stretchr/testify/require"
)
func TestResolveDashboard_ResolveSuccess(t *testing.T) {
m := mocks.NewMockWorkspaceClient(t)
api := m.GetMockDashboardsAPI()
api.EXPECT().
GetByName(mock.Anything, "dashboard").
Return(&sql.Dashboard{
Id: "1234",
}, nil)
ctx := context.Background()
l := resolveDashboard{name: "dashboard"}
result, err := l.Resolve(ctx, m.WorkspaceClient)
require.NoError(t, err)
assert.Equal(t, "1234", result)
}
func TestResolveDashboard_ResolveNotFound(t *testing.T) {
m := mocks.NewMockWorkspaceClient(t)
api := m.GetMockDashboardsAPI()
api.EXPECT().
GetByName(mock.Anything, "dashboard").
Return(nil, &apierr.APIError{StatusCode: 404})
ctx := context.Background()
l := resolveDashboard{name: "dashboard"}
_, err := l.Resolve(ctx, m.WorkspaceClient)
require.ErrorIs(t, err, apierr.ErrNotFound)
}
func TestResolveDashboard_String(t *testing.T) {
l := resolveDashboard{name: "name"}
assert.Equal(t, "dashboard: name", l.String())
}

View File

@ -0,0 +1,24 @@
package variable
import (
"context"
"fmt"
"github.com/databricks/databricks-sdk-go"
)
type resolveInstancePool struct {
name string
}
func (l resolveInstancePool) Resolve(ctx context.Context, w *databricks.WorkspaceClient) (string, error) {
entity, err := w.InstancePools.GetByInstancePoolName(ctx, l.name)
if err != nil {
return "", err
}
return fmt.Sprint(entity.InstancePoolId), nil
}
func (l resolveInstancePool) String() string {
return fmt.Sprintf("instance-pool: %s", l.name)
}

View File

@ -0,0 +1,49 @@
package variable
import (
"context"
"testing"
"github.com/databricks/databricks-sdk-go/apierr"
"github.com/databricks/databricks-sdk-go/experimental/mocks"
"github.com/databricks/databricks-sdk-go/service/compute"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/mock"
"github.com/stretchr/testify/require"
)
func TestResolveInstancePool_ResolveSuccess(t *testing.T) {
m := mocks.NewMockWorkspaceClient(t)
api := m.GetMockInstancePoolsAPI()
api.EXPECT().
GetByInstancePoolName(mock.Anything, "instance_pool").
Return(&compute.InstancePoolAndStats{
InstancePoolId: "5678",
}, nil)
ctx := context.Background()
l := resolveInstancePool{name: "instance_pool"}
result, err := l.Resolve(ctx, m.WorkspaceClient)
require.NoError(t, err)
assert.Equal(t, "5678", result)
}
func TestResolveInstancePool_ResolveNotFound(t *testing.T) {
m := mocks.NewMockWorkspaceClient(t)
api := m.GetMockInstancePoolsAPI()
api.EXPECT().
GetByInstancePoolName(mock.Anything, "instance_pool").
Return(nil, &apierr.APIError{StatusCode: 404})
ctx := context.Background()
l := resolveInstancePool{name: "instance_pool"}
_, err := l.Resolve(ctx, m.WorkspaceClient)
require.ErrorIs(t, err, apierr.ErrNotFound)
}
func TestResolveInstancePool_String(t *testing.T) {
l := resolveInstancePool{name: "name"}
assert.Equal(t, "instance-pool: name", l.String())
}

View File

@ -0,0 +1,24 @@
package variable
import (
"context"
"fmt"
"github.com/databricks/databricks-sdk-go"
)
type resolveJob struct {
name string
}
func (l resolveJob) Resolve(ctx context.Context, w *databricks.WorkspaceClient) (string, error) {
entity, err := w.Jobs.GetBySettingsName(ctx, l.name)
if err != nil {
return "", err
}
return fmt.Sprint(entity.JobId), nil
}
func (l resolveJob) String() string {
return fmt.Sprintf("job: %s", l.name)
}

View File

@ -0,0 +1,49 @@
package variable
import (
"context"
"testing"
"github.com/databricks/databricks-sdk-go/apierr"
"github.com/databricks/databricks-sdk-go/experimental/mocks"
"github.com/databricks/databricks-sdk-go/service/jobs"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/mock"
"github.com/stretchr/testify/require"
)
func TestResolveJob_ResolveSuccess(t *testing.T) {
m := mocks.NewMockWorkspaceClient(t)
api := m.GetMockJobsAPI()
api.EXPECT().
GetBySettingsName(mock.Anything, "job").
Return(&jobs.BaseJob{
JobId: 5678,
}, nil)
ctx := context.Background()
l := resolveJob{name: "job"}
result, err := l.Resolve(ctx, m.WorkspaceClient)
require.NoError(t, err)
assert.Equal(t, "5678", result)
}
func TestResolveJob_ResolveNotFound(t *testing.T) {
m := mocks.NewMockWorkspaceClient(t)
api := m.GetMockJobsAPI()
api.EXPECT().
GetBySettingsName(mock.Anything, "job").
Return(nil, &apierr.APIError{StatusCode: 404})
ctx := context.Background()
l := resolveJob{name: "job"}
_, err := l.Resolve(ctx, m.WorkspaceClient)
require.ErrorIs(t, err, apierr.ErrNotFound)
}
func TestResolveJob_String(t *testing.T) {
l := resolveJob{name: "name"}
assert.Equal(t, "job: name", l.String())
}

View File

@ -0,0 +1,24 @@
package variable
import (
"context"
"fmt"
"github.com/databricks/databricks-sdk-go"
)
type resolveMetastore struct {
name string
}
func (l resolveMetastore) Resolve(ctx context.Context, w *databricks.WorkspaceClient) (string, error) {
entity, err := w.Metastores.GetByName(ctx, l.name)
if err != nil {
return "", err
}
return fmt.Sprint(entity.MetastoreId), nil
}
func (l resolveMetastore) String() string {
return fmt.Sprintf("metastore: %s", l.name)
}

View File

@ -0,0 +1,49 @@
package variable
import (
"context"
"testing"
"github.com/databricks/databricks-sdk-go/apierr"
"github.com/databricks/databricks-sdk-go/experimental/mocks"
"github.com/databricks/databricks-sdk-go/service/catalog"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/mock"
"github.com/stretchr/testify/require"
)
func TestResolveMetastore_ResolveSuccess(t *testing.T) {
m := mocks.NewMockWorkspaceClient(t)
api := m.GetMockMetastoresAPI()
api.EXPECT().
GetByName(mock.Anything, "metastore").
Return(&catalog.MetastoreInfo{
MetastoreId: "abcd",
}, nil)
ctx := context.Background()
l := resolveMetastore{name: "metastore"}
result, err := l.Resolve(ctx, m.WorkspaceClient)
require.NoError(t, err)
assert.Equal(t, "abcd", result)
}
func TestResolveMetastore_ResolveNotFound(t *testing.T) {
m := mocks.NewMockWorkspaceClient(t)
api := m.GetMockMetastoresAPI()
api.EXPECT().
GetByName(mock.Anything, "metastore").
Return(nil, &apierr.APIError{StatusCode: 404})
ctx := context.Background()
l := resolveMetastore{name: "metastore"}
_, err := l.Resolve(ctx, m.WorkspaceClient)
require.ErrorIs(t, err, apierr.ErrNotFound)
}
func TestResolveMetastore_String(t *testing.T) {
l := resolveMetastore{name: "name"}
assert.Equal(t, "metastore: name", l.String())
}

View File

@ -0,0 +1,24 @@
package variable
import (
"context"
"fmt"
"github.com/databricks/databricks-sdk-go"
)
type resolvePipeline struct {
name string
}
func (l resolvePipeline) Resolve(ctx context.Context, w *databricks.WorkspaceClient) (string, error) {
entity, err := w.Pipelines.GetByName(ctx, l.name)
if err != nil {
return "", err
}
return fmt.Sprint(entity.PipelineId), nil
}
func (l resolvePipeline) String() string {
return fmt.Sprintf("pipeline: %s", l.name)
}

View File

@ -0,0 +1,49 @@
package variable
import (
"context"
"testing"
"github.com/databricks/databricks-sdk-go/apierr"
"github.com/databricks/databricks-sdk-go/experimental/mocks"
"github.com/databricks/databricks-sdk-go/service/pipelines"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/mock"
"github.com/stretchr/testify/require"
)
func TestResolvePipeline_ResolveSuccess(t *testing.T) {
m := mocks.NewMockWorkspaceClient(t)
api := m.GetMockPipelinesAPI()
api.EXPECT().
GetByName(mock.Anything, "pipeline").
Return(&pipelines.PipelineStateInfo{
PipelineId: "abcd",
}, nil)
ctx := context.Background()
l := resolvePipeline{name: "pipeline"}
result, err := l.Resolve(ctx, m.WorkspaceClient)
require.NoError(t, err)
assert.Equal(t, "abcd", result)
}
func TestResolvePipeline_ResolveNotFound(t *testing.T) {
m := mocks.NewMockWorkspaceClient(t)
api := m.GetMockPipelinesAPI()
api.EXPECT().
GetByName(mock.Anything, "pipeline").
Return(nil, &apierr.APIError{StatusCode: 404})
ctx := context.Background()
l := resolvePipeline{name: "pipeline"}
_, err := l.Resolve(ctx, m.WorkspaceClient)
require.ErrorIs(t, err, apierr.ErrNotFound)
}
func TestResolvePipeline_String(t *testing.T) {
l := resolvePipeline{name: "name"}
assert.Equal(t, "pipeline: name", l.String())
}

View File

@ -0,0 +1,24 @@
package variable
import (
"context"
"fmt"
"github.com/databricks/databricks-sdk-go"
)
type resolveQuery struct {
name string
}
func (l resolveQuery) Resolve(ctx context.Context, w *databricks.WorkspaceClient) (string, error) {
entity, err := w.Queries.GetByDisplayName(ctx, l.name)
if err != nil {
return "", err
}
return fmt.Sprint(entity.Id), nil
}
func (l resolveQuery) String() string {
return fmt.Sprintf("query: %s", l.name)
}

View File

@ -0,0 +1,49 @@
package variable
import (
"context"
"testing"
"github.com/databricks/databricks-sdk-go/apierr"
"github.com/databricks/databricks-sdk-go/experimental/mocks"
"github.com/databricks/databricks-sdk-go/service/sql"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/mock"
"github.com/stretchr/testify/require"
)
func TestResolveQuery_ResolveSuccess(t *testing.T) {
m := mocks.NewMockWorkspaceClient(t)
api := m.GetMockQueriesAPI()
api.EXPECT().
GetByDisplayName(mock.Anything, "query").
Return(&sql.ListQueryObjectsResponseQuery{
Id: "1234",
}, nil)
ctx := context.Background()
l := resolveQuery{name: "query"}
result, err := l.Resolve(ctx, m.WorkspaceClient)
require.NoError(t, err)
assert.Equal(t, "1234", result)
}
func TestResolveQuery_ResolveNotFound(t *testing.T) {
m := mocks.NewMockWorkspaceClient(t)
api := m.GetMockQueriesAPI()
api.EXPECT().
GetByDisplayName(mock.Anything, "query").
Return(nil, &apierr.APIError{StatusCode: 404})
ctx := context.Background()
l := resolveQuery{name: "query"}
_, err := l.Resolve(ctx, m.WorkspaceClient)
require.ErrorIs(t, err, apierr.ErrNotFound)
}
func TestResolveQuery_String(t *testing.T) {
l := resolveQuery{name: "name"}
assert.Equal(t, "query: name", l.String())
}

View File

@ -0,0 +1,24 @@
package variable
import (
"context"
"fmt"
"github.com/databricks/databricks-sdk-go"
)
type resolveServicePrincipal struct {
name string
}
func (l resolveServicePrincipal) Resolve(ctx context.Context, w *databricks.WorkspaceClient) (string, error) {
entity, err := w.ServicePrincipals.GetByDisplayName(ctx, l.name)
if err != nil {
return "", err
}
return fmt.Sprint(entity.ApplicationId), nil
}
func (l resolveServicePrincipal) String() string {
return fmt.Sprintf("service-principal: %s", l.name)
}

View File

@ -0,0 +1,49 @@
package variable
import (
"context"
"testing"
"github.com/databricks/databricks-sdk-go/apierr"
"github.com/databricks/databricks-sdk-go/experimental/mocks"
"github.com/databricks/databricks-sdk-go/service/iam"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/mock"
"github.com/stretchr/testify/require"
)
func TestResolveServicePrincipal_ResolveSuccess(t *testing.T) {
m := mocks.NewMockWorkspaceClient(t)
api := m.GetMockServicePrincipalsAPI()
api.EXPECT().
GetByDisplayName(mock.Anything, "service-principal").
Return(&iam.ServicePrincipal{
ApplicationId: "5678",
}, nil)
ctx := context.Background()
l := resolveServicePrincipal{name: "service-principal"}
result, err := l.Resolve(ctx, m.WorkspaceClient)
require.NoError(t, err)
assert.Equal(t, "5678", result)
}
func TestResolveServicePrincipal_ResolveNotFound(t *testing.T) {
m := mocks.NewMockWorkspaceClient(t)
api := m.GetMockServicePrincipalsAPI()
api.EXPECT().
GetByDisplayName(mock.Anything, "service-principal").
Return(nil, &apierr.APIError{StatusCode: 404})
ctx := context.Background()
l := resolveServicePrincipal{name: "service-principal"}
_, err := l.Resolve(ctx, m.WorkspaceClient)
require.ErrorIs(t, err, apierr.ErrNotFound)
}
func TestResolveServicePrincipal_String(t *testing.T) {
l := resolveServicePrincipal{name: "name"}
assert.Equal(t, "service-principal: name", l.String())
}

View File

@ -0,0 +1,24 @@
package variable
import (
"context"
"fmt"
"github.com/databricks/databricks-sdk-go"
)
type resolveWarehouse struct {
name string
}
func (l resolveWarehouse) Resolve(ctx context.Context, w *databricks.WorkspaceClient) (string, error) {
entity, err := w.Warehouses.GetByName(ctx, l.name)
if err != nil {
return "", err
}
return fmt.Sprint(entity.Id), nil
}
func (l resolveWarehouse) String() string {
return fmt.Sprintf("warehouse: %s", l.name)
}

View File

@ -0,0 +1,49 @@
package variable
import (
"context"
"testing"
"github.com/databricks/databricks-sdk-go/apierr"
"github.com/databricks/databricks-sdk-go/experimental/mocks"
"github.com/databricks/databricks-sdk-go/service/sql"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/mock"
"github.com/stretchr/testify/require"
)
func TestResolveWarehouse_ResolveSuccess(t *testing.T) {
m := mocks.NewMockWorkspaceClient(t)
api := m.GetMockWarehousesAPI()
api.EXPECT().
GetByName(mock.Anything, "warehouse").
Return(&sql.EndpointInfo{
Id: "abcd",
}, nil)
ctx := context.Background()
l := resolveWarehouse{name: "warehouse"}
result, err := l.Resolve(ctx, m.WorkspaceClient)
require.NoError(t, err)
assert.Equal(t, "abcd", result)
}
func TestResolveWarehouse_ResolveNotFound(t *testing.T) {
m := mocks.NewMockWorkspaceClient(t)
api := m.GetMockWarehousesAPI()
api.EXPECT().
GetByName(mock.Anything, "warehouse").
Return(nil, &apierr.APIError{StatusCode: 404})
ctx := context.Background()
l := resolveWarehouse{name: "warehouse"}
_, err := l.Resolve(ctx, m.WorkspaceClient)
require.ErrorIs(t, err, apierr.ErrNotFound)
}
func TestResolveWarehouse_String(t *testing.T) {
l := resolveWarehouse{name: "name"}
assert.Equal(t, "warehouse: name", l.String())
}

View File

@ -97,7 +97,7 @@ func TestAccBundleInitOnMlopsStacks(t *testing.T) {
require.NoError(t, err)
job, err := w.Jobs.GetByJobId(context.Background(), batchJobId)
assert.NoError(t, err)
assert.Equal(t, fmt.Sprintf("dev-%s-batch-inference-job", projectName), job.Settings.Name)
assert.Contains(t, job.Settings.Name, fmt.Sprintf("dev-%s-batch-inference-job", projectName))
}
func TestAccBundleInitHelpers(t *testing.T) {