support canceling long-running backup/restore operations

authzed · Nov 6, 2023 · 13089c8 · 13089c8
1 parent 7225cfb
commit 13089c8
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 0 deletions.
diff --git a/internal/cmd/backup.go b/internal/cmd/backup.go
@@ -94,6 +94,10 @@ func backupCmdFunc(cmd *cobra.Command, args []string) error {
 
 	var stored uint
 	for {
+		if err := ctx.Err(); err != nil {
+			return fmt.Errorf("aborted backup: %w", err)
+		}
+
 		relsResp, err := relationshipStream.Recv()
 		if err != nil {
 			if !errors.Is(err, io.EOF) {

diff --git a/internal/cmd/restore.go b/internal/cmd/restore.go
@@ -117,6 +117,10 @@ func restoreCmdFunc(cmd *cobra.Command, args []string) error {
 	var written uint64
 	var batchesWritten int
 	for rel, err := decoder.Next(); rel != nil && err == nil; rel, err = decoder.Next() {
+		if err := ctx.Err(); err != nil {
+			return fmt.Errorf("aborted restore: %w", err)
+		}
+
 		batch = append(batch, rel)
 
 		if len(batch)%batchSize == 0 {