Merge branch 'master' into fix_command_counter_increment

README.md

@@ -34,10 +34,18 @@ To learn more, visit [citusdata.com](https://www.citusdata.com) and join
 the [mailing list](https://groups.google.com/forum/#!forum/citus-users) to
 stay on top of the latest developments.
 
-### Quickstart
+### Getting started with Citus
+
+The fastest way to get up and running is to create a Citus Cloud account. You can also setup a local Citus cluster with Docker.
+
+#### Citus Cloud
+
+Citus Cloud runs on top of AWS as a fully managed database as a service and has development plans available for getting started. You can provision a Citus Cloud account at [https://console.citusdata.com](https://console.citusdata.com) and get started with just a few clicks.
 
 #### Local Citus Cluster
 
+If you're looking to get started locally, you can follow the following steps to get up and running.
+
 * Install docker-compose: [Mac][mac_install] | [Linux][linux_install]
 * (Mac only) connect to Docker VM
   ```bash
@@ -100,7 +108,7 @@ stay on top of the latest developments.
 <tr>
   <td>Training and Support</td>
   <td>See our <a
-  href="https://www.citusdata.com/citus-products/citus-data-pricing">support
+  href="https://www.citusdata.com/support">support
   page</a> for training and dedicated support options.</td>
 </tr>
 </table>
@@ -140,7 +148,7 @@ Video](https://www.youtube.com/watch?v=NVl9_6J1G60&list=PLixnExCn6lRpP10ZlpJwx6A
 
 ___
 
-Copyright © 2012–2016 Citus Data, Inc.
+Copyright © 2012–2017 Citus Data, Inc.
 
 [faq]: https://www.citusdata.com/frequently-asked-questions
 [linux_install]: https://www.digitalocean.com/community/tutorials/how-to-install-and-use-docker-compose-on-ubuntu-14-04

src/backend/distributed/master/master_node_protocol.c

@@ -64,7 +64,7 @@
 
 /* Shard related configuration */
 int ShardCount = 32;
-int ShardReplicationFactor = 2; /* desired replication factor for shards */
+int ShardReplicationFactor = 1; /* desired replication factor for shards */
 int ShardMaxSize = 1048576;     /* maximum size in KB one shard can grow to */
 int ShardPlacementPolicy = SHARD_PLACEMENT_ROUND_ROBIN;
 

src/backend/distributed/planner/multi_explain.c

@@ -52,8 +52,6 @@
 
 
 /* Config variables that enable printing distributed query plans */
-bool ExplainMultiLogicalPlan = false;
-bool ExplainMultiPhysicalPlan = false;
 bool ExplainDistributedQueries = true;
 bool ExplainAllTasks = false;
 
@@ -79,6 +77,9 @@ static void ExplainTask(Task *task, int placementIndex, List *explainOutputList,
 static void ExplainTaskPlacement(ShardPlacement *taskPlacement, List *explainOutputList,
 								 ExplainState *es);
 static StringInfo BuildRemoteExplainQuery(char *queryString, ExplainState *es);
+static void MultiExplainOnePlan(PlannedStmt *plan, IntoClause *into,
+								ExplainState *es, const char *queryString,
+								ParamListInfo params, const instr_time *planDuration);
 
 /* Static Explain functions copied from explain.c */
 static void ExplainOpenGroup(const char *objtype, const char *labelname,
@@ -100,17 +101,10 @@ void
 MultiExplainOneQuery(Query *query, IntoClause *into, ExplainState *es,
 					 const char *queryString, ParamListInfo params)
 {
-	MultiPlan *multiPlan = NULL;
-	CmdType commandType = CMD_UNKNOWN;
-	PlannedStmt *initialPlan = NULL;
-	Job *workerJob = NULL;
-	bool routerExecutablePlan = false;
 	instr_time planStart;
 	instr_time planDuration;
-	Query *originalQuery = NULL;
-	RelationRestrictionContext *restrictionContext = NULL;
-	bool localQuery = !NeedsDistributedPlanning(query);
 	int cursorOptions = 0;
+	PlannedStmt *plan = NULL;
 
 #if PG_VERSION_NUM >= 90600
 
@@ -124,85 +118,54 @@ MultiExplainOneQuery(Query *query, IntoClause *into, ExplainState *es,
 	}
 #endif
 
-	/* handle local queries in the same way as ExplainOneQuery */
+	/* plan query, just like ExplainOneQuery does */
-	if (localQuery)
-	{
-		PlannedStmt *plan = NULL;
-
-		INSTR_TIME_SET_CURRENT(planStart);
-
-		/* plan the query */
-		plan = pg_plan_query(query, cursorOptions, params);
-
-		INSTR_TIME_SET_CURRENT(planDuration);
-		INSTR_TIME_SUBTRACT(planDuration, planStart);
-
-		/* run it (if needed) and produce output */
-		ExplainOnePlan(plan, into, es, queryString, params, &planDuration);
-
-		return;
-	}
-
-	/*
-	 * standard_planner scribbles on it's input, but for deparsing we need the
-	 * unmodified form. So copy once we're sure it's a distributed query.
-	 */
-	originalQuery = copyObject(query);
-
-	/* measure the full planning time to display in EXPLAIN ANALYZE */
 	INSTR_TIME_SET_CURRENT(planStart);
 
-	restrictionContext = CreateAndPushRestrictionContext();
+	/* plan the query */
-
+	plan = pg_plan_query(query, cursorOptions, params);
-	PG_TRY();
-	{
-		/* call standard planner to modify the query structure before multi planning */
-		initialPlan = standard_planner(query, cursorOptions, params);
-
-		commandType = initialPlan->commandType;
-		if (commandType == CMD_INSERT || commandType == CMD_UPDATE ||
-			commandType == CMD_DELETE)
-		{
-			if (es->analyze)
-			{
-				ereport(ERROR, (errcode(ERRCODE_FEATURE_NOT_SUPPORTED),
-								errmsg("Using ANALYZE for INSERT/UPDATE/DELETE on "
-									   "distributed tables is not supported.")));
-			}
-		}
-
-		multiPlan = CreatePhysicalPlan(originalQuery, query, restrictionContext);
-	}
-	PG_CATCH();
-	{
-		PopRestrictionContext();
-		PG_RE_THROW();
-	}
-	PG_END_TRY();
-
-	PopRestrictionContext();
 
 	INSTR_TIME_SET_CURRENT(planDuration);
 	INSTR_TIME_SUBTRACT(planDuration, planStart);
 
-	if (ExplainMultiLogicalPlan)
+	/* if not a distributed query, use plain explain infrastructure */
+	if (!HasCitusToplevelNode(plan))
 	{
-		MultiTreeRoot *multiTree = MultiLogicalPlanCreate(query);
+		/* run it (if needed) and produce output */
-		char *logicalPlanString = CitusNodeToString(multiTree);
+		ExplainOnePlan(plan, into, es, queryString, params, &planDuration);
-		char *formattedPlanString = pretty_format_node_dump(logicalPlanString);
+	}
+	else
+	{
+		MultiExplainOnePlan(plan, into, es, queryString, params, &planDuration);
+	}
+}
 
-		appendStringInfo(es->str, "logical plan:\n");
+
-		appendStringInfo(es->str, "%s\n", formattedPlanString);
+/*
+ * MultiExplainOnePlan explains the plan for an individual distributed query.
+ */
+static void
+MultiExplainOnePlan(PlannedStmt *plan, IntoClause *into,
+					ExplainState *es, const char *queryString,
+					ParamListInfo params, const instr_time *planDuration)
+{
+	MultiPlan *multiPlan = NULL;
+	CmdType commandType = CMD_UNKNOWN;
+	Job *workerJob = NULL;
+	bool routerExecutablePlan = false;
+
+	commandType = plan->commandType;
+	if (commandType == CMD_INSERT || commandType == CMD_UPDATE ||
+		commandType == CMD_DELETE)
+	{
+		if (es->analyze)
+		{
+			ereport(ERROR, (errcode(ERRCODE_FEATURE_NOT_SUPPORTED),
+							errmsg("Using ANALYZE for INSERT/UPDATE/DELETE on "
+								   "distributed tables is not supported.")));
+		}
 	}
 
-	if (ExplainMultiPhysicalPlan)
+	multiPlan = GetMultiPlan(plan);
-	{
-		char *physicalPlanString = CitusNodeToString(multiPlan);
-		char *formattedPlanString = pretty_format_node_dump(physicalPlanString);
-
-		appendStringInfo(es->str, "physical plan:\n");
-		appendStringInfo(es->str, "%s\n", formattedPlanString);
-	}
 
 	if (!ExplainDistributedQueries)
 	{
@@ -268,8 +231,6 @@ MultiExplainOneQuery(Query *query, IntoClause *into, ExplainState *es,
 
 	if (!routerExecutablePlan)
 	{
-		PlannedStmt *masterPlan = MultiQueryContainerNode(initialPlan, multiPlan);
-
 		if (es->format == EXPLAIN_FORMAT_TEXT)
 		{
 			appendStringInfoSpaces(es->str, es->indent * 2);
@@ -279,7 +240,7 @@ MultiExplainOneQuery(Query *query, IntoClause *into, ExplainState *es,
 
 		ExplainOpenGroup("Master Query", "Master Query", false, es);
 
-		ExplainMasterPlan(masterPlan, into, es, queryString, params, &planDuration);
+		ExplainMasterPlan(plan, into, es, queryString, params, planDuration);
 
 		ExplainCloseGroup("Master Query", "Master Query", false, es);
 

src/backend/distributed/planner/multi_planner.c

@@ -36,10 +36,15 @@ static List *relationRestrictionContextList = NIL;
 
 /* local function forward declarations */
 static void CheckNodeIsDumpable(Node *node);
-
-
-/* local function forward declarations */
 static char * GetMultiPlanString(PlannedStmt *result);
+static PlannedStmt * MultiQueryContainerNode(PlannedStmt *result,
+											 struct MultiPlan *multiPlan);
+static struct MultiPlan * CreatePhysicalPlan(Query *originalQuery, Query *query,
+											 RelationRestrictionContext *
+											 restrictionContext);
+static RelationRestrictionContext * CreateAndPushRestrictionContext(void);
+static RelationRestrictionContext * CurrentRestrictionContext(void);
+static void PopRestrictionContext(void);
 
 
 /* Distributed planner hook */
@@ -123,7 +128,7 @@ multi_planner(Query *parse, int cursorOptions, ParamListInfo boundParams)
  * target shards. SELECT queries go through the full logical plan/optimize/
  * physical plan process needed to produce distributed query plans.
  */
-MultiPlan *
+static MultiPlan *
 CreatePhysicalPlan(Query *originalQuery, Query *query,
 				   RelationRestrictionContext *restrictionContext)
 {

src/backend/distributed/shared_library_init.c

@@ -295,30 +295,6 @@ RegisterCitusConfigVariables(void)
 		0,
 		NULL, NULL, NULL);
 
-	DefineCustomBoolVariable(
-		"citus.explain_multi_logical_plan",
-		gettext_noop("Enables Explain to print out distributed logical plans."),
-		gettext_noop("We use this private configuration entry as a debugging aid. "
-					 "If enabled, the Explain command prints out the optimized "
-					 "logical plan for distributed queries."),
-		&ExplainMultiLogicalPlan,
-		false,
-		PGC_USERSET,
-		GUC_NO_SHOW_ALL,
-		NULL, NULL, NULL);
-
-	DefineCustomBoolVariable(
-		"citus.explain_multi_physical_plan",
-		gettext_noop("Enables Explain to print out distributed physical plans."),
-		gettext_noop("We use this private configuration entry as a debugging aid. "
-					 "If enabled, the Explain command prints out the physical "
-					 "plan for distributed queries."),
-		&ExplainMultiPhysicalPlan,
-		false,
-		PGC_USERSET,
-		GUC_NO_SHOW_ALL,
-		NULL, NULL, NULL);
-
 	DefineCustomBoolVariable(
 		"citus.explain_distributed_queries",
 		gettext_noop("Enables Explain for distributed queries."),
@@ -393,7 +369,7 @@ RegisterCitusConfigVariables(void)
 					 "configuration value at sharded table creation time, "
 					 "and later reuse the initially read value."),
 		&ShardReplicationFactor,
-		2, 1, 100,
+		1, 1, 100,
 		PGC_USERSET,
 		0,
 		NULL, NULL, NULL);

src/backend/distributed/transaction/multi_shard_transaction.c

@@ -202,12 +202,12 @@ GetShardHashConnections(HTAB *connectionHash, int64 shardId, bool *connectionsFo
 
 
 /*
- * ConnectionList flattens the connection hash to a list of placement connections.
+ * ShardConnectionList returns the list of ShardConnections in connectionHash.
  */
 List *
-ConnectionList(HTAB *connectionHash)
+ShardConnectionList(HTAB *connectionHash)
 {
-	List *connectionList = NIL;
+	List *shardConnectionsList = NIL;
 	HASH_SEQ_STATUS status;
 	ShardConnections *shardConnections = NULL;
 
@@ -221,13 +221,12 @@ ConnectionList(HTAB *connectionHash)
 	shardConnections = (ShardConnections *) hash_seq_search(&status);
 	while (shardConnections != NULL)
 	{
-		List *shardConnectionsList = list_copy(shardConnections->connectionList);
+		shardConnectionsList = lappend(shardConnectionsList, shardConnections);
-		connectionList = list_concat(connectionList, shardConnectionsList);
 
 		shardConnections = (ShardConnections *) hash_seq_search(&status);
 	}
 
-	return connectionList;
+	return shardConnectionsList;
 }
 
 

src/include/distributed/multi_explain.h

@@ -13,8 +13,6 @@
 #include "executor/executor.h"
 
 /* Config variables managed via guc.c to explain distributed query plans */
-extern bool ExplainMultiLogicalPlan;
-extern bool ExplainMultiPhysicalPlan;
 extern bool ExplainDistributedQueries;
 extern bool ExplainAllTasks;
 

src/include/distributed/multi_planner.h

@@ -53,16 +53,8 @@ extern PlannedStmt * multi_planner(Query *parse, int cursorOptions,
 
 extern bool HasCitusToplevelNode(PlannedStmt *planStatement);
 struct MultiPlan;
-extern struct MultiPlan * CreatePhysicalPlan(Query *originalQuery, Query *query,
-											 RelationRestrictionContext *
-											 restrictionContext);
 extern struct MultiPlan * GetMultiPlan(PlannedStmt *planStatement);
-extern PlannedStmt * MultiQueryContainerNode(PlannedStmt *result,
-											 struct MultiPlan *multiPlan);
 extern void multi_relation_restriction_hook(PlannerInfo *root, RelOptInfo *relOptInfo,
 											Index index, RangeTblEntry *rte);
-extern RelationRestrictionContext * CreateAndPushRestrictionContext(void);
-extern RelationRestrictionContext * CurrentRestrictionContext(void);
-extern void PopRestrictionContext(void);
 
 #endif /* MULTI_PLANNER_H */

src/include/distributed/multi_shard_transaction.h

@@ -22,11 +22,7 @@ typedef struct ShardConnections
 {
 	int64 shardId;
 
-	/*
+	/* list of MultiConnection structs */
-	 * XXX: this list contains MultiConnection for multi-shard transactions
-	 * or TransactionConnection for COPY, the latter should be converted to
-	 * use MultiConnection as well.
-	 */
 	List *connectionList;
 } ShardConnections;
 
@@ -36,7 +32,7 @@ extern HTAB * CreateShardConnectionHash(MemoryContext memoryContext);
 extern ShardConnections * GetShardConnections(int64 shardId, bool *shardConnectionsFound);
 extern ShardConnections * GetShardHashConnections(HTAB *connectionHash, int64 shardId,
 												  bool *connectionsFound);
-extern List * ConnectionList(HTAB *connectionHash);
+extern List * ShardConnectionList(HTAB *connectionHash);
 extern void CloseConnections(List *connectionList);
 extern void ResetShardPlacementTransactionState(void);
 

src/test/regress/expected/multi_insert_select.out

@@ -1456,17 +1456,15 @@ INSERT INTO raw_events_first SELECT * FROM raw_events_second WHERE user_id = 100
 ERROR:  single-shard DML commands must not appear in transaction blocks which contain multi-shard data modifications
 ROLLBACK;
 -- Insert after copy is currently disallowed because of the way the 
--- transaction modification state is currently handled. Copy still
+-- transaction modification state is currently handled. Copy is also
--- goes through despite rollback.
+-- rolled back.
 BEGIN;
 COPY raw_events_second (user_id, value_1) FROM STDIN DELIMITER ',';
 INSERT INTO raw_events_first SELECT * FROM raw_events_second;
 ERROR:  multi-shard data modifications must not appear in transaction blocks which contain single-shard DML commands
 ROLLBACK;
 -- Insert after copy is currently allowed for single-shard operation.
--- Since the COPY commits immediately, the result is visible in the
+-- Both insert and copy are rolled back successfully.
--- next operation. Copy goes through despite rollback, while insert
--- rolls back.
 BEGIN;
 COPY raw_events_second (user_id, value_1) FROM STDIN DELIMITER ',';
 INSERT INTO raw_events_first SELECT * FROM raw_events_second WHERE user_id = 101;
@@ -1477,9 +1475,7 @@ SELECT user_id FROM raw_events_first WHERE user_id = 101;
 (1 row)
 
 ROLLBACK;
--- Copy after insert is disallowed since the insert is not immediately
+-- Copy after insert is currently disallowed.
--- committed and the copy uses different connections that will not yet
--- see the result of the insert.
 BEGIN;
 INSERT INTO raw_events_first SELECT * FROM raw_events_second;
 COPY raw_events_first (user_id, value_1) FROM STDIN DELIMITER ',';
@@ -1489,8 +1485,6 @@ ROLLBACK;
 BEGIN;
 INSERT INTO raw_events_first SELECT * FROM raw_events_second WHERE user_id = 100;
 COPY raw_events_first (user_id, value_1) FROM STDIN DELIMITER ',';
-ERROR:  distributed copy operations must not appear in transaction blocks containing other distributed modifications
-CONTEXT:  COPY raw_events_first, line 1: "103,103"
 ROLLBACK;
 -- selecting from views works
 CREATE VIEW test_view AS SELECT * FROM raw_events_first;
@@ -1499,7 +1493,7 @@ INSERT INTO raw_events_first (user_id, time, value_1, value_2, value_3, value_4)
 SELECT count(*) FROM raw_events_second;
  count 
 -------
-    11
+     9
 (1 row)
 
 INSERT INTO raw_events_second SELECT * FROM test_view;
@@ -1509,7 +1503,7 @@ INSERT INTO raw_events_second SELECT * FROM test_view WHERE user_id = 17 GROUP B
 SELECT count(*) FROM raw_events_second;
  count 
 -------
-    13
+    11
 (1 row)
 
 -- inserting into views does not

src/test/regress/expected/multi_modifying_xacts.out

@@ -192,37 +192,178 @@ SELECT * FROM labs WHERE id = 6;
 ----+------
 (0 rows)
 
--- COPY can't happen second,
+-- COPY can happen after single row INSERT
 BEGIN;
 INSERT INTO labs VALUES (6, 'Bell Labs');
 \copy labs from stdin delimiter ','
-ERROR:  distributed copy operations must not appear in transaction blocks containing other distributed modifications
-CONTEXT:  COPY labs, line 1: "10,Weyland-Yutani"
 COMMIT;
--- though it will work if before any modifications
+-- COPY can happen before single row INSERT
 BEGIN;
 \copy labs from stdin delimiter ','
 SELECT name FROM labs WHERE id = 10;
       name      
 ----------------
  Weyland-Yutani
-(1 row)
+ Weyland-Yutani
+(2 rows)
 
 INSERT INTO labs VALUES (6, 'Bell Labs');
 COMMIT;
--- but a double-copy isn't allowed (the first will persist)
+-- two consecutive COPYs in a transaction are allowed
 BEGIN;
 \copy labs from stdin delimiter ','
 \copy labs from stdin delimiter ','
-ERROR:  distributed copy operations must not appear in transaction blocks containing other distributed modifications
-CONTEXT:  COPY labs, line 1: "12,fsociety"
 COMMIT;
-SELECT name FROM labs WHERE id = 11;
+SELECT name FROM labs WHERE id = 11 OR id = 12 ORDER BY id;
       name      
 ----------------
  Planet Express
+ fsociety
+(2 rows)
+
+-- 1pc failure test
+SELECT recover_prepared_transactions();
+ recover_prepared_transactions 
+-------------------------------
+                             0
 (1 row)
 
+-- copy with unique index violation
+BEGIN;
+\copy researchers FROM STDIN delimiter ','
+\copy researchers FROM STDIN delimiter ','
+ERROR:  duplicate key value violates unique constraint "avoid_name_confusion_idx_1200001"
+DETAIL:  Key (lab_id, name)=(6,  'Bjarne Stroustrup') already exists.
+COMMIT;
+-- verify rollback
+SELECT * FROM researchers WHERE lab_id = 6;
+ id | lab_id | name 
+----+--------+------
+(0 rows)
+
+SELECT count(*) FROM pg_dist_transaction;
+ count 
+-------
+     0
+(1 row)
+
+-- 2pc failure and success tests
+SET citus.multi_shard_commit_protocol TO '2pc';
+SELECT recover_prepared_transactions();
+ recover_prepared_transactions 
+-------------------------------
+                             0
+(1 row)
+
+-- copy with unique index violation
+BEGIN;
+\copy researchers FROM STDIN delimiter ','
+\copy researchers FROM STDIN delimiter ','
+ERROR:  duplicate key value violates unique constraint "avoid_name_confusion_idx_1200001"
+DETAIL:  Key (lab_id, name)=(6,  'Bjarne Stroustrup') already exists.
+COMMIT;
+-- verify rollback
+SELECT * FROM researchers WHERE lab_id = 6;
+ id | lab_id | name 
+----+--------+------
+(0 rows)
+
+SELECT count(*) FROM pg_dist_transaction;
+ count 
+-------
+     0
+(1 row)
+
+BEGIN;
+\copy researchers FROM STDIN delimiter ','
+\copy researchers FROM STDIN delimiter ','
+COMMIT;
+-- verify success
+SELECT * FROM researchers WHERE lab_id = 6;
+ id | lab_id |         name         
+----+--------+----------------------
+ 17 |      6 |  'Bjarne Stroustrup'
+ 18 |      6 |  'Dennis Ritchie'
+(2 rows)
+
+-- verify 2pc
+SELECT count(*) FROM pg_dist_transaction;
+ count 
+-------
+     2
+(1 row)
+
+RESET citus.multi_shard_commit_protocol;
+-- create a check function
+SELECT * from run_command_on_workers('CREATE FUNCTION reject_large_id() RETURNS trigger AS $rli$
+    BEGIN
+        IF (NEW.id > 30) THEN
+            RAISE ''illegal value'';
+        END IF;
+
+        RETURN NEW;
+    END;
+$rli$ LANGUAGE plpgsql;')
+ORDER BY nodeport;
+ nodename  | nodeport | success |     result      
+-----------+----------+---------+-----------------
+ localhost |    57637 | t       | CREATE FUNCTION
+ localhost |    57638 | t       | CREATE FUNCTION
+(2 rows)
+
+-- register after insert trigger
+SELECT * FROM run_command_on_placements('researchers', 'CREATE CONSTRAINT TRIGGER reject_large_researcher_id AFTER INSERT ON %s DEFERRABLE INITIALLY DEFERRED FOR EACH ROW EXECUTE PROCEDURE  reject_large_id()')
+ORDER BY nodeport, shardid;
+ nodename  | nodeport | shardid | success |     result     
+-----------+----------+---------+---------+----------------
+ localhost |    57637 | 1200000 | t       | CREATE TRIGGER
+ localhost |    57637 | 1200001 | t       | CREATE TRIGGER
+ localhost |    57638 | 1200000 | t       | CREATE TRIGGER
+ localhost |    57638 | 1200001 | t       | CREATE TRIGGER
+(4 rows)
+
+-- hide postgresql version dependend messages for next test only
+\set VERBOSITY terse
+-- deferred check should abort the transaction
+BEGIN;
+DELETE FROM researchers WHERE lab_id = 6;
+\copy researchers FROM STDIN delimiter ','
+\copy researchers FROM STDIN delimiter ','
+COMMIT;
+WARNING:  illegal value
+WARNING:  failed to commit transaction on localhost:57638
+WARNING:  illegal value
+WARNING:  failed to commit transaction on localhost:57637
+WARNING:  could not commit transaction for shard 1200001 on any active node
+ERROR:  could not commit transaction on any active node
+\unset VERBOSITY
+-- verify everyhing including delete is rolled back
+SELECT * FROM researchers WHERE lab_id = 6;
+ id | lab_id |         name         
+----+--------+----------------------
+ 17 |      6 |  'Bjarne Stroustrup'
+ 18 |      6 |  'Dennis Ritchie'
+(2 rows)
+
+-- cleanup triggers and the function
+SELECT * from run_command_on_placements('researchers', 'drop trigger reject_large_researcher_id on %s')
+ORDER BY nodeport, shardid;
+ nodename  | nodeport | shardid | success |    result    
+-----------+----------+---------+---------+--------------
+ localhost |    57637 | 1200000 | t       | DROP TRIGGER
+ localhost |    57637 | 1200001 | t       | DROP TRIGGER
+ localhost |    57638 | 1200000 | t       | DROP TRIGGER
+ localhost |    57638 | 1200001 | t       | DROP TRIGGER
+(4 rows)
+
+SELECT * FROM run_command_on_workers('drop function reject_large_id()')
+ORDER BY nodeport;
+ nodename  | nodeport | success |    result     
+-----------+----------+---------+---------------
+ localhost |    57637 | t       | DROP FUNCTION
+ localhost |    57638 | t       | DROP FUNCTION
+(2 rows)
+
 -- finally, ALTER and copy aren't compatible
 BEGIN;
 ALTER TABLE labs ADD COLUMN motto text;
@@ -238,11 +379,6 @@ COMMIT;
  id     | bigint | not null
  name   | text   | not null
 
-SELECT * FROM labs WHERE id = 12;
- id | name 
-----+------
-(0 rows)
-
 -- and if the copy is before the ALTER...
 BEGIN;
 \copy labs from stdin delimiter ','
@@ -269,7 +405,7 @@ ALTER TABLE labs ADD COLUMN motto text;
 SELECT master_modify_multiple_shards('DELETE FROM labs');
  master_modify_multiple_shards 
 -------------------------------
-                             5
+                             7
 (1 row)
 
 ALTER TABLE labs ADD COLUMN score float;

src/test/regress/output/multi_copy.source

@@ -707,3 +707,4 @@ HINT:  No function matches the given name and argument types. You might need to
 CONTEXT:  while executing command on localhost:57638
 WARNING:  could not get statistics for shard public.composite_partition_column_table_560164
 DETAIL:  Setting shard statistics to NULL
+ERROR:  failure on connection marked as essential: localhost:57637

src/test/regress/pg_regress_multi.pl

@@ -219,6 +219,7 @@ push(@pgOptions, '-c', "citus.max_running_tasks_per_node=4");
 push(@pgOptions, '-c', "citus.expire_cached_shards=on");
 push(@pgOptions, '-c', "citus.task_tracker_delay=10ms");
 push(@pgOptions, '-c', "citus.remote_task_check_interval=1ms");
+push(@pgOptions, '-c', "citus.shard_replication_factor=2");
 
 # Add externally added options last, so they overwrite the default ones above
 for my $option (@userPgOptions)

src/test/regress/sql/multi_insert_select.sql

@@ -706,8 +706,8 @@ INSERT INTO raw_events_first SELECT * FROM raw_events_second WHERE user_id = 100
 ROLLBACK;
 
 -- Insert after copy is currently disallowed because of the way the 
--- transaction modification state is currently handled. Copy still
+-- transaction modification state is currently handled. Copy is also
--- goes through despite rollback.
+-- rolled back.
 BEGIN;
 COPY raw_events_second (user_id, value_1) FROM STDIN DELIMITER ',';
 100,100
@@ -716,9 +716,7 @@ INSERT INTO raw_events_first SELECT * FROM raw_events_second;
 ROLLBACK;
 
 -- Insert after copy is currently allowed for single-shard operation.
--- Since the COPY commits immediately, the result is visible in the
+-- Both insert and copy are rolled back successfully.
--- next operation. Copy goes through despite rollback, while insert
--- rolls back.
 BEGIN;
 COPY raw_events_second (user_id, value_1) FROM STDIN DELIMITER ',';
 101,101
@@ -727,9 +725,7 @@ INSERT INTO raw_events_first SELECT * FROM raw_events_second WHERE user_id = 101
 SELECT user_id FROM raw_events_first WHERE user_id = 101;
 ROLLBACK;
 
--- Copy after insert is disallowed since the insert is not immediately
+-- Copy after insert is currently disallowed.
--- committed and the copy uses different connections that will not yet
--- see the result of the insert.
 BEGIN;
 INSERT INTO raw_events_first SELECT * FROM raw_events_second;
 COPY raw_events_first (user_id, value_1) FROM STDIN DELIMITER ',';

src/test/regress/sql/multi_modifying_xacts.sql

@@ -148,7 +148,7 @@ COMMIT;
 \d labs
 SELECT * FROM labs WHERE id = 6;
 
--- COPY can't happen second,
+-- COPY can happen after single row INSERT
 BEGIN;
 INSERT INTO labs VALUES (6, 'Bell Labs');
 \copy labs from stdin delimiter ','
@@ -156,7 +156,7 @@ INSERT INTO labs VALUES (6, 'Bell Labs');
 \.
 COMMIT;
 
--- though it will work if before any modifications
+-- COPY can happen before single row INSERT
 BEGIN;
 \copy labs from stdin delimiter ','
 10,Weyland-Yutani
@@ -165,7 +165,7 @@ SELECT name FROM labs WHERE id = 10;
 INSERT INTO labs VALUES (6, 'Bell Labs');
 COMMIT;
 
--- but a double-copy isn't allowed (the first will persist)
+-- two consecutive COPYs in a transaction are allowed
 BEGIN;
 \copy labs from stdin delimiter ','
 11,Planet Express
@@ -175,7 +175,93 @@ BEGIN;
 \.
 COMMIT;
 
-SELECT name FROM labs WHERE id = 11;
+SELECT name FROM labs WHERE id = 11 OR id = 12 ORDER BY id;
+
+-- 1pc failure test
+SELECT recover_prepared_transactions();
+-- copy with unique index violation
+BEGIN;
+\copy researchers FROM STDIN delimiter ','
+17, 6, 'Bjarne Stroustrup'
+\.
+\copy researchers FROM STDIN delimiter ','
+18, 6, 'Bjarne Stroustrup'
+\.
+COMMIT;
+-- verify rollback
+SELECT * FROM researchers WHERE lab_id = 6;
+SELECT count(*) FROM pg_dist_transaction;
+
+-- 2pc failure and success tests
+SET citus.multi_shard_commit_protocol TO '2pc';
+SELECT recover_prepared_transactions();
+-- copy with unique index violation
+BEGIN;
+\copy researchers FROM STDIN delimiter ','
+17, 6, 'Bjarne Stroustrup'
+\.
+\copy researchers FROM STDIN delimiter ','
+18, 6, 'Bjarne Stroustrup'
+\.
+COMMIT;
+-- verify rollback
+SELECT * FROM researchers WHERE lab_id = 6;
+SELECT count(*) FROM pg_dist_transaction;
+
+BEGIN;
+\copy researchers FROM STDIN delimiter ','
+17, 6, 'Bjarne Stroustrup'
+\.
+\copy researchers FROM STDIN delimiter ','
+18, 6, 'Dennis Ritchie'
+\.
+COMMIT;
+-- verify success
+SELECT * FROM researchers WHERE lab_id = 6;
+-- verify 2pc
+SELECT count(*) FROM pg_dist_transaction;
+
+RESET citus.multi_shard_commit_protocol;
+
+-- create a check function
+SELECT * from run_command_on_workers('CREATE FUNCTION reject_large_id() RETURNS trigger AS $rli$
+    BEGIN
+        IF (NEW.id > 30) THEN
+            RAISE ''illegal value'';
+        END IF;
+
+        RETURN NEW;
+    END;
+$rli$ LANGUAGE plpgsql;')
+ORDER BY nodeport;
+
+-- register after insert trigger
+SELECT * FROM run_command_on_placements('researchers', 'CREATE CONSTRAINT TRIGGER reject_large_researcher_id AFTER INSERT ON %s DEFERRABLE INITIALLY DEFERRED FOR EACH ROW EXECUTE PROCEDURE  reject_large_id()')
+ORDER BY nodeport, shardid;
+
+-- hide postgresql version dependend messages for next test only
+\set VERBOSITY terse
+-- deferred check should abort the transaction
+BEGIN;
+DELETE FROM researchers WHERE lab_id = 6;
+\copy researchers FROM STDIN delimiter ','
+31, 6, 'Bjarne Stroustrup'
+\.
+\copy researchers FROM STDIN delimiter ','
+30, 6, 'Dennis Ritchie'
+\.
+COMMIT;
+\unset VERBOSITY
+
+-- verify everyhing including delete is rolled back
+SELECT * FROM researchers WHERE lab_id = 6;
+
+-- cleanup triggers and the function
+SELECT * from run_command_on_placements('researchers', 'drop trigger reject_large_researcher_id on %s')
+ORDER BY nodeport, shardid;
+
+SELECT * FROM run_command_on_workers('drop function reject_large_id()')
+ORDER BY nodeport;
 
 -- finally, ALTER and copy aren't compatible
 BEGIN;
@@ -187,7 +273,6 @@ COMMIT;
 
 -- but the DDL should correctly roll back
 \d labs
-SELECT * FROM labs WHERE id = 12;
 
 -- and if the copy is before the ALTER...
 BEGIN;