De Banderas a Subcomandos: Reescribiendo Aws-Doctor Con Cobra

Introducción

En el artículo anterior te estuve compartiendo mi experiencia creando aws-doctor, una herramienta open source para optimizar costos en AWS desde la terminal. En este artículo te hablaré un poco sobre la versión 2 de aws-doctor y cuales fueron los breaking changes relacionados que provocaron que tuviera que lanzar una nueva versión mayor. Desde ya te comento que estaremos hablando bastante sobre cobra, el framework de Golang para crear CLIs.

Contexto

En la versión v1.10.2 de aws-doctor, la herramienta era bastante estable y cumplía su objetivo. Contaba con tres funcionalidades principales, las cuales eran:

• flujo por defecto: para comparar los gastos del mes en curso con el mes pasado en el mismo período de tiempo.
• --trend flag: para mostrar en la consola la tendencia de costos de los últimos 6 meses.
• --waste flag: para detectar recursos zombies en AWS.

Además de esto, existían otras banderas globales aplicables a cualquier comando:

• --profile para especificar el perfil de AWS CLI a usar.
• --region para especificar la región de AWS a usar.
• --output para especificar el formato de salida (json, table).

Motivación

Uno de los features que estaba disponible en la versión v1.10.2 era la capacidad de especificar para el flujo --waste las revisiones a ejecutar, de esta forma, si el usuario deseaba solo ejecutar la revisión sobre los recursos ec2 y s3, podía ejecutar el siguiente comando:

aws-doctor --waste ec2,s3

Y esto funcionaba perfectamente, pero a partir de ese momento me di cuenta que algo no estaba bien con la herramienta, y ahora te muestro porque. Hasta ese momento, en el proyecto no se estaba utilizando ningún framework para crear la CLI, solamente tenía un servico llamado FlagService que te muestro a continuación:

package flag

import "github.com/elC0mpa/aws-doctor/model"

type service struct{}

// Service is the interface for CLI flag service.
type Service interface {
	GetParsedFlags(args []string) (model.Flags, error)
}

// Package flag provides a service for parsing CLI flags.
package flag

import (
	"flag"
	"strings"

	"github.com/elC0mpa/aws-doctor/model"
)

// NewService creates a new Flag service.
func NewService() Service {
	return &service{}
}

func (s *service) GetParsedFlags(args []string) (model.Flags, error) {
	fs := flag.NewFlagSet("aws-doctor", flag.ContinueOnError)

	region := fs.String("region", "", "AWS region (defaults to AWS_REGION, AWS_DEFAULT_REGION, or ~/.aws/config)")
	profile := fs.String("profile", "", "AWS profile configuration")
	trend := fs.Bool("trend", false, "Display a trend report for the last 6 months")
	waste := fs.Bool("waste", false, "Display AWS waste report (e.g., --waste ec2,s3)")
	output := fs.String("output", "table", "Output format: table or json")
	version := fs.Bool("version", false, "Display version information")
	update := fs.Bool("update", false, "Update aws-doctor to the latest version")

	var wasteChecks []string

	filteredArgs := make([]string, 0, len(args))

	for i := 0; i < len(args); i++ {
		arg := args[i]
		filteredArgs = append(filteredArgs, arg)

		// If the current argument is the waste flag and the next argument is not a flag,
		// we treat the next argument as the list of specific checks to run.
		if (arg == "--waste" || arg == "-waste") && i+1 < len(args) && !strings.HasPrefix(args[i+1], "-") {
			wasteChecks = strings.Split(args[i+1], ",")
			i++ // Consume the next argument so it's not parsed as a positional argument
		}
	}

	if err := fs.Parse(filteredArgs); err != nil {
		return model.Flags{}, err
	}

	return model.Flags{
		Region:      *region,
		Profile:     *profile,
		Trend:       *trend,
		Waste:       *waste,
		WasteChecks: wasteChecks,
		Output:      *output,
		Version:     *version,
		Update:      *update,
	}, nil
}

