From 5750f8712d3f14c6eb44096472549e9040d3e48c Mon Sep 17 00:00:00 2001
From: Omar Ramos <omar@planetscale.com>
Date: Fri, 25 Oct 2024 14:19:00 -0700
Subject: [PATCH] Added context to longer-running functions for quicker
 cancellation

Currently, when using either the `pscale database dump` or `pscale database restore-dump` commands there may be times where you would like to cancel out of the operation because it is taking too long, a mistake was made, etc.

At the moment the behavior for both commands is not ideal in these cases since you can attempt to send a `Ctrl+C` command but it won't be very responsive / immediate as you expect it to be.

Instead, depending on the length of time the operation has remaining, it may actually continue to completion and/or the renewal of the temporary 5 minute credential won't occur and it will stop / error then.

There is likely a more comprehensive approach to using contexts properly here that my limited Go knowledge is unaware of, but it seemed like passing in the current context into the longer-running functions that do the bulk of the data exporting / loading was enough to make `Ctrl+C` responsive enough for my expectations.
---
 internal/dumper/dumper.go | 9 +++++++--
 internal/dumper/loader.go | 9 +++++++--
 2 files changed, 14 insertions(+), 4 deletions(-)

diff --git a/internal/dumper/dumper.go b/internal/dumper/dumper.go
index d651d61e..ad04c378 100644
--- a/internal/dumper/dumper.go
+++ b/internal/dumper/dumper.go
@@ -157,7 +157,7 @@ func (d *Dumper) Run(ctx context.Context) error {
 					zap.Int("thread_conn_id", conn.ID),
 				)
 
-				err := d.dumpTable(conn, database, table, d.cfg.UseReplica, d.cfg.UseRdonly)
+				err := d.dumpTable(ctx, conn, database, table, d.cfg.UseReplica, d.cfg.UseRdonly)
 				if err != nil {
 					d.log.Error("error dumping table", zap.Error(err))
 				}
@@ -219,7 +219,7 @@ func (d *Dumper) dumpTableSchema(conn *Connection, database string, table string
 }
 
 // Dump a table in "MySQL" (multi-inserts) format
-func (d *Dumper) dumpTable(conn *Connection, database string, table string, useReplica, useRdonly bool) error {
+func (d *Dumper) dumpTable(ctx context.Context, conn *Connection, database string, table string, useReplica, useRdonly bool) error {
 	var allBytes uint64
 	var allRows uint64
 	var where string
@@ -278,6 +278,11 @@ func (d *Dumper) dumpTable(conn *Connection, database string, table string, useR
 			return err
 		}
 
+		// Allows for quicker exit when using Ctrl+C at the Terminal:
+		if ctx.Err() != nil {
+			return ctx.Err()
+		}
+
 		values := make([]string, 0, 16)
 		for _, v := range row {
 			if v.Raw() == nil {
diff --git a/internal/dumper/loader.go b/internal/dumper/loader.go
index 3298dcc1..66939c8c 100644
--- a/internal/dumper/loader.go
+++ b/internal/dumper/loader.go
@@ -85,7 +85,7 @@ func (l *Loader) Run(ctx context.Context) error {
 		eg.Go(func() error {
 			defer pool.Put(conn)
 
-			r, err := l.restoreTable(table, conn)
+			r, err := l.restoreTable(ctx, table, conn)
 			if err != nil {
 				return err
 			}
@@ -234,7 +234,7 @@ func (l *Loader) restoreTableSchema(overwrite bool, tables []string, conn *Conne
 	return nil
 }
 
-func (l *Loader) restoreTable(table string, conn *Connection) (int, error) {
+func (l *Loader) restoreTable(ctx context.Context, table string, conn *Connection) (int, error) {
 	bytes := 0
 	part := "0"
 	base := filepath.Base(table)
@@ -277,6 +277,11 @@ func (l *Loader) restoreTable(table string, conn *Connection) (int, error) {
 	querys := strings.Split(query1, ";\n")
 	bytes = len(query1)
 	for _, query := range querys {
+		// Allows for quicker exit when using Ctrl+C at the Terminal:
+		if ctx.Err() != nil {
+			return 0, ctx.Err()
+		}
+
 		if !strings.HasPrefix(query, "/*") && query != "" {
 			err = conn.Execute(query)
 			if err != nil {