plpgsql: add support for unnamed cursors

This patch adds support for opening an "unnamed" cursor in a PLpgSQL routine. A PLpgSQL cursor is unnamed when the value for the cursor variable is `NULL`. When an unnamed cursor is opened, a name will be generated for it like `<unnamed portal 1>` that is guaranteed not to conflict with an existing cursor or portal name. The PLpgSQL variable that represents the cursor's name is updated to reflect the generated name. Informs #109709 Release note (sql change): Added support for unnamed PLpgSQL cursors, which generate a unique name when no cursor name was specified.
cockroachdb · Oct 2, 2023 · d2d9656 · d2d9656
1 parent 1dc2ac2
commit d2d9656
Show file tree

Hide file tree

Showing 4 changed files with 478 additions and 137 deletions.
diff --git a/pkg/sql/logictest/testdata/logic_test/plpgsql_cursor b/pkg/sql/logictest/testdata/logic_test/plpgsql_cursor
@@ -138,25 +138,6 @@ FETCH FORWARD 3 FROM foo;
 ----
 1
 
-# Cursor with empty-string name.
-statement ok
-ABORT;
-CREATE OR REPLACE FUNCTION f() RETURNS INT AS $$
-  DECLARE
-    curs STRING := '';
-  BEGIN
-    OPEN curs FOR SELECT 1;
-    RETURN 0;
-  END
-$$ LANGUAGE PLpgSQL;
-BEGIN;
-SELECT f();
-
-query I
-FETCH FORWARD 3 FROM "";
-----
-1
-
 # Multiple cursors.
 statement ok
 ABORT;
@@ -232,14 +213,28 @@ SELECT * FROM xy;
 1  2
 3  4
 
+# The empty string conflicts with the unnamed portal, which always exists.
 statement ok
 ABORT;
+CREATE OR REPLACE FUNCTION f() RETURNS INT AS $$
+  DECLARE
+    curs STRING := '';
+  BEGIN
+    OPEN curs FOR SELECT 1;
+    RETURN 0;
+  END
+$$ LANGUAGE PLpgSQL;
+BEGIN;
+
+statement error pgcode 42P03 pq: cursor \"\" already in use
+SELECT f();
 
 # It is possible to use the OPEN statement in an implicit transaction, but the
 # cursor is closed at the end of the transaction when the statement execution
 # finishes. So, until FETCH is implemented, we can't actually read from the
 # cursor.
 statement ok
+ABORT;
 CREATE OR REPLACE FUNCTION f() RETURNS INT AS $$
   DECLARE
     curs STRING := 'foo';
@@ -307,23 +302,6 @@ CREATE OR REPLACE FUNCTION f() RETURNS INT AS $$
   END
 $$ LANGUAGE PLpgSQL;
 
-statement ok
-CREATE OR REPLACE FUNCTION f() RETURNS INT AS $$
-  DECLARE
-    curs STRING;
-  BEGIN
-    OPEN curs FOR SELECT 1;
-    RETURN 0;
-  END
-$$ LANGUAGE PLpgSQL;
-BEGIN;
-
-statement error pgcode 0A000 pq: unimplemented: opening an unnamed cursor is not yet supported
-SELECT f();
-
-statement ok
-ABORT;
-
 statement error pgcode 0A000 pq: unimplemented: opening a cursor in a routine with an exception block is not yet supported
 CREATE OR REPLACE FUNCTION f() RETURNS INT AS $$
   DECLARE
@@ -636,23 +614,6 @@ SELECT count(*) FROM pg_cursors;
 ----
 0
 
-# It is currently necessary to assign the cursor a name, since one is not
-# automatically generated.
-statement ok
-ABORT;
-CREATE OR REPLACE FUNCTION f() RETURNS INT AS $$
-  DECLARE
-    curs CURSOR FOR SELECT 1;
-  BEGIN
-    OPEN curs;
-    RETURN 0;
-  END
-$$ LANGUAGE PLpgSQL;
-BEGIN;
-
-statement error pgcode 0A000 pq: unimplemented: opening an unnamed cursor is not yet supported
-SELECT f();
-
 statement ok
 ABORT;
 
@@ -691,3 +652,210 @@ CREATE OR REPLACE FUNCTION f() RETURNS INT AS $$
     RETURN 0;
   END
 $$ LANGUAGE PLpgSQL;
