Merge #144014 #144068

144014: sql: ensure correct cursor closing behavior r=yuzefovich a=DrewKimball

This commit adds cleanup logic so that the container for a persisted cursor is closed if executing the cursor's query resulted in an error. The closing logic is also modified to correctly roll back all cursors in the even of an error during closing.

Fixes #142114

Release note: None

144068: roachtest: fix follower-reads/survival=region r=fqazi a=fqazi

Previously, this test allocated 6 nodes: 3 in us-east1, 2 in us-west1, and 1 in europe-west2. When using survival regions, we would have 5 replicas and 2 would be allowed to be virtually anywhere without constraints. However, this proved problematic in serverless environments, as the additional regions could fall outside the user's configured regions. We later modified this logic to use the same constraints as super regions, which now forces one region to have more replicas. This patch adjusts the test so that the Europe region has one more replica to account for this logic.

Fixes: #143333
Fixes: #143199

Release note: None

Co-authored-by: Drew Kimball <[email protected]>
Co-authored-by: Faizan Qazi <[email protected]>

Author: 3 people

pkg/cmd/roachtest/tests/follower_reads.go

@@ -56,7 +56,7 @@ func registerFollowerReads(r registry.Registry) {
 				6, /* nodeCount */
 				spec.CPU(4),
 				spec.Geo(),
-				spec.GCEZones("us-east1-b,us-east1-b,us-east1-b,us-west1-b,us-west1-b,europe-west2-b"),
+				spec.GCEZones("us-east1-b,us-east1-b,us-east1-b,us-west1-b,europe-west2-b,europe-west2-b"),
 			),
 			CompatibleClouds: registry.OnlyGCE,
 			Suites:           registry.Suites(registry.Nightly),

pkg/sql/logictest/testdata/logic_test/cursor

@@ -872,3 +872,51 @@ statement ok
 RESET autocommit_before_ddl
 
 subtest end
+
+# Regression test for #142114 - close the holdable cursor container if there is
+# an error while executing the query.
+subtest regression_142114
+
+statement error pgcode 22012 pq: sql-cursor: division by zero
+DECLARE foo CURSOR WITH HOLD FOR SELECT 1 // 0;
+
+statement ok
+DECLARE curs CURSOR WITH HOLD FOR SELECT 100;
+
+statement ok
+BEGIN;
+
+statement ok
+DECLARE foo CURSOR FOR SELECT 1;
+
+statement ok
+DECLARE bar CURSOR WITH HOLD FOR SELECT 2;
+
+statement ok
+DECLARE baz CURSOR WITH HOLD FOR SELECT 1 // 0;
+
+statement ok
+DECLARE bar2 CURSOR WITH HOLD FOR SELECT 3;
+
+statement ok
+INSERT INTO a VALUES (-1, -2);
+
+statement error pgcode 22012 pq: sql-cursor: division by zero
+COMMIT;
+
+# All cursors created in the aborted transaction should be closed.
+query T rowsort
+SELECT name FROM pg_cursors;
+----
+curs
+
+# The insert should not have been committed.
+query III
+SELECT * FROM a ORDER BY a LIMIT 1;
+----
+1  2  NULL
+
+statement ok
+CLOSE curs;
+
+subtest end

pkg/sql/sql_cursor.go

