Use placement connection to drop shards instead of node connection

src/backend/distributed/master/master_delete_protocol.c

@@ -340,14 +340,6 @@ DropShards(Oid relationId, char *schemaName, char *relationName,
 	ListCell *shardIntervalCell = NULL;
 	int droppedShardCount = 0;
 
-	if (XactModificationLevel != XACT_MODIFICATION_NONE)
-	{
-		ereport(ERROR, (errcode(ERRCODE_ACTIVE_SQL_TRANSACTION),
-						errmsg("shard drop operations must not appear in "
-							   "transaction blocks containing other distributed "
-							   "modifications")));
-	}
-
 	BeginOrContinueCoordinatedTransaction();
 
 	/* At this point we intentionally decided to not use 2PC for reference tables */
@@ -396,8 +388,8 @@ DropShards(Oid relationId, char *schemaName, char *relationName,
 								 quotedShardName);
 			}
 
-			connection = GetNodeUserDatabaseConnection(connectionFlags, workerName,
+			connection = GetPlacementConnection(connectionFlags, shardPlacement,
-													   workerPort, extensionOwner, NULL);
+												extensionOwner);
 
 			RemoteTransactionBeginIfNecessary(connection);
 

src/test/regress/expected/multi_create_table.out

@@ -438,6 +438,14 @@ WHERE col1 = 132;
 
 DROP TABLE data_load_test1, data_load_test2;
 END;
+-- There should be no table on the worker node
+\c - - - :worker_1_port
+SELECT relname FROM pg_class WHERE relname LIKE 'data_load_test%';
+ relname 
+---------
+(0 rows)
+
+\c - - - :master_port
 -- creating an index after loading data works
 BEGIN;
 CREATE TABLE data_load_test (col1 int, col2 text, col3 serial);
@@ -450,8 +458,10 @@ NOTICE:  Copying data from local table...
 (1 row)
 
 CREATE INDEX data_load_test_idx ON data_load_test (col2);
-END;
+NOTICE:  using one-phase commit for distributed DDL commands
+HINT:  You can enable two-phase commit for extra safety with: SET citus.multi_shard_commit_protocol TO '2pc'
 DROP TABLE data_load_test;
+END;
 -- popping in and out of existence in the same transaction works
 BEGIN;
 CREATE TABLE data_load_test (col1 int, col2 text, col3 serial);
@@ -478,9 +488,6 @@ NOTICE:  Copying data from local table...
 
 INSERT INTO data_load_test VALUES (243, 'world');
 DROP TABLE data_load_test;
-ERROR:  shard drop operations must not appear in transaction blocks containing other distributed modifications
-CONTEXT:  SQL statement "SELECT master_drop_all_shards(v_obj.objid, v_obj.schema_name, v_obj.object_name)"
-PL/pgSQL function citus_drop_trigger() line 21 at PERFORM
 END;
 -- Test data loading after dropping a column
 CREATE TABLE data_load_test (col1 int, col2 text, col3 text);

src/test/regress/expected/multi_transactional_drop_shards.out

@@ -197,9 +197,6 @@ SELECT master_create_empty_shard('transactional_drop_shards');
 BEGIN;
 INSERT INTO transactional_drop_shards VALUES (1);
 DROP TABLE transactional_drop_shards;
-ERROR:  shard drop operations must not appear in transaction blocks containing other distributed modifications
-CONTEXT:  SQL statement "SELECT master_drop_all_shards(v_obj.objid, v_obj.schema_name, v_obj.object_name)"
-PL/pgSQL function citus_drop_trigger() line 21 at PERFORM
 ROLLBACK;
 -- verify metadata is not deleted
 SELECT shardid FROM pg_dist_shard WHERE logicalrelid = 'transactional_drop_shards'::regclass ORDER BY shardid;
@@ -242,7 +239,11 @@ Table "public.transactional_drop_shards_1410005"
 BEGIN;
 INSERT INTO transactional_drop_shards VALUES (1);
 SELECT master_apply_delete_command('DELETE FROM transactional_drop_shards');
-ERROR:  shard drop operations must not appear in transaction blocks containing other distributed modifications
+ master_apply_delete_command 
+-----------------------------
+                           1
+(1 row)
+
 ROLLBACK;
 -- verify metadata is not deleted
 SELECT shardid FROM pg_dist_shard WHERE logicalrelid = 'transactional_drop_shards'::regclass ORDER BY shardid;

