sq/cli/cmd_x.go
Neil O'Toole 452a91547c
Better handling of download cache refresh ()
* A failed download cache refresh no longer destroys the previous cache.
2024-01-26 15:18:38 -07:00

222 lines
5.2 KiB
Go

package cli
import (
"bufio"
"fmt"
"net/url"
"os"
"time"
"github.com/spf13/cobra"
"github.com/neilotoole/sq/cli/flag"
"github.com/neilotoole/sq/cli/run"
"github.com/neilotoole/sq/libsq/core/errz"
"github.com/neilotoole/sq/libsq/core/ioz/checksum"
"github.com/neilotoole/sq/libsq/core/ioz/httpz"
"github.com/neilotoole/sq/libsq/core/lg"
"github.com/neilotoole/sq/libsq/core/progress"
"github.com/neilotoole/sq/libsq/files"
"github.com/neilotoole/sq/libsq/files/downloader"
"github.com/neilotoole/sq/libsq/source"
)
// newXCmd returns the root "x" command, which is the container
// for a set of hidden commands that are useful for development.
// The x commands are not intended for end users.
func newXCmd() *cobra.Command {
cmd := &cobra.Command{
Use: "x",
Short: "Run hidden dev/test commands",
Hidden: true,
}
return cmd
}
func newXLockSrcCmd() *cobra.Command {
cmd := &cobra.Command{
Use: "lock-src-cache @src",
Short: "Test source cache locking",
Hidden: true,
Args: cobra.ExactArgs(1),
ValidArgsFunction: completeHandle(1),
RunE: execXLockSrcCmd,
Example: ` $ sq x lock-src-cache @sakila`,
}
return cmd
}
func execXLockSrcCmd(cmd *cobra.Command, args []string) error {
ctx := cmd.Context()
ru := run.FromContext(ctx)
src, err := ru.Config.Collection.Get(args[0])
if err != nil {
return err
}
timeout := time.Minute * 20
ru.Config.Options[files.OptCacheLockTimeout.Key()] = timeout
unlock, err := ru.Files.CacheLockAcquire(ctx, src)
if err != nil {
return err
}
fmt.Fprintf(ru.Out, "Cache lock acquired for %s\n", src.Handle)
select {
case <-pressEnter():
fmt.Fprintln(ru.Out, "\nENTER received, releasing lock")
case <-ctx.Done():
fmt.Fprintln(ru.Out, "\nContext done, releasing lock")
}
fmt.Fprintf(ru.Out, "Releasing cache lock for %s\n", src.Handle)
unlock()
fmt.Fprintf(ru.Out, "Cache lock released for %s\n", src.Handle)
return nil
}
func newXProgressCmd() *cobra.Command {
cmd := &cobra.Command{
Use: "progress",
Short: "Execute progress test code",
Hidden: true,
RunE: execXProgress,
Example: ` $ sq x progress`,
}
return cmd
}
func execXProgress(cmd *cobra.Command, _ []string) error {
ctx := cmd.Context()
log := lg.FromContext(ctx)
ru := run.FromContext(ctx)
d := time.Second * 5
pb := progress.FromContext(ctx)
bar := pb.NewTimeoutWaiter("Locking @sakila", time.Now().Add(d))
defer bar.Stop()
select {
case <-pressEnter():
bar.Stop()
pb.Stop()
fmt.Fprintln(ru.Out, "\nENTER received")
case <-ctx.Done():
bar.Stop()
pb.Stop()
fmt.Fprintln(ru.Out, "Context done")
case <-time.After(d + time.Second*5):
bar.Stop()
log.Warn("timed out, about to print something")
fmt.Fprintln(ru.Out, "Really timed out")
log.Warn("done printing")
}
fmt.Fprintln(ru.Out, "exiting")
return ctx.Err()
}
func newXDownloadCmd() *cobra.Command {
cmd := &cobra.Command{
Use: "download URL",
Short: "Download a file",
Hidden: true,
Args: cobra.ExactArgs(1),
RunE: execXDownloadCmd,
Example: ` $ sq x download https://sq.io/testdata/actor.csv
# Download a big-ass file
$ sq x download https://sqio-public.s3.amazonaws.com/testdata/payment-large.gen.csv
`,
}
cmd.Flags().BoolP(flag.JSON, flag.JSONShort, false, flag.JSONUsage)
return cmd
}
func execXDownloadCmd(cmd *cobra.Command, args []string) error {
ctx := cmd.Context()
log := lg.FromContext(ctx)
ru := run.FromContext(ctx)
u, err := url.ParseRequestURI(args[0])
if err != nil {
return err
}
sum := checksum.Sum([]byte(u.String()))
fakeSrc := &source.Source{Handle: "@download_" + sum}
cacheDir, err := ru.Files.CacheDirFor(fakeSrc)
if err != nil {
return err
}
c := httpz.NewClient(
httpz.DefaultUserAgent,
httpz.OptResponseTimeout(time.Second*15),
httpz.OptRequestDelay(time.Second*5),
)
dl, err := downloader.New(fakeSrc.Handle, c, u.String(), cacheDir)
if err != nil {
return err
}
h := downloader.NewSinkHandler(log.With("origin", "handler"))
dl.Get(ctx, h.Handler)
switch {
case len(h.Errors) > 0:
err1 := errz.Err(h.Errors[0])
return err1
case len(h.Downloaded) > 0:
fmt.Fprintf(ru.Out, "Cached: %s\n", h.Downloaded[0])
return nil
case len(h.Streams) > 0:
fmt.Fprintf(ru.Out, "Uncached: %d bytes\n", h.Streams[0].Size())
}
return nil
}
func pressEnter() <-chan struct{} {
done := make(chan struct{})
go func() {
buf := bufio.NewReader(os.Stdin)
fmt.Fprintf(os.Stdout, "\nPress [ENTER] to continue\n\n > ")
_, _ = buf.ReadBytes('\n')
close(done)
}()
return done
}
func newXLockConfigCmd() *cobra.Command {
cmd := &cobra.Command{
Use: "lock-config",
Short: "Test config lock",
Hidden: true,
Args: cobra.NoArgs,
ValidArgsFunction: completeHandle(1),
RunE: func(cmd *cobra.Command, args []string) error {
ru := run.FromContext(cmd.Context())
fmt.Fprintf(ru.Out, "Locking config (pid %d)\n", os.Getpid())
unlock, err := lockReloadConfig(cmd)
if err != nil {
return err
}
fmt.Fprintln(ru.Out, "Config locked; ctrl-c to exit")
<-cmd.Context().Done()
unlock()
return nil
},
Example: ` $ sq x lock-config`,
}
return cmd
}