From 282521568617a704df3e09279f568f31b1871411 Mon Sep 17 00:00:00 2001
From: Nikita Glukhov <n.gluhov@postgrespro.ru>
Date: Sat, 1 Apr 2023 16:21:20 +0300
Subject: [PATCH v8 2/7] Pass field accessors to generic subscripting

---
 src/backend/parser/parse_expr.c   | 66 ++++++++++++++++++++-----------
 src/backend/parser/parse_node.c   | 41 +++++++++++++++++--
 src/backend/parser/parse_target.c |  3 +-
 src/backend/utils/adt/arraysubs.c | 13 ++++--
 src/backend/utils/adt/jsonbsubs.c | 11 +++++-
 src/include/parser/parse_node.h   |  3 +-
 6 files changed, 105 insertions(+), 32 deletions(-)

diff --git a/src/backend/parser/parse_expr.c b/src/backend/parser/parse_expr.c
index 2c0f4a50b21..8ea51176196 100644
--- a/src/backend/parser/parse_expr.c
+++ b/src/backend/parser/parse_expr.c
@@ -442,8 +442,9 @@ transformIndirection(ParseState *pstate, A_Indirection *ind)
 	ListCell   *i;
 
 	/*
-	 * We have to split any field-selection operations apart from
-	 * subscripting.  Adjacent A_Indices nodes have to be treated as a single
+	 * Combine field names and subscripts into a single indirection list, as
+	 * some subscripting containers, such as jsonb, support field access using
+	 * dot notation. Adjacent A_Indices nodes have to be treated as a single
 	 * multidimensional subscript operation.
 	 */
 	foreach(i, ind->indirection)
@@ -461,19 +462,43 @@ transformIndirection(ParseState *pstate, A_Indirection *ind)
 		}
 		else
 		{
-			Node	   *newresult;
-
 			Assert(IsA(n, String));
+			subscripts = lappend(subscripts, n);
+		}
+	}
+
+	while (subscripts)
+	{
+		/* try processing container subscripts first */
+		Node	   *newresult = (Node *)
+			transformContainerSubscripts(pstate,
+										 result,
+										 exprType(result),
+										 exprTypmod(result),
+										 &subscripts,
+										 false,
+										 true);
+
+		if (!newresult)
+		{
+			/*
+			 * generic subscripting failed; falling back to function call or
+			 * field selection for a composite type.
+			 */
+			Node	   *n;
+
+			Assert(subscripts);
 
-			/* process subscripts before this field selection */
-			while (subscripts)
-				result = (Node *) transformContainerSubscripts(pstate,
-															   result,
-															   exprType(result),
-															   exprTypmod(result),
-															   &subscripts,
-															   false);
+			n = linitial(subscripts);
+
+			if (!IsA(n, String))
+				ereport(ERROR,
+						(errcode(ERRCODE_DATATYPE_MISMATCH),
+						 errmsg("cannot subscript type %s because it does not support subscripting",
+								format_type_be(exprType(result))),
+						 parser_errposition(pstate, exprLocation(result))));
 
+			/* try to find function for field selection */
 			newresult = ParseFuncOrColumn(pstate,
 										  list_make1(n),
 										  list_make1(result),
@@ -481,19 +506,16 @@ transformIndirection(ParseState *pstate, A_Indirection *ind)
 										  NULL,
 										  false,
 										  location);
-			if (newresult == NULL)
+
+			if (!newresult)
 				unknown_attribute(pstate, result, strVal(n), location);
-			result = newresult;
+
+			/* consume field select */
+			subscripts = list_delete_first(subscripts);
 		}
+
+		result = newresult;
 	}
-	/* process trailing subscripts, if any */
-	while (subscripts)
-		result = (Node *) transformContainerSubscripts(pstate,
-													   result,
-													   exprType(result),
-													   exprTypmod(result),
-													   &subscripts,
-													   false);
 
 	return result;
 }
