Adds additional alter database propagation support (#7253)

DESCRIPTION: Adds database connection limit, rename and set tablespace propagation In this PR, below statement propagations are added alter database <database_name> with allow_connections = <boolean_value>; alter database <database_name> rename to <database_name2>; alter database <database_name> set TABLESPACE <table_space_name> --------- Co-authored-by: Jelte Fennema-Nio <[email protected]> Co-authored-by: Jelte Fennema-Nio <[email protected]> Co-authored-by: Onur Tirtir <[email protected]>
citusdata · Dec 26, 2023 · 181b8ab · 181b8ab
1 parent b877d60
commit 181b8ab
Show file tree

Hide file tree

Showing 9 changed files with 323 additions and 87 deletions.
diff --git a/src/backend/distributed/commands/database.c b/src/backend/distributed/commands/database.c
@@ -192,6 +192,25 @@ PreprocessGrantOnDatabaseStmt(Node *node, const char *queryString,
 }
 
 
+/*
+ * IsSetTablespaceStatement returns true if the statement is a SET TABLESPACE statement,
+ * false otherwise.
+ */
+static bool
+IsSetTablespaceStatement(AlterDatabaseStmt *stmt)
+{
+	DefElem *def = NULL;
+	foreach_ptr(def, stmt->options)
+	{
+		if (strcmp(def->defname, "tablespace") == 0)
+		{
+			return true;
+		}
+	}
+	return false;
+}
+
+
 /*
  * PreprocessAlterDatabaseStmt is executed before the statement is applied to the local
  * postgres instance.
@@ -203,22 +222,38 @@ List *
 PreprocessAlterDatabaseStmt(Node *node, const char *queryString,
 							ProcessUtilityContext processUtilityContext)
 {
-	if (!ShouldPropagate())
+	bool missingOk = false;
+	AlterDatabaseStmt *stmt = castNode(AlterDatabaseStmt, node);
+	ObjectAddress *dbAddress = GetDatabaseAddressFromDatabaseName(stmt->dbname,
+																  missingOk);
+
+	if (!ShouldPropagate() || !IsAnyObjectDistributed(list_make1(dbAddress)))
 	{
 		return NIL;
 	}
 
-	AlterDatabaseStmt *stmt = castNode(AlterDatabaseStmt, node);
-
 	EnsureCoordinator();
 
 	char *sql = DeparseTreeNode((Node *) stmt);
 
 	List *commands = list_make3(DISABLE_DDL_PROPAGATION,
-								(void *) sql,
+								sql,
 								ENABLE_DDL_PROPAGATION);
 
-	return NodeDDLTaskList(NON_COORDINATOR_NODES, commands);
+	if (IsSetTablespaceStatement(stmt))
+	{
+		/*
+		 * Set tablespace does not work inside a transaction.Therefore, we need to use
+		 * NontransactionalNodeDDLTask to run the command on the workers outside
+		 * the transaction block.
+		 */
+
+		return NontransactionalNodeDDLTaskList(NON_COORDINATOR_NODES, commands);
+	}
+	else
+	{
+		return NodeDDLTaskList(NON_COORDINATOR_NODES, commands);
+	}
 }
 
 
