Re-implement pl/pgsql's expression and assignment parsing.

Invent new RawParseModes that allow the core grammar to handle
pl/pgsql expressions and assignments directly, and thereby get rid
of a lot of hackery in pl/pgsql's parser.  This moves a good deal
of knowledge about pl/pgsql into the core code: notably, we have to
invent a CoercionContext that matches pl/pgsql's (rather dubious)
historical behavior for assignment coercions.  That's getting away
from the original idea of pl/pgsql as an arm's-length extension of
the core, but really we crossed that bridge a long time ago.

The main advantage of doing this is that we can now use the core
parser to generate FieldStore and/or SubscriptingRef nodes to handle
assignments to pl/pgsql variables that are records or arrays.  That
fixes a number of cases that had never been implemented in pl/pgsql
assignment, such as nested records and array slicing, and it allows
pl/pgsql assignment to support the datatype-specific subscripting
behaviors introduced in commit c7aba7c.

There are cosmetic benefits too: when a syntax error occurs in a
pl/pgsql expression, the error report no longer includes the confusing
"SELECT" keyword that used to get prefixed to the expression text.
Also, there seem to be some small speed gains.

Discussion: https://postgr.es/m/[email protected]

Author: tglsfdc

contrib/hstore/expected/hstore.out

@@ -1583,6 +1583,10 @@ select f2 from test_json_agg;
  "d"=>NULL, "x"=>"xyzzy"
 (3 rows)
 
+-- Test subscripting in plpgsql
+do $$ declare h hstore;
+begin h['a'] := 'b'; raise notice 'h = %, h[a] = %', h, h['a']; end $$;
+NOTICE:  h = "a"=>"b", h[a] = b
 -- Check the hstore_hash() and hstore_hash_extended() function explicitly.
 SELECT v as value, hstore_hash(v)::bit(32) as standard,
        hstore_hash_extended(v, 0)::bit(32) as extended0,

contrib/hstore/sql/hstore.sql

@@ -372,6 +372,10 @@ select f2['d':'e'] from test_json_agg;  -- error
 update test_json_agg set f2['d'] = f2['e'], f2['x'] = 'xyzzy';
 select f2 from test_json_agg;
 
+-- Test subscripting in plpgsql
+do $$ declare h hstore;
+begin h['a'] := 'b'; raise notice 'h = %, h[a] = %', h, h['a']; end $$;
+
 -- Check the hstore_hash() and hstore_hash_extended() function explicitly.
 SELECT v as value, hstore_hash(v)::bit(32) as standard,
        hstore_hash_extended(v, 0)::bit(32) as extended0,

doc/src/sgml/plpgsql.sgml

@@ -946,8 +946,8 @@ PREPARE <replaceable>statement_name</replaceable>(integer, integer) AS SELECT $1
      database engine.  The expression must yield a single value (possibly
      a row value, if the variable is a row or record variable).  The target
      variable can be a simple variable (optionally qualified with a block
-     name), a field of a row or record variable, or an element of an array
-     that is a simple variable or field.  Equal (<literal>=</literal>) can be
+     name), a field of a row or record target, or an element or slice of
+     an array target.  Equal (<literal>=</literal>) can be
      used instead of PL/SQL-compliant <literal>:=</literal>.
     </para>
 
@@ -968,8 +968,25 @@ PREPARE <replaceable>statement_name</replaceable>(integer, integer) AS SELECT $1
 <programlisting>
 tax := subtotal * 0.06;
 my_record.user_id := 20;
+my_array[j] := 20;
+my_array[1:3] := array[1,2,3];
+complex_array[n].realpart = 12.3;
 </programlisting>
     </para>
+
+    <para>
+     It's useful to know that what follows the assignment operator is
+     essentially treated as a <literal>SELECT</literal> command; as long
+     as it returns a single row and column, it will work.  Thus for example
+     one can write something like
+<programlisting>
+total_sales := sum(quantity) from sales;
+</programlisting>
+     This provides an effect similar to the single-row <literal>SELECT
+     ... INTO</literal> syntax described in
+     <xref linkend="plpgsql-statements-sql-onerow"/>.  However, that syntax
+     is more portable.
+    </para>
    </sect2>
 
    <sect2 id="plpgsql-statements-sql-noresult">

src/backend/commands/functioncmds.c

@@ -1628,6 +1628,7 @@ CreateCast(CreateCastStmt *stmt)
 		case COERCION_ASSIGNMENT:
 			castcontext = COERCION_CODE_ASSIGNMENT;
 			break;
+			/* COERCION_PLPGSQL is intentionally not covered here */
 		case COERCION_EXPLICIT:
 			castcontext = COERCION_CODE_EXPLICIT;
 			break;

src/backend/executor/spi.c

@@ -51,6 +51,12 @@ static _SPI_connection *_SPI_current = NULL;
 static int	_SPI_stack_depth = 0;	/* allocated size of _SPI_stack */
 static int	_SPI_connected = -1;	/* current stack index */
 