Como puedes ver, este servicio se encargaba de parsear las banderas de la aplicación y luego retornarlas en un struct. Pero el problema comenzó a partir de la lógica rara que puedes ver escrita en el ciclo for, que básicamente se encargaba de detectar si el usuario, luego de pasar la bandera --waste, había pasado un argumento adicional con las revisiones a ejecutar, y de ser así, parsear ese argumento y guardarlo en el struct de flags. A pesar de que esto funcionaba, mi instinto me dijo que si continuaba por este camino, la herramienta se volvería cada vez más difícil de mantener, y por este motivo empecé a buscar alternativas.

La solución: Cobra

Luego de hacer una pequeña investigación, llegué a la conclusión de que la forma de resolver este problema y que se mantuviera estable la herramienta era utilizando el framework Cobra, que entre otras cosas, es usado por CLIs muy famosos, como el de GitHub y el de GitLab.

Cambio de paradigma

Anteriormente, como ya leíste, la herramienta usaba banderas para todo, incluso para ejecutar funcionalidades específicas, como el caso de --waste y --trend. Luego de estudiar un poco sobre Cobra, aprendí que estos no deberían ser banderas, sino subcomandos, y que las banderas deberían ser usadas para modificar el comportamiento de los comandos. Por ejemplo, en lugar de ejecutar aws-doctor --waste, ahora el comando para ejecutar la funcionalidad de waste sería aws-doctor waste, y como cada subcomando tiene argumentos propios, en el caso de este subcomando, los argumentos serían las revisiones a ejecutar, por lo que el comando completo para ejecutar la revisión de ec2 y s3 sería aws-doctor waste ec2 s3.

Este es el motivo principal por que decidí lanzar una nueva versión mayor, ya que este cambio de paradigma rompía la forma en la que el usuario final usaba la herramienta.

Implementación

Ahora, la implementación es mucho más estructurada y escalable. Se recomienda tener una carpeta cmd en la raíz del proyecto, donde haya un fichero por cada subcomando, y un fichero root.go que se encargue de definir el comando raíz y las banderas globales. En cada uno de los ficheros de los subcomandos, se define el comando específico y sus argumentos, y luego se implementa la lógica correspondiente en la función Run.

Este es el código en el fichero cmd/root.go:

package cmd

import (
	"context"
	"fmt"

	"github.com/elC0mpa/aws-doctor/model"
	awsconfig "github.com/elC0mpa/aws-doctor/service/aws_config"
	"github.com/elC0mpa/aws-doctor/service/cloudwatchlogs"
	"github.com/elC0mpa/aws-doctor/service/cloudwatchmetrics"
	awscostexplorer "github.com/elC0mpa/aws-doctor/service/costexplorer"
	awsec2 "github.com/elC0mpa/aws-doctor/service/ec2"
	"github.com/elC0mpa/aws-doctor/service/elb"
	"github.com/elC0mpa/aws-doctor/service/orchestrator"
	"github.com/elC0mpa/aws-doctor/service/output"
	"github.com/elC0mpa/aws-doctor/service/rds"
	"github.com/elC0mpa/aws-doctor/service/s3"
	awssts "github.com/elC0mpa/aws-doctor/service/sts"
	"github.com/elC0mpa/aws-doctor/service/update"
	"github.com/elC0mpa/aws-doctor/utils/banner"
	"github.com/elC0mpa/aws-doctor/utils/spinner"
	"github.com/spf13/cobra"
)

var (
	region              string
	profile             string
	outputFormat        string
	versionInfo         model.VersionInfo
	orchestratorBuilder = buildOrchestrator
)