@@ -445,54 +445,67 @@ const (
 )
 
 func (c *cursorMap) closeAll(p *planner, reason cursorCloseReason) error {
-	for n, curs := range c.cursors {
-		switch reason {
-		case cursorCloseForTxnCommit:
-			if curs.withHold {
-				// Cursors declared using WITH HOLD are not closed at transaction
-				// commit, and become the responsibility of the session.
-				curs.committed = true
-				if !curs.persisted {
-					// Execute the cursor's query to completion and persist the result so
-					// that it can survive the transaction's commit.
-					if err := persistCursor(p, curs); err != nil {
-						return err
-					}
+	if reason == cursorCloseForTxnCommit {
+		for _, curs := range c.cursors {
+			if curs.withHold && !curs.persisted {
+				// Execute the cursor's query to completion and persist the result so
+				// that it can survive the transaction's commit. The cursor will become
+				// the responsibility of the session.
+				if err := persistCursor(p, curs); err != nil {
+					return errors.CombineErrors(err, c.closeAll(p, cursorCloseForTxnRollback))
 				}
-				continue
-			}
-		case cursorCloseForTxnRollback:
-			if curs.committed {
-				// Transaction rollback should only remove cursors that were created in
-				// the current transaction.
-				continue
-			}
-		case cursorCloseForTxnPrepare:
-			if curs.withHold && !curs.committed {
-				// Disallow preparing a transaction that has created a cursor WITH HOLD.
-				// It's fine for previous transactions to have created holdable cursors.
-				// NOTE: Postgres also disallows this.
-				//
-				// Make sure to close all cursors before returning the error.
-				err := c.closeAll(p, cursorCloseForTxnRollback)
-				return errors.CombineErrors(
-					pgerror.New(pgcode.FeatureNotSupported,
-						"cannot PREPARE a transaction that has created a cursor WITH HOLD"),
-					err,
-				)
 			}
 		}
-		if err := curs.Close(); err != nil {
-			return err
+	}
+	// Close the cursor iterators/containers. Make sure to continue closing even
+	// if one of them fails.
+	var retErr error
+	for _, curs := range c.cursors {
+		if reason != cursorCloseForExplicitClose && curs.committed {
+			// A holdable cursor from a previously committed transaction should remain
+			// open, except for a session close.
+			continue
+		}
+		if reason == cursorCloseForTxnCommit && curs.withHold {
+			// A holdable cursor remains open after transaction commit.
+			continue
 		}
-		delete(c.cursors, n)
+		if reason == cursorCloseForTxnPrepare && curs.withHold && !curs.committed {
+			// Disallow preparing a transaction that has created a cursor WITH HOLD.
+			// It's fine for previous transactions to have created holdable cursors.
+			// NOTE: Postgres also disallows this.
+			//
+			// Make sure to close all cursors before returning the error.
+			err := c.closeAll(p, cursorCloseForTxnRollback)
+			return errors.CombineErrors(
+				pgerror.New(pgcode.FeatureNotSupported,
+					"cannot PREPARE a transaction that has created a cursor WITH HOLD"),
+				err,
+			)
+		}
+		retErr = errors.CombineErrors(retErr, curs.Close())
+	}
+	if reason == cursorCloseForTxnCommit && retErr != nil {
+		return errors.CombineErrors(retErr, c.closeAll(p, cursorCloseForTxnRollback))
 	}
+	// Remove closed cursors from the map. Mark holdable cursors as committed
+	// for a transaction commit.
 	if reason == cursorCloseForExplicitClose {
 		// All cursors are closed for explicit close, so we can lose the reference
 		// to the map.
 		c.cursors = nil
+	} else {
+		for n, curs := range c.cursors {
+			if reason == cursorCloseForTxnCommit && curs.withHold {
+				curs.committed = true
+			}
+			if curs.committed {
+				continue
+			}
+			delete(c.cursors, n)
+		}
 	}
-	return nil
+	return retErr
 }
 
 func (c *cursorMap) closeCursor(s tree.Name) error {
@@ -583,7 +596,7 @@ func (p *planner) checkNoConflictingCursors(stmt tree.Statement) error {
 
 // persistCursor runs the given cursor to completion and stores the result in a
 // row container that can outlive the cursor's transaction.
-func persistCursor(p *planner, cursor *sqlCursor) error {
+func persistCursor(p *planner, cursor *sqlCursor) (retErr error) {
 	// Use context.Background() because the cursor can outlive the context in
 	// which it was created.
 	helper := persistedCursorHelper{
@@ -602,6 +615,12 @@ func persistCursor(p *planner, cursor *sqlCursor) error {
 		p.ExtendedEvalContextCopy(),
 		"persisted_cursor", /* opName */
 	)
+	defer func() {
+		if retErr != nil {
+			// Close the container if we encountered an error.
+			helper.container.Close(helper.ctx)
+		}
+	}()
 	for {
 		ok, err := cursor.Next(helper.ctx)
 		if err != nil {
@@ -619,6 +638,7 @@ func persistCursor(p *planner, cursor *sqlCursor) error {
 		return err
 	}
 	cursor.Rows = &helper
+	cursor.persisted = true
 	return nil
 }