From 29d523bd4f33f11a408dc6eb4ecb10b205d2f0fb Mon Sep 17 00:00:00 2001 From: krau <71133316+krau@users.noreply.github.com> Date: Fri, 19 Dec 2025 14:20:00 +0800 Subject: [PATCH] feat: refactor upload command and implement progress tracking --- cmd/root.go | 2 + cmd/{upload.go => upload/cmd.go} | 41 +++++- cmd/upload/progress_stub.go | 35 +++++ cmd/upload/progress_tea.go | 178 +++++++++++++++++++++++++ common/utils/ioutil/progress_reader.go | 65 +++++++++ go.mod | 11 +- go.sum | 15 +++ 7 files changed, 341 insertions(+), 6 deletions(-) rename cmd/{upload.go => upload/cmd.go} (72%) create mode 100644 cmd/upload/progress_stub.go create mode 100644 cmd/upload/progress_tea.go create mode 100644 common/utils/ioutil/progress_reader.go diff --git a/cmd/root.go b/cmd/root.go index 99957a3..e11f616 100644 --- 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) { diff --git a/cmd/upload.go b/cmd/upload/cmd.go similarity index 72% rename from cmd/upload.go rename to cmd/upload/cmd.go index 082173a..e6676c4 100644 --- a/cmd/upload.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 } diff --git a/cmd/upload/progress_stub.go b/cmd/upload/progress_stub.go new file mode 100644 index 0000000..c3ebf24 --- /dev/null +++ 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() {} diff --git a/cmd/upload/progress_tea.go b/cmd/upload/progress_tea.go new file mode 100644 index 0000000..8990011 --- /dev/null +++ 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() +} diff --git a/common/utils/ioutil/progress_reader.go b/common/utils/ioutil/progress_reader.go new file mode 100644 index 0000000..e58e20e --- /dev/null +++ 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() +} diff --git a/go.mod b/go.mod index 316b79b..3caffa1 100644 --- 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 diff --git a/go.sum b/go.sum index 782a34e..c193831 100644 --- 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=