refactor

Signed-off-by: Wantong Jiang <[email protected]>

Author: jwtty

tools/fleet/drain.go tools/fleet/cmd/drain/drain.gotools/fleet/drain.go renamed to tools/fleet/cmd/drain/drain.go

@@ -14,15 +14,17 @@ See the License for the specific language governing permissions and
 limitations under the License.
 */
 
-package main
+package drain
 
 import (
 	"context"
 	"fmt"
 	"log"
 
+	"github.com/spf13/cobra"
 	k8errors "k8s.io/apimachinery/pkg/api/errors"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
+	"k8s.io/apimachinery/pkg/runtime"
 	"k8s.io/apimachinery/pkg/types"
 	"k8s.io/apimachinery/pkg/util/uuid"
 	"k8s.io/apimachinery/pkg/util/validation"
@@ -43,30 +45,54 @@ const (
 	resourceIdentifierKeyFormat = "%s/%s/%s/%s/%s"
 )
 
+var (
+	hubClusterContext string
+	clusterName       string
+)
+
+// setupClient creates and configures the Kubernetes client
+func setupClient() (*runtime.Scheme, client.Client, error) {
+	scheme := runtime.NewScheme()
+
+	if err := clusterv1beta1.AddToScheme(scheme); err != nil {
+		log.Fatalf("failed to add custom APIs (cluster) to the runtime scheme: %v", err)
+	}
+	if err := placementv1beta1.AddToScheme(scheme); err != nil {
+		log.Fatalf("failed to add custom APIs (placement) to the runtime scheme: %v", err)
+	}
+
+	hubClient, err := toolsutils.GetClusterClientFromClusterContext(hubClusterContext, scheme)
+	if err != nil {
+		log.Fatalf("failed to create hub cluster client: %v", err)
+	}
+
+	return scheme, hubClient, nil
+}
+
 type drainHelper struct {
-	hubClient   client.Client
-	clusterName string
+	HubClient   client.Client
+	ClusterName string
 }
 
 func (h *drainHelper) Drain(ctx context.Context) (bool, error) {
 	if err := h.cordon(ctx); err != nil {
-		return false, fmt.Errorf("failed to cordon member cluster %s: %w", h.clusterName, err)
+		return false, fmt.Errorf("failed to cordon member cluster %s: %w", h.ClusterName, err)
 	}
-	log.Printf("Successfully cordoned member cluster %s by adding cordon taint", h.clusterName)
+	log.Printf("Successfully cordoned member cluster %s by adding cordon taint", h.ClusterName)
 
 	crpNameMap, err := h.fetchClusterResourcePlacementNamesToEvict(ctx)
 	if err != nil {
 		return false, err
 	}
 
 	if len(crpNameMap) == 0 {
-		log.Printf("There are currently no resources propagated to %s from fleet using ClusterResourcePlacement resources", h.clusterName)
+		log.Printf("There are currently no resources propagated to %s from fleet using ClusterResourcePlacement resources", h.ClusterName)
 		return true, nil
 	}
 
 	isDrainSuccessful := true
 	for crpName := range crpNameMap {
-		evictionName, err := generateDrainEvictionName(crpName, h.clusterName)
+		evictionName, err := generateDrainEvictionName(crpName, h.ClusterName)
 		if err != nil {
 			return false, err
 		}
@@ -80,54 +106,54 @@ func (h *drainHelper) Drain(ctx context.Context) (bool, error) {
 				},
 				Spec: placementv1beta1.PlacementEvictionSpec{
 					PlacementName: crpName,
-					ClusterName:   h.clusterName,
+					ClusterName:   h.ClusterName,
 				},
 			}
-			return h.hubClient.Create(ctx, &eviction)
+			return h.HubClient.Create(ctx, &eviction)
 		})
 
 		if err != nil {
-			return false, fmt.Errorf("failed to create eviction %s for CRP %s targeting member cluster %s: %w", evictionName, crpName, h.clusterName, err)
+			return false, fmt.Errorf("failed to create eviction %s for CRP %s targeting member cluster %s: %w", evictionName, crpName, h.ClusterName, err)
 		}
 
