Skip to content

Commit

Permalink
Added filtering flags for cluster list commands (#1703)
Browse files Browse the repository at this point in the history
## Changes
Fixes #1701 

## Tests
```
Usage:
  databricks clusters list [flags]

Flags:
      --cluster-sources []string   Filter clusters by source
      --cluster-states []string    Filter clusters by states
  -h, --help                       help for list
      --is-pinned                  Filter clusters by pinned status
      --page-size int              Use this field to specify the maximum number of results to be returned by the server.
      --page-token string          Use next_page_token or prev_page_token returned from the previous request to list the next or previous page of clusters respectively.
      --policy-id string           Filter clusters by policy id
```
  • Loading branch information
andrewnester authored Aug 21, 2024
1 parent f5df211 commit 6f34529
Showing 1 changed file with 67 additions and 1 deletion.
68 changes: 67 additions & 1 deletion cmd/workspace/clusters/overrides.go
Original file line number Diff line number Diff line change
@@ -1,17 +1,83 @@
package clusters

import (
"strings"

"github.com/databricks/cli/libs/cmdio"
"github.com/databricks/databricks-sdk-go/service/compute"
"github.com/spf13/cobra"
)

func listOverride(listCmd *cobra.Command, _ *compute.ListClustersRequest) {
// Below we add overrides for filter flags for cluster list command to allow for custom filtering
// Auto generating such flags is not yet supported by the CLI generator
func listOverride(listCmd *cobra.Command, listReq *compute.ListClustersRequest) {
listCmd.Annotations["headerTemplate"] = cmdio.Heredoc(`
{{header "ID"}} {{header "Name"}} {{header "State"}}`)
listCmd.Annotations["template"] = cmdio.Heredoc(`
{{range .}}{{.ClusterId | green}} {{.ClusterName | cyan}} {{if eq .State "RUNNING"}}{{green "%s" .State}}{{else if eq .State "TERMINATED"}}{{red "%s" .State}}{{else}}{{blue "%s" .State}}{{end}}
{{end}}`)

listReq.FilterBy = &compute.ListClustersFilterBy{}
listCmd.Flags().BoolVar(&listReq.FilterBy.IsPinned, "is-pinned", false, "Filter clusters by pinned status")
listCmd.Flags().StringVar(&listReq.FilterBy.PolicyId, "policy-id", "", "Filter clusters by policy id")

sources := &clusterSources{source: &listReq.FilterBy.ClusterSources}
listCmd.Flags().Var(sources, "cluster-sources", "Filter clusters by source")

states := &clusterStates{state: &listReq.FilterBy.ClusterStates}
listCmd.Flags().Var(states, "cluster-states", "Filter clusters by states")
}

type clusterSources struct {
source *[]compute.ClusterSource
}

func (c *clusterSources) String() string {
s := make([]string, len(*c.source))
for i, source := range *c.source {
s[i] = string(source)
}

return strings.Join(s, ",")
}

func (c *clusterSources) Set(value string) error {
splits := strings.Split(value, ",")
for _, split := range splits {
*c.source = append(*c.source, compute.ClusterSource(split))
}

return nil
}

func (c *clusterSources) Type() string {
return "[]string"
}

type clusterStates struct {
state *[]compute.State
}

func (c *clusterStates) String() string {
s := make([]string, len(*c.state))
for i, source := range *c.state {
s[i] = string(source)
}

return strings.Join(s, ",")
}

func (c *clusterStates) Set(value string) error {
splits := strings.Split(value, ",")
for _, split := range splits {
*c.state = append(*c.state, compute.State(split))
}

return nil
}

func (c *clusterStates) Type() string {
return "[]string"
}

func listNodeTypesOverride(listNodeTypesCmd *cobra.Command) {
Expand Down

0 comments on commit 6f34529

Please sign in to comment.