diff --git a/src/backend/parser/parse_node.c b/src/backend/parser/parse_node.c
index 19a6b678e67..b3e476eb181 100644
--- a/src/backend/parser/parse_node.c
+++ b/src/backend/parser/parse_node.c
@@ -238,6 +238,8 @@ transformContainerType(Oid *containerType, int32 *containerTypmod)
  * containerTypMod	typmod for the container
  * indirection		Untransformed list of subscripts (must not be NIL)
  * isAssignment		True if this will become a container assignment.
+ * noError			True for return NULL with no error, if the container type
+ * 					is not subscriptable.
  */
 SubscriptingRef *
 transformContainerSubscripts(ParseState *pstate,
@@ -245,13 +247,15 @@ transformContainerSubscripts(ParseState *pstate,
 							 Oid containerType,
 							 int32 containerTypMod,
 							 List **indirection,
-							 bool isAssignment)
+							 bool isAssignment,
+							 bool noError)
 {
 	SubscriptingRef *sbsref;
 	const SubscriptRoutines *sbsroutines;
 	Oid			elementType;
 	bool		isSlice = false;
 	ListCell   *idx;
+	int			indirection_length = list_length(*indirection);
 
 	/*
 	 * Determine the actual container type, smashing any domain.  In the
@@ -267,11 +271,16 @@ transformContainerSubscripts(ParseState *pstate,
 	 */
 	sbsroutines = getSubscriptingRoutines(containerType, &elementType);
 	if (!sbsroutines)
+	{
+		if (noError)
+			return NULL;
+
 		ereport(ERROR,
 				(errcode(ERRCODE_DATATYPE_MISMATCH),
 				 errmsg("cannot subscript type %s because it does not support subscripting",
 						format_type_be(containerType)),
 				 parser_errposition(pstate, exprLocation(containerBase))));
+	}
 
 	/*
 	 * Detect whether any of the indirection items are slice specifiers.
@@ -282,9 +291,9 @@ transformContainerSubscripts(ParseState *pstate,
 	 */
 	foreach(idx, *indirection)
 	{
-		A_Indices  *ai = lfirst_node(A_Indices, idx);
+		Node	   *ai = lfirst(idx);
 
-		if (ai->is_slice)
+		if (IsA(ai, A_Indices) && castNode(A_Indices, ai)->is_slice)
 		{
 			isSlice = true;
 			break;
@@ -312,6 +321,32 @@ transformContainerSubscripts(ParseState *pstate,
 	sbsroutines->transform(sbsref, indirection, pstate,
 						   isSlice, isAssignment);
 
+	/*
+	 * Error out, if datatype failed to consume any indirection elements.
+	 */
+	if (list_length(*indirection) == indirection_length)
+	{
+		Node	   *ind = linitial(*indirection);
+
+		if (noError)
+			return NULL;
+
+		if (IsA(ind, String))
+			ereport(ERROR,
+					(errcode(ERRCODE_DATATYPE_MISMATCH),
+					 errmsg("type %s does not support dot notation",
+							format_type_be(containerType)),
+					 parser_errposition(pstate, exprLocation(containerBase))));
+		else if (IsA(ind, A_Indices))
+			ereport(ERROR,
+					(errcode(ERRCODE_DATATYPE_MISMATCH),
+					 errmsg("type %s does not support array subscripting",
+							format_type_be(containerType)),
+					 parser_errposition(pstate, exprLocation(containerBase))));
+		else
+			elog(ERROR, "invalid indirection operation: %d", nodeTag(ind));
+	}
+
 	/*
 	 * Verify we got a valid type (this defends, for example, against someone
 	 * using array_subscript_handler as typsubscript without setting typelem).
diff --git a/src/backend/parser/parse_target.c b/src/backend/parser/parse_target.c
index 4675a523045..3ef5897f2eb 100644
--- a/src/backend/parser/parse_target.c
+++ b/src/backend/parser/parse_target.c
@@ -937,7 +937,8 @@ transformAssignmentSubscripts(ParseState *pstate,
 										  containerType,
 										  containerTypMod,
 										  &subscripts,
-										  true);
+										  true,
+										  false);
 
 	typeNeeded = sbsref->refrestype;
 	typmodNeeded = sbsref->reftypmod;
diff --git a/src/backend/utils/adt/arraysubs.c b/src/backend/utils/adt/arraysubs.c
index 234c2c278c1..d03d3519dfd 100644
--- a/src/backend/utils/adt/arraysubs.c
+++ b/src/backend/utils/adt/arraysubs.c
@@ -62,6 +62,7 @@ array_subscript_transform(SubscriptingRef *sbsref,
 	List	   *upperIndexpr = NIL;
 	List	   *lowerIndexpr = NIL;
 	ListCell   *idx;
+	int			ndim;
 
 	/*
 	 * Transform the subscript expressions, and separate upper and lower
@@ -73,9 +74,14 @@ array_subscript_transform(SubscriptingRef *sbsref,
 	 */
 	foreach(idx, *indirection)
 	{
-		A_Indices  *ai = lfirst_node(A_Indices, idx);
+		A_Indices  *ai;
 		Node	   *subexpr;
 
+		if (!IsA(lfirst(idx), A_Indices))
+			break;
+
+		ai = lfirst_node(A_Indices, idx);
+
 		if (isSlice)
 		{
 			if (ai->lidx)
@@ -145,14 +151,15 @@ array_subscript_transform(SubscriptingRef *sbsref,
 	sbsref->reflowerindexpr = lowerIndexpr;
 
 	/* Verify subscript list lengths are within implementation limit */
-	if (list_length(upperIndexpr) > MAXDIM)
+	ndim = list_length(upperIndexpr);
+	if (ndim > MAXDIM)
 		ereport(ERROR,
 				(errcode(ERRCODE_PROGRAM_LIMIT_EXCEEDED),
 				 errmsg("number of array dimensions (%d) exceeds the maximum allowed (%d)",
 						list_length(upperIndexpr), MAXDIM)));
 	/* We need not check lowerIndexpr separately */
 
-	*indirection = NIL;
+	*indirection = list_delete_first_n(*indirection, ndim);
 
 	/*
 	 * Determine the result type of the subscripting operation.  It's the same
diff --git a/src/backend/utils/adt/jsonbsubs.c b/src/backend/utils/adt/jsonbsubs.c
index 8ad6aa1ad4f..a0d38a0fd80 100644
--- a/src/backend/utils/adt/jsonbsubs.c
+++ b/src/backend/utils/adt/jsonbsubs.c
@@ -55,9 +55,14 @@ jsonb_subscript_transform(SubscriptingRef *sbsref,
 	 */
 	foreach(idx, *indirection)
 	{
-		A_Indices  *ai = lfirst_node(A_Indices, idx);
+		A_Indices  *ai;
 		Node	   *subExpr;
 
+		if (!IsA(lfirst(idx), A_Indices))
+			break;
+
+		ai = lfirst_node(A_Indices, idx);
+
 		if (isSlice)
 		{
 			Node	   *expr = ai->uidx ? ai->uidx : ai->lidx;
@@ -160,7 +165,9 @@ jsonb_subscript_transform(SubscriptingRef *sbsref,
 	sbsref->refrestype = JSONBOID;
 	sbsref->reftypmod = -1;
 
-	*indirection = NIL;
+	/* Remove processed elements */
+	if (upperIndexpr)
+		*indirection = list_delete_first_n(*indirection, list_length(upperIndexpr));
 }
 
 /*
diff --git a/src/include/parser/parse_node.h b/src/include/parser/parse_node.h
index 5ae11ccec33..71b04bd503c 100644
--- a/src/include/parser/parse_node.h
+++ b/src/include/parser/parse_node.h
@@ -378,7 +378,8 @@ extern SubscriptingRef *transformContainerSubscripts(ParseState *pstate,
 													 Oid containerType,
 													 int32 containerTypMod,
 													 List **indirection,
-													 bool isAssignment);
+													 bool isAssignment,
+													 bool noError);
 extern Const *make_const(ParseState *pstate, A_Const *aconst);
 
 #endif							/* PARSE_NODE_H */
-- 
2.39.5 (Apple Git-154)

