1 files changed, 44 insertions, 1 deletions

diff --git a/src/backend/commands/trigger.c b/src/backend/commands/trigger.c
index 7a5ffe32f60..97c087929f3 100644
--- a/src/backend/commands/trigger.c
+++ b/src/backend/commands/trigger.c
@@ -43,6 +43,7 @@
 #include "parser/parse_relation.h"
 #include "partitioning/partdesc.h"
 #include "pgstat.h"
+#include "rewrite/rewriteHandler.h"
 #include "rewrite/rewriteManip.h"
 #include "storage/lmgr.h"
 #include "utils/acl.h"
@@ -101,6 +102,7 @@ static void AfterTriggerSaveEvent(EState *estate, ResultRelInfo *relinfo,
 								  bool is_crosspart_update);
 static void AfterTriggerEnlargeQueryState(void);
 static bool before_stmt_triggers_fired(Oid relid, CmdType cmdType);
+static HeapTuple check_modified_virtual_generated(TupleDesc tupdesc, HeapTuple tuple);
 
 
 /*
@@ -641,7 +643,8 @@ CreateTriggerFiringOn(CreateTrigStmt *stmt, const char *queryString,
 					if (TRIGGER_FOR_BEFORE(tgtype) &&
 						var->varattno == 0 &&
 						RelationGetDescr(rel)->constr &&
-						RelationGetDescr(rel)->constr->has_generated_stored)
+						(RelationGetDescr(rel)->constr->has_generated_stored ||
+						 RelationGetDescr(rel)->constr->has_generated_virtual))
 						ereport(ERROR,
 								(errcode(ERRCODE_INVALID_OBJECT_DEFINITION),
 								 errmsg("BEFORE trigger's WHEN condition cannot reference NEW generated columns"),
@@ -2504,6 +2507,8 @@ ExecBRInsertTriggers(EState *estate, ResultRelInfo *relinfo,
 		}
 		else if (newtuple != oldtuple)
 		{
+			newtuple = check_modified_virtual_generated(RelationGetDescr(relinfo->ri_RelationDesc), newtuple);
+
 			ExecForceStoreHeapTuple(newtuple, slot, false);
 
 			/*
@@ -3061,6 +3066,8 @@ ExecBRUpdateTriggers(EState *estate, EPQState *epqstate,
 		}
 		else if (newtuple != oldtuple)
 		{
+			newtuple = check_modified_virtual_generated(RelationGetDescr(relinfo->ri_RelationDesc), newtuple);
+
 			ExecForceStoreHeapTuple(newtuple, newslot, false);
 
 			/*
@@ -3491,6 +3498,8 @@ TriggerEnabled(EState *estate, ResultRelInfo *relinfo,
 
 			oldContext = MemoryContextSwitchTo(estate->es_query_cxt);
 			tgqual = stringToNode(trigger->tgqual);
+			tgqual = expand_generated_columns_in_expr(tgqual, relinfo->ri_RelationDesc, PRS2_OLD_VARNO);
+			tgqual = expand_generated_columns_in_expr(tgqual, relinfo->ri_RelationDesc, PRS2_NEW_VARNO);
 			/* Change references to OLD and NEW to INNER_VAR and OUTER_VAR */
 			ChangeVarNodes(tgqual, PRS2_OLD_VARNO, INNER_VAR, 0);
 			ChangeVarNodes(tgqual, PRS2_NEW_VARNO, OUTER_VAR, 0);
@@ -6621,3 +6630,37 @@ pg_trigger_depth(PG_FUNCTION_ARGS)
 {
 	PG_RETURN_INT32(MyTriggerDepth);
 }
+
+/*
+ * Check whether a trigger modified a virtual generated column and replace the
+ * value with null if so.
+ *
+ * We need to check this so that we don't end up storing a non-null value in a
+ * virtual generated column.
+ *
+ * We don't need to check for stored generated columns, since those will be
+ * overwritten later anyway.
+ */
+static HeapTuple
+check_modified_virtual_generated(TupleDesc tupdesc, HeapTuple tuple)
+{
+	if (!(tupdesc->constr && tupdesc->constr->has_generated_virtual))
+		return tuple;
+
+	for (int i = 0; i < tupdesc->natts; i++)
+	{
+		if (TupleDescAttr(tupdesc, i)->attgenerated == ATTRIBUTE_GENERATED_VIRTUAL)
+		{
+			if (!heap_attisnull(tuple, i + 1, tupdesc))
+			{
+				int			replCol = i + 1;
+				Datum		replValue = 0;
+				bool		replIsnull = true;
+
+				tuple = heap_modify_tuple_by_cols(tuple, tupdesc, 1, &replCol, &replValue, &replIsnull);
+			}
+		}
+	}
+
+	return tuple;
+}