+
+statement ok
+DELETE FROM xy WHERE x <> 1 AND x <> 3;
+
+# Testing unnamed cursors.
+statement ok
+DROP FUNCTION f();
+CREATE OR REPLACE FUNCTION f() RETURNS STRING AS $$
+  DECLARE
+    curs STRING;
+  BEGIN
+    OPEN curs FOR SELECT 1;
+    RETURN curs;
+  END
+$$ LANGUAGE PLpgSQL;
+
+statement ok
+BEGIN;
+
+query T rowsort
+SELECT name FROM pg_cursors;
+----
+
+query T
+SELECT f();
+----
+<unnamed portal 1>
+
+query I
+FETCH FORWARD 3 FROM "<unnamed portal 1>";
+----
+1
+
+query T rowsort
+SELECT name FROM pg_cursors;
+----
+<unnamed portal 1>
+
+query T
+SELECT f();
+----
+<unnamed portal 2>
+
+query I
+FETCH FORWARD 3 FROM "<unnamed portal 2>";
+----
+1
+
+query T
+SELECT f();
+----
+<unnamed portal 3>
+
+query I
+FETCH FORWARD 3 FROM "<unnamed portal 3>";
+----
+1
+
+query T rowsort
+SELECT name FROM pg_cursors;
+----
+<unnamed portal 1>
+<unnamed portal 2>
+<unnamed portal 3>
+
+# The generated name does not "fill in gaps".
+statement ok
+CLOSE "<unnamed portal 2>";
+CLOSE "<unnamed portal 1>";
+
+query T
+SELECT f();
+----
+<unnamed portal 4>
+
+query T
+SELECT f();
+----
+<unnamed portal 5>
+
+query T rowsort
+SELECT name FROM pg_cursors;
+----
+<unnamed portal 4>
+<unnamed portal 5>
+<unnamed portal 3>
+
+statement ok
+ABORT;
+BEGIN;
+
+query T
+SELECT f();
+----
+<unnamed portal 6>
+
+# The counter for the generated name keeps incrementing as long as the session
+# is open.
+query T rowsort
+SELECT name FROM pg_cursors;
+----
+<unnamed portal 6>
+
+# The generated name will not conflict with manually created cursors.
+statement ok
+DECLARE "<unnamed portal 7>" CURSOR FOR SELECT 1;
+DECLARE "<unnamed portal 8>" CURSOR FOR SELECT 1;
+
+query T rowsort
+SELECT name FROM pg_cursors;
+----
+<unnamed portal 7>
+<unnamed portal 8>
+<unnamed portal 6>
+
+query T
+SELECT f();
+----
+<unnamed portal 9>
+
+query T rowsort
+SELECT name FROM pg_cursors;
+----
+<unnamed portal 6>
+<unnamed portal 7>
+<unnamed portal 8>
+<unnamed portal 9>
+
+# Do not generate a new name if one was supplied.
+statement ok
+ABORT;
+CREATE OR REPLACE FUNCTION f() RETURNS STRING AS $$
+  DECLARE
+    curs STRING := 'foo';
+  BEGIN
+    OPEN curs FOR SELECT 1;
+    RETURN curs;
+  END
+$$ LANGUAGE PLpgSQL;
+BEGIN;
+
+query T rowsort
+SELECT name FROM pg_cursors;
+----
+
+query T
+SELECT f();
+----
+foo
+
+query T rowsort
+SELECT name FROM pg_cursors;
+----
+foo
+
+# The unnamed portal counter shouldn't have incremented for the named cursor,
+# since no name was generated.
+statement ok
+CREATE OR REPLACE FUNCTION f_unnamed() RETURNS STRING AS $$
+  DECLARE
+    curs STRING;
+  BEGIN
+    OPEN curs FOR SELECT 2;
+    RETURN curs;
+  END
+$$ LANGUAGE PLpgSQL;
+
+query T
+SELECT f_unnamed();
+----
+<unnamed portal 10>
+
+query T rowsort
+SELECT name FROM pg_cursors;
+----
+<unnamed portal 10>
+foo
+
+query I
+FETCH FORWARD 3 FROM "<unnamed portal 10>";
+----
+2
+
+# A bound, unnamed cursor.
+statement ok
+ABORT;
+DROP FUNCTION f();
+CREATE OR REPLACE FUNCTION f() RETURNS INT AS $$
+  DECLARE
+    curs CURSOR FOR SELECT 100;
+  BEGIN
+    OPEN curs;
+    RETURN 0;
+  END
+$$ LANGUAGE PLpgSQL;
+BEGIN;
+
+statement ok
+SELECT f();
+
+query I
+FETCH FORWARD 3 FROM "<unnamed portal 11>";
+----
+100
+
+statement ok
+ABORT;
diff --git a/pkg/sql/opt/optbuilder/plpgsql.go b/pkg/sql/opt/optbuilder/plpgsql.go
@@ -539,6 +539,12 @@ func (b *plpgsqlBuilder) buildPLpgSQLStatements(stmts []ast.Statement, s *scope)
 				}
 				panic(err)
 			}
