feat: refactor upload command and implement progress tracking

2026-05-11 05:19:40 +08:00 · 2025-12-19 14:20:00 +08:00
parent 3f9c3f2a73
commit 29d523bd4f
7 changed files with 341 additions and 6 deletions
--- a/cmd/root.go
+++ b/cmd/root.go
@@ -4,6 +4,7 @@ import (
 	"context"
 	"fmt"

+	"github.com/krau/SaveAny-Bot/cmd/upload"
 	"github.com/krau/SaveAny-Bot/config"
 	"github.com/spf13/cobra"
 )
@@ -16,6 +17,7 @@ var rootCmd = &cobra.Command{

 func init() {
 	config.RegisterFlags(rootCmd)
+	upload.Register(rootCmd)
 }

 func Execute(ctx context.Context) {
--- a/cmd/upload/cmd.go
+++ b/cmd/upload/cmd.go
@@ -1,8 +1,9 @@
-package cmd
+package upload

 import (
 	"context"
 	"fmt"
+	"io"
 	"os"
 	"path"
 	"path/filepath"
@@ -10,6 +11,7 @@ import (
 	"github.com/charmbracelet/log"
 	"github.com/krau/SaveAny-Bot/client/bot"
 	"github.com/krau/SaveAny-Bot/common/cache"
+	"github.com/krau/SaveAny-Bot/common/utils/ioutil"
 	"github.com/krau/SaveAny-Bot/common/utils/tgutil"
 	"github.com/krau/SaveAny-Bot/config"
 	"github.com/krau/SaveAny-Bot/database"
@@ -25,13 +27,14 @@ var uploadCmd = &cobra.Command{
 	RunE:  Upload,
 }

-func init() {
+func Register(root *cobra.Command) {
 	uploadCmd.Flags().StringP("file", "f", "", "file path to upload")
 	uploadCmd.MarkFlagRequired("file")
 	uploadCmd.Flags().StringP("storage", "s", "", "storage name to upload to")
 	uploadCmd.MarkFlagRequired("storage")
 	uploadCmd.Flags().StringP("dir", "d", "", "storage dir to upload to, default is the base_path of the storage")
-	rootCmd.AddCommand(uploadCmd)
+	uploadCmd.Flags().Bool("no-progress", false, "disable progress bar")
+	root.AddCommand(uploadCmd)
 }

 func Upload(cmd *cobra.Command, args []string) error {
@@ -47,6 +50,10 @@ func Upload(cmd *cobra.Command, args []string) error {
 	if err != nil {
 		return err
 	}
+	noProgress, err := cmd.Flags().GetBool("no-progress")
+	if err != nil {
+		return err
+	}

 	ctx := cmd.Context()
 	log := log.FromContext(ctx)
@@ -86,10 +93,36 @@ func Upload(cmd *cobra.Command, args []string) error {
 	ctx = context.WithValue(ctx, ctxkey.ContentLength, fileSize)
 	ctx = tgutil.ExtWithContext(ctx, bot.ExtContext())

+	// Create progress reader and UI
+	var reader io.Reader
+	var progressUI *UploadProgress
 	log.Info("Uploading file...", "file", fp, "to", storname, "as", uploadPath)
-	if err := stor.Save(ctx, file, uploadPath); err != nil {
+
+	if !noProgress && fileSize > 0 {
+		progressUI = NewUploadProgress(ctx, fileName, fileSize)
+		progressUI.Start()
+
+		reader = ioutil.NewProgressReader(file, fileSize, func(read int64, total int64) {
+			if total > 0 {
+				progressUI.UpdateProgress(float64(read) / float64(total))
+			}
+		})
+	} else {
+		reader = file
+	}
+
+	if err := stor.Save(ctx, reader, uploadPath); err != nil {
+		if progressUI != nil {
+			progressUI.SetError(err)
+			progressUI.Wait()
+		}
 		log.Fatal("Failed to upload file", "error", err)
 	}
+
+	if progressUI != nil {
+		progressUI.Done()
+		progressUI.Wait()
+	}
 	log.Info("File uploaded successfully")
 	return nil
 }
--- a/cmd/upload/progress_stub.go
+++ b/cmd/upload/progress_stub.go
@@ -0,0 +1,35 @@
+//go:build no_bubbletea
+
+package upload
+
+import "context"
+
+type uploadModel struct {
+}
+
+// UploadProgress manages the progress UI for uploads
+type UploadProgress struct {
+}
+
+// NewUploadProgress creates a new upload progress tracker
+func NewUploadProgress(ctx context.Context, fileName string, fileSize int64) *UploadProgress {
+	return &UploadProgress{}
+}
+
+// Start starts the progress UI in a goroutine and returns immediately
+func (up *UploadProgress) Start() {}
+
+// UpdateProgress updates the progress bar with a new percentage (0.0 - 1.0)
+func (up *UploadProgress) UpdateProgress(percent float64) {}
+
+// SetError sets an error and quits the progress UI
+func (up *UploadProgress) SetError(err error) {}
+
+// Done signals that the upload is complete
+func (up *UploadProgress) Done() {}
+
+// Wait waits for the progress UI to finish
+func (up *UploadProgress) Wait() {}
+
+// Quit quits the progress UI
+func (up *UploadProgress) Quit() {}
--- a/cmd/upload/progress_tea.go
+++ b/cmd/upload/progress_tea.go
@@ -0,0 +1,178 @@
+//go:build !no_bubbletea
+
+package upload
+
+import (
+	"context"
+	"fmt"
+	"strings"
+
+	"github.com/charmbracelet/bubbles/progress"
+	tea "github.com/charmbracelet/bubbletea"
+	"github.com/charmbracelet/lipgloss"
+	"github.com/dustin/go-humanize"
+)
+
+var (
+	helpStyle = lipgloss.NewStyle().Foreground(lipgloss.Color("#626262"))
+)
+
+// progressMsg is sent to update the progress bar
+type progressMsg float64
+
+// progressErrMsg is sent when an error occurs
+type progressErrMsg struct{ err error }
+
+// progressDoneMsg is sent when the upload is complete
+type progressDoneMsg struct{}
+
+// uploadModel is the bubbletea model for the upload progress UI
+type uploadModel struct {
+	progress  progress.Model
+	fileName  string
+	fileSize  int64
+	bytesRead int64
+	err       error
+	done      bool
+	quitting  bool
+	width     int
+}
+
+func newUploadModel(fileName string, fileSize int64) uploadModel {
+	p := progress.New(
+		progress.WithDefaultGradient(),
+		progress.WithWidth(50),
+	)
+	return uploadModel{
+		progress: p,
+		fileName: fileName,
+		fileSize: fileSize,
+	}
+}
+
+func (m uploadModel) Init() tea.Cmd {
+	return nil
+}
+
+func (m uploadModel) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
+	switch msg := msg.(type) {
+	case tea.WindowSizeMsg:
+		m.width = msg.Width
+		m.progress.Width = min(msg.Width-10, 80)
+		return m, nil
+
+	case progressMsg:
+		var cmds []tea.Cmd
+		percent := float64(msg)
+		m.bytesRead = int64(percent * float64(m.fileSize))
+
+		cmds = append(cmds, m.progress.SetPercent(percent))
+		return m, tea.Batch(cmds...)
+
+	case progressErrMsg:
+		m.err = msg.err
+		return m, tea.Quit
+
+	case progressDoneMsg:
+		m.done = true
+		m.progress.SetPercent(1.0)
+		return m, tea.Quit
+
+	case progress.FrameMsg:
+		// Don't process frame messages if we're done or quitting
+		if m.done || m.quitting {
+			return m, nil
+		}
+		progressModel, cmd := m.progress.Update(msg)
+		m.progress = progressModel.(progress.Model)
+		return m, cmd
+	}
+
+	return m, nil
+}
+
+func (m uploadModel) View() string {
+	if m.err != nil {
+		return fmt.Sprintf("\n  ❌ Error: %s\n\n", m.err.Error())
+	}
+
+	var sb strings.Builder
+	sb.WriteString("\n")
+
+	// File info
+	sb.WriteString(fmt.Sprintf("  📁 %s\n", m.fileName))
+	sb.WriteString(fmt.Sprintf("  📊 %s / %s\n\n",
+		humanize.Bytes(uint64(m.bytesRead)),
+		humanize.Bytes(uint64(m.fileSize)),
+	))
+
+	// Progress bar
+	sb.WriteString("  ")
+	sb.WriteString(m.progress.View())
+	sb.WriteString("\n\n")
+
+	if m.done {
+		sb.WriteString("  √ Upload complete!\n\n")
+	} else {
+		sb.WriteString(helpStyle.Render("  Press Ctrl+C to cancel"))
+		sb.WriteString("\n\n")
+	}
+
+	return sb.String()
+}
+
+// UploadProgress manages the progress UI for uploads
+type UploadProgress struct {
+	program *tea.Program
+	ctx     context.Context
+	cancel  context.CancelFunc
+}
+
+// NewUploadProgress creates a new upload progress tracker
+func NewUploadProgress(ctx context.Context, fileName string, fileSize int64) *UploadProgress {
+	model := newUploadModel(fileName, fileSize)
+	ctx, cancel := context.WithCancel(ctx)
+	p := tea.NewProgram(
+		model,
+		tea.WithoutSignalHandler(),
+		tea.WithContext(ctx),
+		tea.WithInput(nil), // Disable keyboard input, rely on context cancellation
+	)
+	return &UploadProgress{
+		program: p,
+		ctx:     ctx,
+		cancel:  cancel,
+	}
+}
+
+// Start starts the progress UI in a goroutine and returns immediately
+func (up *UploadProgress) Start() {
+	go func() {
+		up.program.Run()
+	}()
+}
+
+// UpdateProgress updates the progress bar with a new percentage (0.0 - 1.0)
+func (up *UploadProgress) UpdateProgress(percent float64) {
+	up.program.Send(progressMsg(percent))
+}
+
+// SetError sets an error and quits the progress UI
+func (up *UploadProgress) SetError(err error) {
+	up.program.Send(progressErrMsg{err: err})
+}
+
+// Done signals that the upload is complete
+func (up *UploadProgress) Done() {
+	up.program.Send(progressDoneMsg{})
+}
+
+// Wait waits for the progress UI to finish
+func (up *UploadProgress) Wait() {
+	up.program.Wait()
+}
+
+// Quit quits the progress UI
+func (up *UploadProgress) Quit() {
+	up.program.Quit()
+}
--- a/common/utils/ioutil/progress_reader.go
+++ b/common/utils/ioutil/progress_reader.go
@@ -0,0 +1,65 @@
+package ioutil
+
+import (
+	"io"
+	"sync/atomic"
+)
+
+var _ io.ReadSeeker = (*ProgressReadSeeker)(nil)
+
+// ProgressReadSeeker wraps an io.ReadSeeker and tracks read progress
+type ProgressReadSeeker struct {
+	reader     io.ReadSeeker
+	total      atomic.Int64
+	read       atomic.Int64
+	onProgress func(read int64, total int64)
+}
+
+// Seek implements io.ReadSeeker.
+func (pr *ProgressReadSeeker) Seek(offset int64, whence int) (int64, error) {
+	return pr.reader.Seek(offset, whence)
+}
+
+// NewProgressReader creates a new ProgressReader
+func NewProgressReader(rs io.ReadSeeker, total int64, onProgress func(read int64, total int64)) *ProgressReadSeeker {
+	prs := &ProgressReadSeeker{
+		reader:     rs,
+		total:      atomic.Int64{},
+		read:       atomic.Int64{},
+		onProgress: onProgress,
+	}
+	prs.total.Store(total)
+	return prs
+}
+
+// Read implements io.Reader
+func (pr *ProgressReadSeeker) Read(p []byte) (int, error) {
+	n, err := pr.reader.Read(p)
+	if n > 0 {
+		pr.read.Add(int64(n))
+		read := pr.read.Load()
+
+		if pr.onProgress != nil {
+			pr.onProgress(read, pr.total.Load())
+		}
+	}
+	return n, err
+}
+
+// Progress returns the current progress as a float64 between 0 and 1
+func (pr *ProgressReadSeeker) Progress() float64 {
+	if pr.total.Load() <= 0 {
+		return 0
+	}
+	return float64(pr.read.Load()) / float64(pr.total.Load())
+}
+
+// Read returns the number of bytes read so far
+func (pr *ProgressReadSeeker) BytesRead() int64 {
+	return pr.read.Load()
+}
+
+// Total returns the total number of bytes
+func (pr *ProgressReadSeeker) Total() int64 {
+	return pr.total.Load()
+}
--- a/go.mod
+++ b/go.mod
@@ -6,7 +6,11 @@ require (
 	github.com/blang/semver v3.5.1+incompatible
 	github.com/celestix/gotgproto v1.0.0-beta22
 	github.com/cenkalti/backoff/v4 v4.3.0
+	github.com/charmbracelet/bubbles v0.21.0
+	github.com/charmbracelet/bubbletea v1.3.10
+	github.com/charmbracelet/lipgloss v1.1.0
 	github.com/charmbracelet/log v0.4.2
+	github.com/dustin/go-humanize v1.0.1
 	github.com/gabriel-vasile/mimetype v1.4.10
 	github.com/goccy/go-yaml v1.18.0
 	github.com/gotd/contrib v0.21.1
@@ -31,7 +35,7 @@ require (
 	github.com/aymanbagabas/go-osc52/v2 v2.0.1 // indirect
 	github.com/cespare/xxhash/v2 v2.3.0 // indirect
 	github.com/charmbracelet/colorprofile v0.3.2 // indirect
-	github.com/charmbracelet/lipgloss v1.1.0 // indirect
+	github.com/charmbracelet/harmonica v0.2.0 // indirect
 	github.com/charmbracelet/x/ansi v0.10.2 // indirect
 	github.com/charmbracelet/x/cellbuf v0.0.13 // indirect
 	github.com/charmbracelet/x/term v0.2.1 // indirect
@@ -39,7 +43,7 @@ require (
 	github.com/coder/websocket v1.8.14 // indirect
 	github.com/deckarep/golang-set/v2 v2.7.0 // indirect
 	github.com/dlclark/regexp2 v1.11.5 // indirect
-	github.com/dustin/go-humanize v1.0.1 // indirect
+	github.com/erikgeiser/coninput v0.0.0-20211004153227-1c3628e74d0f // indirect
 	github.com/fatih/color v1.18.0 // indirect
 	github.com/ghodss/yaml v1.0.0 // indirect
 	github.com/glebarez/go-sqlite v1.22.0 // indirect
@@ -67,9 +71,12 @@ require (
 	github.com/lucasb-eyer/go-colorful v1.3.0 // indirect
 	github.com/mattn/go-colorable v0.1.14 // indirect
 	github.com/mattn/go-isatty v0.0.20 // indirect
+	github.com/mattn/go-localereader v0.0.1 // indirect
 	github.com/mattn/go-runewidth v0.0.19 // indirect
 	github.com/minio/crc64nvme v1.1.1 // indirect
 	github.com/minio/md5-simd v1.1.2 // indirect
+	github.com/muesli/ansi v0.0.0-20230316100256-276c6243b2f6 // indirect
+	github.com/muesli/cancelreader v0.2.2 // indirect
 	github.com/muesli/termenv v0.16.0 // indirect
 	github.com/ncruces/go-strftime v1.0.0 // indirect
 	github.com/ncruces/julianday v1.0.0 // indirect
--- a/go.sum
+++ b/go.sum
@@ -42,8 +42,14 @@ github.com/cespare/xxhash/v2 v2.3.0 h1:UL815xU9SqsFlibzuggzjXhog7bL6oX9BbNZnL2UF
 github.com/cespare/xxhash/v2 v2.3.0/go.mod h1:VGX0DQ3Q6kWi7AoAeZDth3/j3BFtOZR5XLFGgcrjCOs=
 github.com/cevatbarisyilmaz/ara v0.0.4 h1:SGH10hXpBJhhTlObuZzTuFn1rrdmjQImITXnZVPSodc=
 github.com/cevatbarisyilmaz/ara v0.0.4/go.mod h1:BfFOxnUd6Mj6xmcvRxHN3Sr21Z1T3U2MYkYOmoQe4Ts=
+github.com/charmbracelet/bubbles v0.21.0 h1:9TdC97SdRVg/1aaXNVWfFH3nnLAwOXr8Fn6u6mfQdFs=
+github.com/charmbracelet/bubbles v0.21.0/go.mod h1:HF+v6QUR4HkEpz62dx7ym2xc71/KBHg+zKwJtMw+qtg=
+github.com/charmbracelet/bubbletea v1.3.10 h1:otUDHWMMzQSB0Pkc87rm691KZ3SWa4KUlvF9nRvCICw=
+github.com/charmbracelet/bubbletea v1.3.10/go.mod h1:ORQfo0fk8U+po9VaNvnV95UPWA1BitP1E0N6xJPlHr4=
 github.com/charmbracelet/colorprofile v0.3.2 h1:9J27WdztfJQVAQKX2WOlSSRB+5gaKqqITmrvb1uTIiI=
 github.com/charmbracelet/colorprofile v0.3.2/go.mod h1:mTD5XzNeWHj8oqHb+S1bssQb7vIHbepiebQ2kPKVKbI=
+github.com/charmbracelet/harmonica v0.2.0 h1:8NxJWRWg/bzKqqEaaeFNipOu77YR5t8aSwG4pgaUBiQ=
+github.com/charmbracelet/harmonica v0.2.0/go.mod h1:KSri/1RMQOZLbw7AHqgcBycp8pgJnQMYYT8QZRqZ1Ao=
 github.com/charmbracelet/lipgloss v1.1.0 h1:vYXsiLHVkK7fp74RkV7b2kq9+zDLoEU4MZoFqR/noCY=
 github.com/charmbracelet/lipgloss v1.1.0/go.mod h1:/6Q8FR2o+kj8rz4Dq0zQc3vYf7X+B0binUUBwA0aL30=
 github.com/charmbracelet/log v0.4.2 h1:hYt8Qj6a8yLnvR+h7MwsJv/XvmBJXiueUcI3cIxsyig=
@@ -76,6 +82,8 @@ github.com/duke-git/lancet/v2 v2.3.7 h1:nnNBA9KyoqwbPm4nFmEFVIbXeAmpqf6IDCH45+HH
 github.com/duke-git/lancet/v2 v2.3.7/go.mod h1:zGa2R4xswg6EG9I6WnyubDbFO/+A/RROxIbXcwryTsc=
 github.com/dustin/go-humanize v1.0.1 h1:GzkhY7T5VNhEkwH0PVJgjz+fX1rhBrR7pRT3mDkpeCY=
 github.com/dustin/go-humanize v1.0.1/go.mod h1:Mu1zIs6XwVuF/gI1OepvI0qD18qycQx+mFykh5fBlto=
+github.com/erikgeiser/coninput v0.0.0-20211004153227-1c3628e74d0f h1:Y/CXytFA4m6baUTXGLOoWe4PQhGxaX0KpnayAqC48p4=
+github.com/erikgeiser/coninput v0.0.0-20211004153227-1c3628e74d0f/go.mod h1:vw97MGsxSvLiUE2X8qFplwetxpGLQrlU1Q9AUEIzCaM=
 github.com/fatih/color v1.18.0 h1:S8gINlzdQ840/4pfAwic/ZE0djQEH3wM94VfqLTZcOM=
 github.com/fatih/color v1.18.0/go.mod h1:4FelSpRwEGDpQ12mAdzqdOukCy4u8WUtOY6lkT/6HfU=
 github.com/frankban/quicktest v1.14.6 h1:7Xjx+VpznH+oBnejlPUj8oUpdxnVs4f8XU8WnHkI4W8=
@@ -168,6 +176,8 @@ github.com/mattn/go-colorable v0.1.14 h1:9A9LHSqF/7dyVVX6g0U9cwm9pG3kP9gSzcuIPHP
 github.com/mattn/go-colorable v0.1.14/go.mod h1:6LmQG8QLFO4G5z1gPvYEzlUgJ2wF+stgPZH1UqBm1s8=
 github.com/mattn/go-isatty v0.0.20 h1:xfD0iDuEKnDkl03q4limB+vH+GxLEtL/jb4xVJSWWEY=
 github.com/mattn/go-isatty v0.0.20/go.mod h1:W+V8PltTTMOvKvAeJH7IuucS94S2C6jfK/D7dTCTo3Y=
+github.com/mattn/go-localereader v0.0.1 h1:ygSAOl7ZXTx4RdPYinUpg6W99U8jWvWi9Ye2JC/oIi4=
+github.com/mattn/go-localereader v0.0.1/go.mod h1:8fBrzywKY7BI3czFoHkuzRoWE9C+EiG4R1k4Cjx5p88=
 github.com/mattn/go-runewidth v0.0.19 h1:v++JhqYnZuu5jSKrk9RbgF5v4CGUjqRfBm05byFGLdw=
 github.com/mattn/go-runewidth v0.0.19/go.mod h1:XBkDxAl56ILZc9knddidhrOlY5R/pDhgLpndooCuJAs=
 github.com/minio/crc64nvme v1.1.1 h1:8dwx/Pz49suywbO+auHCBpCtlW1OfpcLN7wYgVR6wAI=
@@ -180,6 +190,10 @@ github.com/mitchellh/go-ps v1.0.0 h1:i6ampVEEF4wQFF+bkYfwYgY+F/uYJDktmvLPf7qIgjc
 github.com/mitchellh/go-ps v1.0.0/go.mod h1:J4lOc8z8yJs6vUwklHw2XEIiT4z4C40KtWVN3nvg8Pg=
 github.com/mitchellh/mapstructure v1.5.0 h1:jeMsZIYE/09sWLaz43PL7Gy6RuMjD2eJVyuac5Z2hdY=
 github.com/mitchellh/mapstructure v1.5.0/go.mod h1:bFUtVrKA4DC2yAKiSyO/QUcy7e+RRV2QTWOzhPopBRo=
+github.com/muesli/ansi v0.0.0-20230316100256-276c6243b2f6 h1:ZK8zHtRHOkbHy6Mmr5D264iyp3TiX5OmNcI5cIARiQI=
+github.com/muesli/ansi v0.0.0-20230316100256-276c6243b2f6/go.mod h1:CJlz5H+gyd6CUWT45Oy4q24RdLyn7Md9Vj2/ldJBSIo=
+github.com/muesli/cancelreader v0.2.2 h1:3I4Kt4BQjOR54NavqnDogx/MIoWBFa0StPA8ELUXHmA=
+github.com/muesli/cancelreader v0.2.2/go.mod h1:3XuTXfFS2VjM+HTLZY9Ak0l6eUKfijIfMUZ4EgX0QYo=
 github.com/muesli/termenv v0.16.0 h1:S5AlUN9dENB57rsbnkPyfdGuWIlkmzJjbFf0Tf5FWUc=
 github.com/muesli/termenv v0.16.0/go.mod h1:ZRfOIKPFDYQoDFF4Olj7/QJbW60Ol/kL1pU3VfY/Cnk=
 github.com/ncruces/go-sqlite3 v0.30.1 h1:pHC3YsyRdJv4pCMB4MO1Q2BXw/CAa+Hoj7GSaKtVk+g=
@@ -303,6 +317,7 @@ golang.org/x/sync v0.18.0/go.mod h1:9KTHXmSnoGruLpwFjVSX0lNNA75CykiMECbovNTZqGI=
 golang.org/x/sys v0.0.0-20190215142949-d0b11bdaac8a/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY=
 golang.org/x/sys v0.0.0-20201119102817-f84b799fce68/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
 golang.org/x/sys v0.0.0-20210615035016-665e8c7367d1/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
+golang.org/x/sys v0.0.0-20210809222454-d867a43fc93e/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
 golang.org/x/sys v0.0.0-20220520151302-bc2c85ada10a/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
 golang.org/x/sys v0.0.0-20220722155257-8c9f86f7a55f/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
 golang.org/x/sys v0.5.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=