Merge cockroachdb#112779

112779: plpgql: implement special OTHERS exception-handling branch r=DrewKimball a=DrewKimball

PLpgSQL exception-handling allows use of the special `OTHERS` condition that matches any error code apart from `query_canceled` and `assert_failure` (although these can be caught explicitly). Note that transaction-rollback (Class 40) errors still cannot be caught, either explicitly or by `OTHERS`.

Informs cockroachdb#105253

Release note (sql change): Added support for the special `OTHERS` condition in PLpgSQL exception blocks, which allows matching any error code apart from `query_canceled` and `assert_failure`. Note that Class 40 errors (40000, 40001, 40003, 40002, and 40P01) cannot be caught either, tracked in cockroachdb#111446.

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

pkg/sql/logictest/testdata/logic_test/udf_plpgsql

@@ -2216,3 +2216,108 @@ SELECT * FROM xy;
 ----
 1  2
 3  4
+
+# Testing the special OTHERS condition for exception handling.
+subtest others
+
+statement ok
+DROP FUNCTION f(INT);
+CREATE OR REPLACE FUNCTION f(n INT) RETURNS TEXT AS $$
+  BEGIN
+    IF n = 0 THEN
+      RAISE division_by_zero;
+    ELSIF n = 1 THEN
+      RAISE data_exception;
+    ELSIF n = 2 THEN
+      RAISE not_null_violation;
+    ELSIF n = 3 THEN
+      RAISE query_canceled;
+    ELSE
+      RAISE assert_failure;
+    END IF;
+  EXCEPTION
+    WHEN not_null_violation THEN
+      RETURN 'not_null_violation';
+    WHEN OTHERS THEN
+      RETURN 'others';
+    WHEN assert_failure THEN
+      RETURN 'assert_failure';
+    WHEN data_exception THEN
+      RETURN 'data_exception';
+  END
+$$ LANGUAGE PLpgSQL;
+
+# OTHERS should catch division_by_zero.
+query T
+SELECT f(0);
+----
+others
+
+# OTHERS should catch data_exception despite the later branch that
+# explicitly catches it.
+query T
+SELECT f(1);
+----
+others
+
+# OTHERS can catch not_null_violation, but does not in this instance because
+# the branch that explicitly catches it comes first.
+query T
+SELECT f(2);
+----
+not_null_violation
+
+# OTHERS cannot catch query_canceled or assert_failure; however, they can be
+# explicitly caught.
+statement error pgcode 57014 pq: query_canceled
+SELECT f(3);
+
+query T
+SELECT f(4);
+----
+assert_failure
+
+# Postgres allows multiple OTHERS branches.
+statement ok
+DROP FUNCTION f();
+CREATE OR REPLACE FUNCTION f() RETURNS TEXT AS $$
+  BEGIN
+    RAISE division_by_zero;
+  EXCEPTION
+    WHEN OTHERS THEN
+      RETURN 'first branch';
+    WHEN OTHERS THEN
+      RETURN 'second branch';
+    WHEN division_by_zero THEN
+      RETURN 'third branch';
+    WHEN OTHERS THEN
+      RETURN 'fourth branch';
+  END
+$$ LANGUAGE PLpgSQL;
+
+query T
+SELECT f();
+----
+first branch
+
+# The first OTHERS branch is taken even though it raises its own exception.
+statement ok
+CREATE OR REPLACE FUNCTION f() RETURNS TEXT AS $$
+  BEGIN
+    RAISE division_by_zero;
+  EXCEPTION
+    WHEN OTHERS THEN
+      RAISE null_value_not_allowed;
+    WHEN OTHERS THEN
+      RETURN 'second branch';
+    WHEN division_by_zero THEN
+      RETURN 'third branch';
+    WHEN OTHERS THEN
+      RETURN 'fourth branch';
+  END
+$$ LANGUAGE PLpgSQL;
+
+statement error pgcode 22004 pq: null_value_not_allowed
+SELECT f();
+
+subtest end

pkg/sql/opt/memo/expr.go

