diff --git a/src/backend/distributed/planner/deparse_shard_query.c b/src/backend/distributed/planner/deparse_shard_query.c index c83170eac..43a493a89 100644 --- a/src/backend/distributed/planner/deparse_shard_query.c +++ b/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); diff --git a/src/backend/distributed/planner/local_plan_cache.c b/src/backend/distributed/planner/local_plan_cache.c index f0c30a01e..9bb9e51f1 100644 --- a/src/backend/distributed/planner/local_plan_cache.c +++ b/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 diff --git a/src/include/distributed/deparse_shard_query.h b/src/include/distributed/deparse_shard_query.h index bdd1eb600..9370e51e2 100644 --- a/src/include/distributed/deparse_shard_query.h +++ b/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 */ diff --git a/src/test/regress/expected/distributed_planning.out b/src/test/regress/expected/distributed_planning.out index 060ba0773..42dac188d 100644 --- a/src/test/regress/expected/distributed_planning.out +++ b/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 diff --git a/src/test/regress/expected/distributed_planning_create_load.out b/src/test/regress/expected/distributed_planning_create_load.out index b83735e71..adfed245e 100644 --- a/src/test/regress/expected/distributed_planning_create_load.out +++ b/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; diff --git a/src/test/regress/sql/distributed_planning.sql b/src/test/regress/sql/distributed_planning.sql index 4de7097f5..b388ad390 100644 --- a/src/test/regress/sql/distributed_planning.sql +++ b/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; + diff --git a/src/test/regress/sql/distributed_planning_create_load.sql b/src/test/regress/sql/distributed_planning_create_load.sql index bed1374e8..d1530cc34 100644 --- a/src/test/regress/sql/distributed_planning_create_load.sql +++ b/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;