Merge pull request #2248 from citusdata/partition_key_extraction_fix

Strip implicit coercions when determining partition key value of an INSERT

src/backend/distributed/planner/multi_router_planner.c

@@ -2661,6 +2661,7 @@ ExtractInsertPartitionKeyValue(Query *query)
 	Var *partitionColumn = NULL;
 	TargetEntry *targetEntry = NULL;
 	Const *singlePartitionValueConst = NULL;
+	Node *targetExpression = NULL;
 
 	char partitionMethod = PartitionMethod(distributedTableId);
 	if (partitionMethod == DISTRIBUTE_BY_NONE)
@@ -2676,13 +2677,15 @@ ExtractInsertPartitionKeyValue(Query *query)
 		return NULL;
 	}
 
+	targetExpression = strip_implicit_coercions((Node *) targetEntry->expr);
+
 	/*
 	 * Multi-row INSERTs have a Var in the target list that points to
 	 * an RTE_VALUES.
 	 */
-	if (IsA(targetEntry->expr, Var))
+	if (IsA(targetExpression, Var))
 	{
-		Var *partitionVar = (Var *) targetEntry->expr;
+		Var *partitionVar = (Var *) targetExpression;
 		RangeTblEntry *referencedRTE = NULL;
 		ListCell *valuesListCell = NULL;
 
@@ -2691,7 +2694,9 @@ ExtractInsertPartitionKeyValue(Query *query)
 		foreach(valuesListCell, referencedRTE->values_lists)
 		{
 			List *rowValues = (List *) lfirst(valuesListCell);
-			Expr *partitionValueExpr = list_nth(rowValues, partitionVar->varattno - 1);
+			Node *partitionValueNode = list_nth(rowValues, partitionVar->varattno - 1);
+			Expr *partitionValueExpr = (Expr *) strip_implicit_coercions(
+				partitionValueNode);
 			Const *partitionValueConst = NULL;
 
 			if (!IsA(partitionValueExpr, Const))
@@ -2720,10 +2725,10 @@ ExtractInsertPartitionKeyValue(Query *query)
 			}
 		}
 	}
-	else if (IsA(targetEntry->expr, Const))
+	else if (IsA(targetExpression, Const))
 	{
 		/* single-row INSERT with a constant partition column value */
-		singlePartitionValueConst = (Const *) targetEntry->expr;
+		singlePartitionValueConst = (Const *) targetExpression;
 	}
 	else
 	{

src/test/regress/expected/multi_mx_router_planner.out

@@ -62,6 +62,7 @@ SET client_min_messages TO 'DEBUG2';
 INSERT INTO articles_single_shard_hash_mx VALUES (50, 10, 'anjanette', 19519);
 DEBUG:  Creating router plan
 DEBUG:  Plan is router executable
+DETAIL:  distribution column value: 10
 -- single-shard tests
 -- test simple select for a single row
 SELECT * FROM articles_hash_mx WHERE author_id = 10 AND id = 50;
@@ -1460,6 +1461,7 @@ CREATE MATERIALIZED VIEW mv_articles_hash_mx_error AS
 INSERT INTO articles_hash_mx VALUES (51,  1, 'amateus', 1814);
 DEBUG:  Creating router plan
 DEBUG:  Plan is router executable
+DETAIL:  distribution column value: 1
 -- verify insert is successfull (not router plannable and executable)
 SELECT id
 	FROM articles_hash_mx

src/test/regress/expected/multi_router_planner.out

@@ -122,6 +122,7 @@ SET client_min_messages TO 'DEBUG2';
 INSERT INTO articles_single_shard_hash VALUES (50, 10, 'anjanette', 19519);
 DEBUG:  Creating router plan
 DEBUG:  Plan is router executable
+DETAIL:  distribution column value: 10
 -- single-shard tests
 -- test simple select for a single row
 SELECT * FROM articles_hash WHERE author_id = 10 AND id = 50;
@@ -455,6 +456,7 @@ DEBUG:  data-modifying statements are not supported in the WITH clauses of distr
 DEBUG:  generating subplan 81_1 for CTE new_article: INSERT INTO public.articles_hash (id, author_id, title, word_count) VALUES (1, 1, 'arsenous'::character varying, 9) RETURNING id, author_id, title, word_count
 DEBUG:  Creating router plan
 DEBUG:  Plan is router executable
+DETAIL:  distribution column value: 1
 DEBUG:  Plan 81 query after replacing subqueries and CTEs: SELECT id, author_id, title, word_count FROM (SELECT intermediate_result.id, intermediate_result.author_id, intermediate_result.title, intermediate_result.word_count FROM read_intermediate_result('81_1'::text, 'binary'::citus_copy_format) intermediate_result(id bigint, author_id bigint, title character varying(20), word_count integer)) new_article
 DEBUG:  Creating router plan
 DEBUG:  Plan is router executable
@@ -2305,7 +2307,7 @@ DETAIL:  distribution column value: 1
 (5 rows)
 
 -- insert query is router plannable even under task-tracker
-INSERT INTO articles_hash VALUES (51,  1, 'amateus', 1814);
+INSERT INTO articles_hash VALUES (51, 1, 'amateus', 1814), (52, 1, 'second amateus', 2824);
 DEBUG:  Creating router plan
 DEBUG:  Plan is router executable
 -- verify insert is successfull (not router plannable and executable)
@@ -2323,7 +2325,8 @@ DETAIL:  distribution column value: 1
  31
  41
  51
-(6 rows)
+ 52
+(7 rows)
 
 SET client_min_messages to 'NOTICE';
 -- test that a connection failure marks placements invalid

src/test/regress/sql/multi_router_planner.sql

@@ -1043,7 +1043,7 @@ SELECT id
 	WHERE author_id = 1;
 
 -- insert query is router plannable even under task-tracker
-INSERT INTO articles_hash VALUES (51,  1, 'amateus', 1814);
+INSERT INTO articles_hash VALUES (51, 1, 'amateus', 1814), (52, 1, 'second amateus', 2824);
 
 -- verify insert is successfull (not router plannable and executable)
 SELECT id