@@ -744,7 +744,8 @@ type UDFDefinition struct {
 type ExceptionBlock struct {
 	// Codes is a list of pgcode strings (see pgcode/codes.go). When the body of a
 	// routine with an ExceptionBlock returns an error, the code of that error is
-	// compared against the Codes slice for a match.
+	// compared against the Codes slice for a match. As a special case, the code
+	// may be "OTHERS", indicating that (almost) any error code should be matched.
 	Codes []pgcode.Code
 
 	// Actions contains routine definitions that represent exception handlers for

pkg/sql/opt/memo/expr_format.go

@@ -984,7 +984,12 @@ func (f *ExprFmtCtx) formatScalarWithLabel(
 			for i := range udf.Def.ExceptionBlock.Codes {
 				code := udf.Def.ExceptionBlock.Codes[i]
 				body := udf.Def.ExceptionBlock.Actions[i].Body
-				branch := n.Childf("SQLSTATE '%s'", code)
+				var branch treeprinter.Node
+				if code.String() == "OTHERS" {
+					branch = n.Child("OTHERS")
+				} else {
+					branch = n.Childf("SQLSTATE '%s'", code)
+				}
 				for j := range body {
 					f.formatExpr(body[j], branch)
 				}

pkg/sql/opt/optbuilder/plpgsql.go

@@ -1030,6 +1030,11 @@ func (b *plpgsqlBuilder) buildExceptions(block *ast.Block) *memo.ExceptionBlock
 				continue
 			}
 			// The match condition was supplied by name instead of code.
+			if strings.ToUpper(cond.SqlErrName) == "OTHERS" {
+				// The special "OTHERS" condition matches (almost) any error code.
+				addHandler("OTHERS" /* codeStr */, handlerCon.def)
+				continue
+			}
 			branchCodes, ok := pgcode.PLpgSQLConditionNameToCode[cond.SqlErrName]
 			if !ok {
 				panic(pgerror.Newf(

pkg/sql/opt/optbuilder/testdata/udf_plpgsql

@@ -5484,3 +5484,70 @@ project
                      │                                                                └── projections
                      │                                                                     └── const: 0 [as=stmt_return_4:10]
                      └── const: 1
+
+# Correctly display the OTHERS branch.
+exec-ddl
+CREATE OR REPLACE FUNCTION f() RETURNS INT AS $$
+  BEGIN
+    RETURN 1 // 0;
+  EXCEPTION
+    WHEN assert_failure THEN
+      RETURN -1;
+    WHEN OTHERS THEN
+      RETURN -2;
+    WHEN query_canceled THEN
+      RETURN -3;
+  END
+$$ LANGUAGE PLpgSQL;
+----
+
+build format=show-scalars
+SELECT f();
+----
+project
+ ├── columns: f:6
+ ├── values
+ │    └── tuple
+ └── projections
+      └── udf: f [as=f:6]
+           └── body
+                └── limit
+                     ├── columns: exception_block_7:5
+                     ├── project
+                     │    ├── columns: exception_block_7:5
+                     │    ├── values
+                     │    │    └── tuple
+                     │    └── projections
+                     │         └── udf: exception_block_7 [as=exception_block_7:5]
+                     │              ├── body
+                     │              │    └── project
+                     │              │         ├── columns: stmt_return_8:4!null
+                     │              │         ├── values
+                     │              │         │    └── tuple
+                     │              │         └── projections
+                     │              │              └── floor-div [as=stmt_return_8:4]
+                     │              │                   ├── const: 1
+                     │              │                   └── const: 0
+                     │              └── exception-handler
+                     │                   ├── SQLSTATE 'P0004'
+                     │                   │    └── project
+                     │                   │         ├── columns: stmt_return_2:1!null
+                     │                   │         ├── values
+                     │                   │         │    └── tuple
+                     │                   │         └── projections
+                     │                   │              └── const: -1 [as=stmt_return_2:1]
+                     │                   ├── OTHERS
+                     │                   │    └── project
+                     │                   │         ├── columns: stmt_return_4:2!null
+                     │                   │         ├── values
+                     │                   │         │    └── tuple
+                     │                   │         └── projections
+                     │                   │              └── const: -2 [as=stmt_return_4:2]
+                     │                   └── SQLSTATE '57014'
+                     │                        └── project
+                     │                             ├── columns: stmt_return_6:3!null
+                     │                             ├── values
+                     │                             │    └── tuple
+                     │                             └── projections
+                     │                                  └── const: -3 [as=stmt_return_6:3]
+                     └── const: 1

pkg/sql/routine.go

@@ -325,7 +325,14 @@ func (g *routineGenerator) handleException(ctx context.Context, err error) error
 	exceptionHandler := blockState.ExceptionHandler
 	blockState.ExceptionHandler = nil
 	for i, code := range exceptionHandler.Codes {
-		if code == caughtCode {
+		caughtException := code == caughtCode
+		if code.String() == "OTHERS" {
+			// The special OTHERS condition matches any error code apart from
+			// query_canceled and assert_failure (though they can still be caught
+			// explicitly).
+			caughtException = caughtCode != pgcode.QueryCanceled && caughtCode != pgcode.AssertFailure
+		}
+		if caughtException {
 			cursErr := g.closeCursors(blockState)
 			if cursErr != nil {
 				return errors.CombineErrors(err, cursErr)

pkg/sql/sem/tree/routine.go

@@ -191,7 +191,8 @@ func (node *RoutineExpr) Walk(v Visitor) Expr {
 // RoutineExceptionHandler encapsulates the information needed to match and
 // handle errors for the exception block of a routine defined with PLpgSQL.
 type RoutineExceptionHandler struct {
-	// Codes is a list of pgcode strings used to match exceptions.
+	// Codes is a list of pgcode strings used to match exceptions. Note that as a
+	// special case, the code may be "OTHERS", which matches most error codes.
 	Codes []pgcode.Code
 
 	// Actions contains a routine to handle each error code.