feat: support foreign keys that reference partitioned tables

java/yb-pgsql/src/test/java/org/yb/pgsql/TestPgRegressForeignKey.java

@@ -18,7 +18,7 @@
 import org.yb.YBTestRunner;
 
 @RunWith(value=YBTestRunner.class)
-public class TestPgRegressForeignKey extends BasePgRegressTest {
+public class TestPgRegressForeignKey extends BasePgRegressTestPorted {
   @Override
   protected Map<String, String> getTServerFlags() {
     Map<String, String> flagMap = super.getTServerFlags();

src/postgres/src/backend/commands/tablecmds.c

@@ -9767,16 +9767,6 @@ ATAddForeignKeyConstraint(List **wqueue, AlteredTableInfo *tab, Relation rel,
 	 * Validity checks (permission checks wait till we have the column
 	 * numbers)
 	 */
-	/*
-	 * YB_TODO(feat): begin: Remove after adding support for foreign keys that reference
-	 * partitioned tables
-	 */
-	if (pkrel->rd_rel->relkind == RELKIND_PARTITIONED_TABLE)
-		ereport(ERROR, (errcode(ERRCODE_WRONG_OBJECT_TYPE),
-						errmsg("cannot reference partitioned table \"%s\"",
-							   RelationGetRelationName(pkrel))));
-	/* YB_TODO(feat): end */
-
 	if (rel->rd_rel->relkind == RELKIND_PARTITIONED_TABLE)
 	{
 		if (!recurse)
@@ -12282,6 +12272,7 @@ typedef struct YbFKTriggerScanDescData
 	int buffered_tuples_size;
 	int current_tuple_idx;
 	bool all_tuples_processed;
+	EState* estate;
 	TupleTableSlot* buffered_tuples[];
 } YbFKTriggerScanDescData;
 
@@ -12328,7 +12319,7 @@ YbGetNext(YbFKTriggerScanDesc desc, TupleTableSlot *slot)
 				ExecDropSingleTupleTableSlot(new_slot);
 				break;
 			}
-			YbAddTriggerFKReferenceIntent(desc->trigger, desc->fk_rel, new_slot);
+			YbAddTriggerFKReferenceIntent(desc->trigger, desc->fk_rel, new_slot, desc->estate);
 			desc->buffered_tuples[desc->buffered_tuples_size++] = new_slot;
 		}
 	}
@@ -12367,6 +12358,8 @@ YbFKTriggerScanBegin(TableScanDesc scan,
 					  &YbFKTriggerScanVTableIsYugaByteEnabled :
 					  &YbFKTriggerScanVTableNotYugaByteEnabled;
 	descr->per_batch_cxt = per_batch_cxt;
+	descr->estate = CreateExecutorState();
+	// elog(INFO, "descr->estate %p", descr->estate);
 	return descr;
 }
 
@@ -12457,6 +12450,7 @@ validateForeignKeyConstraint(char *conname,
 	{
 		LOCAL_FCINFO(fcinfo, 0);
 		TriggerData trigdata = {0};
+		// elog(INFO, "TriggerData in validateForeignKeyConstraint");
 
 		CHECK_FOR_INTERRUPTS();
 
@@ -12476,6 +12470,7 @@ validateForeignKeyConstraint(char *conname,
 		trigdata.tg_trigtuple = ExecFetchSlotHeapTuple(ybSlot, false, NULL);
 		trigdata.tg_trigslot = ybSlot;
 		trigdata.tg_trigger = &trig;
+		trigdata.estate = fk_scan->estate;
 
 		fcinfo->context = (Node *) &trigdata;
 
@@ -12490,6 +12485,7 @@ validateForeignKeyConstraint(char *conname,
 	MemoryContextSwitchTo(oldcxt);
 	MemoryContextDelete(perTupCxt);
 	table_endscan(scan);
+	FreeExecutorState(fk_scan->estate);
 	pfree(fk_scan);
 	UnregisterSnapshot(snapshot);
 	if (!IsYBRelation(rel))

src/postgres/src/backend/commands/trigger.c

@@ -2453,6 +2453,7 @@ ExecCallTriggerFunc(TriggerData *trigdata,
 void
 ExecBSInsertTriggers(EState *estate, ResultRelInfo *relinfo)
 {
+	// elog(INFO, "ExecBSInsertTriggers");
 	TriggerDesc *trigdesc;
 	int			i;
 	TriggerData LocTriggerData = {0};
@@ -2473,6 +2474,7 @@ ExecBSInsertTriggers(EState *estate, ResultRelInfo *relinfo)
 	LocTriggerData.tg_event = TRIGGER_EVENT_INSERT |
 		TRIGGER_EVENT_BEFORE;
 	LocTriggerData.tg_relation = relinfo->ri_RelationDesc;
+	LocTriggerData.estate = estate;
 	for (i = 0; i < trigdesc->numtriggers; i++)
 	{
 		Trigger    *trigger = &trigdesc->triggers[i];
@@ -2518,6 +2520,7 @@ bool
 ExecBRInsertTriggers(EState *estate, ResultRelInfo *relinfo,
 					 TupleTableSlot *slot)
 {
+	//elog(INFO, "ExecBRInsertTriggers");
 	TriggerDesc *trigdesc = relinfo->ri_TrigDesc;
 	HeapTuple	newtuple = NULL;
 	bool		should_free;
@@ -2529,6 +2532,7 @@ ExecBRInsertTriggers(EState *estate, ResultRelInfo *relinfo,
 		TRIGGER_EVENT_ROW |
 		TRIGGER_EVENT_BEFORE;
 	LocTriggerData.tg_relation = relinfo->ri_RelationDesc;
+	LocTriggerData.estate = estate;
 	for (i = 0; i < trigdesc->numtriggers; i++)
 	{
 		Trigger    *trigger = &trigdesc->triggers[i];
@@ -2611,6 +2615,7 @@ bool
 ExecIRInsertTriggers(EState *estate, ResultRelInfo *relinfo,
 					 TupleTableSlot *slot)
 {
+	//elog(INFO, "ExecIRInsertTriggers");
 	TriggerDesc *trigdesc = relinfo->ri_TrigDesc;
 	HeapTuple	newtuple = NULL;
 	bool		should_free;
@@ -2622,6 +2627,7 @@ ExecIRInsertTriggers(EState *estate, ResultRelInfo *relinfo,
 		TRIGGER_EVENT_ROW |
 		TRIGGER_EVENT_INSTEAD;
 	LocTriggerData.tg_relation = relinfo->ri_RelationDesc;
+	LocTriggerData.estate = estate;
 	for (i = 0; i < trigdesc->numtriggers; i++)
 	{
 		Trigger    *trigger = &trigdesc->triggers[i];
@@ -2691,6 +2697,7 @@ ExecBSDeleteTriggers(EState *estate, ResultRelInfo *relinfo)
 	LocTriggerData.tg_event = TRIGGER_EVENT_DELETE |
 		TRIGGER_EVENT_BEFORE;
 	LocTriggerData.tg_relation = relinfo->ri_RelationDesc;
+	LocTriggerData.estate = estate;
 	for (i = 0; i < trigdesc->numtriggers; i++)
 	{
 		Trigger    *trigger = &trigdesc->triggers[i];
@@ -2788,6 +2795,7 @@ ExecBRDeleteTriggers(EState *estate, EPQState *epqstate,
 		TRIGGER_EVENT_ROW |
 		TRIGGER_EVENT_BEFORE;
 	LocTriggerData.tg_relation = relinfo->ri_RelationDesc;
+	LocTriggerData.estate = estate;
 	for (i = 0; i < trigdesc->numtriggers; i++)
 	{
 		HeapTuple	newtuple;
@@ -2878,6 +2886,7 @@ ExecIRDeleteTriggers(EState *estate, ResultRelInfo *relinfo,
 		TRIGGER_EVENT_ROW |
 		TRIGGER_EVENT_INSTEAD;
 	LocTriggerData.tg_relation = relinfo->ri_RelationDesc;
+	LocTriggerData.estate = estate;
 
 	ExecForceStoreHeapTuple(trigtuple, slot, false);
 
@@ -2940,6 +2949,7 @@ ExecBSUpdateTriggers(EState *estate, ResultRelInfo *relinfo)
 	LocTriggerData.tg_event = TRIGGER_EVENT_UPDATE |
 		TRIGGER_EVENT_BEFORE;
 	LocTriggerData.tg_relation = relinfo->ri_RelationDesc;
+	LocTriggerData.estate = estate;
 	LocTriggerData.tg_updatedcols = updatedCols;
 	for (i = 0; i < trigdesc->numtriggers; i++)
 	{
@@ -3057,6 +3067,7 @@ ExecBRUpdateTriggers(EState *estate, EPQState *epqstate,
 		TRIGGER_EVENT_ROW |
 		TRIGGER_EVENT_BEFORE;
 	LocTriggerData.tg_relation = relinfo->ri_RelationDesc;
+	LocTriggerData.estate = estate;
 	updatedCols = ExecGetAllUpdatedCols(relinfo, estate);
 	LocTriggerData.tg_updatedcols = updatedCols;
 	for (i = 0; i < trigdesc->numtriggers; i++)
@@ -3205,6 +3216,7 @@ ExecIRUpdateTriggers(EState *estate, ResultRelInfo *relinfo,
 		TRIGGER_EVENT_ROW |
 		TRIGGER_EVENT_INSTEAD;
 	LocTriggerData.tg_relation = relinfo->ri_RelationDesc;
+	LocTriggerData.estate = estate;
 
 	ExecForceStoreHeapTuple(trigtuple, oldslot, false);
 
@@ -3273,6 +3285,7 @@ ExecBSTruncateTriggers(EState *estate, ResultRelInfo *relinfo)
 	LocTriggerData.tg_event = TRIGGER_EVENT_TRUNCATE |
 		TRIGGER_EVENT_BEFORE;
 	LocTriggerData.tg_relation = relinfo->ri_RelationDesc;
+	LocTriggerData.estate = estate;
 
 	for (i = 0; i < trigdesc->numtriggers; i++)
 	{
@@ -4510,6 +4523,8 @@ AfterTriggerExecute(EState *estate,
 	LocTriggerData.tg_event =
 		evtshared->ats_event & (TRIGGER_EVENT_OPMASK | TRIGGER_EVENT_ROW);
 	LocTriggerData.tg_relation = rel;
+	LocTriggerData.estate = estate;
+
 	if (TRIGGER_FOR_UPDATE(LocTriggerData.tg_trigger->tgtype))
 		LocTriggerData.tg_updatedcols = evtshared->ats_modifiedcols;
 
@@ -6126,6 +6141,13 @@ AfterTriggerSaveEvent(EState *estate, ResultRelInfo *relinfo,
 					  bool is_crosspart_update)
 {
 	Relation	rel = relinfo->ri_RelationDesc;
+	// elog(INFO, "AfterTriggerSaveEvent relation: %s, event %d", RelationGetRelationName(rel), event);
+
+	// Relation srcrel = NULL;
+	// if (src_partinfo)
+	// 	srcrel = src_partinfo->ri_RelationDesc;
+	// elog(INFO, "AfterTriggerSaveEvent rel: %s partrel: %s, src_partinfo: %p", RelationGetRelationName(rel), srcrel ? RelationGetRelationName(rel) : "NULL", src_partinfo);
+
 	TriggerDesc *trigdesc = relinfo->ri_TrigDesc;
 	AfterTriggerEventData new_event;
 	AfterTriggerSharedData new_shared;
@@ -6419,13 +6441,17 @@ AfterTriggerSaveEvent(EState *estate, ResultRelInfo *relinfo,
 					if (is_crosspart_update &&
 						TRIGGER_FIRED_BY_DELETE(event) &&
 						trigger->tgisclone)
+					{
+						// elog(INFO, "Skipping this RI_TRIGGER_PK event (delete, cross partition)");
 						continue;
+					}
 
 					/* Update or delete on trigger's PK table */
 					if (!RI_FKey_pk_upd_check_required(trigger, rel,
 													   oldslot, newslot,
 													   &estate->yb_skip_entities))
 					{
+						// elog(INFO, "Skipping this RI_TRIGGER_PK event");
 						/* skip queuing this event */
 						continue;
 					}
@@ -6449,6 +6475,7 @@ AfterTriggerSaveEvent(EState *estate, ResultRelInfo *relinfo,
 													   oldslot, newslot,
 													   &estate->yb_skip_entities))
 					{
+						// elog(INFO, "Skipping this RI_TRIGGER_FK event");
 						/* skip queuing this event */
 						continue;
 					}
@@ -6547,7 +6574,7 @@ AfterTriggerSaveEvent(EState *estate, ResultRelInfo *relinfo,
 		}
 
 		if (IsYBBackedRelation(rel) && RI_FKey_trigger_type(trigger->tgfoid) == RI_TRIGGER_FK)
-			YbAddTriggerFKReferenceIntent(trigger, rel, newslot);
+			YbAddTriggerFKReferenceIntent(trigger, rel, newslot, estate);
 
 		afterTriggerAddEvent(&afterTriggers.query_stack[afterTriggers.query_depth].events,
 							 &new_event, &new_shared);

src/postgres/src/backend/executor/execMain.c

@@ -1832,6 +1832,7 @@ ExecPartitionCheck(ResultRelInfo *resultRelInfo, TupleTableSlot *slot,
 		List	   *qual = RelationGetPartitionQual(resultRelInfo->ri_RelationDesc);
 
 		resultRelInfo->ri_PartitionCheckExpr = ExecPrepareCheck(qual, estate);
+		// elog(INFO, "resultRelInfo->ri_PartitionCheckExpr was null, now updated to %p", resultRelInfo->ri_PartitionCheckExpr);
 		MemoryContextSwitchTo(oldcxt);
 	}
 

src/postgres/src/backend/executor/execPartition.c

@@ -489,6 +489,129 @@ ExecFindPartition(ModifyTableState *mtstate,
 	return rri;
 }
 
+Oid
+FindLeafPartitionOid(ResultRelInfo *rootResultRelInfo,
+					 PartitionTupleRouting *proute, TupleTableSlot *slot,
+					 EState *estate)
+{
+	PartitionDispatch *pd = proute->partition_dispatch_info;
+	Datum values[PARTITION_MAX_KEYS];
+	bool isnull[PARTITION_MAX_KEYS];
+	// Relation rel;
+	ExprContext *ecxt = GetPerTupleExprContext(estate);
+	TupleTableSlot *ecxt_scantuple_saved = ecxt->ecxt_scantuple;
+	PartitionDispatch dispatch;
+	PartitionDesc partdesc;
+	TupleTableSlot *myslot = NULL;
+	MemoryContext oldcxt;
+	Oid resultOid = InvalidOid;
+
+	/* use per-tuple context here to avoid leaking memory */
+	oldcxt = MemoryContextSwitchTo(GetPerTupleMemoryContext(estate));
+
+	/*
+	 * First check the root table's partition constraint, if any.  No point in
+	 * routing the tuple if it doesn't belong in the root table itself.
+	 */
+	if (rootResultRelInfo->ri_RelationDesc->rd_rel->relispartition)
+		ExecPartitionCheck(rootResultRelInfo, slot, estate, true);
+
+	/* start with the root partitioned table */
+	dispatch = pd[0];
+	while (dispatch != NULL)
+	{
+		int partidx = -1;
+
+		CHECK_FOR_INTERRUPTS();
+
+		// rel = dispatch->reldesc;
+		partdesc = dispatch->partdesc;
+
+		/*
+		 * Extract partition key from tuple. Expression evaluation machinery
+		 * that FormPartitionKeyDatum() invokes expects ecxt_scantuple to
+		 * point to the correct tuple slot.  The slot might have changed from
+		 * what was used for the parent table if the table of the current
+		 * partitioning level has different tuple descriptor from the parent.
+		 * So update ecxt_scantuple accordingly.
+		 */
+		ecxt->ecxt_scantuple = slot;
+		FormPartitionKeyDatum(dispatch, slot, estate, values, isnull);
+
+		/*
+		 * If this partitioned table has no partitions or no partition for
+		 * these values, error out.
+		 */
+		if (partdesc->nparts == 0 ||
+			(partidx = get_partition_for_tuple(dispatch, values, isnull)) < 0)
+			break;
+
+		if (partdesc->is_leaf[partidx])
+		{
+			resultOid = partdesc->oids[partidx];
+			break;
+		}
+		else
+		{
+			/*
+			 * Partition is a sub-partitioned table; get the PartitionDispatch
+			 */
+			if (likely(dispatch->indexes[partidx] >= 0))
+			{
+				/* Already built. */
+				Assert(dispatch->indexes[partidx] < proute->num_dispatch);
+				/*
+				 * Move down to the next partition level and search again
+				 * until we find a leaf partition that matches this tuple
+				 */
+				dispatch = pd[dispatch->indexes[partidx]];
+			}
+			else
+			{
+				/* Not yet built. Do that now. */
+				PartitionDispatch subdispatch;
+
+				/*
+				 * Create the new PartitionDispatch.  We pass the current one
+				 * in as the parent PartitionDispatch
+				 */
+				subdispatch = ExecInitPartitionDispatchInfo(
+					estate, proute, partdesc->oids[partidx], dispatch, partidx,
+					NULL);
+				Assert(dispatch->indexes[partidx] >= 0 &&
+					   dispatch->indexes[partidx] < proute->num_dispatch);
+
+				dispatch = subdispatch;
+			}
+
+			/*
+			 * Convert the tuple to the new parent's layout, if different from
+			 * the previous parent.
+			 */
+			if (dispatch->tupslot)
+			{
+				AttrMap *map = dispatch->tupmap;
+				TupleTableSlot *tempslot = myslot;
+
+				myslot = dispatch->tupslot;
+				slot = execute_attr_map_slot(map, slot, myslot);
+
+				if (tempslot != NULL)
+					ExecClearTuple(tempslot);
+			}
+		}
+	}
+
+	/* Release the tuple in the lowest parent's dedicated slot. */
+	if (myslot != NULL)
+		ExecClearTuple(myslot);
+	/* and restore ecxt's scantuple */
+	ecxt->ecxt_scantuple = ecxt_scantuple_saved;
+	MemoryContextSwitchTo(oldcxt);
+	// ExecPartitionCheckEmitError(rootResultRelInfo, slot, estate);
+	return resultOid;
+}
+
 /*
  * ExecInitPartitionInfo
  *		Lock the partition and initialize ResultRelInfo.  Also setup other
@@ -1312,6 +1435,8 @@ FormPartitionKeyDatum(PartitionDispatch pd,
 	ListCell   *partexpr_item;
 	int			i;
 
+	// bool is_virtual = TTS_IS_VIRTUAL(slot);
+	// elog(INFO, "FormPartitionKeyDatum is_virtual %d", is_virtual);
 	if (pd->key->partexprs != NIL && pd->keystate == NIL)
 	{
 		/* Check caller has set up context correctly */
@@ -1332,6 +1457,12 @@ FormPartitionKeyDatum(PartitionDispatch pd,
 		if (keycol != 0)
 		{
 			/* Plain column; get the value directly from the heap tuple */
+			// if (is_virtual)
+			// {
+			// 	datum = slot->tts_values[keycol-1];
+			// 	isNull = slot->tts_isnull[keycol-1];
+			// }
+			// else
 			datum = slot_getattr(slot, keycol, &isNull);
 		}
 		else