-		log.Printf("Created eviction %s for CRP %s targeting member cluster %s", evictionName, crpName, h.clusterName)
+		log.Printf("Created eviction %s for CRP %s targeting member cluster %s", evictionName, crpName, h.ClusterName)
 
 		var eviction placementv1beta1.ClusterResourcePlacementEviction
 		err = wait.ExponentialBackoffWithContext(ctx, retry.DefaultBackoff, func(ctx context.Context) (bool, error) {
-			if err := h.hubClient.Get(ctx, types.NamespacedName{Name: evictionName}, &eviction); err != nil {
-				return false, fmt.Errorf("failed to get eviction %s for CRP %s targeting member cluster %s: %w", evictionName, crpName, h.clusterName, err)
+			if err := h.HubClient.Get(ctx, types.NamespacedName{Name: evictionName}, &eviction); err != nil {
+				return false, fmt.Errorf("failed to get eviction %s for CRP %s targeting member cluster %s: %w", evictionName, crpName, h.ClusterName, err)
 			}
 			return evictionutils.IsEvictionInTerminalState(&eviction), nil
 		})
 
 		if err != nil {
-			return false, fmt.Errorf("failed to wait for eviction %s for CRP %s targeting member cluster %s to reach terminal state: %w", evictionName, crpName, h.clusterName, err)
+			return false, fmt.Errorf("failed to wait for eviction %s for CRP %s targeting member cluster %s to reach terminal state: %w", evictionName, crpName, h.ClusterName, err)
 		}
 
 		validCondition := eviction.GetCondition(string(placementv1beta1.PlacementEvictionConditionTypeValid))
 		if validCondition != nil && validCondition.Status == metav1.ConditionFalse {
 			if validCondition.Reason == condition.EvictionInvalidMissingCRPMessage ||
 				validCondition.Reason == condition.EvictionInvalidDeletingCRPMessage ||
 				validCondition.Reason == condition.EvictionInvalidMissingCRBMessage {
-				log.Printf("eviction %s is invalid with reason %s for CRP %s targeting member cluster %s, but drain will succeed", evictionName, validCondition.Reason, crpName, h.clusterName)
+				log.Printf("eviction %s is invalid with reason %s for CRP %s targeting member cluster %s, but drain will succeed", evictionName, validCondition.Reason, crpName, h.ClusterName)
 				continue
 			}
 		}
 		executedCondition := eviction.GetCondition(string(placementv1beta1.PlacementEvictionConditionTypeExecuted))
 		if executedCondition == nil || executedCondition.Status == metav1.ConditionFalse {
 			isDrainSuccessful = false
-			log.Printf("eviction %s was not executed successfully for CRP %s targeting member cluster %s", evictionName, crpName, h.clusterName)
+			log.Printf("eviction %s was not executed successfully for CRP %s targeting member cluster %s", evictionName, crpName, h.ClusterName)
 			continue
 		}
-		log.Printf("eviction %s was executed successfully for CRP %s targeting member cluster %s", evictionName, crpName, h.clusterName)
-		
+		log.Printf("eviction %s was executed successfully for CRP %s targeting member cluster %s", evictionName, crpName, h.ClusterName)
+
 		clusterScopedResourceIdentifiers, err := h.collectClusterScopedResourcesSelectedByCRP(ctx, crpName)
 		if err != nil {
 			log.Printf("failed to collect cluster scoped resources selected by CRP %s: %v", crpName, err)
 			continue
 		}
 		for _, resourceIdentifier := range clusterScopedResourceIdentifiers {
-			log.Printf("evicted resource %s propagated by CRP %s targeting member cluster %s", generateResourceIdentifierKey(resourceIdentifier), crpName, h.clusterName)
+			log.Printf("evicted resource %s propagated by CRP %s targeting member cluster %s", generateResourceIdentifierKey(resourceIdentifier), crpName, h.ClusterName)
 		}
 	}
 