@@ -256,6 +291,36 @@ PreprocessAlterDatabaseRefreshCollStmt(Node *node, const char *queryString,
 
 #endif
 
+/*
+ * PreprocessAlterDatabaseRenameStmt is executed before the statement is applied to the local
+ * postgres instance. In this stage we prepare ALTER DATABASE RENAME statement to be run on
+ * all workers.
+ */
+List *
+PostprocessAlterDatabaseRenameStmt(Node *node, const char *queryString)
+{
+	bool missingOk = false;
+	RenameStmt *stmt = castNode(RenameStmt, node);
+	ObjectAddress *dbAddress = GetDatabaseAddressFromDatabaseName(stmt->newname,
+																  missingOk);
+
+	if (!ShouldPropagate() || !IsAnyObjectDistributed(list_make1(dbAddress)))
+	{
+		return NIL;
+	}
+
+	EnsureCoordinator();
+
+	char *sql = DeparseTreeNode((Node *) stmt);
+
+	List *commands = list_make3(DISABLE_DDL_PROPAGATION,
+								(void *) sql,
+								ENABLE_DDL_PROPAGATION);
+
+	return NodeDDLTaskList(NON_COORDINATOR_NODES, commands);
+}
+
+
 /*
  * PreprocessAlterDatabaseSetStmt is executed before the statement is applied to the local
  * postgres instance.

diff --git a/src/backend/distributed/commands/distribute_object_ops.c b/src/backend/distributed/commands/distribute_object_ops.c
@@ -522,6 +522,16 @@ static DistributeObjectOps Database_Set = {
 	.markDistributed = false,
 };
 
+static DistributeObjectOps Database_Rename = {
+	.deparse = DeparseAlterDatabaseRenameStmt,
+	.qualify = NULL,
+	.preprocess = NULL,
+	.postprocess = PostprocessAlterDatabaseRenameStmt,
+	.objectType = OBJECT_DATABASE,
+	.operationType = DIST_OPS_ALTER,
+	.address = NULL,
+	.markDistributed = false,
+};
 
 static DistributeObjectOps Domain_Alter = {
 	.deparse = DeparseAlterDomainStmt,
@@ -2087,6 +2097,11 @@ GetDistributeObjectOps(Node *node)
 					return &Collation_Rename;
 				}
 
+				case OBJECT_DATABASE:
+				{
+					return &Database_Rename;
+				}
+
 				case OBJECT_DOMAIN:
 				{
 					return &Domain_Rename;

diff --git a/src/backend/distributed/deparser/deparse_database_stmts.c b/src/backend/distributed/deparser/deparse_database_stmts.c
@@ -30,12 +30,14 @@
 static void AppendAlterDatabaseOwnerStmt(StringInfo buf, AlterOwnerStmt *stmt);
 static void AppendAlterDatabaseSetStmt(StringInfo buf, AlterDatabaseSetStmt *stmt);
 static void AppendAlterDatabaseStmt(StringInfo buf, AlterDatabaseStmt *stmt);
-static void AppendDefElemConnLimit(StringInfo buf, DefElem *def);
 static void AppendCreateDatabaseStmt(StringInfo buf, CreatedbStmt *stmt);
 static void AppendDropDatabaseStmt(StringInfo buf, DropdbStmt *stmt);
 static void AppendGrantOnDatabaseStmt(StringInfo buf, GrantStmt *stmt);
+static void AppendBasicAlterDatabaseOptions(StringInfo buf, AlterDatabaseStmt *stmt);
+static void AppendGrantDatabases(StringInfo buf, GrantStmt *stmt);
+static void AppendAlterDatabaseSetTablespace(StringInfo buf, DefElem *def, char *dbname);
 
-const DefElemOptionFormat create_database_option_formats[] = {
+const DefElemOptionFormat createDatabaseOptionFormats[] = {
 	{ "owner", " OWNER %s", OPTION_FORMAT_STRING },
 	{ "template", " TEMPLATE %s", OPTION_FORMAT_STRING },
 	{ "encoding", " ENCODING %s", OPTION_FORMAT_LITERAL_CSTR },
@@ -53,6 +55,14 @@ const DefElemOptionFormat create_database_option_formats[] = {
 	{ "is_template", " IS_TEMPLATE %s", OPTION_FORMAT_BOOLEAN }
 };
 
+
+const DefElemOptionFormat alterDatabaseOptionFormats[] = {
+	{ "is_template", " IS_TEMPLATE %s", OPTION_FORMAT_BOOLEAN },
+	{ "allow_connections", " ALLOW_CONNECTIONS %s", OPTION_FORMAT_BOOLEAN },
+	{ "connection_limit", " CONNECTION LIMIT %d", OPTION_FORMAT_INTEGER },
+};
+
+
 char *
 DeparseAlterDatabaseOwnerStmt(Node *node)
 {
@@ -111,52 +121,67 @@ AppendGrantOnDatabaseStmt(StringInfo buf, GrantStmt *stmt)
 }
 
 
-static void
-AppendDefElemConnLimit(StringInfo buf, DefElem *def)
-{
-	appendStringInfo(buf, " CONNECTION LIMIT %ld", (long int) defGetNumeric(def));
-}
-
-
 static void
 AppendAlterDatabaseStmt(StringInfo buf, AlterDatabaseStmt *stmt)
 {
-	appendStringInfo(buf, "ALTER DATABASE %s ", quote_identifier(stmt->dbname));
+	if (list_length(stmt->options) == 0)
+	{
+		elog(ERROR, "got unexpected number of options for ALTER DATABASE");
+	}
 
 	if (stmt->options)
 	{
-		ListCell *cell = NULL;
-		appendStringInfo(buf, "WITH ");
-		foreach(cell, stmt->options)
+		DefElem *firstOption = linitial(stmt->options);
+		if (strcmp(firstOption->defname, "tablespace") == 0)
 		{
-			DefElem *def = castNode(DefElem, lfirst(cell));
-			if (strcmp(def->defname, "is_template") == 0)
-			{
-				appendStringInfo(buf, "IS_TEMPLATE %s",
-								 quote_literal_cstr(strVal(def->arg)));
-			}
-			else if (strcmp(def->defname, "connection_limit") == 0)
-			{
-				AppendDefElemConnLimit(buf, def);
-			}
-			else if (strcmp(def->defname, "allow_connections") == 0)
-			{
-				ereport(ERROR,
-						errmsg("ALLOW_CONNECTIONS is not supported"));
-			}
-			else
-			{
-				ereport(ERROR,
-						errmsg("unrecognized ALTER DATABASE option: %s",
-							   def->defname));
-			}
+			AppendAlterDatabaseSetTablespace(buf, firstOption, stmt->dbname);
+
+			/* SET tablespace cannot be combined with other options */
+			return;
 		}
