Correctly dump permissions for a Babelfish logical database (#244)

Previously, we were not dumping sysadmin membership for a database user during database level dump due to which users other than sysadmin were unable to drop the restored database. This commit fixes the issue by correctly dumping role memberships for all the database roles. Additionally, refactored the code a little bit in dump_babel_utils.c Task: BABEL-4516 Signed-off-by: Rishabh Tanwar [email protected]
babelfish-for-postgresql · Nov 2, 2023 · 4f6179b · 4f6179b
1 parent ed8d6d1
commit 4f6179b
Show file tree

Hide file tree

Showing 2 changed files with 30 additions and 25 deletions.
diff --git a/src/bin/pg_dump/dump_babel_utils.c b/src/bin/pg_dump/dump_babel_utils.c
@@ -39,7 +39,7 @@ static char *babel_init_user = NULL;
 
 static char *getMinOid(Archive *fout);
 static bool isBabelfishConfigTable(TableInfo *tbinfo);
-static void addFromClauseForLogicalDatabaseDump(PQExpBuffer buf, TableInfo *tbinfo, bool is_builtin_db);
+static void addFromClauseForLogicalDatabaseDump(PQExpBuffer buf, TableInfo *tbinfo);
 static void addFromClauseForPhysicalDatabaseDump(PQExpBuffer buf, TableInfo *tbinfo);
 static int getMbstrlen(const char *mbstr,Archive *fout);
 static bool is_ms_shipped(DumpableObject *dobj, Archive *fout);
@@ -1019,26 +1019,27 @@ setBabelfishDependenciesForLogicalDatabaseDump(Archive *fout)
  * corresponding to specified logical database.
  */
 void
-addFromClauseForLogicalDatabaseDump(PQExpBuffer buf, TableInfo *tbinfo, bool is_builtin_db)
+addFromClauseForLogicalDatabaseDump(PQExpBuffer buf, TableInfo *tbinfo)
 {
 	if (strcmp(tbinfo->dobj.name, "babelfish_sysdatabases") == 0)
 	{
-		appendPQExpBuffer(buf, " FROM ONLY %s a WHERE a.dbid = %d",
-						  fmtQualifiedDumpable(tbinfo), bbf_db_id);
 		/*
-		 * builtin db will already be present in the target server so
-		 * no need to dump catalog entry for it.
+		 * Dump database catalog entry for specified logical database unless
+		 * it's a builtin database (dbid 1, 2, 3 or 4), in which case the db
+		 * will already be present in the target server so no need to dump
+		 * catalog entry for it.
 		 */
-		if (is_builtin_db)
-			appendPQExpBufferStr(buf, " LIMIT 0");
+		appendPQExpBuffer(buf, " FROM ONLY %s a WHERE a.dbid = %d AND a.dbid > 4",
+						  fmtQualifiedDumpable(tbinfo), bbf_db_id);
 	}
 	else if (strcmp(tbinfo->dobj.name, "babelfish_namespace_ext") == 0)
 	{
-		appendPQExpBuffer(buf, " FROM ONLY %s a WHERE a.dbid = %d",
+		appendPQExpBuffer(buf, " FROM ONLY %s a WHERE a.dbid = %d "
+						  "AND a.nspname NOT IN "
+						  "('master_dbo', 'master_guest', "
+						  "'msdb_dbo', 'msdb_guest', "
+						  "'tempdb_dbo', 'tempdb_guest') ",
 						  fmtQualifiedDumpable(tbinfo), bbf_db_id);
-		if (is_builtin_db)
-			appendPQExpBuffer(buf, " AND a.nspname NOT IN ('%s_dbo', '%s_guest')",
-							  escaped_bbf_db_name, escaped_bbf_db_name);
 	}
 	else if (strcmp(tbinfo->dobj.name, "babelfish_view_def") == 0 ||
 			 strcmp(tbinfo->dobj.name, "babelfish_extended_properties") == 0)
@@ -1211,7 +1212,7 @@ fixCursorForBbfCatalogTableData(Archive *fout, TableInfo *tbinfo, PQExpBuffer bu
 	if (bbf_db_name == NULL)
 		addFromClauseForPhysicalDatabaseDump(buf, tbinfo);
 	else
-		addFromClauseForLogicalDatabaseDump(buf, tbinfo, is_builtin_db);
+		addFromClauseForLogicalDatabaseDump(buf, tbinfo);
 }
 
 /*
@@ -1310,7 +1311,7 @@ fixCopyCommand(Archive *fout, PQExpBuffer copyBuf, TableInfo *tbinfo, bool isFro
 		if (bbf_db_name == NULL)
 			addFromClauseForPhysicalDatabaseDump(copyBuf, tbinfo);
 		else
-			addFromClauseForLogicalDatabaseDump(copyBuf, tbinfo, is_builtin_db);
+			addFromClauseForLogicalDatabaseDump(copyBuf, tbinfo);
 		appendPQExpBufferStr(copyBuf, ") TO stdout;");
 	}
 	destroyPQExpBuffer(q);

diff --git a/src/bin/pg_dump/dumpall_babel_utils.c b/src/bin/pg_dump/dumpall_babel_utils.c
@@ -274,7 +274,7 @@ getBabelfishRoleMembershipQuery(PGconn *conn, PQExpBuffer buf,
 
 	resetPQExpBuffer(buf);
 	appendPQExpBufferStr(buf, "WITH bbf_catalog AS (");
-	/* Include logins only in case of Babelfish physical database dump. */
+	/* Include all the logins only in case of Babelfish physical database dump. */
 	if (!bbf_db_name)
 	{
 		char *babel_init_user = getBabelfishInitUser(conn);
@@ -285,6 +285,10 @@ getBabelfishRoleMembershipQuery(PGconn *conn, PQExpBuffer buf,
 						  babel_init_user);
 		pfree(babel_init_user);
 	}
+	/* Just include sysadmin role memberships in case of Babelfish logical database dump. */
+	else
+		appendPQExpBufferStr(buf,
+							 "SELECT 'sysadmin' AS rolname UNION ");
 	appendPQExpBuffer(buf,
 					  "SELECT rolname FROM sys.babelfish_authid_user_ext ");
 	/* Only dump users of the specific logical database we are currently dumping. */
@@ -304,14 +308,14 @@ getBabelfishRoleMembershipQuery(PGconn *conn, PQExpBuffer buf,
 					  "bbf_roles AS (SELECT rc.* FROM %s rc INNER JOIN bbf_catalog bcat "
 					  "ON rc.rolname = bcat.rolname) ", role_catalog);
 
-	appendPQExpBuffer(buf, "SELECT ur.rolname AS roleid, "
-					  "um.rolname AS member, "
-					  "a.admin_option, "
-					  "ug.rolname AS grantor "
-					  "FROM pg_auth_members a "
-					  "LEFT JOIN %s ur on ur.oid = a.roleid "
-					  "INNER JOIN bbf_roles um on um.oid = a.member "
-					  "LEFT JOIN %s ug on ug.oid = a.grantor "
-					  "WHERE NOT (ur.rolname ~ '^pg_' AND um.rolname ~ '^pg_')"
-					  "ORDER BY 1,2,3", role_catalog, role_catalog);
+	appendPQExpBufferStr(buf, "SELECT ur.rolname AS roleid, "
+						 "um.rolname AS member, "
+						 "a.admin_option, "
+						 "ug.rolname AS grantor "
+						 "FROM pg_auth_members a "
+						 "INNER JOIN bbf_roles ur on ur.oid = a.roleid "
+						 "INNER JOIN bbf_roles um on um.oid = a.member "
+						 "LEFT JOIN bbf_roles ug on ug.oid = a.grantor "
+						 "WHERE NOT (ur.rolname ~ '^pg_' AND um.rolname ~ '^pg_') "
+						 "ORDER BY 1,2,3");
 }