From 41539f97f3d2d2a0b045bf1b6f8bfe8b8b2f9204 Mon Sep 17 00:00:00 2001 From: Pieter Noordhuis Date: Wed, 20 Nov 2024 23:57:33 +0100 Subject: [PATCH] Support lookup by name of notification destinations --- bundle/config/variable/lookup.go | 5 ++ .../resolve_notification_destination.go | 46 +++++++++++ .../resolve_notification_destination_test.go | 82 +++++++++++++++++++ 3 files changed, 133 insertions(+) create mode 100644 bundle/config/variable/resolve_notification_destination.go create mode 100644 bundle/config/variable/resolve_notification_destination_test.go diff --git a/bundle/config/variable/lookup.go b/bundle/config/variable/lookup.go index f8cb6719..37e380f1 100755 --- a/bundle/config/variable/lookup.go +++ b/bundle/config/variable/lookup.go @@ -22,6 +22,8 @@ type Lookup struct { Metastore string `json:"metastore,omitempty"` + NotificationDestination string `json:"notification_destination,omitempty"` + Pipeline string `json:"pipeline,omitempty"` Query string `json:"query,omitempty"` @@ -63,6 +65,9 @@ func (l *Lookup) constructResolver() (resolver, error) { if l.Metastore != "" { resolvers = append(resolvers, resolveMetastore{name: l.Metastore}) } + if l.NotificationDestination != "" { + resolvers = append(resolvers, resolveNotificationDestination{name: l.NotificationDestination}) + } if l.Pipeline != "" { resolvers = append(resolvers, resolvePipeline{name: l.Pipeline}) } diff --git a/bundle/config/variable/resolve_notification_destination.go b/bundle/config/variable/resolve_notification_destination.go new file mode 100644 index 00000000..4c4cd892 --- /dev/null +++ b/bundle/config/variable/resolve_notification_destination.go @@ -0,0 +1,46 @@ +package variable + +import ( + "context" + "fmt" + + "github.com/databricks/databricks-sdk-go" + "github.com/databricks/databricks-sdk-go/service/settings" +) + +type resolveNotificationDestination struct { + name string +} + +func (l resolveNotificationDestination) Resolve(ctx context.Context, w *databricks.WorkspaceClient) (string, error) { + result, err := w.NotificationDestinations.ListAll(ctx, settings.ListNotificationDestinationsRequest{ + // The default page size for this API is 20. + // We use a higher value to make fewer API calls. + PageSize: 200, + }) + if err != nil { + return "", err + } + + // Collect all notification destinations with the given name. + var entities []settings.ListNotificationDestinationsResult + for _, entity := range result { + if entity.DisplayName == l.name { + entities = append(entities, entity) + } + } + + // Return the ID of the first matching notification destination. + switch len(entities) { + case 0: + return "", fmt.Errorf("notification destination named %q does not exist", l.name) + case 1: + return entities[0].Id, nil + default: + return "", fmt.Errorf("there are %d instances of clusters named %q", len(entities), l.name) + } +} + +func (l resolveNotificationDestination) String() string { + return fmt.Sprintf("notification-destination: %s", l.name) +} diff --git a/bundle/config/variable/resolve_notification_destination_test.go b/bundle/config/variable/resolve_notification_destination_test.go new file mode 100644 index 00000000..2b8201d1 --- /dev/null +++ b/bundle/config/variable/resolve_notification_destination_test.go @@ -0,0 +1,82 @@ +package variable + +import ( + "context" + "fmt" + "testing" + + "github.com/databricks/databricks-sdk-go/experimental/mocks" + "github.com/databricks/databricks-sdk-go/service/settings" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/mock" + "github.com/stretchr/testify/require" +) + +func TestResolveNotificationDestination_ResolveSuccess(t *testing.T) { + m := mocks.NewMockWorkspaceClient(t) + + api := m.GetMockNotificationDestinationsAPI() + api.EXPECT(). + ListAll(mock.Anything, mock.Anything). + Return([]settings.ListNotificationDestinationsResult{ + {Id: "1234", DisplayName: "destination"}, + }, nil) + + ctx := context.Background() + l := resolveNotificationDestination{name: "destination"} + result, err := l.Resolve(ctx, m.WorkspaceClient) + require.NoError(t, err) + assert.Equal(t, "1234", result) +} + +func TestResolveNotificationDestination_ResolveError(t *testing.T) { + m := mocks.NewMockWorkspaceClient(t) + + api := m.GetMockNotificationDestinationsAPI() + api.EXPECT(). + ListAll(mock.Anything, mock.Anything). + Return(nil, fmt.Errorf("bad")) + + ctx := context.Background() + l := resolveNotificationDestination{name: "destination"} + _, err := l.Resolve(ctx, m.WorkspaceClient) + assert.ErrorContains(t, err, "bad") +} + +func TestResolveNotificationDestination_ResolveNotFound(t *testing.T) { + m := mocks.NewMockWorkspaceClient(t) + + api := m.GetMockNotificationDestinationsAPI() + api.EXPECT(). + ListAll(mock.Anything, mock.Anything). + Return([]settings.ListNotificationDestinationsResult{}, nil) + + ctx := context.Background() + l := resolveNotificationDestination{name: "destination"} + _, err := l.Resolve(ctx, m.WorkspaceClient) + require.Error(t, err) + assert.ErrorContains(t, err, `notification destination named "destination" does not exist`) +} + +func TestResolveNotificationDestination_ResolveMultiple(t *testing.T) { + m := mocks.NewMockWorkspaceClient(t) + + api := m.GetMockNotificationDestinationsAPI() + api.EXPECT(). + ListAll(mock.Anything, mock.Anything). + Return([]settings.ListNotificationDestinationsResult{ + {Id: "1234", DisplayName: "destination"}, + {Id: "5678", DisplayName: "destination"}, + }, nil) + + ctx := context.Background() + l := resolveNotificationDestination{name: "destination"} + _, err := l.Resolve(ctx, m.WorkspaceClient) + require.Error(t, err) + assert.ErrorContains(t, err, `there are 2 instances of clusters named "destination"`) +} + +func TestResolveNotificationDestination_String(t *testing.T) { + l := resolveNotificationDestination{name: "name"} + assert.Equal(t, "notification-destination: name", l.String()) +}