Support adding users to T-SQL db_owner database role (#481)

This commit adds support in Babelfish dump utility to dump the linked roles associated with members of db_owner role that are not present in Babelfish catalogs. Task: BABEL-4899 Signed-off-by: Sharu Goel <[email protected]>
babelfish-for-postgresql · Dec 11, 2024 · dec4c87 · dec4c87
1 parent 93d3921
commit dec4c87
Showing 1 changed file with 64 additions and 0 deletions.
diff --git a/src/bin/pg_dump/dumpall_babel_utils.c b/src/bin/pg_dump/dumpall_babel_utils.c
@@ -241,6 +241,38 @@ getBabelfishRolesQuery(PGconn *conn, PQExpBuffer buf, char *role_catalog,
 		appendPQExpBuffer(buf, "WHERE database_name = '%s' ", escaped_bbf_db_name);
 		pfree(escaped_bbf_db_name);
 	}
+
+	/* Append roles associated with db_owner members which are not present in bbf catalogs */
+	appendPQExpBuffer(buf,
+					  "UNION "
+					  "SELECT pr.rolname "
+					  "FROM ( "
+					  "SELECT DISTINCT pr.rolname AS member_role "
+					  "FROM sys.babelfish_authid_user_ext ta "
+					  "JOIN pg_roles r ON r.rolname = ta.rolname "
+					  "JOIN pg_auth_members m ON r.oid = m.roleid "
+					  "JOIN pg_roles pr ON pr.oid = m.member "
+					  "WHERE pr.rolname != ta.rolname "
+					  "AND ta.orig_username = 'db_owner' ");
+
+	/* Only dump users of the specific logical database we are currently dumping. */
+	if (bbf_db_name != NULL)
+	{
+		/*
+		 * Get escaped bbf_db_name to handle special characters in it.
+		 * 2*strlen+1 bytes are required for PQescapeString according to the documentation.
+		 */
+		char *escaped_bbf_db_name = pg_malloc(2 * strlen(bbf_db_name) + 1);
+
+		PQescapeString(escaped_bbf_db_name, bbf_db_name, strlen(bbf_db_name));
+		appendPQExpBuffer(buf, "AND ta.database_name = '%s' ", escaped_bbf_db_name);
+		pfree(escaped_bbf_db_name);
+	}
+
+	appendPQExpBuffer(buf,
+					  ") AS db_owner_members "
+					  "JOIN pg_roles pr ON pr.rolname = sys.babelfish_truncate_identifier(db_owner_members.member_role || '_obj') ");
+
 	appendPQExpBuffer(buf, "), "
 					  "bbf_roles AS (SELECT rc.* FROM %s rc INNER JOIN bbf_catalog bcat "
 					  "ON rc.rolname = bcat.rolname) ", role_catalog);
@@ -327,6 +359,38 @@ getBabelfishRoleMembershipQuery(PGconn *conn, PQExpBuffer buf,
 		appendPQExpBuffer(buf, "WHERE database_name = '%s' ", escaped_bbf_db_name);
 		pfree(escaped_bbf_db_name);
 	}
+
+		/* Append roles associated with db_owner members which are not present in bbf catalogs */
+	appendPQExpBuffer(buf,
+					  "UNION "
+					  "SELECT pr.rolname "
+					  "FROM ( "
+					  "SELECT DISTINCT pr.rolname AS member_role "
+					  "FROM sys.babelfish_authid_user_ext ta "
+					  "JOIN pg_roles r ON r.rolname = ta.rolname "
+					  "JOIN pg_auth_members m ON r.oid = m.roleid "
+					  "JOIN pg_roles pr ON pr.oid = m.member "
+					  "WHERE pr.rolname != ta.rolname "
+					  "AND ta.orig_username = 'db_owner' ");
+
+	/* Only dump users of the specific logical database we are currently dumping. */
+	if (bbf_db_name != NULL)
+	{
+		/*
+		 * Get escaped bbf_db_name to handle special characters in it.
+		 * 2*strlen+1 bytes are required for PQescapeString according to the documentation.
+		 */
+		char *escaped_bbf_db_name = pg_malloc(2 * strlen(bbf_db_name) + 1);
+
+		PQescapeString(escaped_bbf_db_name, bbf_db_name, strlen(bbf_db_name));
+		appendPQExpBuffer(buf, "AND ta.database_name = '%s' ", escaped_bbf_db_name);
+		pfree(escaped_bbf_db_name);
+	}
+
+	appendPQExpBuffer(buf,
+					  ") AS db_owner_members "
+					  "JOIN pg_roles pr ON pr.rolname = sys.babelfish_truncate_identifier(db_owner_members.member_role || '_obj') ");
+
 	appendPQExpBuffer(buf, "), "
 					  "bbf_roles AS (SELECT rc.* FROM %s rc INNER JOIN bbf_catalog bcat "
 					  "ON rc.rolname = bcat.rolname) ", role_catalog);