endpoint

Author: alyssa-db

cmd/pipelines/run.go

@@ -27,7 +27,6 @@ import (
 	"github.com/databricks/databricks-sdk-go/service/pipelines"
 	"github.com/spf13/cobra"
 	"golang.org/x/exp/maps"
-	"golang.org/x/exp/slices"
 )
 
 type PipelineUpdateData struct {
@@ -37,7 +36,7 @@ type PipelineUpdateData struct {
 	LastEventTime       string
 }
 
-const pipelineUpdateTemplate = `Update {{ .Update.UpdateId }} for pipeline {{- if .Update.Config }}{{ .Update.Config.Name }}{{ end }} {{- if .Update.Config }}{{ .Update.Config.Id }}{{ end }} completed successfully.
+const pipelineUpdateTemplate = `Update {{ .Update.UpdateId }} for pipeline {{- if .Update.Config }} {{ .Update.Config.Name }}{{ end }} {{- if .Update.Config }} {{ .Update.Config.Id }}{{ end }} completed successfully.
 {{- if .Update.Cause }}
 Cause: {{ .Update.Cause }}
 {{- end }}
@@ -90,30 +89,6 @@ func getRefreshSelectionString(update pipelines.UpdateInfo) string {
 	return "default refresh-all"
 }
 
-func fetchUpdateProgressEventsForUpdateAscending(ctx context.Context, bundle *bundle.Bundle, pipelineId, updateId string) ([]pipelines.PipelineEvent, error) {
-	w := bundle.WorkspaceClient()
-
-	req := pipelines.ListPipelineEventsRequest{
-		PipelineId: pipelineId,
-		Filter:     fmt.Sprintf("update_id='%s' AND event_type='update_progress'", updateId),
-		// OrderBy:    []string{"timestamp asc"}, TODO: Add this back in when the API is fixed
-	}
-
-	iterator := w.Pipelines.ListPipelineEvents(ctx, req)
-	var events []pipelines.PipelineEvent
-
-	for iterator.HasNext(ctx) {
-		event, err := iterator.Next(ctx)
-		if err != nil {
-			return nil, err
-		}
-		events = append(events, event)
-	}
-	slices.Reverse(events)
-
-	return events, nil
-}
-
 func fetchAndDisplayPipelineUpdate(ctx context.Context, bundle *bundle.Bundle, ref bundleresources.Reference, updateId string) error {
 	w := bundle.WorkspaceClient()
 
@@ -137,12 +112,17 @@ func fetchAndDisplayPipelineUpdate(ctx context.Context, bundle *bundle.Bundle, r
 
 	latestUpdate := *getUpdateResponse.Update
 
-	if latestUpdate.State == pipelines.UpdateInfoStateCompleted {
-		events, err := fetchUpdateProgressEventsForUpdateAscending(ctx, bundle, pipelineID, updateId)
-		if err != nil {
-			return err
-		}
+	params := &PipelineEventsQueryParams{
+		Filter:  fmt.Sprintf("update_id='%s' AND event_type='update_progress'", updateId),
+		OrderBy: "timestamp asc",
+	}
 
+	events, err := fetchAllPipelineEvents(ctx, w, pipelineID, params)
+	if err != nil {
+		return err
+	}
+
+	if latestUpdate.State == pipelines.UpdateInfoStateCompleted {
 		err = displayPipelineUpdate(ctx, latestUpdate, pipelineID, events)
 		if err != nil {
 			return err
@@ -165,7 +145,6 @@ func getLastEventTime(events []pipelines.PipelineEvent) string {
 	return parsedTime.Format("2006-01-02T15:04:05Z")
 }
 
-// displayPipelineUpdate displays pipeline update information
 func displayPipelineUpdate(ctx context.Context, update pipelines.UpdateInfo, pipelineID string, events []pipelines.PipelineEvent) error {
 	data := PipelineUpdateData{
 		PipelineId:          pipelineID,

cmd/pipelines/utils.go

@@ -4,12 +4,16 @@ import (
 	"context"
 	"errors"
 	"fmt"
+	"strconv"
 
 	"github.com/databricks/cli/bundle"
 	configresources "github.com/databricks/cli/bundle/config/resources"
 	"github.com/databricks/cli/bundle/resources"
 	"github.com/databricks/cli/bundle/run"
 	"github.com/databricks/cli/libs/cmdio"
+	"github.com/databricks/databricks-sdk-go"
+	"github.com/databricks/databricks-sdk-go/client"
+	"github.com/databricks/databricks-sdk-go/service/pipelines"
 )
 
 // Copied from cmd/bundle/run.go
@@ -88,3 +92,52 @@ func keyToRunner(b *bundle.Bundle, arg string) (run.Runner, error) {
 
 	return runner, nil
 }
+
+type PipelineEventsResponse struct {
+	Events        []pipelines.PipelineEvent `json:"events"`
+	NextPageToken string                    `json:"next_page_token,omitempty"`
+}
+
+type PipelineEventsQueryParams struct {
+	Filter     string `json:"filter,omitempty"`
+	MaxResults int    `json:"max_results,omitempty"`
+	PageToken  string `json:"page_token,omitempty"`
+	OrderBy    string `json:"order_by,omitempty"`
+}
+
+func fetchAllPipelineEvents(ctx context.Context, w *databricks.WorkspaceClient, pipelineID string, params *PipelineEventsQueryParams) ([]pipelines.PipelineEvent, error) {
+	apiClient, err := client.New(w.Config)
+	if err != nil {
+		return nil, fmt.Errorf("failed to create API client: %w", err)
+	}
+
+	path := fmt.Sprintf("/api/2.0/pipelines/%s/events", pipelineID)
+
+	queryParams := map[string]string{}
+	if params.Filter != "" {
+		queryParams["filter"] = params.Filter
+	}
+	if params.MaxResults > 0 {
+		queryParams["max_results"] = strconv.Itoa(params.MaxResults)
+	}
+
+	if params.OrderBy != "" {
+		queryParams["order_by"] = params.OrderBy
+	}
+
+	var response PipelineEventsResponse
+	err = apiClient.Do(
+		ctx,
+		"GET",
+		path,
+		nil,
+		nil,
+		queryParams,
+		&response,
+	)
+	if err != nil {
+		return nil, fmt.Errorf("failed to fetch pipeline events: %w", err)
+	}
+
+	return response.Events, nil
+}