Address comments

Signed-off-by: Harsh Lunagariya <[email protected]>

contrib/babelfishpg_tsql/src/pl_handler.c

@@ -3285,27 +3285,22 @@ bbf_ProcessUtility(PlannedStmt *pstmt,
 								{
 									RoleSpec   *rolspec = lfirst(item);
 									char	   *user_name;
-									char       *db_principal;
+									const char *db_principal_type = drop_user ? "user" : "role";
 									const char *db_owner_name;
 									int	       role_oid;
 									int        rolename_len;
 
 									user_name = get_physical_user_name(db_name, rolspec->rolename, false);
-									db_owner_name = get_db_owner_name(get_cur_db_name());
+									db_owner_name = get_db_owner_name(db_name);
 									role_oid = get_role_oid(user_name, true);
 									rolename_len = strlen(rolspec->rolename);
 
-									if (drop_user)
-										db_principal = "user";
-									else
-										db_principal = "role";
-
 									/* If user is dbo or role is db_owner, restrict dropping */
 									if ((drop_user && rolename_len == 3 && strncmp(rolspec->rolename, "dbo", 3) == 0) ||
 										(drop_role && rolename_len == 8 && strncmp(rolspec->rolename, "db_owner", 8) == 0))
 										ereport(ERROR,
 												(errcode(ERRCODE_CHECK_VIOLATION),
-												errmsg("Cannot drop the %s '%s'.", db_principal, rolspec->rolename)));
+												 errmsg("Cannot drop the %s '%s'.", db_principal_type, rolspec->rolename)));
 
 									/* 
 									 * Check for current_user's privileges 
@@ -3315,7 +3310,7 @@ bbf_ProcessUtility(PlannedStmt *pstmt,
 										!is_member_of_role(GetUserId(), get_role_oid(db_owner_name, false)))
 										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, rolspec->rolename)));
+												 errmsg("Cannot drop the %s '%s', because it does not exist or you do not have permission.", db_principal_type, rolspec->rolename)));
 
 									/*
 									 * If a role has members, do not drop it.

test/JDBC/expected/restrict_drop_user_role.out

@@ -11,6 +11,15 @@ go
 create user no_priv_user2 for login no_priv_login2
 go
 
+create login no_priv_login3 with password = '12345678'
+go
+
+create user no_priv_user3 for login no_priv_login3
+go
+
+create role no_priv_role3;
+go
+
 create database restrict_user_db1
 go
 
@@ -123,7 +132,7 @@ go
 ~~ERROR (Message: Cannot drop the role 'fake_role', because it does not exist or you do not have permission.)~~
 
 
-drop user fake_user
+drop user fake_user;
 go
 ~~ERROR (Code: 33557097)~~
 
@@ -193,6 +202,30 @@ go
 drop role dont_drop_role
 go
 
+-- 3.3 - when drop user/role has IF EXISTS clause
+-- Should not throw error for the case where IF EXISTS clause is included
+drop role if exists fake_role;
+go
+
+drop user if exists fake_user;
+go
+
+-- both of the following statements should be executed successfully since user/role exists
+drop user if exists no_priv_user3;
+go
+
+drop user if exists no_priv_role3;
+go
+
+select count(*) from sys.database_principals where name like '%no_priv_%3';
+go
+~~START~~
+int
+0
+~~END~~
+
+
+
 -- psql
 -- Need to terminate active session before cleaning up the login
 SELECT pg_terminate_backend(pid) FROM pg_stat_get_activity(NULL)
@@ -224,5 +257,8 @@ t
 drop login no_priv_login2
 go
 
+drop login no_priv_login3
+go
+
 drop database restrict_user_db1
 go

test/JDBC/input/restrict_drop_user_role.mix

@@ -11,6 +11,15 @@ go
 create user no_priv_user2 for login no_priv_login2
 go
 
+create login no_priv_login3 with password = '12345678'
+go
+
+create user no_priv_user3 for login no_priv_login3
+go
+
+create role no_priv_role3;
+go
+
 create database restrict_user_db1
 go
 
@@ -83,7 +92,7 @@ go
 drop role fake_role
 go
 
-drop user fake_user
+drop user fake_user;
 go
 
 drop user db
@@ -125,6 +134,25 @@ go
 drop role dont_drop_role
 go
 
+-- 3.3 - when drop user/role has IF EXISTS clause
+-- Should not throw error for the case where IF EXISTS clause is included
+drop role if exists fake_role;
+go
+
+drop user if exists fake_user;
+go
+
+-- both of the following statements should be executed successfully since user/role exists
+drop user if exists no_priv_user3;
+go
+
+drop user if exists no_priv_role3;
+go
+
+select count(*) from sys.database_principals where name like '%no_priv_%3';
+go
+
+
 -- psql
 -- Need to terminate active session before cleaning up the login
 SELECT pg_terminate_backend(pid) FROM pg_stat_get_activity(NULL)
@@ -146,5 +174,8 @@ go
 drop login no_priv_login2
 go
 
+drop login no_priv_login3
+go
+
 drop database restrict_user_db1
 go