From 3958bc58e6ff8f56552b520fc9126dbef333d77c Mon Sep 17 00:00:00 2001
From: Metin Doslu <metin@citusdata.com>
Date: Wed, 21 Jun 2017 16:43:13 +0300
Subject: [PATCH] fix 9.5

---
 src/backend/distributed/utils/ruleutils_95.c | 35 ++++++++++++++++----
 1 file changed, 29 insertions(+), 6 deletions(-)

diff --git a/src/backend/distributed/utils/ruleutils_95.c b/src/backend/distributed/utils/ruleutils_95.c
index f8a60fe7d..8c4771886 100644
--- a/src/backend/distributed/utils/ruleutils_95.c
+++ b/src/backend/distributed/utils/ruleutils_95.c
@@ -35,6 +35,7 @@
 #include "catalog/pg_type.h"
 #include "distributed/citus_nodefuncs.h"
 #include "distributed/citus_ruleutils.h"
+#include "distributed/multi_router_planner.h"
 #include "commands/defrem.h"
 #include "commands/extension.h"
 #include "foreign/foreign.h"
@@ -3144,17 +3145,39 @@ get_update_query_def(Query *query, deparse_context *context)
 	 * Start the query with UPDATE relname SET
 	 */
 	rte = rt_fetch(query->resultRelation, query->rtable);
-	Assert(rte->rtekind == RTE_RELATION);
+	Assert(rte->rtekind == RTE_RELATION || UpdateFromQuery(query));
 	if (PRETTY_INDENT(context))
 	{
 		appendStringInfoChar(buf, ' ');
 		context->indentLevel += PRETTYINDENT_STD;
 	}
-	appendStringInfo(buf, "UPDATE %s%s",
-					 only_marker(rte),
-					 generate_relation_or_shard_name(rte->relid,
-													 context->distrelid,
-													 context->shardid, NIL));
+
+	/* if it's a shard, do differently */
+	if (GetRangeTblKind(rte) == CITUS_RTE_SHARD)
+	{
+		char *fragmentSchemaName = NULL;
+		char *fragmentTableName = NULL;
+
+		ExtractRangeTblExtraData(rte, NULL, &fragmentSchemaName, &fragmentTableName, NULL);
+
+		/* Use schema and table name from the remote alias */
+		appendStringInfo(buf, "UPDATE %s%s",
+						 only_marker(rte),
+						 generate_fragment_name(fragmentSchemaName, fragmentTableName));
+
+		if(rte->eref != NULL)
+			appendStringInfo(buf, " %s",
+					quote_identifier(rte->eref->aliasname));
+	}
+	else
+	{
+		appendStringInfo(buf, "UPDATE %s%s",
+						 only_marker(rte),
+						 generate_relation_or_shard_name(rte->relid,
+														 context->distrelid,
+														 context->shardid, NIL));
+	}
+
 	if (rte->alias != NULL)
 		appendStringInfo(buf, " %s",
 						 quote_identifier(rte->alias->aliasname));