From 716bf17e322a0a441b5941eaaf465892b1f637f4 Mon Sep 17 00:00:00 2001 From: Morgan Fainberg Date: Tue, 5 Mar 2024 15:00:30 -0800 Subject: [PATCH] Cleanup Aliasing Aliases in Cobra and Viper are relatively complex and viper.RegisterAlias does not do what would be expected as ENV handling is a very special case. This commit fixes up the handling of aliases specifically for chunk size and minimuim chunk size. * CLI Option Set * NEW ENV VAR * Deprecated ENV Var * Default Value For ChunkSize and MinimumChunkSize we now explicitly disallow setting both --chunk-size and --minimum-chunk-size as CLI options. If PGET_CHUNK_SIZE is set, we utilize that value except if the cli option(s) are set. If PGET_MINIMUM_CHUNK_SIZE is set * If PGET_CHUNK_SIZE is set to a non-default value a warning is emitted, and PGET_CHUNK_SIZE is used.. * If PGET_CHUNK_SIZE is the default value, PGET_MINIMUM_CHUNK_SIZE is set as PGET_CHUNK_SIZE and a deprecation warning is emitted. Additionally, viper calls now all reference config.OptChunkSize --- cmd/multifile/multifile.go | 2 +- cmd/root/root.go | 64 ++++++++++++++++++++++++-------------- pkg/config/config.go | 10 ------ 3 files changed, 42 insertions(+), 34 deletions(-) diff --git a/cmd/multifile/multifile.go b/cmd/multifile/multifile.go index 0fc5d0b..4da0f8f 100644 --- a/cmd/multifile/multifile.go +++ b/cmd/multifile/multifile.go @@ -97,7 +97,7 @@ func maxConcurrentFiles() int { } func multifileExecute(ctx context.Context, manifest pget.Manifest) error { - chunkSize, err := humanize.ParseBytes(viper.GetString(config.OptMinimumChunkSize)) + chunkSize, err := humanize.ParseBytes(viper.GetString(config.OptChunkSize)) if err != nil { return err } diff --git a/cmd/root/root.go b/cmd/root/root.go index ca2882b..4cb3d19 100644 --- a/cmd/root/root.go +++ b/cmd/root/root.go @@ -3,6 +3,7 @@ package root import ( "context" "fmt" + "github.com/replicate/pget/pkg/logging" "os" "runtime" "time" @@ -44,6 +45,8 @@ var concurrency int var pidFile *cli.PIDFile var chunkSize string +const chunkSizeDefault = "125M" + func GetCommand() *cobra.Command { cmd := &cobra.Command{ Use: "pget [flags] ", @@ -51,7 +54,6 @@ func GetCommand() *cobra.Command { Long: rootLongDesc, PersistentPreRunE: rootPersistentPreRunEFunc, PersistentPostRunE: rootPersistentPostRunEFunc, - PreRun: rootCmdPreRun, RunE: runRootCMD, Args: validateArgs, Example: ` pget https://example.com/file.tar ./target-dir`, @@ -68,6 +70,7 @@ func GetCommand() *cobra.Command { fmt.Println(err) os.Exit(1) } + err = viper.BindPFlags(cmd.Flags()) if err != nil { fmt.Println(err) @@ -102,6 +105,7 @@ func pidFlock(pidFilePath string) error { } func rootPersistentPreRunEFunc(cmd *cobra.Command, args []string) error { + logger := logging.GetLogger() if err := config.PersistentStartupProcessFlags(); err != nil { return err } @@ -110,6 +114,38 @@ func rootPersistentPreRunEFunc(cmd *cobra.Command, args []string) error { return err } } + + // Handle chunk size flags (deprecation and overwriting where needed) + // + // Expected Behavior for chunk size flags: + // * If either cli option is set, use that value + // * If both are set, emit an error + // * If neither are set, use ENV values + // ** If PGET_CHUNK_SIZE is set, use that value + // ** If PGET_CHUNK_SIZE is not set, use PGET_MINIMUM_CHUNK_SIZE if set + // NOTE: PGET_MINIMUM_CHUNK_SIZE value is just set over the key for PGET_CHUNK_SIZE + // Warning message will be emitted + // ** If both PGET_CHUNK_SIZE and PGET_MINIMUM_CHUNK_SIZE are set, use PGET_CHUNK_SIZE + // Warning message will be emitted + // * If neither are set, use the default value + + changedMin := cmd.PersistentFlags().Changed(config.OptMinimumChunkSize) + changedChunk := cmd.PersistentFlags().Changed(config.OptChunkSize) + if changedMin && changedChunk { + return fmt.Errorf("--minimum-chunk-size and --chunk-size cannot be used at the same time, use --chunk-size instead") + } else if !(changedMin && changedChunk) { + minChunkSizeEnv := viper.GetString(config.OptMinimumChunkSize) + chunkSizeEnv := viper.GetString(config.OptChunkSize) + if minChunkSizeEnv != chunkSizeDefault { + if chunkSizeEnv == chunkSizeDefault { + logger.Warn().Msg("Using PGET_MINIMUM_CHUNK_SIZE is deprecated, use PGET_CHUNK_SIZE instead") + viper.Set(config.OptChunkSize, minChunkSizeEnv) + } else { + logger.Warn().Msg("Both PGET_MINIMUM_CHUNK_SIZE and PGET_CHUNK_SIZE are set, using PGET_CHUNK_SIZE") + } + } + } + return nil } @@ -125,8 +161,8 @@ func persistentFlags(cmd *cobra.Command) error { cmd.PersistentFlags().IntVarP(&concurrency, config.OptConcurrency, "c", runtime.GOMAXPROCS(0)*4, "Maximum number of concurrent downloads/maximum number of chunks for a given file") cmd.PersistentFlags().IntVar(&concurrency, config.OptMaxChunks, runtime.GOMAXPROCS(0)*4, "Maximum number of chunks for a given file") cmd.PersistentFlags().Duration(config.OptConnTimeout, 5*time.Second, "Timeout for establishing a connection, format is , e.g. 10s") - cmd.PersistentFlags().StringVarP(&chunkSize, config.OptChunkSize, "m", "125M", "Chunk size (in bytes) to use when downloading a file (e.g. 10M)") - cmd.PersistentFlags().StringVar(&chunkSize, config.OptMinimumChunkSize, "125M", "Minimum chunk size (in bytes) to use when downloading a file (e.g. 10M)") + cmd.PersistentFlags().StringVarP(&chunkSize, config.OptChunkSize, "m", chunkSizeDefault, "Chunk size (in bytes) to use when downloading a file (e.g. 10M)") + cmd.PersistentFlags().StringVar(&chunkSize, config.OptMinimumChunkSize, chunkSizeDefault, "Minimum chunk size (in bytes) to use when downloading a file (e.g. 10M)") cmd.PersistentFlags().BoolP(config.OptForce, "f", false, "OptForce download, overwriting existing file") cmd.PersistentFlags().StringSlice(config.OptResolve, []string{}, "OptResolve hostnames to specific IPs") cmd.PersistentFlags().IntP(config.OptRetries, "r", 5, "Number of retries when attempting to retrieve a file") @@ -137,10 +173,6 @@ func persistentFlags(cmd *cobra.Command) error { cmd.PersistentFlags().StringP(config.OptOutputConsumer, "o", "file", "Output Consumer (file, tar, null)") cmd.PersistentFlags().String(config.OptPIDFile, defaultPidFilePath(), "PID file path") - if err := config.AddFlagAlias(cmd, config.OptConcurrency, config.OptMaxChunks); err != nil { - return err - } - if err := hideAndDeprecateFlags(cmd); err != nil { return err } @@ -166,12 +198,6 @@ func hideAndDeprecateFlags(cmd *cobra.Command) error { } -func rootCmdPreRun(cmd *cobra.Command, args []string) { - if viper.GetBool(config.OptExtract) { - viper.Set(config.OptOutputConsumer, config.ConsumerTarExtractor) - } -} - func runRootCMD(cmd *cobra.Command, args []string) error { // After we run through the PreRun functions we want to silence usage from being printed // on all errors @@ -182,7 +208,7 @@ func runRootCMD(cmd *cobra.Command, args []string) error { log.Info().Str("url", urlString). Str("dest", dest). - Str("minimum_chunk_size", viper.GetString(config.OptMinimumChunkSize)). + Str("chunk_size", viper.GetString(config.OptChunkSize)). Msg("Initiating") // OMG BODGE FIX THIS @@ -202,18 +228,10 @@ func runRootCMD(cmd *cobra.Command, args []string) error { // rootExecute is the main function of the program and encapsulates the general logic // returns any/all errors to the caller. func rootExecute(ctx context.Context, urlString, dest string) error { - chunkSize, err := humanize.ParseBytes(viper.GetString(config.OptMinimumChunkSize)) + chunkSize, err := humanize.ParseBytes(viper.GetString(config.OptChunkSize)) if err != nil { return fmt.Errorf("error parsing chunk size: %w", err) } - minChunkSize, err := humanize.ParseBytes(viper.GetString(config.OptMinimumChunkSize)) - if err != nil { - return fmt.Errorf("error parsing minimum chunk size: %w", err) - } - if chunkSize < minChunkSize { - // minChunkSize is deprecated but we should still respect it - chunkSize = minChunkSize - } resolveOverrides, err := config.ResolveOverridesToMap(viper.GetStringSlice(config.OptResolve)) if err != nil { diff --git a/pkg/config/config.go b/pkg/config/config.go index 57d3f53..b4d36c8 100644 --- a/pkg/config/config.go +++ b/pkg/config/config.go @@ -78,16 +78,6 @@ func DeprecateFlags(cmd *cobra.Command, deprecations ...DeprecatedFlag) error { return nil } -func AddFlagAlias(cmd *cobra.Command, alias, flag string) error { - f := cmd.Flag(flag) - if f == nil { - return fmt.Errorf("flag %s does not exist", flag) - } - - viper.RegisterAlias(alias, flag) - return nil -} - func ViperInit() { viper.SetEnvPrefix(viperEnvPrefix) viper.SetEnvKeyReplacer(strings.NewReplacer("-", "_"))