Fixed usage and progress display

This commit is contained in:
Leon Mika 2025-01-04 14:21:10 +11:00
parent 12f82e106e
commit 3403deb463
5 changed files with 123 additions and 26 deletions

6
.idea/vcs.xml Normal file
View file

@ -0,0 +1,6 @@
<?xml version="1.0" encoding="UTF-8"?>
<project version="4">
<component name="VcsDirectoryMappings">
<mapping directory="$PROJECT_DIR$" vcs="Git" />
</component>
</project>

View file

@ -17,13 +17,17 @@ type chunker struct {
gc *gokapiClient gc *gokapiClient
parallelChunks int parallelChunks int
chunkSize int chunkSize int
maxDownloads int
maxDays int
} }
func newChunker(gc *gokapiClient, parallelChunks, chunkSize int) *chunker { func newChunker(gc *gokapiClient, config Config) *chunker {
return &chunker{ return &chunker{
gc: gc, gc: gc,
parallelChunks: parallelChunks, parallelChunks: config.ParallelChunks,
chunkSize: chunkSize, chunkSize: config.ChunkSize,
maxDownloads: config.MaxDownloads,
maxDays: config.MaxDays,
} }
} }
@ -45,8 +49,8 @@ func (c *chunker) UploadFile(ctx context.Context, filename string, progress func
filename: fname, filename: fname,
totalSize: fstat.Size(), totalSize: fstat.Size(),
contentType: mime.TypeByExtension(filepath.Ext(fname)), contentType: mime.TypeByExtension(filepath.Ext(fname)),
allowedDownloads: 5, allowedDownloads: c.maxDownloads,
expiryDays: 7, expiryDays: c.maxDays,
password: "", password: "",
} }
@ -61,6 +65,10 @@ func (c *chunker) upload(ctx context.Context, fi uploadInfo, r io.ReaderAt, prog
} }
chunks := int(fi.totalSize/int64(c.chunkSize) + 1) chunks := int(fi.totalSize/int64(c.chunkSize) + 1)
uploaders := c.parallelChunks
if chunks < uploaders {
uploaders = 1
}
chunkUploaded := make(chan uploadedChunk) chunkUploaded := make(chan uploadedChunk)
doneChunkReport := make(chan struct{}) doneChunkReport := make(chan struct{})
@ -91,7 +99,7 @@ func (c *chunker) upload(ctx context.Context, fi uploadInfo, r io.ReaderAt, prog
}() }()
errGroup, egctx := errgroup.WithContext(ctx) errGroup, egctx := errgroup.WithContext(ctx)
errGroup.SetLimit(c.parallelChunks) errGroup.SetLimit(uploaders)
for i := 0; i < chunks; i++ { for i := 0; i < chunks; i++ {
errGroup.Go(func() error { errGroup.Go(func() error {

View file

@ -5,4 +5,6 @@ type Config struct {
APIKey string APIKey string
ParallelChunks int ParallelChunks int
ChunkSize int ChunkSize int
MaxDownloads int
MaxDays int
} }

64
display.go Normal file
View file

@ -0,0 +1,64 @@
package main
import (
"flag"
"fmt"
"github.com/schollz/progressbar/v3"
"os"
"path/filepath"
"strings"
)
var (
ansiMoveUp = []byte{keyEscape, '[', '1', 'A'}
ansiClearLine = []byte{keyEscape, '[', '2', 'K'}
ansiLineFeed = []byte{keyEscape, '[', '1', 'G'}
)
func envVarUsage() {
fmt.Fprintf(flag.CommandLine.Output(), " %-20s Gokapi URL (e.g. https://gokapi.example.com/)\n", envVarURL)
fmt.Fprintf(flag.CommandLine.Output(), " %-20s Gokapi API key\n", envVarAPIKey)
}
func init() {
flag.Usage = func() {
fmt.Fprintf(flag.CommandLine.Output(), "Usage: %s [FLAGS] FILES...\n", os.Args[0])
flag.PrintDefaults()
fmt.Fprintf(flag.CommandLine.Output(), "Environment variables:\n")
envVarUsage()
}
}
func progressBarReport(filename string, maxFileLen int) func(cr ChunkReport) {
var pr *progressbar.ProgressBar
displayedFileName := formatFilename(filename, maxFileLen)
return func(cr ChunkReport) {
if cr.UploadedChunks == 0 && pr == nil {
pr = progressbar.DefaultBytes(cr.TotalSize, displayedFileName)
}
pr.Set(int(cr.UploadedBytes))
}
}
func displayCompletion(filename string, maxFileLen int, ur UploadResponse) {
os.Stdout.Write(ansiMoveUp)
os.Stdout.Write(ansiClearLine)
os.Stdout.Write(ansiLineFeed)
displayedFileName := formatFilename(filename, maxFileLen)
fmt.Printf("%v: done - %v\n", displayedFileName, ur.FileInfo.UrlHotlink)
}
func displayFailedCompletion(filename string, maxFileLen int, err error) {
os.Stdout.Write(ansiClearLine)
os.Stdout.Write(ansiLineFeed)
displayedFileName := formatFilename(filename, maxFileLen)
fmt.Printf("%v: error - %v\n", displayedFileName, err)
}
func formatFilename(filename string, maxFileLen int) string {
baseFileName := filepath.Base(filename)
return strings.Repeat(" ", maxFileLen-len(baseFileName)) + baseFileName
}

57
main.go
View file

@ -3,21 +3,43 @@ package main
import ( import (
"context" "context"
"flag" "flag"
"github.com/schollz/progressbar/v3" "fmt"
"log" "log"
"net/url" "net/url"
"os" "os"
"path/filepath" "path/filepath"
) )
const (
envVarURL = "SEND2GOKAPI_URL"
envVarAPIKey = "SEND2GOKAPI_API_KEY"
keyEscape = 27
)
func main() { func main() {
flagDownloads := flag.Int("d", 5, "Max number of downloads")
flagDays := flag.Int("t", 7, "Amount of time to keep file, in days")
flag.Parse()
if flag.NArg() < 1 {
flag.Usage()
os.Exit(2)
}
config := Config{ config := Config{
Hostname: os.Getenv("GOKAPI_HOSTNAME"), Hostname: os.Getenv(envVarURL),
APIKey: os.Getenv("GOKAPI_API_KEY"), APIKey: os.Getenv(envVarAPIKey),
ParallelChunks: 4, ParallelChunks: 4,
ChunkSize: 1024 * 100, ChunkSize: 1024 * 100,
MaxDownloads: max(*flagDownloads, 0),
MaxDays: max(*flagDays, 0),
}
if config.Hostname == "" || config.APIKey == "" {
fmt.Fprintf(flag.CommandLine.Output(), "Please set the following environment variables:\n")
envVarUsage()
os.Exit(2)
} }
flag.Parse()
hostUrl, err := url.Parse(config.Hostname) hostUrl, err := url.Parse(config.Hostname)
if err != nil { if err != nil {
@ -27,24 +49,19 @@ func main() {
ctx := context.Background() ctx := context.Background()
client := newGokapiClient(hostUrl, config.APIKey) client := newGokapiClient(hostUrl, config.APIKey)
cnkr := newChunker(client, config.ParallelChunks, config.ChunkSize) cnkr := newChunker(client, config)
maxFileLen := 0
for _, file := range flag.Args() {
maxFileLen = max(maxFileLen, len(filepath.Base(file)))
}
for _, file := range flag.Args() { for _, file := range flag.Args() {
_, err := cnkr.UploadFile(ctx, file, progressBarReport(file)) ur, err := cnkr.UploadFile(ctx, file, progressBarReport(file, maxFileLen))
if err != nil { if err == nil {
log.Fatal(err) displayCompletion(file, maxFileLen, ur)
} else {
displayFailedCompletion(file, maxFileLen, err)
} }
} }
}
func progressBarReport(filename string) func(cr ChunkReport) {
var pr *progressbar.ProgressBar
return func(cr ChunkReport) {
if cr.UploadedChunks == 0 && pr == nil {
pr = progressbar.DefaultBytes(cr.TotalSize, filepath.Base(filename))
}
pr.Set(int(cr.UploadedBytes))
}
} }