From adfe3f5771f933bcb6012d24aa53a436370be677 Mon Sep 17 00:00:00 2001
From: Philipp Wolfer <ph.wolfer@gmail.com>
Date: Thu, 22 May 2025 08:52:52 +0200
Subject: [PATCH] Use the transfer context also for the progress bars

---
 internal/cli/progress.go |  6 ++++--
 internal/cli/transfer.go | 14 ++++++++++----
 2 files changed, 14 insertions(+), 6 deletions(-)

diff --git a/internal/cli/progress.go b/internal/cli/progress.go
index e93ec18..db862a1 100644
--- a/internal/cli/progress.go
+++ b/internal/cli/progress.go
@@ -18,6 +18,7 @@ Scotty. If not, see <https://www.gnu.org/licenses/>.
 package cli
 
 import (
+	"context"
 	"sync"
 	"time"
 
@@ -39,9 +40,10 @@ type progressBarUpdater struct {
 	importedItems    int
 }
 
-func setupProgressBars(updateChan chan models.TransferProgress) progressBarUpdater {
+func setupProgressBars(ctx context.Context, updateChan chan models.TransferProgress) progressBarUpdater {
 	wg := &sync.WaitGroup{}
-	p := mpb.New(
+	p := mpb.NewWithContext(
+		ctx,
 		mpb.WithWaitGroup(wg),
 		mpb.WithOutput(color.Output),
 		// mpb.WithWidth(64),
diff --git a/internal/cli/transfer.go b/internal/cli/transfer.go
index 79be3f0..3aabb4b 100644
--- a/internal/cli/transfer.go
+++ b/internal/cli/transfer.go
@@ -110,11 +110,13 @@ func (c *TransferCmd[E, I, R]) Transfer(exp backends.ExportProcessor[R], imp bac
 	}
 	printTimestamp("From timestamp: %v (%v)", timestamp)
 
+	// Use a context with cancel to abort the transfer
+	ctx, cancel := context.WithCancel(context.Background())
+
 	// Prepare progress bars
 	progressChan := make(chan models.TransferProgress)
-	progress := setupProgressBars(progressChan)
+	progress := setupProgressBars(ctx, progressChan)
 
-	ctx, cancel := context.WithCancel(context.Background())
 	wg := &sync.WaitGroup{}
 
 	// Export from source
@@ -126,8 +128,12 @@ func (c *TransferCmd[E, I, R]) Transfer(exp backends.ExportProcessor[R], imp bac
 	go imp.Process(ctx, wg, exportChan, resultChan, progressChan)
 	result := <-resultChan
 
-	// Once import is done, the context can be cancelled
-	cancel()
+	// If the import has errored, the context can be cancelled immediately
+	if result.Error != nil {
+		cancel()
+	} else {
+		defer cancel()
+	}
 
 	// Wait for all goroutines to finish
 	wg.Wait()