+typedef struct SPICallbackArg
+{
+	const char *query;
+	RawParseMode mode;
+} SPICallbackArg;
+
 static Portal SPI_cursor_open_internal(const char *name, SPIPlanPtr plan,
 									   ParamListInfo paramLI, bool read_only);
 
@@ -1479,6 +1485,7 @@ SPI_cursor_open_internal(const char *name, SPIPlanPtr plan,
 	Snapshot	snapshot;
 	MemoryContext oldcontext;
 	Portal		portal;
+	SPICallbackArg spicallbackarg;
 	ErrorContextCallback spierrcontext;
 
 	/*
@@ -1533,8 +1540,10 @@ SPI_cursor_open_internal(const char *name, SPIPlanPtr plan,
 	 * Setup error traceback support for ereport(), in case GetCachedPlan
 	 * throws an error.
 	 */
+	spicallbackarg.query = plansource->query_string;
+	spicallbackarg.mode = plan->parse_mode;
 	spierrcontext.callback = _SPI_error_callback;
-	spierrcontext.arg = unconstify(char *, plansource->query_string);
+	spierrcontext.arg = &spicallbackarg;
 	spierrcontext.previous = error_context_stack;
 	error_context_stack = &spierrcontext;
 
@@ -1952,6 +1961,7 @@ SPI_plan_get_cached_plan(SPIPlanPtr plan)
 {
 	CachedPlanSource *plansource;
 	CachedPlan *cplan;
+	SPICallbackArg spicallbackarg;
 	ErrorContextCallback spierrcontext;
 
 	Assert(plan->magic == _SPI_PLAN_MAGIC);
@@ -1966,8 +1976,10 @@ SPI_plan_get_cached_plan(SPIPlanPtr plan)
 	plansource = (CachedPlanSource *) linitial(plan->plancache_list);
 
 	/* Setup error traceback support for ereport() */
+	spicallbackarg.query = plansource->query_string;
+	spicallbackarg.mode = plan->parse_mode;
 	spierrcontext.callback = _SPI_error_callback;
-	spierrcontext.arg = unconstify(char *, plansource->query_string);
+	spierrcontext.arg = &spicallbackarg;
 	spierrcontext.previous = error_context_stack;
 	error_context_stack = &spierrcontext;
 
@@ -2094,13 +2106,16 @@ _SPI_prepare_plan(const char *src, SPIPlanPtr plan)
 	List	   *raw_parsetree_list;
 	List	   *plancache_list;
 	ListCell   *list_item;
+	SPICallbackArg spicallbackarg;
 	ErrorContextCallback spierrcontext;
 
 	/*
 	 * Setup error traceback support for ereport()
 	 */
+	spicallbackarg.query = src;
+	spicallbackarg.mode = plan->parse_mode;
 	spierrcontext.callback = _SPI_error_callback;
-	spierrcontext.arg = unconstify(char *, src);
+	spierrcontext.arg = &spicallbackarg;
 	spierrcontext.previous = error_context_stack;
 	error_context_stack = &spierrcontext;
 
@@ -2199,13 +2214,16 @@ _SPI_prepare_oneshot_plan(const char *src, SPIPlanPtr plan)
 	List	   *raw_parsetree_list;
 	List	   *plancache_list;
 	ListCell   *list_item;
+	SPICallbackArg spicallbackarg;
 	ErrorContextCallback spierrcontext;
 
 	/*
 	 * Setup error traceback support for ereport()
 	 */
+	spicallbackarg.query = src;
+	spicallbackarg.mode = plan->parse_mode;
 	spierrcontext.callback = _SPI_error_callback;
-	spierrcontext.arg = unconstify(char *, src);
+	spierrcontext.arg = &spicallbackarg;
 	spierrcontext.previous = error_context_stack;
 	error_context_stack = &spierrcontext;
 
@@ -2263,15 +2281,18 @@ _SPI_execute_plan(SPIPlanPtr plan, ParamListInfo paramLI,
 	SPITupleTable *my_tuptable = NULL;
 	int			res = 0;
 	bool		pushed_active_snap = false;
+	SPICallbackArg spicallbackarg;
 	ErrorContextCallback spierrcontext;
 	CachedPlan *cplan = NULL;
 	ListCell   *lc1;
 
 	/*
 	 * Setup error traceback support for ereport()
 	 */
+	spicallbackarg.query = NULL;	/* we'll fill this below */
+	spicallbackarg.mode = plan->parse_mode;
 	spierrcontext.callback = _SPI_error_callback;
-	spierrcontext.arg = NULL;	/* we'll fill this below */
+	spierrcontext.arg = &spicallbackarg;
 	spierrcontext.previous = error_context_stack;
 	error_context_stack = &spierrcontext;
 
@@ -2318,7 +2339,7 @@ _SPI_execute_plan(SPIPlanPtr plan, ParamListInfo paramLI,
 		List	   *stmt_list;
 		ListCell   *lc2;
 
-		spierrcontext.arg = unconstify(char *, plansource->query_string);
+		spicallbackarg.query = plansource->query_string;
 
 		/*
 		 * If this is a one-shot plan, we still need to do parse analysis.
@@ -2722,7 +2743,8 @@ _SPI_pquery(QueryDesc *queryDesc, bool fire_triggers, uint64 tcount)
 static void
 _SPI_error_callback(void *arg)
 {
-	const char *query = (const char *) arg;
+	SPICallbackArg *carg = (SPICallbackArg *) arg;
+	const char *query = carg->query;
 	int			syntaxerrposition;
 
 	if (query == NULL)			/* in case arg wasn't set yet */
@@ -2740,7 +2762,23 @@ _SPI_error_callback(void *arg)
 		internalerrquery(query);
 	}
 	else
-		errcontext("SQL statement \"%s\"", query);
+	{
+		/* Use the parse mode to decide how to describe the query */
+		switch (carg->mode)
+		{
+			case RAW_PARSE_PLPGSQL_EXPR:
+				errcontext("SQL expression \"%s\"", query);
+				break;
+			case RAW_PARSE_PLPGSQL_ASSIGN1:
+			case RAW_PARSE_PLPGSQL_ASSIGN2:
+			case RAW_PARSE_PLPGSQL_ASSIGN3:
+				errcontext("PL/pgSQL assignment \"%s\"", query);
+				break;
+			default:
+				errcontext("SQL statement \"%s\"", query);
+				break;
+		}
+	}
 }
 
 /*

src/backend/nodes/copyfuncs.c

@@ -3199,6 +3199,20 @@ _copySetOperationStmt(const SetOperationStmt *from)
 	return newnode;
 }
 
+static PLAssignStmt *
+_copyPLAssignStmt(const PLAssignStmt *from)
+{
+	PLAssignStmt *newnode = makeNode(PLAssignStmt);
+
+	COPY_STRING_FIELD(name);
+	COPY_NODE_FIELD(indirection);
+	COPY_SCALAR_FIELD(nnames);
+	COPY_NODE_FIELD(val);
+	COPY_LOCATION_FIELD(location);
+
+	return newnode;
+}
+
 static AlterTableStmt *
 _copyAlterTableStmt(const AlterTableStmt *from)
 {
@@ -5220,6 +5234,9 @@ copyObjectImpl(const void *from)
 		case T_SetOperationStmt:
 			retval = _copySetOperationStmt(from);
 			break;
+		case T_PLAssignStmt:
+			retval = _copyPLAssignStmt(from);
+			break;
 		case T_AlterTableStmt:
 			retval = _copyAlterTableStmt(from);
 			break;

src/backend/nodes/equalfuncs.c

@@ -1085,6 +1085,18 @@ _equalSetOperationStmt(const SetOperationStmt *a, const SetOperationStmt *b)
 	return true;
 }
 
+static bool
+_equalPLAssignStmt(const PLAssignStmt *a, const PLAssignStmt *b)
+{
+	COMPARE_STRING_FIELD(name);
+	COMPARE_NODE_FIELD(indirection);
+	COMPARE_SCALAR_FIELD(nnames);
+	COMPARE_NODE_FIELD(val);
+	COMPARE_LOCATION_FIELD(location);
+
+	return true;
+}
+
 static bool
 _equalAlterTableStmt(const AlterTableStmt *a, const AlterTableStmt *b)
 {
@@ -3275,6 +3287,9 @@ equal(const void *a, const void *b)
 		case T_SetOperationStmt:
 			retval = _equalSetOperationStmt(a, b);
 			break;
+		case T_PLAssignStmt:
+			retval = _equalPLAssignStmt(a, b);
+			break;
 		case T_AlterTableStmt:
 			retval = _equalAlterTableStmt(a, b);
 			break;

src/backend/nodes/nodeFuncs.c

@@ -3669,6 +3669,16 @@ raw_expression_tree_walker(Node *node,
 					return true;
 			}
 			break;
+		case T_PLAssignStmt:
+			{
+				PLAssignStmt *stmt = (PLAssignStmt *) node;
+
+				if (walker(stmt->indirection, context))
+					return true;
+				if (walker(stmt->val, context))
+					return true;
+			}
+			break;
 		case T_A_Expr:
 			{
 				A_Expr	   *expr = (A_Expr *) node;

src/backend/nodes/outfuncs.c

@@ -2775,6 +2775,18 @@ _outSelectStmt(StringInfo str, const SelectStmt *node)
 	WRITE_NODE_FIELD(rarg);
 }
 
+static void
+_outPLAssignStmt(StringInfo str, const PLAssignStmt *node)
+{
+	WRITE_NODE_TYPE("PLASSIGN");
+
+	WRITE_STRING_FIELD(name);
+	WRITE_NODE_FIELD(indirection);
+	WRITE_INT_FIELD(nnames);
+	WRITE_NODE_FIELD(val);
+	WRITE_LOCATION_FIELD(location);
+}
+
 static void
 _outFuncCall(StringInfo str, const FuncCall *node)
 {
@@ -4211,6 +4223,9 @@ outNode(StringInfo str, const void *obj)
 			case T_SelectStmt:
 				_outSelectStmt(str, obj);
 				break;
+			case T_PLAssignStmt:
+				_outPLAssignStmt(str, obj);
+				break;
 			case T_ColumnDef:
 				_outColumnDef(str, obj);
 				break;