plpgsql: add support for CONSTANT variable declarations

It is possible to declare variables as CONSTANT in PLpgSQL. This will cause
any attempt to assign to the variable to result in a compile-time error.
This patch implements CONSTANT variables by adding a check whenever an
assignment is built that the assigned variable is non-constant.

Fixes cockroachdb#105241

Release note (sql change): Added support for CONSTANT variable
declarations in PLpgSQL routines. Any assignment to a variable declared
with the CONSTANT keyword will raise a compile-time error.

Author: DrewKimball

pkg/sql/logictest/testdata/logic_test/udf_plpgsql

@@ -708,3 +708,103 @@ CREATE OR REPLACE FUNCTION f() RETURNS INT AS $$
     RETURN i;
   END
 $$ LANGUAGE PLpgSQL;
+
+statement ok
+CREATE OR REPLACE FUNCTION f() RETURNS INT AS $$
+  DECLARE
+    i CONSTANT INT;
+  BEGIN
+    RETURN i;
+  END
+$$ LANGUAGE PLpgSQL;
+
+query I
+SELECT f();
+----
+NULL
+
+statement ok
+CREATE OR REPLACE FUNCTION f() RETURNS INT AS $$
+  DECLARE
+    i CONSTANT INT := 0;
+  BEGIN
+    RETURN i;
+  END
+$$ LANGUAGE PLpgSQL;
+
+query I
+SELECT f();
+----
+0
+
+statement ok
+CREATE OR REPLACE FUNCTION f() RETURNS INT AS $$
+  DECLARE
+    i CONSTANT INT := (SELECT x FROM xy ORDER BY x LIMIT 1);
+  BEGIN
+    RETURN i;
+  END
+$$ LANGUAGE PLpgSQL;
+
+query I
+SELECT f();
+----
+1
+
+statement ok
+CREATE OR REPLACE FUNCTION f(n INT) RETURNS INT AS $$
+  DECLARE
+    i CONSTANT INT := n;
+  BEGIN
+    RETURN i;
+  END
+$$ LANGUAGE PLpgSQL;
+
+query IIIIII
+SELECT f(-100), f(-1), f(0), f(1), f(100), f(NULL);
+----
+-100  -1  0  1  100  NULL
+
+statement error pgcode 22005 pq: variable \"i\" is declared CONSTANT
+CREATE OR REPLACE FUNCTION f() RETURNS INT AS $$
+  DECLARE
+    i CONSTANT INT;
+  BEGIN
+    i := i + 1;
+    RETURN i;
+  END
+$$ LANGUAGE PLpgSQL;
+
+statement error pgcode 22005 pq: variable \"i\" is declared CONSTANT
+CREATE OR REPLACE FUNCTION f() RETURNS INT AS $$
+  DECLARE
+    i CONSTANT INT := 0;
+  BEGIN
+    i := i + 1;
+    RETURN i;
+  END
+$$ LANGUAGE PLpgSQL;
+
+statement error pgcode 22005 pq: variable \"i\" is declared CONSTANT
+CREATE OR REPLACE FUNCTION f(n INT) RETURNS INT AS $$
+  DECLARE
+    i CONSTANT INT := 0;
+  BEGIN
+    IF n > 0 THEN
+      i := i + 1;
+    END IF;
+    RETURN i;
+  END
+$$ LANGUAGE PLpgSQL;
+
+statement error pgcode 22005 pq: variable \"i\" is declared CONSTANT
+CREATE OR REPLACE FUNCTION f() RETURNS INT AS $$
+  DECLARE
+    i CONSTANT INT := 0;
+  BEGIN
+    LOOP IF i >= 10 THEN EXIT; END IF;
+      i := i + 1;
+    END LOOP;
+    RETURN i;
+  END
+$$ LANGUAGE PLpgSQL;

pkg/sql/opt/optbuilder/plpgsql.go

@@ -121,6 +121,9 @@ type plpgsqlBuilder struct {
 	// varTypes maps from the name of each variable to its type.
 	varTypes map[tree.Name]*types.T
 
+	// constants tracks the variables that were declared as constant.
+	constants map[tree.Name]struct{}
+
 	// returnType is the return type of the PL/pgSQL function.
 	returnType *types.T
 
@@ -164,12 +167,6 @@ func (b *plpgsqlBuilder) init(
 				"not-null PL/pgSQL variables are not yet supported",
 			))
 		}
-		if dec.Constant {
-			panic(unimplemented.NewWithIssueDetail(105241,
-				"constant variable",
-				"constant PL/pgSQL variables are not yet supported",
-			))
-		}
 		if dec.Collate != "" {
 			panic(unimplemented.NewWithIssueDetail(105245,
 				"variable collation",
@@ -185,14 +182,20 @@ func (b *plpgsqlBuilder) build(block *plpgsqltree.PLpgSQLStmtBlock, s *scope) *s
 	s = s.push()
 	b.ensureScopeHasExpr(s)
 
-	// Some variable declarations initialize the variable.
+	b.constants = make(map[tree.Name]struct{})
 	for _, dec := range b.decls {
 		if dec.Expr != nil {
+			// Some variable declarations initialize the variable.
 			s = b.addPLpgSQLAssign(s, dec.Var, dec.Expr)
 		} else {
 			// Uninitialized variables are null.
 			s = b.addPLpgSQLAssign(s, dec.Var, &tree.CastExpr{Expr: tree.DNull, Type: dec.Typ})
 		}
+		if dec.Constant {
+			// Add to the constants map after initializing the variable, since
+			// constant variables only prevent assignment, not initialization.
+			b.constants[dec.Var] = struct{}{}
+		}
 	}
 	if s = b.buildPLpgSQLStatements(block.Body, s); s != nil {
 		return s
@@ -399,6 +402,11 @@ func (b *plpgsqlBuilder) buildPLpgSQLStatements(
 func (b *plpgsqlBuilder) addPLpgSQLAssign(
 	inScope *scope, ident plpgsqltree.PLpgSQLVariable, val plpgsqltree.PLpgSQLExpr,
 ) *scope {
+	if b.constants != nil {
+		if _, ok := b.constants[ident]; ok {
+			panic(pgerror.Newf(pgcode.ErrorInAssignment, "variable \"%s\" is declared CONSTANT", ident))
+		}
+	}
 	typ, ok := b.varTypes[ident]
 	if !ok {
 		panic(pgerror.Newf(pgcode.Syntax, "\"%s\" is not a known variable", ident))

pkg/sql/sem/plpgsqltree/statements.go

@@ -60,7 +60,6 @@ type PLpgSQLStmtBlock struct {
 	Decls      []PLpgSQLDecl
 	Body       []PLpgSQLStatement
 	Exceptions *PLpgSQLExceptionBlock
-	Scope      VariableScope
 }
 
 // TODO(drewk): format Label and Exceptions fields.

pkg/sql/sem/plpgsqltree/variable.go

@@ -13,9 +13,3 @@ package plpgsqltree
 import "github.com/cockroachdb/cockroach/pkg/sql/sem/tree"
 
 type PLpgSQLVariable = tree.Name
-
-// Scope contains all the variables defined in the DECLARE section of current statement block.
-type VariableScope struct {
-	Variables    []*PLpgSQLVariable
-	VarNameToIdx map[string]int // mapping from variable
-}