@@ -137,7 +163,7 @@ func (h *drainHelper) Drain(ctx context.Context) (bool, error) {
 func (h *drainHelper) cordon(ctx context.Context) error {
 	return retry.RetryOnConflict(retry.DefaultRetry, func() error {
 		var mc clusterv1beta1.MemberCluster
-		if err := h.hubClient.Get(ctx, types.NamespacedName{Name: h.clusterName}, &mc); err != nil {
+		if err := h.HubClient.Get(ctx, types.NamespacedName{Name: h.ClusterName}, &mc); err != nil {
 			return err
 		}
 
@@ -149,20 +175,20 @@ func (h *drainHelper) cordon(ctx context.Context) error {
 
 		mc.Spec.Taints = append(mc.Spec.Taints, toolsutils.CordonTaint)
 
-		return h.hubClient.Update(ctx, &mc)
+		return h.HubClient.Update(ctx, &mc)
 	})
 }
 
 func (h *drainHelper) fetchClusterResourcePlacementNamesToEvict(ctx context.Context) (map[string]bool, error) {
 	var crbList placementv1beta1.ClusterResourceBindingList
-	if err := h.hubClient.List(ctx, &crbList); err != nil {
+	if err := h.HubClient.List(ctx, &crbList); err != nil {
 		return map[string]bool{}, fmt.Errorf("failed to list cluster resource bindings: %w", err)
 	}
 
 	crpNameMap := make(map[string]bool)
 	for i := range crbList.Items {
 		crb := crbList.Items[i]
-		if crb.Spec.TargetCluster == h.clusterName && crb.DeletionTimestamp == nil {
+		if crb.Spec.TargetCluster == h.ClusterName && crb.DeletionTimestamp == nil {
 			crpName, ok := crb.GetLabels()[placementv1beta1.PlacementTrackingLabel]
 			if !ok {
 				return map[string]bool{}, fmt.Errorf("failed to get CRP name from binding %s", crb.Name)
@@ -176,7 +202,7 @@ func (h *drainHelper) fetchClusterResourcePlacementNamesToEvict(ctx context.Cont
 
 func (h *drainHelper) collectClusterScopedResourcesSelectedByCRP(ctx context.Context, crpName string) ([]placementv1beta1.ResourceIdentifier, error) {
 	var crp placementv1beta1.ClusterResourcePlacement
-	if err := h.hubClient.Get(ctx, types.NamespacedName{Name: crpName}, &crp); err != nil {
+	if err := h.HubClient.Get(ctx, types.NamespacedName{Name: crpName}, &crp); err != nil {
 		return nil, fmt.Errorf("failed to get ClusterResourcePlacement %s: %w", crpName, err)
 	}
 
@@ -209,4 +235,64 @@ func generateResourceIdentifierKey(r placementv1beta1.ResourceIdentifier) string
 		return fmt.Sprintf(resourceIdentifierKeyFormat, r.Group, r.Version, r.Kind, "''", r.Name)
 	}
 	return fmt.Sprintf(resourceIdentifierKeyFormat, r.Group, r.Version, r.Kind, r.Namespace, r.Name)
-}
+}
+
+// NewCmdDrain creates a new drain command
+func NewCmdDrain() *cobra.Command {
+	cmd := &cobra.Command{
+		Use:   "drain",
+		Short: "Drain a member cluster",
+		Long:  "Drain a member cluster by cordoning it and removing propagated resources",
+		RunE: func(command *cobra.Command, args []string) error {
+			return runDrain()
+		},
+	}
+
+	// Add flags specific to drain command
+	cmd.Flags().StringVar(&hubClusterContext, "hubClusterContext", "", "kubectl context for the hub cluster (required)")
+	cmd.Flags().StringVar(&clusterName, "clusterName", "", "name of the member cluster (required)")
+
+	// Mark required flags
+	_ = cmd.MarkFlagRequired("hubClusterContext")
+	_ = cmd.MarkFlagRequired("clusterName")
+
+	return cmd
+}
+
+func runDrain() error {
+	_, hubClient, err := setupClient()
+	if err != nil {
+		return err
+	}
+
+	ctx := context.Background()
+	drainHelper := &drainHelper{
+		HubClient:   hubClient,
+		ClusterName: clusterName,
+	}
+
+	isDrainSuccessful, err := drainHelper.Drain(ctx)
+	if err != nil {
+		log.Fatalf("failed to drain member cluster %s: %v", clusterName, err)
+	}
+
+	if isDrainSuccessful {
+		log.Printf("drain was successful for cluster %s", clusterName)
+	} else {
+		log.Printf("drain was not successful for cluster %s", clusterName)
+	}
+
+	log.Printf("retrying drain to ensure all resources propagated from hub cluster are evicted")
+	isDrainRetrySuccessful, err := drainHelper.Drain(ctx)
+	if err != nil {
+		log.Fatalf("failed to drain cluster on retry %s: %v", clusterName, err)
+	}
+	if isDrainRetrySuccessful {
+		log.Printf("drain retry was successful for cluster %s", clusterName)
+	} else {
+		log.Printf("drain retry was not successful for cluster %s", clusterName)
+	}
+
+	log.Printf("reminder: uncordon the cluster %s to remove cordon taint if needed", clusterName)
+	return nil
+}

tools/fleet/drain_test.go tools/fleet/cmd/drain/drain_test.gotools/fleet/drain_test.go renamed to tools/fleet/cmd/drain/drain_test.go

@@ -14,7 +14,7 @@ See the License for the specific language governing permissions and
 limitations under the License.
 */
 
-package main
+package drain
 
 import (
 	"context"
@@ -157,8 +157,8 @@ func TestFetchClusterResourcePlacementNamesToEvict(t *testing.T) {
 				WithObjects(objects...).
 				Build()
 			h := drainHelper{
-				hubClient:   fakeClient,
-				clusterName: tc.targetCluster,
+				HubClient:   fakeClient,
+				ClusterName: tc.targetCluster,
 			}
 			gotMap, gotErr := h.fetchClusterResourcePlacementNamesToEvict(context.Background())
 			if tc.wantErr == nil {
@@ -274,7 +274,7 @@ func TestCollectClusterScopedResourcesSelectedByCRP(t *testing.T) {
 				Build()
 
 			h := drainHelper{
-				hubClient: fakeClient,
+				HubClient: fakeClient,
 			}
 
 			gotResources, gotErr := h.collectClusterScopedResourcesSelectedByCRP(context.Background(), tc.crpName)
@@ -480,8 +480,8 @@ func TestCordon(t *testing.T) {
 				Build()
 
 			h := drainHelper{
-				hubClient:   fakeClient,
-				clusterName: "test-cluster",
+				HubClient:   fakeClient,
+				ClusterName: "test-cluster",
 			}
 
 			gotErr := h.cordon(context.Background())