Added more testcases

Signed-off-by: ANJU BHARTI <[email protected]>
amazon-aurora · Aug 16, 2024 · 4330cdd · 4330cdd
1 parent 675f72e
commit 4330cdd
Show file tree

Hide file tree

Showing 3 changed files with 49 additions and 11 deletions.
diff --git a/contrib/babelfishpg_tsql/src/pl_handler.c b/contrib/babelfishpg_tsql/src/pl_handler.c
@@ -3287,27 +3287,25 @@ bbf_ProcessUtility(PlannedStmt *pstmt,
 									char	   *user_name;
 									char       *db_principal;
 									const char *db_owner_name;
-									int	      	role_oid;
-									int       	rolename_len;
-									char	   *logical_role_name = NULL;
+									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());
 									role_oid = get_role_oid(user_name, true);
-									logical_role_name = rolspec->rolename;
-									rolename_len = strlen(logical_role_name);
+									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 && strncmp(logical_role_name, "dbo", rolename_len) == 0) ||
-										(drop_role && strncmp(logical_role_name, "db_owner", rolename_len) == 0))
+									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, logical_role_name)));
+												errmsg("Cannot drop the %s '%s'.", db_principal, rolspec->rolename)));
 
 									/* 
 									 * Check for current_user's privileges 
@@ -3317,7 +3315,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, logical_role_name)));
+												errmsg("Cannot drop the %s '%s', because it does not exist or you do not have permission.", db_principal, rolspec->rolename)));
 
 									/*
 									 * If a role has members, do not drop it.
@@ -3337,7 +3335,7 @@ bbf_ProcessUtility(PlannedStmt *pstmt,
 									 * if enabled. 3. Otherwise throw an
 									 * error.
 									 */
-									if (drop_user && strcmp(logical_role_name, "guest") == 0)
+									if (drop_user && strcmp(rolspec->rolename, "guest") == 0)
 									{
 										if (guest_has_dbaccess(db_name))
 										{
@@ -3346,7 +3344,7 @@ bbf_ProcessUtility(PlannedStmt *pstmt,
 														(errcode(ERRCODE_CHECK_VIOLATION),
 														 errmsg("Cannot disable access to the guest user in master or tempdb.")));
 
-											alter_user_can_connect(false, logical_role_name, db_name);
+											alter_user_can_connect(false, rolspec->rolename, db_name);
 											return;
 										}
 										else

diff --git a/test/JDBC/expected/restrict_drop_user_role.out b/test/JDBC/expected/restrict_drop_user_role.out
@@ -130,6 +130,34 @@ go
 ~~ERROR (Message: Cannot drop the user 'fake_user', because it does not exist or you do not have permission.)~~
 
 
+drop user db
+go
+~~ERROR (Code: 33557097)~~
+
+~~ERROR (Message: Cannot drop the user 'db', because it does not exist or you do not have permission.)~~
+
+
+drop role db_own
+go
+~~ERROR (Code: 33557097)~~
+
+~~ERROR (Message: Cannot drop the role 'db_own', because it does not exist or you do not have permission.)~~
+
+
+drop user dbo_u1
+go
+~~ERROR (Code: 33557097)~~
+
+~~ERROR (Message: Cannot drop the user 'dbo_u1', because it does not exist or you do not have permission.)~~
+
+
+drop role db_owner_r1
+go
+~~ERROR (Code: 33557097)~~
+
+~~ERROR (Message: Cannot drop the role 'db_owner_r1', because it does not exist or you do not have permission.)~~
+
+
 -- should deny
 -- try to drop dbo user, db_owner role
 drop role db_owner

diff --git a/test/JDBC/input/restrict_drop_user_role.mix b/test/JDBC/input/restrict_drop_user_role.mix
@@ -86,6 +86,18 @@ go
 drop user fake_user
 go
 
+drop user db
+go
+
+drop role db_own
+go
+
+drop user dbo_u1
+go
+
+drop role db_owner_r1
+go
+
 -- should deny
 -- try to drop dbo user, db_owner role
 drop role db_owner