Fix missing from entry

src/backend/distributed/planner/deparse_shard_query.c

@@ -39,7 +39,6 @@
 #include "utils/syscache.h"
 
 
-static void AddInsertAliasIfNeeded(Query *query);
 static void UpdateTaskQueryString(Query *query, Task *task);
 static RelationShard * FindRelationShard(Oid inputRelationId, List *relationShardList);
 static void ConvertRteToSubqueryWithEmptyResult(RangeTblEntry *rte);
@@ -159,7 +158,7 @@ RebuildQueryStrings(Job *workerJob)
  * deparsing issues (e.g. RETURNING might reference the original table name,
  * which has been replaced by a shard name).
  */
-static void
+void
 AddInsertAliasIfNeeded(Query *query)
 {
 	Assert(query->commandType == CMD_INSERT);

src/backend/distributed/planner/local_plan_cache.c

@@ -174,6 +174,8 @@ DeparseLocalShardQuery(Query *jobQuery, List *relationShardList, Oid
 		 */
 		Assert(!CheckInsertSelectQuery(jobQuery));
 
+		AddInsertAliasIfNeeded(jobQuery);
+
 		/*
 		 * For INSERT queries we cannot use pg_get_query_def. Mainly because we
 		 * cannot run UpdateRelationToShardNames on an INSERT query. This is

src/include/distributed/deparse_shard_query.h

@@ -29,6 +29,7 @@ extern void SetTaskQueryStringList(Task *task, List *queryStringList);
 extern char * TaskQueryString(Task *task);
 extern char * TaskQueryStringAtIndex(Task *task, int index);
 extern int GetTaskQueryType(Task *task);
+extern void AddInsertAliasIfNeeded(Query *query);
 
 
 #endif /* DEPARSE_SHARD_QUERY_H */

src/test/regress/expected/distributed_planning.out

@@ -537,6 +537,233 @@ ORDER BY 1;
 (4 rows)
 
 ROLLBACK;
+prepare p1 as INSERT INTO t1(a,c) VALUES (15, 15) ON CONFLICT (c) DO UPDATE SET a=EXCLUDED.a + 10 RETURNING *;
+execute p1(5);
+ a  | c
+---------------------------------------------------------------------
+ 15 | 15
+(1 row)
+
+execute p1(5);
+ a  | c
+---------------------------------------------------------------------
+ 25 | 15
+(1 row)
+
+execute p1(5);
+ a  | c
+---------------------------------------------------------------------
+ 25 | 15
+(1 row)
+
+execute p1(5);
+ a  | c
+---------------------------------------------------------------------
+ 25 | 15
+(1 row)
+
+execute p1(5);
+ a  | c
+---------------------------------------------------------------------
+ 25 | 15
+(1 row)
+
+execute p1(5);
+ a  | c
+---------------------------------------------------------------------
+ 25 | 15
+(1 row)
+
+execute p1(5);
+ a  | c
+---------------------------------------------------------------------
+ 25 | 15
+(1 row)
+
+prepare p5(int) as INSERT INTO t1(a,c) VALUES (15, $1) ON CONFLICT (c) DO UPDATE SET a=EXCLUDED.a + 10 RETURNING *;
+execute p5(5);
+ a  | c
+---------------------------------------------------------------------
+ 15 | 5
+(1 row)
+
+execute p5(5);
+ a  | c
+---------------------------------------------------------------------
+ 25 | 5
+(1 row)
+
+execute p5(5);
+ a  | c
+---------------------------------------------------------------------
+ 25 | 5
+(1 row)
+
+execute p5(5);
+ a  | c
+---------------------------------------------------------------------
+ 25 | 5
+(1 row)
+
+execute p5(5);
+ a  | c
+---------------------------------------------------------------------
+ 25 | 5
+(1 row)
+
+execute p5(5);
+ a  | c
+---------------------------------------------------------------------
+ 25 | 5
+(1 row)
+
+execute p5(5);
+ a  | c
+---------------------------------------------------------------------
+ 25 | 5
+(1 row)
+
+INSERT INTO "companies" ("id","meta","name","created_at","updated_at","deleted_at") VALUES (1,'{"test":123}','Name','2016-11-07 17:34:22.101807','2021-05-20 22:16:55.424521',NULL) ON CONFLICT (id) DO UPDATE SET "meta"=EXCLUDED."meta"  RETURNING "id";
+ id
+---------------------------------------------------------------------
+  1
+(1 row)
+
+INSERT INTO "companies" ("id","meta","name","created_at","updated_at","deleted_at") VALUES (1,'{"test":123}','Name','2016-11-07 17:34:22.101807','2021-05-20 22:16:55.424521',NULL) ON CONFLICT (id) DO UPDATE SET "meta"=EXCLUDED."meta"  RETURNING "id";
+ id
+---------------------------------------------------------------------
+  1
+(1 row)
+
+PREPARE p6 AS INSERT INTO "companies" ("id","meta","name","created_at","updated_at","deleted_at") VALUES (1,'{"test":123}','Name','2016-11-07 17:34:22.101807','2021-05-20 22:16:55.424521',NULL) ON CONFLICT (id) DO UPDATE SET "meta"=EXCLUDED."meta"  RETURNING "id";;
+EXECUTE p6;
+ id
+---------------------------------------------------------------------
+  1
+(1 row)
+
+EXECUTE p6;
+ id
+---------------------------------------------------------------------
+  1
+(1 row)
+
+EXECUTE p6;
+ id
+---------------------------------------------------------------------
+  1
+(1 row)
+
+EXECUTE p6;
+ id
+---------------------------------------------------------------------
+  1
+(1 row)
+
+EXECUTE p6;
+ id
+---------------------------------------------------------------------
+  1
+(1 row)
+
+EXECUTE p6;
+ id
+---------------------------------------------------------------------
+  1
+(1 row)
+
+EXECUTE p6;
+ id
+---------------------------------------------------------------------
+  1
+(1 row)
+
+prepare insert_select(int) as insert into companies SELECT * FROM companies WHERE id >= $1 ON CONFLICT(id) DO UPDATE SET "meta"=EXCLUDED."meta" RETURNING "id";;
+EXECUTE insert_select(1);
+ id
+---------------------------------------------------------------------
+  1
+(1 row)
+
+EXECUTE insert_select(1);
+ id
+---------------------------------------------------------------------
+  1
+(1 row)
+
+EXECUTE insert_select(1);
+ id
+---------------------------------------------------------------------
+  1
+(1 row)
+
+EXECUTE insert_select(1);
+ id
+---------------------------------------------------------------------
+  1
+(1 row)
+
+EXECUTE insert_select(1);
+ id
+---------------------------------------------------------------------
+  1
+(1 row)
+
+EXECUTE insert_select(1);
+ id
+---------------------------------------------------------------------
+  1
+(1 row)
+
+EXECUTE insert_select(1);
+ id
+---------------------------------------------------------------------
+  1
+(1 row)
+
+prepare insert_select_1 as insert into companies SELECT * FROM companies WHERE id >= 1 ON CONFLICT(id) DO UPDATE SET "meta"=EXCLUDED."meta" RETURNING "id";;
+EXECUTE insert_select_1;
+ id
+---------------------------------------------------------------------
+  1
+(1 row)
+
+EXECUTE insert_select_1;
+ id
+---------------------------------------------------------------------
+  1
+(1 row)
+
+EXECUTE insert_select_1;
+ id
+---------------------------------------------------------------------
+  1
+(1 row)
+
+EXECUTE insert_select_1;
+ id
+---------------------------------------------------------------------
+  1
+(1 row)
+
+EXECUTE insert_select_1;
+ id
+---------------------------------------------------------------------
+  1
+(1 row)
+
+EXECUTE insert_select_1;
+ id
+---------------------------------------------------------------------
+  1
+(1 row)
+
+EXECUTE insert_select_1;
+ id
+---------------------------------------------------------------------
+  1
+(1 row)
+
 -- query fails on the shards should be handled
 -- nicely
 \set VERBOSITY terse

src/test/regress/expected/distributed_planning_create_load.out

@@ -90,3 +90,36 @@ INSERT INTO upsert_test (part_key, other_col) VALUES (1, 1), (2, 2) RETURNING *;
         2 |         2 |
 (2 rows)
 
+create table t1(a int, b int, c int primary key, d int);
+ALTER TABLE t1 DROP COLUMN b;
+select create_distributed_table('t1','c');
+ create_distributed_table
+---------------------------------------------------------------------
+
+(1 row)
+
+ALTER TABLE t1 DROP COLUMN d;
+CREATE TABLE companies (
+    id bigint NOT NULL,
+    name character varying NOT NULL,
+    created_at timestamp without time zone NOT NULL,
+    a int,
+    b int,
+    updated_at timestamp without time zone NOT NULL,
+    fmi boolean DEFAULT false NOT NULL,
+    meta jsonb DEFAULT '{}'::jsonb NOT NULL,
+    deleted_at timestamp without time zone,
+    c int,
+    flex boolean DEFAULT false
+);
+ALTER TABLE ONLY companies
+    ADD CONSTRAINT companies_pkey PRIMARY KEY (id);
+alter table companies drop column a;
+alter table companies drop column b;
+SELECT create_reference_table('companies');
+ create_reference_table
+---------------------------------------------------------------------
+
+(1 row)
+
+alter table companies drop column c;

src/test/regress/sql/distributed_planning.sql

@@ -298,6 +298,53 @@ ORDER BY 1;
 
 ROLLBACK;
 
+prepare p1 as INSERT INTO t1(a,c) VALUES (15, 15) ON CONFLICT (c) DO UPDATE SET a=EXCLUDED.a + 10 RETURNING *;
+execute p1(5);
+execute p1(5);
+execute p1(5);
+execute p1(5);
+execute p1(5);
+execute p1(5);
+execute p1(5);
+
+prepare p5(int) as INSERT INTO t1(a,c) VALUES (15, $1) ON CONFLICT (c) DO UPDATE SET a=EXCLUDED.a + 10 RETURNING *;
+execute p5(5);
+execute p5(5);
+execute p5(5);
+execute p5(5);
+execute p5(5);
+execute p5(5);
+execute p5(5);
+
+INSERT INTO "companies" ("id","meta","name","created_at","updated_at","deleted_at") VALUES (1,'{"test":123}','Name','2016-11-07 17:34:22.101807','2021-05-20 22:16:55.424521',NULL) ON CONFLICT (id) DO UPDATE SET "meta"=EXCLUDED."meta"  RETURNING "id";
+INSERT INTO "companies" ("id","meta","name","created_at","updated_at","deleted_at") VALUES (1,'{"test":123}','Name','2016-11-07 17:34:22.101807','2021-05-20 22:16:55.424521',NULL) ON CONFLICT (id) DO UPDATE SET "meta"=EXCLUDED."meta"  RETURNING "id";
+PREPARE p6 AS INSERT INTO "companies" ("id","meta","name","created_at","updated_at","deleted_at") VALUES (1,'{"test":123}','Name','2016-11-07 17:34:22.101807','2021-05-20 22:16:55.424521',NULL) ON CONFLICT (id) DO UPDATE SET "meta"=EXCLUDED."meta"  RETURNING "id";;
+EXECUTE p6;
+EXECUTE p6;
+EXECUTE p6;
+EXECUTE p6;
+EXECUTE p6;
+EXECUTE p6;
+EXECUTE p6;
+
+prepare insert_select(int) as insert into companies SELECT * FROM companies WHERE id >= $1 ON CONFLICT(id) DO UPDATE SET "meta"=EXCLUDED."meta" RETURNING "id";;
+EXECUTE insert_select(1);
+EXECUTE insert_select(1);
+EXECUTE insert_select(1);
+EXECUTE insert_select(1);
+EXECUTE insert_select(1);
+EXECUTE insert_select(1);
+EXECUTE insert_select(1);
+
+prepare insert_select_1 as insert into companies SELECT * FROM companies WHERE id >= 1 ON CONFLICT(id) DO UPDATE SET "meta"=EXCLUDED."meta" RETURNING "id";;
+EXECUTE insert_select_1;
+EXECUTE insert_select_1;
+EXECUTE insert_select_1;
+EXECUTE insert_select_1;
+EXECUTE insert_select_1;
+EXECUTE insert_select_1;
+EXECUTE insert_select_1;
+
 -- query fails on the shards should be handled
 -- nicely
 \set VERBOSITY terse
@@ -311,3 +358,4 @@ SELECT count(*), t1.event FROM date_part_table t1 JOIN date_part_table USING (us
 SELECT count(*), event FROM date_part_table WHERE event_time > '2020-01-05' GROUP BY event ORDER BY count(*) DESC, event DESC LIMIT 5;
 SELECT count(*), event FROM date_part_table WHERE user_id = 12 AND event_time = '2020-01-12 12:00:00' GROUP BY event ORDER BY count(*) DESC, event DESC LIMIT 5;
 SELECT count(*), t1.event FROM date_part_table t1 JOIN date_part_table t2 USING (user_id) WHERE t1.user_id = 1 AND t2.event_time > '2020-01-03' GROUP BY t1.event ORDER BY count(*) DESC, t1.event DESC LIMIT 5;
+

src/test/regress/sql/distributed_planning_create_load.sql

@@ -53,3 +53,32 @@ SELECT create_distributed_table('upsert_test', 'part_key');
 
 -- do a regular insert
 INSERT INTO upsert_test (part_key, other_col) VALUES (1, 1), (2, 2) RETURNING *;
+
+create table t1(a int, b int, c int primary key, d int);
+ALTER TABLE t1 DROP COLUMN b;
+select create_distributed_table('t1','c');
+ALTER TABLE t1 DROP COLUMN d;
+
+CREATE TABLE companies (
+    id bigint NOT NULL,
+    name character varying NOT NULL,
+    created_at timestamp without time zone NOT NULL,
+    a int,
+    b int,
+    updated_at timestamp without time zone NOT NULL,
+    fmi boolean DEFAULT false NOT NULL,
+    meta jsonb DEFAULT '{}'::jsonb NOT NULL,
+    deleted_at timestamp without time zone,
+    c int,
+    flex boolean DEFAULT false
+);
+
+ALTER TABLE ONLY companies
+    ADD CONSTRAINT companies_pkey PRIMARY KEY (id);
+
+alter table companies drop column a;
+alter table companies drop column b;
+
+SELECT create_reference_table('companies');
+
+alter table companies drop column c;