func buildOrchestrator(needsAWS bool) (orchestrator.Service, error) {
	outputService := output.NewService(outputFormat)
	updateService := update.NewService()

	if !needsAWS {
		return orchestrator.NewService(nil, nil, nil, nil, nil, nil, nil, outputService, updateService, versionInfo), nil
	}

	banner.DrawBannerTitle()

	cfgService := awsconfig.NewService()

	awsCfg, err := cfgService.GetAWSCfg(context.Background(), region, profile)
	if err != nil {
		return nil, fmt.Errorf("failed to load AWS config: %w", err)
	}

	spinner.StartSpinner()

	costService := awscostexplorer.NewService(awsCfg)
	stsService := awssts.NewService(awsCfg)
	ec2Service := awsec2.NewService(awsCfg)
	elbService := elb.NewService(awsCfg)
	s3Service := s3.NewService(awsCfg)
	cloudwatchlogsService := cloudwatchlogs.NewService(awsCfg)
	cwMetricsService := cloudwatchmetrics.NewService(awsCfg)
	rdsService := rds.NewService(awsCfg, cwMetricsService)

	return orchestrator.NewService(stsService, costService, ec2Service, elbService, s3Service, cloudwatchlogsService, rdsService, outputService, updateService, versionInfo), nil
}

var rootCmd = &cobra.Command{
	Use:   "aws-doctor",
	Short: "A comprehensive health check for your AWS accounts",
}

// Execute adds all child commands to the root command and sets flags appropriately.
func Execute(version, commit, date string) error {
	versionInfo = model.VersionInfo{
		Version: version,
		Commit:  commit,
		Date:    date,
	}

	return rootCmd.Execute()
}

func init() {
	rootCmd.PersistentFlags().StringVar(&region, "region", "", "AWS region (defaults to AWS_REGION, AWS_DEFAULT_REGION, or ~/.aws/config)")
	rootCmd.PersistentFlags().StringVar(&profile, "profile", "", "AWS profile configuration")
	rootCmd.PersistentFlags().StringVar(&outputFormat, "output", "table", "Output format: table, json or csv")
}

En este fichero se hacen todas las inicializaciones necesarias para la herramienta, para esto se usa el método buildOrchestrator, que se encarga de inicializar todos los servicios necesarios para ejecutar las funcionalidades de la herramienta. Este método se usa desde cada uno de los otros subcomandos, y se le pasa un booleano indicando si el subcomando necesita o no acceso a AWS, de esta forma, si el subcomando no necesita acceso a AWS, no se inicializan los servicios relacionados, lo que hace que la herramienta sea más rápida y eficiente.

A continuación, te muestro el fichero cmd/waste.go, que es el encargado de ejecutar la funcionalidad waste:

package cmd

import (
	"strings"

	"github.com/elC0mpa/aws-doctor/model"
	"github.com/spf13/cobra"
)

var wasteCmd = &cobra.Command{
	Use:   "waste [checks...]",
	Short: "Display AWS waste report (e.g., ec2 s3)",
	RunE: func(cmd *cobra.Command, args []string) error {
		orch, err := orchestratorBuilder(true)
		if err != nil {
			return err
		}

		var parsedChecks []string

		for _, arg := range args {
			parsedChecks = append(parsedChecks, strings.Split(arg, ",")...)
		}

		flags := model.Flags{
			Region:      region,
			Profile:     profile,
			Output:      outputFormat,
			Waste:       true,
			WasteChecks: parsedChecks,
		}

		return orch.Orchestrate(flags)
	},
}

func init() {
	rootCmd.AddCommand(wasteCmd)
}

Como puedes ver, básicamente se construye el orquestrador con acceso a AWS, se parsean los argumentos para obtener las revisiones a ejecutar, y luego se llama al método Orchestrate del orquestador, pasando las flags correspondientes.

Es justo en este ejemplo donde se puede ver claramente la ventaja de usar Cobra, ya que la implementación es mucho más limpia y fácil de entender, además de que se sigue el paradigma recomendado para CLIs, lo que hace que la herramienta sea más intuitiva para los usuarios.

Conclusiones

Haber tomado la decisión de migrar a Cobra fue un gran paso para mejorar la escalabilidad de esta CLI así como su legibilidad, y aunque implicó hacer un cambio de paradigma y lanzar una nueva versión mayor, creo que fue la decisión correcta a largo plazo. Ahora la herramienta está mucho mejor estructurada, es más fácil de mantener y de extender con nuevas funcionalidades en el futuro. Además, al seguir el paradigma recomendado para CLIs, creo que la experiencia del usuario final también ha mejorado significativamente.