From 30b5b17d040eda6b28757e8f3085278828c36c8e Mon Sep 17 00:00:00 2001
From: Aaron Son <aaron@dolthub.com>
Date: Wed, 5 Mar 2025 20:02:06 -0800
Subject: [PATCH] server/doltgres_handler.go: Ensure we run queries with the
 Context that comes back from ProcessList.BeginQuery, so that we can cancel
 their contexts through the process list.

---
 server/doltgres_handler.go | 17 ++++-------------
 1 file changed, 4 insertions(+), 13 deletions(-)

diff --git a/server/doltgres_handler.go b/server/doltgres_handler.go
index 619d878b69..82b952bc34 100644
--- a/server/doltgres_handler.go
+++ b/server/doltgres_handler.go
@@ -301,16 +301,10 @@ func (h *DoltgresHandler) doQuery(ctx context.Context, c *mysql.Conn, query stri
 	sqlCtx.GetLogger().Debugf("Starting query")
 	sqlCtx.GetLogger().Tracef("beginning execution")
 
-	oCtx := ctx
-
 	// TODO: it would be nice to put this logic in the engine, not the handler, but we don't want the process to be
 	//  marked done until we're done spooling rows over the wire
-	ctx, err = sqlCtx.ProcessList.BeginQuery(sqlCtx, query)
-	defer func() {
-		if err != nil && ctx != nil {
-			sqlCtx.ProcessList.EndQuery(sqlCtx)
-		}
-	}()
+	sqlCtx, err = sqlCtx.ProcessList.BeginQuery(sqlCtx, query)
+	defer sqlCtx.ProcessList.EndQuery(sqlCtx)
 
 	schema, rowIter, qFlags, err := queryExec(sqlCtx, query, parsed, analyzedPlan)
 	if err != nil {
@@ -341,9 +335,6 @@ func (h *DoltgresHandler) doQuery(ctx context.Context, c *mysql.Conn, query stri
 		return err
 	}
 
-	// errGroup context is now canceled
-	ctx = oCtx
-
 	sqlCtx.GetLogger().Debugf("Query finished in %d ms", time.Since(start).Milliseconds())
 
 	// processedAtLeastOneBatch means we already called callback() at least
@@ -528,7 +519,7 @@ func (h *DoltgresHandler) resultForDefaultIter(ctx *sql.Context, schema sql.Sche
 		for {
 			select {
 			case <-ctx.Done():
-				return nil
+				return context.Cause(ctx)
 			default:
 				row, err := iter.Next(ctx)
 				if err == io.EOF {
@@ -578,7 +569,7 @@ func (h *DoltgresHandler) resultForDefaultIter(ctx *sql.Context, schema sql.Sche
 
 			select {
 			case <-ctx.Done():
-				return nil
+				return context.Cause(ctx)
 			case row, ok := <-rowChan:
 				if !ok {
 					return nil