Attempt to address planner context crashes

Both of these are a bit of a shot in the dark. In one case, we noticed
a stack trace where a caller received a null pointer and attempted to
dereference the memory context field (at 0x010). In the other, I saw
that any error thrown from within AdjustParseTree could keep the stack
from being cleaned up (presumably if we push we should always pop).

Both stack traces were collected during times of high memory pressure
and locally reproducing the problem locally or otherwise has been very
tricky (i.e. it hasn't been reproduced reliably at all).

src/backend/distributed/planner/distributed_planner.c

@@ -157,6 +157,10 @@ distributed_planner(Query *parse, int cursorOptions, ParamListInfo boundParams)
 
 			result = CreateDistributedPlannedStmt(planId, result, originalQuery, parse,
 												  boundParams, plannerRestrictionContext);
+
+			setPartitionedTablesInherited = true;
+			AdjustPartitioningForDistributedPlanning(parse,
+													 setPartitionedTablesInherited);
 		}
 	}
 	PG_CATCH();
@@ -166,13 +170,6 @@ distributed_planner(Query *parse, int cursorOptions, ParamListInfo boundParams)
 	}
 	PG_END_TRY();
 
-	if (needsDistributedPlanning)
-	{
-		setPartitionedTablesInherited = true;
-
-		AdjustPartitioningForDistributedPlanning(parse, setPartitionedTablesInherited);
-	}
-
 	/* remove the context from the context list */
 	PopPlannerRestrictionContext();
 
@@ -1517,6 +1514,13 @@ CurrentPlannerRestrictionContext(void)
 	plannerRestrictionContext =
 		(PlannerRestrictionContext *) linitial(plannerRestrictionContextList);
 
+	if (plannerRestrictionContext == NULL)
+	{
+		ereport(ERROR, (errcode(ERRCODE_INTERNAL_ERROR),
+						errmsg("planner restriction context stack was empty"),
+						errdetail("Please report this to the Citus core team.")));
+	}
+
 	return plannerRestrictionContext;
 }