From 6f345293b1e5f4febcd702da8a362b15b606ebd9 Mon Sep 17 00:00:00 2001
From: Andrew Nester <andrew.nester@databricks.com>
Date: Wed, 21 Aug 2024 17:05:49 +0200
Subject: [PATCH] Added filtering flags for cluster list commands (#1703)

## 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
```
---
 cmd/workspace/clusters/overrides.go | 68 ++++++++++++++++++++++++++++-
 1 file changed, 67 insertions(+), 1 deletion(-)

diff --git a/cmd/workspace/clusters/overrides.go b/cmd/workspace/clusters/overrides.go
index 55976d40..6038978a 100644
--- a/cmd/workspace/clusters/overrides.go
+++ b/cmd/workspace/clusters/overrides.go
@@ -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) {