Fix DROP CONSTRAINT in command string with other commands (#7012)

Co-authored-by: Marco Slot <marco.slot@gmail.com>

src/backend/distributed/commands/table.c

@@ -1382,6 +1382,16 @@ PreprocessAlterTableStmt(Node *node, const char *alterTableCommand,
 															   constraintName, missingOk);
 				rightRelationId = GetReferencedTableId(foreignKeyId);
 			}
+
+			/*
+			 * We support deparsing for DROP CONSTRAINT, but currently deparsing is only
+			 * possible if all subcommands are supported.
+			 */
+			if (list_length(commandList) == 1 &&
+				alterTableStatement->objtype == OBJECT_TABLE)
+			{
+				deparseAT = true;
+			}
 		}
 		else if (alterTableType == AT_AddColumn)
 		{
@@ -1589,11 +1599,10 @@ PreprocessAlterTableStmt(Node *node, const char *alterTableCommand,
 	DDLJob *ddlJob = palloc0(sizeof(DDLJob));
 	ObjectAddressSet(ddlJob->targetObjectAddress, RelationRelationId, leftRelationId);
 
-	const char *sqlForTaskList = alterTableCommand;
 	if (deparseAT)
 	{
 		newStmt->cmds = list_make1(newCmd);
-		sqlForTaskList = DeparseTreeNode((Node *) newStmt);
+		alterTableCommand = DeparseTreeNode((Node *) newStmt);
 	}
 
 	ddlJob->metadataSyncCommand = useInitialDDLCommandString ? alterTableCommand : NULL;
@@ -1609,13 +1618,13 @@ PreprocessAlterTableStmt(Node *node, const char *alterTableCommand,
 		{
 			/* if foreign key or attaching partition index related, use specialized task list function ... */
 			ddlJob->taskList = InterShardDDLTaskList(leftRelationId, rightRelationId,
-													 sqlForTaskList);
+													 alterTableCommand);
 		}
 	}
 	else
 	{
 		/* ... otherwise use standard DDL task list function */
-		ddlJob->taskList = DDLTaskList(leftRelationId, sqlForTaskList);
+		ddlJob->taskList = DDLTaskList(leftRelationId, alterTableCommand);
 		if (!propagateCommandToWorkers)
 		{
 			ddlJob->taskList = NIL;

src/backend/distributed/deparser/deparse_table_stmts.c

@@ -31,6 +31,8 @@ static void AppendAlterTableStmt(StringInfo buf, AlterTableStmt *stmt);
 static void AppendAlterTableCmd(StringInfo buf, AlterTableCmd *alterTableCmd,
 								AlterTableStmt *stmt);
 static void AppendAlterTableCmdAddColumn(StringInfo buf, AlterTableCmd *alterTableCmd);
+static void AppendAlterTableCmdDropConstraint(StringInfo buf,
+											  AlterTableCmd *alterTableCmd);
 
 char *
 DeparseAlterTableSchemaStmt(Node *node)
@@ -411,6 +413,12 @@ AppendAlterTableCmd(StringInfo buf, AlterTableCmd *alterTableCmd, AlterTableStmt
 			break;
 		}
 
+		case AT_DropConstraint:
+		{
+			AppendAlterTableCmdDropConstraint(buf, alterTableCmd);
+			break;
+		}
+
 		case AT_AddConstraint:
 		{
 			Constraint *constraint = (Constraint *) alterTableCmd->def;
@@ -488,3 +496,27 @@ AppendAlterTableCmdAddColumn(StringInfo buf, AlterTableCmd *alterTableCmd)
 		appendStringInfo(buf, " COLLATE %s", identifier);
 	}
 }
+
+
+/*
+ * AppendAlterTableCmdDropConstraint builds and appends to the given buffer an
+ * AT_DropConstraint command from given AlterTableCmd object in the form
+ * DROP CONSTRAINT ...
+ */
+static void
+AppendAlterTableCmdDropConstraint(StringInfo buf, AlterTableCmd *alterTableCmd)
+{
+	appendStringInfoString(buf, " DROP CONSTRAINT");
+
+	if (alterTableCmd->missing_ok)
+	{
+		appendStringInfoString(buf, " IF EXISTS");
+	}
+
+	appendStringInfo(buf, " %s", quote_identifier(alterTableCmd->name));
+
+	if (alterTableCmd->behavior == DROP_CASCADE)
+	{
+		appendStringInfoString(buf, " CASCADE");
+	}
+}

src/test/regress/expected/create_single_shard_table.out

@@ -989,7 +989,24 @@ ALTER TABLE null_key_dist ADD CONSTRAINT fkey_add_test_4 FOREIGN KEY(a)
 ALTER TABLE null_key_dist DROP CONSTRAINT fkey_add_test_1;
 ALTER TABLE null_key_dist DROP CONSTRAINT fkey_add_test_2;
 ALTER TABLE null_key_dist DROP CONSTRAINT fkey_add_test_3;
+-- mimic a conjoined django command
+CREATE FUNCTION drop_fkey_4()
+  RETURNS void
+  LANGUAGE plpgsql
+AS $function$
+BEGIN
+  EXECUTE $$
+    SET CONSTRAINTS fkey_add_test_4 IMMEDIATE;
     ALTER TABLE null_key_dist DROP CONSTRAINT fkey_add_test_4;
+  $$;
+END
+$function$;
+SELECT drop_fkey_4();
+ drop_fkey_4
+---------------------------------------------------------------------
+
+(1 row)
+
 ALTER TABLE "NULL_!_dist_key"."nullKeyTable.1!?!9012345678901234567890123456789012345678901234567890123456789" DROP CONSTRAINT fkey_to_dummy_dist;
 -- create a view that depends on the single-shard table
 CREATE VIEW public.v1 AS SELECT * FROM null_key_dist;

src/test/regress/sql/create_single_shard_table.sql

@@ -677,7 +677,21 @@ ALTER TABLE null_key_dist ADD CONSTRAINT fkey_add_test_4 FOREIGN KEY(a)
 ALTER TABLE null_key_dist DROP CONSTRAINT fkey_add_test_1;
 ALTER TABLE null_key_dist DROP CONSTRAINT fkey_add_test_2;
 ALTER TABLE null_key_dist DROP CONSTRAINT fkey_add_test_3;
+
+-- mimic a conjoined django command
+CREATE FUNCTION drop_fkey_4()
+  RETURNS void
+  LANGUAGE plpgsql
+AS $function$
+BEGIN
+  EXECUTE $$
+    SET CONSTRAINTS fkey_add_test_4 IMMEDIATE;
     ALTER TABLE null_key_dist DROP CONSTRAINT fkey_add_test_4;
+  $$;
+END
+$function$;
+SELECT drop_fkey_4();
+
 ALTER TABLE "NULL_!_dist_key"."nullKeyTable.1!?!9012345678901234567890123456789012345678901234567890123456789" DROP CONSTRAINT fkey_to_dummy_dist;
 
 -- create a view that depends on the single-shard table