+			// TODO(drewk): this should check REFCURSOR.
+			if !source.(*scopeColumn).typ.Equivalent(types.String) {
+				panic(pgerror.Newf(pgcode.DatatypeMismatch,
+					"variable \"%s\" must be of type cursor or refcursor", t.CurVar,
+				))
+			}
 			// Initialize the routine with the information needed to pipe the first
 			// body statement into a cursor.
 			query := b.resolveOpenQuery(t)
@@ -558,7 +564,16 @@ func (b *plpgsqlBuilder) buildPLpgSQLStatements(stmts []ast.Statement, s *scope)
 			}
 			b.appendBodyStmt(&openCon, openScope)
 			b.appendPlpgSQLStmts(&openCon, stmts[i+1:])
-			return b.callContinuation(&openCon, s)
+
+			// Build a statement to generate a unique name for the cursor if one
+			// was not supplied. Add this to its own volatile routine to ensure that
+			// the name generation isn't reordered with other operations. Use the
+			// resulting projected column as input to the OPEN continuation.
+			nameCon := b.makeContinuation("_gen_cursor_name")
+			nameCon.def.Volatility = volatility.Volatile
+			nameScope := b.buildCursorNameGen(&nameCon, t.CurVar)
+			b.appendBodyStmt(&nameCon, b.callContinuation(&openCon, nameScope))
+			return b.callContinuation(&nameCon, s)
 
 		default:
 			panic(unimplemented.New(
@@ -609,6 +624,44 @@ func (b *plpgsqlBuilder) resolveOpenQuery(open *ast.Open) tree.Statement {
 	return stmt
 }
 
+// buildCursorNameGen builds a statement that generates a unique name for the
+// cursor if the variable containing the name is unset. The unique name
+// generation is implemented by the crdb_internal.plpgsql_gen_cursor_name
+// builtin function.
+func (b *plpgsqlBuilder) buildCursorNameGen(nameCon *continuation, nameVar ast.Variable) *scope {
+	_, source, _, _ := nameCon.s.FindSourceProvidingColumn(b.ob.ctx, nameVar)
+	const nameFnName = "crdb_internal.plpgsql_gen_cursor_name"
+	props, overloads := builtinsregistry.GetBuiltinProperties(nameFnName)
+	if len(overloads) != 1 {
+		panic(errors.AssertionFailedf("expected one overload for %s", nameFnName))
+	}
+	nameCall := b.ob.factory.ConstructFunction(
+		memo.ScalarListExpr{b.ob.factory.ConstructVariable(source.(*scopeColumn).id)},
+		&memo.FunctionPrivate{
+			Name:       nameFnName,
+			Typ:        types.String,
+			Properties: props,
+			Overload:   &overloads[0],
+		},
+	)
+	// Build an expression that calls the builtin function if the name is unset.
+	scalar := b.ob.factory.ConstructCase(memo.TrueSingleton,
+		memo.ScalarListExpr{
+			b.ob.factory.ConstructWhen(
+				b.ob.factory.ConstructIs(
+					b.ob.factory.ConstructVariable(source.(*scopeColumn).id), memo.NullSingleton,
+				),
+				nameCall,
+			),
+		},
+		b.ob.factory.ConstructVariable(source.(*scopeColumn).id),
+	)
+	nameScope := nameCon.s.push()
+	b.ob.synthesizeColumn(nameScope, scopeColName(nameVar), types.String, nil /* expr */, scalar)
+	b.ob.constructProjectForScope(nameCon.s, nameScope)
+	return nameScope
+}
+
 // addPLpgSQLAssign adds a PL/pgSQL assignment to the current scope as a
 // new column with the variable name that projects the assigned expression.
 // If there is a column with the same name in the previous scope, it will be