Fix tuple routing to initialize batching only for inserts

A cross-partition update on a partitioned table is implemented as a
delete followed by an insert. With foreign partitions, this was however
causing issues, because the FDW and core may disagree on when to enable
batching.  postgres_fdw was only allowing batching for plain inserts
(CMD_INSERT) while core was trying to batch the insert component of the
cross-partition update.  Fix by restricting core to apply batching only
to plain CMD_INSERT queries.

It's possible to allow batching for cross-partition updates, but that
will require more extensive changes, so better to leave that for a
separate patch.

Author: Amit Langote
Reviewed-by: Tomas Vondra, Takayuki Tsunakawa
Discussion: https://postgr.es/m/20200628151002.7x5laxwpgvkyiu3q@development

diff --git a/contrib/postgres_fdw/expected/postgres_fdw.out b/contrib/postgres_fdw/expected/postgres_fdw.out
index 60c7e115d612c1fd2c1134290d572161a2f99727..3326f1b542affdc25186d08a7bdd0d145fad3d2b 100644 (file)
--- a/contrib/postgres_fdw/expected/postgres_fdw.out
+++ b/contrib/postgres_fdw/expected/postgres_fdw.out
@@ -9414,5 +9414,26 @@ SELECT COUNT(*) FROM batch_table;
     66
 (1 row)
 
+-- Check that enabling batched inserts doesn't interfere with cross-partition
+-- updates
+CREATE TABLE batch_cp_upd_test (a int) PARTITION BY LIST (a);
+CREATE TABLE batch_cp_upd_test1 (LIKE batch_cp_upd_test);
+CREATE FOREIGN TABLE batch_cp_upd_test1_f
+   PARTITION OF batch_cp_upd_test
+   FOR VALUES IN (1)
+   SERVER loopback
+   OPTIONS (table_name 'batch_cp_upd_test1', batch_size '10');
+CREATE TABLE batch_cp_up_test1 PARTITION OF batch_cp_upd_test
+   FOR VALUES IN (2);
+INSERT INTO batch_cp_upd_test VALUES (1), (2);
+-- The following moves a row from the local partition to the foreign one
+UPDATE batch_cp_upd_test t SET a = 1 FROM (VALUES (1), (2)) s(a) WHERE t.a = s.a;
+SELECT tableoid::regclass, * FROM batch_cp_upd_test;
+       tableoid       | a 
+----------------------+---
+ batch_cp_upd_test1_f | 1
+ batch_cp_upd_test1_f | 1
+(2 rows)
+
 -- Clean up
-DROP TABLE batch_table CASCADE;
+DROP TABLE batch_table, batch_cp_upd_test CASCADE;

diff --git a/contrib/postgres_fdw/postgres_fdw.c b/contrib/postgres_fdw/postgres_fdw.c
index 368997d9d1f71cc1f48b02e411991f3ba51f37c0..35b48575c59353c930c844830d69c424d715eb1b 100644 (file)
--- a/contrib/postgres_fdw/postgres_fdw.c
+++ b/contrib/postgres_fdw/postgres_fdw.c
@@ -1934,17 +1934,26 @@ static int
 postgresGetForeignModifyBatchSize(ResultRelInfo *resultRelInfo)
 {
    int batch_size;
+   PgFdwModifyState *fmstate = resultRelInfo->ri_FdwState ?
+                           (PgFdwModifyState *) resultRelInfo->ri_FdwState :
+                           NULL;
 
    /* should be called only once */
    Assert(resultRelInfo->ri_BatchSize == 0);
 
+   /*
+    * Should never get called when the insert is being performed as part of
+    * a row movement operation.
+    */
+   Assert(fmstate == NULL || fmstate->aux_fmstate == NULL);
+
    /*
     * In EXPLAIN without ANALYZE, ri_fdwstate is NULL, so we have to lookup
     * the option directly in server/table options. Otherwise just use the
     * value we determined earlier.
     */
-   if (resultRelInfo->ri_FdwState)
-       batch_size = ((PgFdwModifyState *) resultRelInfo->ri_FdwState)->batch_size;
+   if (fmstate)
+       batch_size = fmstate->batch_size;
    else
        batch_size = get_batch_size_option(resultRelInfo->ri_RelationDesc);
 

diff --git a/contrib/postgres_fdw/sql/postgres_fdw.sql b/contrib/postgres_fdw/sql/postgres_fdw.sql
index 151f4f18341eb97290c3380ad07bed5f9b843da1..2b525ea44a88bc5dc432d73a96ce7110a30da1b3 100644 (file)
--- a/contrib/postgres_fdw/sql/postgres_fdw.sql
+++ b/contrib/postgres_fdw/sql/postgres_fdw.sql
@@ -2909,5 +2909,22 @@ CREATE TABLE batch_table_p2
 INSERT INTO batch_table SELECT * FROM generate_series(1, 66) i;
 SELECT COUNT(*) FROM batch_table;
 
+-- Check that enabling batched inserts doesn't interfere with cross-partition
+-- updates
+CREATE TABLE batch_cp_upd_test (a int) PARTITION BY LIST (a);
+CREATE TABLE batch_cp_upd_test1 (LIKE batch_cp_upd_test);
+CREATE FOREIGN TABLE batch_cp_upd_test1_f
+   PARTITION OF batch_cp_upd_test
+   FOR VALUES IN (1)
+   SERVER loopback
+   OPTIONS (table_name 'batch_cp_upd_test1', batch_size '10');
+CREATE TABLE batch_cp_up_test1 PARTITION OF batch_cp_upd_test
+   FOR VALUES IN (2);
+INSERT INTO batch_cp_upd_test VALUES (1), (2);
+
+-- The following moves a row from the local partition to the foreign one
+UPDATE batch_cp_upd_test t SET a = 1 FROM (VALUES (1), (2)) s(a) WHERE t.a = s.a;
+SELECT tableoid::regclass, * FROM batch_cp_upd_test;
+
 -- Clean up
-DROP TABLE batch_table CASCADE;
+DROP TABLE batch_table, batch_cp_upd_test CASCADE;

diff --git a/src/backend/executor/execPartition.c b/src/backend/executor/execPartition.c
index b9e4f2d80b16f85a0f7f93eebb11f5c290ee3235..b8da4c5967d8fc7cd63fb4550d50304f12dc2c43 100644 (file)
--- a/src/backend/executor/execPartition.c
+++ b/src/backend/executor/execPartition.c
@@ -1000,7 +1000,8 @@ ExecInitRoutingInfo(ModifyTableState *mtstate,
     *
     * If the FDW does not support batching, we set the batch size to 1.
     */
-   if (partRelInfo->ri_FdwRoutine != NULL &&
+   if (mtstate->operation == CMD_INSERT &&
+       partRelInfo->ri_FdwRoutine != NULL &&
        partRelInfo->ri_FdwRoutine->GetForeignModifyBatchSize &&
        partRelInfo->ri_FdwRoutine->ExecForeignBatchInsert)
        partRelInfo->ri_BatchSize =