From 8d6d27a241f6b525be982777b87e4293424181f4 Mon Sep 17 00:00:00 2001
From: gindibay <gindibay@microsoft.com>
Date: Thu, 12 Oct 2023 18:31:47 +0300
Subject: [PATCH] Adds set tablespace support

---
 src/backend/distributed/commands/database.c   | 30 +++++++-----------
 .../deparser/deparse_database_stmts.c         |  3 +-
 src/include/distributed/metadata_sync.h       |  2 ++
 .../sql/alter_database_propagation.sql        | 31 ++++++++++++++++---
 4 files changed, 42 insertions(+), 24 deletions(-)

diff --git a/src/backend/distributed/commands/database.c b/src/backend/distributed/commands/database.c
index d8c74c623..637e14cbb 100644
--- a/src/backend/distributed/commands/database.c
+++ b/src/backend/distributed/commands/database.c
@@ -182,30 +182,24 @@ PreprocessAlterDatabaseStmt(Node *node, const char *queryString,
 
 	char *sql = DeparseTreeNode((Node *) stmt);
 
+	List *commands = NULL;
+
 	if (strstr(sql, "SET TABLESPACE") != NULL)
 	{
-		if (IsCoordinator())
-		{
-			ereport(NOTICE, (errmsg(
-								 "Citus partially supports ALTER DATABASE SET TABLESPACE for "
-								 "distributed databases"),
-							 errdetail(
-								 "Citus does not propagate ALTER DATABASE SET TABLESPACE "
-								 "command to workers"),
-							 errhint(
-								 "You can manually alter a tablespace for a database and its "
-								 "extensions on workers.")));
-		}
+		commands = list_make5(DISABLE_DDL_PROPAGATION,
+							  COMMIT_TRANSACTION,
+							  sql,
+							  BEGIN_TRANSACTION,
+							  ENABLE_DDL_PROPAGATION);
 	}
 	else
 	{
-		List *commands = list_make3(DISABLE_DDL_PROPAGATION,
-									(void *) sql,
-									ENABLE_DDL_PROPAGATION);
-
-		return NodeDDLTaskList(NON_COORDINATOR_NODES, commands);
+		commands = list_make3(DISABLE_DDL_PROPAGATION,
+							  (void *) sql,
+							  ENABLE_DDL_PROPAGATION);
 	}
-	return NIL;
+
+	return NodeDDLTaskList(NON_COORDINATOR_NODES, commands);
 }
 
 
diff --git a/src/backend/distributed/deparser/deparse_database_stmts.c b/src/backend/distributed/deparser/deparse_database_stmts.c
index c7eebcd63..dfc46ffb4 100644
--- a/src/backend/distributed/deparser/deparse_database_stmts.c
+++ b/src/backend/distributed/deparser/deparse_database_stmts.c
@@ -131,7 +131,8 @@ AppendBasicAlterDatabaseOptions(StringInfo buf, DefElem *def, bool
 static void
 AppendAlterDatabaseSetTablespace(StringInfo buf, DefElem *def, char *dbname)
 {
-	appendStringInfo(buf, "ALTER DATABASE %s SET TABLESPACE %s",
+	appendStringInfo(buf,
+					 "ALTER DATABASE %s SET TABLESPACE %s",
 					 quote_identifier(dbname), quote_identifier(defGetString(def)));
 }
 
diff --git a/src/include/distributed/metadata_sync.h b/src/include/distributed/metadata_sync.h
index 237df363a..a7fac1e7f 100644
--- a/src/include/distributed/metadata_sync.h
+++ b/src/include/distributed/metadata_sync.h
@@ -196,6 +196,8 @@ extern void SendInterTableRelationshipCommands(MetadataSyncContext *context);
 	"SELECT pg_catalog.worker_drop_sequence_dependency(%s);"
 
 #define DISABLE_DDL_PROPAGATION "SET citus.enable_ddl_propagation TO 'off'"
+#define BEGIN_TRANSACTION "BEGIN"
+#define COMMIT_TRANSACTION "COMMIT"
 #define ENABLE_DDL_PROPAGATION "SET citus.enable_ddl_propagation TO 'on'"
 #define DISABLE_METADATA_SYNC "SET citus.enable_metadata_sync TO 'off'"
 #define ENABLE_METADATA_SYNC "SET citus.enable_metadata_sync TO 'on'"
diff --git a/src/test/regress/sql/alter_database_propagation.sql b/src/test/regress/sql/alter_database_propagation.sql
index 5e7b4c977..693ff8aaa 100644
--- a/src/test/regress/sql/alter_database_propagation.sql
+++ b/src/test/regress/sql/alter_database_propagation.sql
@@ -54,18 +54,25 @@ alter database regression2 with CONNECTION LIMIT 100;
 alter database regression2 with IS_TEMPLATE true CONNECTION LIMIT 50;
 alter database regression2 with IS_TEMPLATE false;
 
-alter database regression2 set TABLESPACE pg_default;
+
+\set alter_db_tablespace :abs_srcdir '/tmp_check/ts3'
+CREATE TABLESPACE alter_db_tablespace LOCATION :'alter_db_tablespace';
 
 \c - - - :worker_1_port
-
-alter database regression2 set TABLESPACE pg_default;
+\set alter_db_tablespace :abs_srcdir '/tmp_check/ts4'
+CREATE TABLESPACE alter_db_tablespace LOCATION :'alter_db_tablespace';
 
 \c - - - :worker_2_port
-
-alter database regression2 set TABLESPACE pg_default;
+\set alter_db_tablespace :abs_srcdir '/tmp_check/ts5'
+CREATE TABLESPACE alter_db_tablespace LOCATION :'alter_db_tablespace';
 
 \c - - - :master_port
 
+set citus.log_remote_commands = true;
+set citus.grep_remote_commands = '%ALTER DATABASE%';
+
+alter database regression2 set TABLESPACE alter_db_tablespace;
+
 set citus.enable_create_database_propagation=on;
 alter database regression2 rename to regression3;
 
@@ -73,3 +80,17 @@ drop database regression3;
 
 set citus.log_remote_commands = false;
 set citus.enable_create_database_propagation=off;
+
+drop tablespace alter_db_tablespace;
+
+\c - - - :worker_1_port
+
+drop tablespace alter_db_tablespace;
+
+\c - - - :worker_2_port
+
+drop tablespace alter_db_tablespace;
+
+\c - - - :master_port
+
+