Skip to content

Commit

Permalink
Cleanup Aliasing
Browse files Browse the repository at this point in the history 
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
  • Loading branch information
@tempusfrangit
tempusfrangit committed Mar 5, 2024
1 parent fae7259 commit e71b6b2
Show file tree
Hide file tree
Showing 3 changed files with 42 additions and 34 deletions.
2 changes: 1 addition & 1 deletion cmd/multifile/multifile.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
}
Expand Down
64 changes: 41 additions & 23 deletions cmd/root/root.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,7 @@ import (
"github.com/replicate/pget/pkg/client"
"github.com/replicate/pget/pkg/config"
"github.com/replicate/pget/pkg/download"
"github.com/replicate/pget/pkg/logging"
)

const rootLongDesc = `
Expand All @@ -44,14 +45,15 @@ var concurrency int
var pidFile *cli.PIDFile
var chunkSize string

const chunkSizeDefault = "125M"

func GetCommand() *cobra.Command {
cmd := &cobra.Command{
Use: "pget [flags] <url> <dest>",
Short: "pget",
Long: rootLongDesc,
PersistentPreRunE: rootPersistentPreRunEFunc,
PersistentPostRunE: rootPersistentPostRunEFunc,
PreRun: rootCmdPreRun,
RunE: runRootCMD,
Args: validateArgs,
Example: ` pget https://example.com/file.tar ./target-dir`,
Expand All @@ -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)
Expand Down Expand Up @@ -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
}
Expand All @@ -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
}

Expand All @@ -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 <number><unit>, 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")
Expand All @@ -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
}
Expand All @@ -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
Expand All @@ -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
Expand All @@ -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 {
Expand Down
10 changes: 0 additions & 10 deletions pkg/config/config.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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("-", "_"))
Expand Down

0 comments on commit e71b6b2

Please sign in to comment.