+
+
+		appendStringInfo(buf, "ALTER DATABASE %s WITH",
+						 quote_identifier(stmt->dbname));
+
+		AppendBasicAlterDatabaseOptions(buf, stmt);
 	}
 
 	appendStringInfo(buf, ";");
 }
 
 
+static void
+AppendAlterDatabaseSetTablespace(StringInfo buf, DefElem *def, char *dbname)
+{
+	appendStringInfo(buf,
+					 "ALTER DATABASE %s SET TABLESPACE %s",
+					 quote_identifier(dbname), quote_identifier(defGetString(def)));
+}
+
+
+/*
+ * AppendBasicAlterDatabaseOptions appends basic ALTER DATABASE options to a string buffer.
+ * Basic options are those that can be appended to the ALTER DATABASE statement
+ * after the "WITH" keyword.(i.e. ALLOW_CONNECTIONS, CONNECTION LIMIT, IS_TEMPLATE)
+ * For example, the tablespace option is not a basic option since it is defined via SET keyword.
+ *
+ * This function takes a string buffer and an AlterDatabaseStmt as input.
+ * It appends the basic options to the string buffer.
+ *
+ */
+static void
+AppendBasicAlterDatabaseOptions(StringInfo buf, AlterDatabaseStmt *stmt)
+{
+	DefElem *def = NULL;
+	foreach_ptr(def, stmt->options)
+	{
+		DefElemOptionToStatement(buf, def, alterDatabaseOptionFormats, lengthof(
+									 alterDatabaseOptionFormats));
+	}
+}
+
+
 char *
 DeparseGrantOnDatabaseStmt(Node *node)
 {
@@ -216,6 +241,22 @@ AppendAlterDatabaseSetStmt(StringInfo buf, AlterDatabaseSetStmt *stmt)
 }
 
 
+char *
+DeparseAlterDatabaseRenameStmt(Node *node)
+{
+	RenameStmt *stmt = (RenameStmt *) node;
+
+	StringInfoData str;
+	initStringInfo(&str);
+
+	appendStringInfo(&str, "ALTER DATABASE %s RENAME TO %s",
+					 quote_identifier(stmt->subname),
+					 quote_identifier(stmt->newname));
+
+	return str.data;
+}
+
+
 char *
 DeparseAlterDatabaseSetStmt(Node *node)
 {
@@ -246,8 +287,8 @@ AppendCreateDatabaseStmt(StringInfo buf, CreatedbStmt *stmt)
 	DefElem *option = NULL;
 	foreach_ptr(option, stmt->options)
 	{
-		DefElemOptionToStatement(buf, option, create_database_option_formats,
-								 lengthof(create_database_option_formats));
+		DefElemOptionToStatement(buf, option, createDatabaseOptionFormats,
+								 lengthof(createDatabaseOptionFormats));
 	}
 }
 

diff --git a/src/include/distributed/commands.h b/src/include/distributed/commands.h
@@ -244,6 +244,8 @@ extern List * DropDatabaseStmtObjectAddress(Node *node, bool missingOk,
 											bool isPostprocess);
 extern List * CreateDatabaseStmtObjectAddress(Node *node, bool missingOk,
 											  bool isPostprocess);
+extern List * GenerateGrantDatabaseCommandList(void);
+extern List * PostprocessAlterDatabaseRenameStmt(Node *node, const char *queryString);
 extern void EnsureSupportedCreateDatabaseCommand(CreatedbStmt *stmt);
 extern char * CreateDatabaseDDLCommand(Oid dbId);
 

diff --git a/src/include/distributed/deparser.h b/src/include/distributed/deparser.h
@@ -251,6 +251,7 @@ extern char * DeparseAlterDatabaseRefreshCollStmt(Node *node);
 extern char * DeparseAlterDatabaseSetStmt(Node *node);
 extern char * DeparseCreateDatabaseStmt(Node *node);
 extern char * DeparseDropDatabaseStmt(Node *node);
+extern char * DeparseAlterDatabaseRenameStmt(Node *node);
 
 
 /* forward declaration for deparse_publication_stmts.c */