src/test/regress/input/multi_alter_table_statements.source

@@ -442,3 +442,16 @@ DROP TABLE lineitem_alter;
 \c - - - :worker_1_port
 SELECT relname FROM pg_class WHERE relname LIKE 'lineitem_alter%';
 \c - - - :master_port
+
+-- Test alter table with drop table in the same transaction
+BEGIN;
+CREATE TABLE test_table_1(id int);
+SELECT create_distributed_table('test_table_1','id');
+ALTER TABLE test_table_1 ADD CONSTRAINT u_key UNIQUE(id);
+DROP TABLE test_table_1;
+END;
+
+-- There should be no test_table_1 shard on workers
+\c - - - :worker_1_port
+SELECT relname FROM pg_class WHERE relname LIKE 'test_table_1%';
+\c - - - :master_port

src/test/regress/input/multi_copy.source

@@ -764,3 +764,19 @@ magnetic disk, test
 SELECT * FROM test_smgr;
 
 DROP TABLE test_smgr;
+
+-- Test drop table with copy in the same transaction
+BEGIN;
+CREATE TABLE tt1(id int);
+SELECT create_distributed_table('tt1','id');
+\copy tt1 from STDIN;
+1
+2
+\.
+DROP TABLE tt1;
+END;
+
+-- There should be no "tt1" shard on the worker nodes
+\c - - - :worker_1_port
+SELECT relname FROM pg_class WHERE relname LIKE 'tt1%';
+\c - - - :master_port

src/test/regress/output/multi_alter_table_statements.source

@@ -873,3 +873,25 @@ SELECT relname FROM pg_class WHERE relname LIKE 'lineitem_alter%';
 (1 row)
 
 \c - - - :master_port
+-- Test alter table with drop table in the same transaction
+BEGIN;
+CREATE TABLE test_table_1(id int);
+SELECT create_distributed_table('test_table_1','id');
+ create_distributed_table 
+--------------------------
+ 
+(1 row)
+
+ALTER TABLE test_table_1 ADD CONSTRAINT u_key UNIQUE(id);
+NOTICE:  using one-phase commit for distributed DDL commands
+HINT:  You can enable two-phase commit for extra safety with: SET citus.multi_shard_commit_protocol TO '2pc'
+DROP TABLE test_table_1;
+END;
+-- There should be no test_table_1 shard on workers
+\c - - - :worker_1_port
+SELECT relname FROM pg_class WHERE relname LIKE 'test_table_1%';
+ relname 
+---------
+(0 rows)
+
+\c - - - :master_port

src/test/regress/output/multi_copy.source

@@ -1035,3 +1035,23 @@ SELECT * FROM test_smgr;
 (1 row)
 
 DROP TABLE test_smgr;
+-- Test drop table with copy in the same transaction
+BEGIN;
+CREATE TABLE tt1(id int);
+SELECT create_distributed_table('tt1','id');
+ create_distributed_table 
+--------------------------
+ 
+(1 row)
+
+\copy tt1 from STDIN;
+DROP TABLE tt1;
+END;
+-- There should be no "tt1" shard on the worker nodes
+\c - - - :worker_1_port
+SELECT relname FROM pg_class WHERE relname LIKE 'tt1%';
+ relname 
+---------
+(0 rows)
+
+\c - - - :master_port

src/test/regress/sql/multi_create_table.sql

@@ -233,14 +233,19 @@ WHERE col1 = 132;
 DROP TABLE data_load_test1, data_load_test2;
 END;
 
+-- There should be no table on the worker node
+\c - - - :worker_1_port
+SELECT relname FROM pg_class WHERE relname LIKE 'data_load_test%';
+\c - - - :master_port
+
 -- creating an index after loading data works
 BEGIN;
 CREATE TABLE data_load_test (col1 int, col2 text, col3 serial);
 INSERT INTO data_load_test VALUES (132, 'hello');
 SELECT create_distributed_table('data_load_test', 'col1');
 CREATE INDEX data_load_test_idx ON data_load_test (col2);
-END;
 DROP TABLE data_load_test;
+END;
 
 -- popping in and out of existence in the same transaction works
 BEGIN;