Address comment

amazon-aurora · Aug 20, 2024 · 54bdf5b · 54bdf5b
1 parent 8d84b3b
commit 54bdf5b
Show file tree

Hide file tree

Showing 3 changed files with 110 additions and 7 deletions.
diff --git a/contrib/babelfishpg_tsql/src/pl_handler.c b/contrib/babelfishpg_tsql/src/pl_handler.c
@@ -2813,6 +2813,7 @@ bbf_ProcessUtility(PlannedStmt *pstmt,
 									int			role_oid;
 									int			rolename_len;
 									bool		is_tsql_db_principal = false;
+									bool		is_psql_db_principal = false;
 									Oid			dbowner;
 
 									user_name = get_physical_user_name(db_name, rolspec->rolename);
@@ -2824,6 +2825,7 @@ bbf_ProcessUtility(PlannedStmt *pstmt,
 									is_tsql_db_principal = OidIsValid(role_oid) &&
 														   ((drop_user && is_user(role_oid)) ||
 															(drop_role && is_role(role_oid)));
+									is_psql_db_principal = OidIsValid(role_oid) && !is_tsql_db_principal;
 
 									/* If user is dbo or role is db_owner, restrict dropping */
 									if ((drop_user && rolename_len == 3 && strncmp(rolspec->rolename, "dbo", 3) == 0) ||
@@ -2838,7 +2840,7 @@ bbf_ProcessUtility(PlannedStmt *pstmt,
 									 */
 									if ((!stmt->missing_ok && !is_tsql_db_principal) ||
 										!is_member_of_role(GetUserId(), dbowner) ||
-										(is_tsql_db_principal && !is_member_of_role(dbowner, role_oid)))
+										(is_tsql_db_principal && !is_member_of_role(dbowner, role_oid)) || is_psql_db_principal)
 										ereport(ERROR,
 												(errcode(ERRCODE_INSUFFICIENT_PRIVILEGE),
 												 errmsg("Cannot drop the %s '%s', because it does not exist or you do not have permission.", db_principal_type, rolspec->rolename)));

diff --git a/test/JDBC/expected/restrict_drop_user_role.out b/test/JDBC/expected/restrict_drop_user_role.out
@@ -211,11 +211,25 @@ drop user if exists fake_user;
 go
 
 -- both of the following statements should be executed successfully since user/role exists
+drop role if exists no_priv_user3;
+go
+~~ERROR (Code: 33557097)~~
+
+~~ERROR (Message: Cannot drop the role 'no_priv_user3', because it does not exist or you do not have permission.)~~
+
+
 drop user if exists no_priv_user3;
 go
 
 drop user if exists no_priv_role3;
 go
+~~ERROR (Code: 33557097)~~
+
+~~ERROR (Message: Cannot drop the user 'no_priv_role3', because it does not exist or you do not have permission.)~~
+
+
+drop role if exists no_priv_role3;
+go
 
 select count(*) from sys.database_principals where name like '%no_priv_%3';
 go
@@ -251,14 +265,14 @@ drop user if exists pguser1;
 go
 ~~ERROR (Code: 33557097)~~
 
-~~ERROR (Message: permission denied to drop role)~~
+~~ERROR (Message: Cannot drop the user 'pguser1', because it does not exist or you do not have permission.)~~
 
 
 drop role if exists pguser1;
 go
 ~~ERROR (Code: 33557097)~~
 
-~~ERROR (Message: permission denied to drop role)~~
+~~ERROR (Message: Cannot drop the role 'pguser1', because it does not exist or you do not have permission.)~~
 
 
 -- psql
@@ -301,18 +315,43 @@ go
 drop database restrict_user_db1
 go
 
--- tsql
+-- psql
 -- 3.5 - Try dropping user created for another database
+ALTER SYSTEM SET babelfishpg_tsql.migration_mode = 'multi-db';
+GO
+SELECT pg_reload_conf();
+GO
+~~START~~
+bool
+t
+~~END~~
+
+-- Wait to sync with another session
+SELECT pg_sleep(1);
+GO
+~~START~~
+void
+
+~~END~~
+
+
+-- tsql
 create database restrict_user_db1;
 go
 
+create database restrict_user_db1_restrict_user_db1;
+go
+
 create login restrict_user_l1 with password = '12345678'
 go
 
+use restrict_user_db1
+go
+
 create user restrict_user_db1_user1 for login restrict_user_l1;
 go
 
-use restrict_user_db1
+use restrict_user_db1_restrict_user_db1
 go
 
 drop user user1;
@@ -324,6 +363,10 @@ go
 
 drop user if exists user1;
 go
+~~ERROR (Code: 33557097)~~
+
+~~ERROR (Message: Cannot drop the user 'user1', because it does not exist or you do not have permission.)~~
+
 
 use master
 go
@@ -339,11 +382,32 @@ bool
 
 
 -- tsql
+use restrict_user_db1
+go
+
 drop user restrict_user_db1_user1;
 go
 
+use master
+go
+
 drop login restrict_user_l1;
 go
 
 drop database restrict_user_db1;
 go
+
+drop database restrict_user_db1_restrict_user_db1;
+go
+
+-- psql
+ALTER SYSTEM SET babelfishpg_tsql.migration_mode = 'single-db';
+GO
+-- Wait to sync with another session
+SELECT pg_sleep(1);
+GO
+~~START~~
+void
+
+~~END~~
+
diff --git a/test/JDBC/input/restrict_drop_user_role.mix b/test/JDBC/input/restrict_drop_user_role.mix
@@ -143,12 +143,18 @@ drop user if exists fake_user;
 go
 
 -- both of the following statements should be executed successfully since user/role exists
+drop role if exists no_priv_user3;
+go
+
 drop user if exists no_priv_user3;
 go
 
 drop user if exists no_priv_role3;
 go
 
+drop role if exists no_priv_role3;
+go
+
 select count(*) from sys.database_principals where name like '%no_priv_%3';
 go
 
@@ -203,17 +209,32 @@ drop database restrict_user_db1
 go
 
 -- 3.5 - Try dropping user created for another database
+-- psql
+ALTER SYSTEM SET babelfishpg_tsql.migration_mode = 'multi-db';
+GO
+SELECT pg_reload_conf();
+GO
+-- Wait to sync with another session
+SELECT pg_sleep(1);
+GO
+
 -- tsql
 create database restrict_user_db1;
 go
 
+create database restrict_user_db1_restrict_user_db1;
+go
+
 create login restrict_user_l1 with password = '12345678'
 go
 
+use restrict_user_db1
+go
+
 create user restrict_user_db1_user1 for login restrict_user_l1;
 go
 
-use restrict_user_db1
+use restrict_user_db1_restrict_user_db1
 go
 
 drop user user1;
@@ -232,11 +253,27 @@ WHERE sys.suser_name(usesysid) = 'restrict_user_l1' AND backend_type = 'client b
 go
 
 -- tsql
+use restrict_user_db1
+go
+
 drop user restrict_user_db1_user1;
 go
 
+use master
+go
+
 drop login restrict_user_l1;
 go
 
 drop database restrict_user_db1;
-go
+go
+
+drop database restrict_user_db1_restrict_user_db1;
+go
+
+-- psql
+ALTER SYSTEM SET babelfishpg_tsql.migration_mode = 'single-db';
+GO
+-- Wait to sync with another session
+SELECT pg_sleep(1);
+GO